TESS 2: Un modelo de lenguaje de difusión generalista a gran escala

TESS 2: A Large-Scale Generalist Diffusion Language Model

February 19, 2025

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

cs.AI

Resumen

Presentamos TESS 2, un modelo de lenguaje de difusión general que sigue instrucciones y supera a los modelos contemporáneos de difusión ajustados para instrucciones, además de igualar y, en ocasiones, superar a los fuertes modelos autoregresivos (AR). Entrenamos TESS 2 adaptando primero un modelo AR robusto mediante un preentrenamiento continuo utilizando la entropía cruzada habitual como pérdida de difusión, y luego realizando un ajuste adicional para instrucciones. Descubrimos que el entrenamiento de adaptación, así como la elección del modelo base, son cruciales para entrenar buenos modelos de difusión que sigan instrucciones. Además, proponemos la guía de recompensa, un procedimiento novedoso y modular de guía en tiempo de inferencia para alinear las salidas del modelo sin necesidad de entrenar el modelo subyacente. Finalmente, demostramos que TESS 2 mejora aún más con un mayor cómputo en tiempo de inferencia, destacando la utilidad de los modelos de lenguaje de difusión al permitir un control detallado sobre la cantidad de cómputo utilizado durante la inferencia. El código y los modelos están disponibles en 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.