多區塊擴散語言模型

摘要

區塊擴散語言模型(BD-LMs)通過KV快取與可變長度生成技術，優化了基於擴散的文本生成。自然的下一步是將其從單塊擴散(SingleBD)擴展至多塊擴散(MultiBD)，其核心機制在於對連續區塊組成的「運行集」進行並行解碼，以實現塊間並行處理。然而，現有BD-LMs大多採用「教師強制」策略訓練——模型僅能觀察到一個帶噪區塊，且該區塊以乾淨前綴為條件。近期提出的「擴散強制」策略雖賦予模型觀察多個帶噪區塊的能力，但其訓練狀態仍與MultiBD推論情境存在差異：在推論階段，解碼過程作用於具有異質性逐槽位雜訊模式的有限運行集。為彌合此鴻溝，我們提出多塊擴散語言模型(MBD-LMs)，其透過「多塊教師強制」(MultiTF)對BD-LMs進行後訓練實現。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.