TESS 2: Um Modelo de Linguagem de Difusão Generalista em Grande Escala

TESS 2: A Large-Scale Generalist Diffusion Language Model

February 19, 2025

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

cs.AI

Resumo

Apresentamos o TESS 2, um modelo de linguagem de difusão geral que segue instruções e supera os modelos contemporâneos de difusão ajustados para instruções, além de igualar e, em alguns casos, exceder modelos autoregressivos (AR) robustos. Treinamos o TESS 2 primeiro adaptando um modelo AR robusto por meio de pré-treinamento contínuo com a entropia cruzada usual como função de perda de difusão, e depois realizando um ajuste adicional para instruções. Descobrimos que o treinamento de adaptação, assim como a escolha do modelo base, é crucial para treinar bons modelos de difusão que seguem instruções. Além disso, propomos a orientação por recompensa, um procedimento modular e inovador de orientação durante a inferência para alinhar as saídas do modelo sem a necessidade de treinar o modelo subjacente. Por fim, demonstramos que o TESS 2 melhora ainda mais com o aumento do poder computacional durante a inferência, destacando a utilidade dos modelos de linguagem de difusão em permitir um controle refinado sobre a quantidade de computação utilizada no momento da inferência. O código e os modelos estão disponíveis em 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.