블록 디퓨전: 자기회귀 모델과 디퓨전 언어 모델 간의 보간

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

초록

확산 언어 모델(Diffusion Language Models)은 병렬 생성 가능성과 제어 가능성 측면에서 자기회귀 모델(Autoregressive Models)에 비해 독특한 장점을 제공하지만, 가능도 모델링(likelihood modeling)에서는 뒤처지며 고정 길이 생성에 제한되는 단점이 있습니다. 본 연구에서는 이산적 노이즈 제거 확산(discrete denoising diffusion)과 자기회귀 모델 사이를 보간하는 블록 확산 언어 모델(Block Diffusion Language Models) 클래스를 소개합니다. 블록 확산은 유연한 길이 생성을 지원하고 KV 캐싱(KV caching) 및 병렬 토큰 샘플링(parallel token sampling)을 통해 추론 효율성을 개선함으로써 두 접근법의 주요 한계를 극복합니다. 우리는 효과적인 블록 확산 모델을 구축하기 위한 레시피를 제안하며, 여기에는 효율적인 훈련 알고리즘, 그래디언트 분산 추정기, 그리고 분산을 최소화하기 위한 데이터 기반 노이즈 스케줄(data-driven noise schedules)이 포함됩니다. 블록 확산은 언어 모델링 벤치마크에서 확산 모델 중 새로운 최첨단 성능을 달성하며, 임의 길이 시퀀스 생성도 가능하게 합니다. 프로젝트 페이지(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/