TESS 2: Un Modello Linguistico Diffusivo Generalista su Larga Scala

TESS 2: A Large-Scale Generalist Diffusion Language Model

February 19, 2025

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

cs.AI

Abstract

Presentiamo TESS 2, un modello linguistico di diffusione generico per il seguimento di istruzioni che supera i modelli di diffusione contemporanei ottimizzati per istruzioni, oltre a eguagliare e talvolta superare i forti modelli autoregressivi (AR). Addestriamo TESS 2 adattando prima un robusto modello AR attraverso un ulteriore pre-addestramento con la consueta cross-entropia come funzione di perdita per la diffusione, per poi eseguire un ulteriore ottimizzazione per il seguimento di istruzioni. Riteniamo che l'addestramento di adattamento, così come la scelta del modello di base, sia cruciale per addestrare modelli di diffusione efficaci nel seguimento di istruzioni. Proponiamo inoltre la guida basata su ricompensa, una nuova e modulare procedura di guida al momento dell'inferenza per allineare gli output del modello senza la necessità di addestrare il modello sottostante. Infine, dimostriamo che TESS 2 migliora ulteriormente con un aumento della potenza di calcolo al momento dell'inferenza, evidenziando l'utilità dei modelli linguistici di diffusione nel consentire un controllo fine sulla quantità di calcolo utilizzata durante l'inferenza. Codice e modelli sono disponibili all'indirizzo 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.