TESS 2: 大規模汎用拡散言語モデル

TESS 2: A Large-Scale Generalist Diffusion Language Model

February 19, 2025

著者: Jaesung Tae, Hamish Ivison, Sachin Kumar, Arman Cohan

cs.AI

要旨

TESS 2を紹介します。これは一般的な指示追従型拡散言語モデルであり、現代の指示チューニングされた拡散モデルを上回り、強力な自己回帰（AR）モデルに匹敵し、時にはそれを超える性能を発揮します。TESS 2のトレーニングは、まず強力なARモデルを拡散損失として通常のクロスエントロピーを用いて継続事前学習により適応させ、その後さらに指示チューニングを行うことで行います。適応トレーニングおよびベースモデルの選択が、優れた指示追従型拡散モデルのトレーニングにおいて重要であることがわかりました。さらに、基盤となるモデルをトレーニングすることなくモデル出力を整合させるための、新規でモジュール型の推論時ガイダンス手順である報酬ガイダンスを提案します。最後に、TESS 2が推論時の計算量の増加に伴ってさらに改善されることを示し、拡散言語モデルが推論時に使用する計算量を細かく制御できる有用性を強調します。コードとモデルは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.