Block Diffusion: Interpolando entre Modelos de Linguagem Autoregressivos e de Difusão

Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

March 12, 2025

Autores: Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, Volodymyr Kuleshov

cs.AI

Resumo

Os modelos de linguagem baseados em difusão oferecem benefícios únicos em relação aos modelos autoregressivos devido ao seu potencial para geração paralelizada e controlabilidade. No entanto, eles ficam atrás na modelagem de verossimilhança e são limitados à geração de sequências de comprimento fixo. Neste trabalho, introduzimos uma classe de modelos de linguagem de difusão em blocos que interpola entre a difusão de ruído discreta e os modelos autoregressivos. A difusão em blocos supera limitações-chave de ambas as abordagens ao suportar geração de comprimento flexível e melhorar a eficiência de inferência com cache de KV e amostragem paralela de tokens. Propomos uma receita para construir modelos de difusão em blocos eficazes, que inclui um algoritmo de treinamento eficiente, estimadores de variância de gradiente e cronogramas de ruído baseados em dados para minimizar a variância. A difusão em blocos estabelece um novo estado da arte entre os modelos de difusão em benchmarks de modelagem de linguagem e permite a geração de sequências de comprimento arbitrário. Disponibilizamos o código, juntamente com os pesos do modelo e um post no blog na página do projeto: 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/