Multi-Block Diffusion Language Models

2026-06-28

AI Summary

This work addresses the inconsistency between training and inference states in existing diffusion language models when generating multiple tokens in parallel, which limits both efficiency and accuracy. To resolve this, the authors propose a Multi-Token Teacher Forcing (MultiTF) training strategy coupled with a Block Buffer decoding mechanism. By integrating a stochastic noise scheduler and reusing KV caches, their approach aligns the training dynamics with multi-token parallel inference. The method substantially improves parallel generation performance: the MBD-LLaDA2-Mini model increases its average number of generated tokens per forward pass from 3.47 to 6.19 and boosts accuracy from 79.95% to 81.03%. When combined with DMax, it achieves a throughput of 9.34 tokens per forward pass with only a 1.02% accuracy drop.

Abstract

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 \textit{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 \textit{running-set} with heterogeneous slot-wise noise patterns. To bridge this gap, we propose \textit{Multi-Block Diffusion Language Models} (MBD-LMs), obtained by post-training BD-LMs with \textit{Multi-block Teacher Forcing} (MultiTF). MultiTF integrates teacher forcing and diffusion forcing by training on bounded \textit{noise-groups} conditioned on clean prefixes, with randomized \textit{noise-schedulers} that better match MultiBD inference states. To make MultiBD practically executable, we further introduce an optimized decoding algorithm based on the \textit{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 \textbf{6.19} and improves average accuracy from 79.95\% to \textbf{81.03\%}; when combined with DMax, MBD-LLaDA2-Mini-DMax reaches an average TPF of \textbf{9.34} with only a 1.02\% accuracy drop on math and code benchmarks.

Problem

• Multi-Block Diffusion
• Diffusion Language Models
• Teacher Forcing
• Noise Scheduling
• Parallel Decoding

Innovation

• Multi-Block Diffusion
• Multi-block Teacher Forcing
• Block Buffer
• KV Caching
• Diffusion Language Models