TESS 2: Ein groß angelegtes generalistisches Diffusionssprachmodell

TESS 2: A Large-Scale Generalist Diffusion Language Model

February 19, 2025

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

cs.AI

Zusammenfassung

Wir stellen TESS 2 vor, ein allgemeines, befehlsfolgendes Diffusions-Sprachmodell, das zeitgenössische, befehlsoptimierte Diffusionsmodelle übertrifft und mit starken autoregressiven (AR) Modellen gleichzieht oder diese manchmal sogar übertrifft. Wir trainieren TESS 2, indem wir zunächst ein leistungsstarkes AR-Modell durch fortgesetztes Vortraining mit der üblichen Kreuzentropie als Diffusionsverlust anpassen und anschließend eine weitere Befehloptimierung durchführen. Wir stellen fest, dass sowohl das Anpassungstraining als auch die Wahl des Basismodells entscheidend für das Training guter, befehlsfolgender Diffusionsmodelle sind. Darüber hinaus schlagen wir Reward Guidance vor, ein neuartiges und modulares Inferenzzeit-Leitverfahren, um Modellausgaben auszurichten, ohne das zugrunde liegende Modell trainieren zu müssen. Schließlich zeigen wir, dass TESS 2 mit erhöhter Inferenzzeit-Rechenleistung weiter verbessert wird, was den Nutzen von Diffusions-Sprachmodellen bei der fein abgestuften Steuerung der zur Inferenzzeit verwendeten Rechenleistung unterstreicht. Code und Modelle sind unter https://github.com/hamishivi/tess-2 verfügbar.

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.