BLOCK DIFFUSION: INTERPOLATING BETWEEN AU-TOREGRESSIVE AND DIFFUSION LANGUAGE MODELS

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 code1, along with the model weights and blog post on the project page: https://m-arriola.com/bd3lms

한 줄 요약

BLOCK DIFFUSION은 확산 모델과 오토회귀 모델의 장점을 결합한 언어 모델로, 가변 길이 생성과 향상된 퍼플렉시티를 달성한다.

핵심 기여도

**BD3-LM**이라는 새로운 언어 모델을 제안하여, 확산과 오토회귀 모델의 한계를 극복함.
**KV 캐싱**과 **병렬 토큰 샘플링**을 통해 추론 효율성을 향상시킴.
**데이터 기반 노이즈 스케줄**과 **그라디언트 분산 추정기**를 도입하여 퍼플렉시티 향상.
**LM1B, OpenWebText** 데이터셋에서 **확산 모델 최고 성능** 기록.

핵심 아이디어

BLOCK DIFFUSION은 토큰 블록 단위로 확산 과정을 수행하는 **semi-autoregressive** 모델로, 블록 내부는 확산, 블록 간은 오토회귀 구조를 갖는다. 이는 기존 확산 모델의 고정 길이 생성 문제를 해결하고, KV 캐싱을 통해 추론 효율성을 높인다. 핵심 아이디어는 **블록 단위 확산**(discrete denoising diffusion over blocks)과 **블록 간 오토회귀 조건**(conditional probability over blocks)의 결합이다.

또한, 확산 모델의 학습 과정에서 발생하는 **그라디언트 분산**(gradient variance)이 성능 저하의 주요 원인임을 밝히고, 이를 줄이기 위해 **데이터 기반 노이즈 스케줄**(data-driven noise schedules)을 제안한다. 이는 확산 모델의 퍼플렉시티 향상에 기여하며, 오토회귀 모델과의 성능 격차를 줄이는 데 효과적이다.

기술적 접근법

**BD3-LM**은 **transformer 기반** 모델로, **블록-인과적 어텐션 마스크**(block-causal attention mask)를 사용하여 블록 간 오토회귀적 관계를 유지.
**KV 캐싱**(Key-Value Caching)을 통해 이전 블록의 계산 결과를 재사용, 추론 효율성 향상.
**하이퍼파라미터**: 블록 크기 $ L' = L $, 최대 850K 그라디언트 스텝으로 사전 학습, 150K 스텝으로 미세 조정.
**노이즈 스케줄 최적화**: $ \beta, \omega $ 파라미터를 **그리드 서치**로 최적화하여 훈련 시 **분산 최소화**.

주요 결과

**LM1B** 데이터셋에서 **퍼플렉시티 20.7**, **OpenWebText**에서 **23.4** 달성.
기존 확산 모델 대비 **+15% 이상 퍼플렉시티 개선**.
**가변 길이 생성** 가능, 훈련 길이를 초과하는 시퀀스도 생성.
**Gaussian diffusion over embeddings** 방식 대비 **10배 적은 생성 단계**로 동일 또는 더 높은 품질 생성.

의의 및 한계

BLOCK DIFFUSION은 확산 모델의 주요 한계인 고정 길이 생성과 퍼플렉시티 저하를 해결하며, 오토회귀 모델과의 성능 격차를 줄이는 데 기여한다. 특히, **KV 캐싱**과 **블록 단위 확산**을 결합한 구조는 추론 효율성과 생성 유연성을 동시에 확보하는 새로운 패러다임을 제시한다.

하지만, **오토회귀 모델 대비 여전히 낮은 퍼플렉시티**가 남아 있으며, **복잡한 블록 크기 조절**이 성능에 미치는 영향은 추가 연구가 필요하다. 또한, **대규모 데이터셋에서의 확장성**과 **실시간 생성 성능**도 검증 대상이다.

실용적 활용

BLOCK DIFFUSION은 **대화 시스템**, **문서 생성**, **생성형 AI** 등에서 유연한 길이의 텍스트 생성이 필요한 상황에 적합하다. 특히, **KV 캐싱을 활용한 추론 효율성**은 대규모 모델의 실시간 응용에 유리하며, **확산 기반 제어 기능**을 활용한 생성물 조절이 가능하다.