TESS 2: Een grootschalig generalistisch diffusie-taalmodel

TESS 2: A Large-Scale Generalist Diffusion Language Model

February 19, 2025

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

cs.AI

Samenvatting

We introduceren TESS 2, een algemeen instructievolgend diffusie-taalmodel dat hedendaagse instructie-afgestemde diffusiemodellen overtreft, en dat evenaart en soms zelfs sterke autoregressieve (AR) modellen overstijgt. We trainen TESS 2 door eerst een sterk AR-model aan te passen via voortgezette vooropleiding met de gebruikelijke kruisentropie als diffusieverlies, en vervolgens verdere instructie-afstemming uit te voeren. We ontdekken dat aanpassingstraining, evenals de keuze van het basismodel, cruciaal zijn voor het trainen van goede instructievolgende diffusiemodellen. We stellen verder beloningsbegeleiding voor, een nieuwe en modulaire begeleidingsprocedure tijdens de inferentie om modeluitvoeringen uit te lijnen zonder het onderliggende model te hoeven trainen. Tot slot tonen we aan dat TESS 2 verder verbetert met meer rekencapaciteit tijdens de inferentie, wat het nut van diffusie-LM's benadrukt bij het hebben van fijnmazige controle over de hoeveelheid rekenkracht die tijdens de inferentie wordt gebruikt. Code en modellen zijn beschikbaar op 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.