Multi-Block Diffusion Language Models

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.

한 줄 요약

Multi-Block Diffusion Language Models (MBD-LMs)는 Multi-Block Diffusion (MultiBD) 추론을 위한 학습-추론 일치성을 확보한 확장형 BD-LM으로, TPF 9.34 달성.

핵심 기여도

Multi-block Teacher Forcing (MultiTF)를 통해 BD-LMs를 MBD-LMs로 post-training하여 MultiBD 추론과 일치하는 학습 상태 구현.
Block Buffer 메커니즘을 도입해 입력 형태를 정적 유지하면서 KV 캐싱과 prefix 캐싱을 보존하고, 병렬성을 실제 속도 향상으로 전환.
MBD-LLaDA2-Mini-DMax 모델에서 TPF 9.34 달성 (SingleBD 대비 +47.1%) 및 정확도 84.59% 유지 (0.08% 하락).

핵심 아이디어

기존 BD-LMs는 SingleBD 방식에서 훈련되며, 이는 하나의 노이즈 블록만 처리하는 teacher forcing 기반으로, MultiBD 추론 시 필요한 복수 블록의 이질적 노이즈 패턴과는 일치하지 않는다. 이를 해결하기 위해, MultiTF는 clean prefix와 함께 여러 노이즈 블록을 결합한 **noise-group**을 훈련 입력으로 사용하며, **Group-Aware Dual-Stream Mask**를 통해 블록 간 상호참조를 허용한다. 또한, **randomized noise-scheduler**를 도입해 실제 MultiBD 추론 시의 노이즈 패턴을 시뮬레이션한다. 이는 MultiBD의 **bounded running-set** 구조와 일치하는 학습 상태를 구축하는 핵심이다.

기술적 접근법

**Multi-block Teacher Forcing (MultiTF)**: clean prefix와 **bounded noise-group**을 결합한 입력으로 훈련.
**Group-Aware Dual-Stream Mask**: 블록 간 상호참조를 허용하는 attention mask 적용.
**Randomized noise-scheduler**: 블록 단위 노이즈 패턴을 랜덤하게 생성하여 실제 MultiBD 추론과 유사한 훈련 상태 구축.
**Block Buffer**: 입력 형태를 정적 유지하면서 KV 캐싱과 prefix 캐싱을 보존하는 추론 메커니즘.
**TPF (Tokens Per Forward pass)**: 평균 3.47 → 6.19 → 9.34로 증가.

주요 결과

**LLaDA2-Mini** 기반에서 MBD-LLaDA2-Mini는 평균 TPF 3.47 → 6.19 (+78.4%), 정확도 79.95% → 81.03%.
**DMax**와 결합한 MBD-LLaDA2-Mini-DMax는 평균 TPF 9.34 달성 (SingleBD 대비 +47.1%), 정확도 84.67% → 84.59% (0.08% 하락).
**HumanEval+ 및 GSM8K**에서 TPF 6.57 → 9.87, AUP 536.89 → 805.34 증가.

의의 및 한계

MBD-LMs는 MultiBD 추론에서의 병렬성을 실제 속도 향상으로 전환하는 첫 사례로, BD-LM의 학습-추론 일치성을 확보한 새로운 훈련 전략을 제시한다. 특히, **Block Buffer**는 추론 효율성과 캐싱 유지 사이의 균형을 맞춘 실용적 기술이다. 그러나, MultiTF는 post-training 기법으로 기존 BD-LM에만 적용 가능하며, 초기 훈련 단계에서 MultiBD를 고려한 모델 아키텍처는 아직 제시되지 않았다.

실용적 활용

MBD-LMs는 수학 및 코드 생성과 같은 고성능 추론이 필요한 분야에서 유용하게 활용될 수 있다. 특히, **TPS (Tokens Per Second)**가 781.50 → 951.41로 증가한 점을 고려할 때, 대규모 언어 모델 기반의 실시간 서비스 및 대화형 시스템에 적합하다.