Block Diffusion: Interpolación entre Modelos de Lenguaje Autoregresivos y de Difusión

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

Resumen

Los modelos de lenguaje basados en difusión ofrecen ventajas únicas sobre los modelos autorregresivos debido a su potencial para la generación paralelizada y la controlabilidad, aunque se quedan atrás en el modelado de verosimilitud y están limitados a la generación de longitud fija. En este trabajo, presentamos una clase de modelos de lenguaje de difusión por bloques que interpolan entre la difusión de eliminación de ruido discreta y los modelos autorregresivos. La difusión por bloques supera limitaciones clave de ambos enfoques al permitir la generación de longitud flexible y mejorar la eficiencia de inferencia con el almacenamiento en caché de KV y el muestreo paralelo de tokens. Proponemos una receta para construir modelos de difusión por bloques efectivos que incluye un algoritmo de entrenamiento eficiente, estimadores de varianza del gradiente y programaciones de ruido basadas en datos para minimizar la varianza. La difusión por bloques establece un nuevo estado del arte en el rendimiento entre los modelos de difusión en benchmarks de modelado de lenguaje y permite la generación de secuencias de longitud arbitraria. Proporcionamos el código, junto con los pesos del modelo y una entrada de blog en la página del proyecto: 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/