提升文本生成效率,通过多块并行解码实现更快的推理速度。
Multi-Block Diffusion Language Models
- 采用多块教师强制训练,模拟真实推理中的噪声状态。
- 解码效率提升,每轮前向传播生成tokens数从3.47增至9.34。
- 适合需要高速生成的场景,如代码和数学推理任务。
块扩散语言模型(BD-LMs)通过键值缓存和灵活长度生成,提升了基于扩散的文本生成效果。为进一步扩展单块到多块并行生成,本文提出多块扩散语言模型(MBD-LMs),通过后训练方式引入多块教师强制(MultiTF)策略。该方法在受限噪声组上训练,结合清洁前缀与随机噪声调度,更贴近多块解码的实际状态。同时设计基于块缓冲机制的优化解码算法,保持前缀缓存复用、输入形状稳定,并将并行度提升转化为实际推理加速。实验表明,MBD-LLaDA2-Mini将平均每轮前向传播生成tokens数(TPF)从3.47提升至6.19,准确率从79.95%升至81.03%;结合DMax后,TPF达9.34,仅在数学与代码基准上下降1.02%准确率。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。