多块扩散语言模型

摘要

块扩散语言模型（BD-LMs）通过KV缓存和灵活长度生成改进了基于扩散的文本生成。其自然的发展方向是从单块扩散（SingleBD）扩展到多块扩散（MultiBD），即通过同时解码连续块的运行集来实现块间并行性。然而，现有的BD-LMs大多在教师强制（teacher forcing）下训练，此时模型仅观察到在一个干净前缀条件下带有噪声的单个块。尽管最近的扩散强制策略引入了多个噪声块之间的可见性，但其训练状态仍与MultiBD推理存在差异——在后者中，解码操作作用于具有异质槽位噪声模式的有限运行集上。为弥合这一差距，我们提出多块扩散语言模型（MBD-LMs），通过对BD-LMs进行多块教师强制（MultiTF）后训练获得。MultiTF将教师强制与扩散强制相结合，通过在干净前缀条件下对有限噪声组进行训练，并采用随机化噪声调度器以更好匹配MultiBD推理状态。为使MultiBD可实际执行，我们进一步引入基于块缓冲区（Block Buffer）机制的优化解码算法，该算法保留了前缀缓存复用，保持输入形状静态，并将增加的解码并行性转化为实际加速。实验表明，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.