TESS 2 : Un modèle de langage à diffusion généraliste à grande échelle

TESS 2: A Large-Scale Generalist Diffusion Language Model

February 19, 2025

Auteurs: Jaesung Tae, Hamish Ivison, Sachin Kumar, Arman Cohan

cs.AI

Résumé

Nous présentons TESS 2, un modèle de langage à diffusion généraliste capable de suivre des instructions, qui surpasse les modèles à diffusion contemporains ajustés aux instructions, et rivalise voire dépasse parfois les modèles autoregressifs (AR) performants. Nous entraînons TESS 2 en adaptant d'abord un modèle AR robuste via un pré-entraînement continu avec la perte d'entropie croisée habituelle comme fonction de diffusion, puis en effectuant un ajustement supplémentaire aux instructions. Nous constatons que l'entraînement d'adaptation ainsi que le choix du modèle de base sont cruciaux pour former de bons modèles à diffusion capables de suivre des instructions. Nous proposons en outre le guidage par récompense, une nouvelle procédure modulaire de guidage au moment de l'inférence pour aligner les sorties du modèle sans avoir besoin de réentraîner le modèle sous-jacent. Enfin, nous montrons que TESS 2 s'améliore davantage avec une augmentation des ressources de calcul lors de l'inférence, mettant en avant l'utilité des modèles de langage à diffusion pour offrir un contrôle précis sur la quantité de calcul utilisée au moment de l'inférence. Le code et les modèles sont disponibles à l'adresse https://github.com/hamishivi/tess-2.

English

We introduce TESS 2, a general instruction-following diffusion language model that outperforms contemporary instruction-tuned diffusion models, as well as matches and sometimes exceeds strong autoregressive (AR) models. We train TESS 2 by first adapting a strong AR model via continued pretraining with the usual cross-entropy as diffusion loss, and then performing further instruction tuning. We find that adaptation training as well as the choice of the base model is crucial for training good instruction-following diffusion models. We further propose reward guidance, a novel and modular inference-time guidance procedure to align model outputs without needing to train the underlying model. Finally, we show that TESS 2 further improves with increased inference-time compute, highlighting the utility of diffusion LMs in having fine-grained controllability over the amount of compute used at inference time. Code and models are available at https://github.com/hamishivi/tess-2.