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.