Блок-диффузия: Интерполяция между авторегрессивными и диффузионными языковыми моделями

Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

March 12, 2025

Авторы: Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, Volodymyr Kuleshov

cs.AI

Аннотация

Диффузионные языковые модели предлагают уникальные преимущества по сравнению с авторегрессивными моделями благодаря их потенциалу для параллелизованной генерации и управляемости, однако они уступают в моделировании правдоподобия и ограничены генерацией фиксированной длины. В данной работе мы представляем класс блочных диффузионных языковых моделей, которые интерполируют между дискретной диффузией с удалением шума и авторегрессивными моделями. Блочная диффузия преодолевает ключевые ограничения обоих подходов, поддерживая генерацию гибкой длины и повышая эффективность вывода за счет кэширования ключей и значений (KV) и параллельной выборки токенов. Мы предлагаем рецепт для создания эффективных блочных диффузионных моделей, включающий эффективный алгоритм обучения, оценщики дисперсии градиента и основанные на данных расписания шума для минимизации дисперсии. Блочная диффузия устанавливает новый уровень производительности среди диффузионных моделей на тестах языкового моделирования и позволяет генерировать последовательности произвольной длины. Мы предоставляем код, веса модели и блог-пост на странице проекта: https://m-arriola.com/bd3lms/

English

Diffusion language models offer unique benefits over autoregressive models due to their potential for parallelized generation and controllability, yet they lag in likelihood modeling and are limited to fixed-length generation. In this work, we introduce a class of block diffusion language models that interpolate between discrete denoising diffusion and autoregressive models. Block diffusion overcomes key limitations of both approaches by supporting flexible-length generation and improving inference efficiency with KV caching and parallel token sampling. We propose a recipe for building effective block diffusion models that includes an efficient training algorithm, estimators of gradient variance, and data-driven noise schedules to minimize the variance. Block diffusion sets a new state-of-the-art performance among diffusion models on language modeling benchmarks and enables generation of arbitrary-length sequences. We provide the code, along with the model weights and blog post on the project page: https://m-arriola.com/bd3lms/