arXiv:2511.05563cs.LGcs.AI2025-11被引 13

通过前瞻式选路提升扩散语言模型解码精度,避免早期错误传播。

Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models

  • 将解码建模为所有可能去掩码顺序的路径选择,无需奖励模型。
  • 仅需2-3条路径即达最佳效果,显著提升数学、编程等任务性能。
  • 适用于基础与微调模型,可与强化学习互补,适合高效推理场景。

掩码扩散模型(MDMs)通过迭代去掩码生成文本,其性能高度依赖解码时的去掩码顺序。现有启发式方法(如基于置信度采样)只考虑局部优化,无法利用额外测试时间计算,且易导致早期错误累积。本文提出前瞻去掩码(Lookahead Unmasking, LookUM),将采样重构为在所有可能去掩码顺序中进行路径选择,无需外部奖励模型。该框架包含:(i) 路径生成器,从去掩码集合池中采样生成候选路径;(ii) 验证器,计算路径不确定性并执行重要性采样以选定最终路径。实验证明,错误去掩码会显著增加序列级不确定性,而本方法有效规避高风险路径。在六项基准测试(包括数学、规划、编码)上均取得一致提升。仅需2至3条路径即可达到峰值性能,效率极高。对LLaDA及微调版LLaDA 1.5的改进尤为显著:使用LookUM的基线模型性能媲美强化学习微调的LLaDA 1.5,且进一步提升后者表现,说明基于不确定性的验证与强化学习具有正交增益,凸显方法普适性。代码将公开发布。

原文摘要 · Abstract (English)

Masked Diffusion Models (MDMs) as language models generate by iteratively unmasking tokens, yet their performance crucially depends on the inference time order of unmasking. Prevailing heuristics, such as confidence based sampling, are myopic: they optimize locally, fail to leverage extra test-time compute, and let early decoding mistakes cascade. We propose Lookahead Unmasking (LookUM), which addresses these concerns by reformulating sampling as path selection over all possible unmasking orders without the need for an external reward model. Our framework couples (i) a path generator that proposes paths by sampling from pools of unmasking sets with (ii) a verifier that computes the uncertainty of the proposed paths and performs importance sampling to subsequently select the final paths. Empirically, erroneous unmasking measurably inflates sequence level uncertainty, and our method exploits this to avoid error-prone trajectories. We validate our framework across six benchmarks, such as mathematics, planning, and coding, and demonstrate consistent performance improvements. LookUM requires only two to three paths to achieve peak performance, demonstrating remarkably efficient path selection. The consistent improvements on both LLaDA and post-trained LLaDA 1.5 are particularly striking: base LLaDA with LookUM rivals the performance of RL-tuned LLaDA 1.5, while LookUM further enhances LLaDA 1.5 itself showing that uncertainty based verification provides orthogonal benefits to reinforcement learning and underscoring the versatility of our framework. Code will be publicly released.

扩散模型解码优化路径选择语言模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。