扩散语言模型解掩码策略研究：856K参数策略提升LLaDA-8B推理效率
What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
LLaDA-8B 上的研究，教你训练一个 856K 参数的小策略，让扩散语言模型解码更快更准，还能迁移到 Dream-7B。
arXiv 上的这篇论文系统研究了扩散语言模型中的可学习解掩码策略，将监督策略视为对解码过程的蒸馏，并分别检验教师标注、状态输入和策略容量等因素的影响。实验发现仅标注与最终输出一致的预测在低 NFE（函数评估次数）下过于宽松，隐藏状态单独作用弱但可与置信度互补，约 10% 的 NFE 降幅即可抵消最高 7.5M 参数策略的开销。作者据此用轨迹收缩目标训练策略，在 LLaDA-8B-Instruct 上以 856K 参数在 GSM8K 上取得最佳准确率-NFE 前沿和最高吞吐，在 MATH-500 和 HumanEval 中等目标下表现有竞争力，并能较大程度迁移到 Dream-7B-Instruct。
What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
Masked diffusion language models generate tokens in parallel, but their efficiency depends on which predictions are committed. Learned unmasking policies improve the speed--accuracy trade-off, yet differ in supervision, inputs, and capacity, obscuring which choices matter. We view supervised policies as distilling a decoding procedure through a teacher and examine independently what the teacher labels as committable, on which states the policy imitates it, what the policy observes, and how much the number of function evaluations (NFE) must decrease to offset policy overhead. We find that labeling predictions matching the final output is overly permissive at low NFE, and hidden states are weak alone but can complement confidence. About a 10% NFE reduction suffices to offset policy overhead for policies up to 7.5M parameters across tested settings. We therefore train a policy on trajectory-contraction targets, which accept only commitments staying on a reference trajectory, using confidence and hidden states. On LLaDA-8B-Instruct, the resulting 856K-parameter policy achieves the best accuracy--NFE frontier and highest throughput on GSM8K, is competitive at moderate targets on MATH-500 and HumanEval, and largely transfers to Dream-7B-Instruct.
