멀티 블록 확산 언어 모델

초록

블록 확산 언어 모델(BD-LM)은 KV 캐싱 및 유연한 길이 생성을 통해 확산 기반 텍스트 생성을 개선한다. 자연스러운 다음 단계는 단일 블록 확산(SingleBD)에서 다중 블록 확산(MultiBD)으로 확장하는 것이며, 여기서 연속된 블록들의 실행 집합이 동시에 디코딩되어 블록 간 병렬성을 확보한다. 그러나 기존 BD-LM은 주로 교사 강제(teacher forcing) 방식으로 훈련되며, 모델은 깨끗한 접두사(clean prefix)를 조건으로 단 하나의 잡음이 있는 블록만 관찰한다. 최근 확산 강제(diffusion forcing) 전략이 여러 잡음 블록 간의 가시성을 도입했지만, 그 훈련 상태는 여전히 MultiBD 추론과 차이를 보이며, 추론 시 디코딩은 이질적인 슬롯별 잡음 패턴을 가진 제한된 실행 집합에서 작동한다. 이러한 격차를 해소하기 위해 우리는 다중 블록 교사 강제(MultiTF)를 통해 BD-LM을 사후 훈련(post-training)하여 얻은 다중 블록 확산 언어 모델(MBD-LM)을 제안한다. MultiTF는 교사 강제와 확산 강제를 통합하며, 깨끗한 접두사를 조건으로 제한된 잡음 그룹에 대해 훈련하고, MultiBD 추론 상태에 더 잘 부합하는 무작위 잡음 스케줄러를 사용한다. MultiBD를 실제로 실행 가능하게 만들기 위해, 우리는 접두사 캐시 재사용을 유지하고 입력 형태를 정적으로 유지하며 증가된 디코딩 병렬성을 실시간 가속으로 변환하는 블록 버퍼 메커니즘에 기반한 최적화된 디코딩 알고리즘을 추가로 도입한다. 실험적으로, MBD-LLaDA2-Mini는 평균 전방향 전달당 토큰 수(TPF)를 3.47에서 6.19로 증가시키고 평균 정확도를 79.95%에서 81.03%로 향상시킨다; DMax와 결합 시, MBD-LLaDA2-Mini-DMax는 평균 TPF 9.34에 도달하며 수학 및 코드 벤치마크에서 1.02%의 정확도 하락만을 보인다.

English

Block Diffusion Language Models (BD-LMs) improve diffusion-based text generation with KV caching and flexible-length generation. A natural next step is to extend them from Single-Block Diffusion (SingleBD) to Multi-Block Diffusion (MultiBD), where a running-set of consecutive blocks is decoded concurrently for inter-block parallelism. However, existing BD-LMs are mostly trained under teacher forcing, where the model observes only one noisy block conditioned on a clean prefix. While the recent diffusion forcing strategy introduces visibility among multiple noisy blocks, its training states still differ from MultiBD inference, where decoding operates on a bounded running-set with heterogeneous slot-wise noise patterns. To bridge this gap, we propose Multi-Block Diffusion Language Models (MBD-LMs), obtained by post-training BD-LMs with Multi-block Teacher Forcing (MultiTF). MultiTF integrates teacher forcing and diffusion forcing by training on bounded noise-groups conditioned on clean prefixes, with randomized noise-schedulers that better match MultiBD inference states. To make MultiBD practically executable, we further introduce an optimized decoding algorithm based on the Block Buffer mechanism that preserves prefix-cache reuse, keeps input shapes static, and translates increased decoding parallelism into wall-clock acceleration. Empirically, MBD-LLaDA2-Mini increases average Tokens Per Forward pass (TPF) from 3.47 to 6.19 and improves average accuracy from 79.95% to 81.03%; when combined with DMax, MBD-LLaDA2-Mini-DMax reaches an average TPF of 9.34 with only a 1.02% accuracy drop on math and code benchmarks.