[Paper] 通过量化不确定性优化 Masked Diffusion Models 中的解码路径

Source: arXiv - 2512.21336v1

Overview

Masked Diffusion Models (MDMs) 已经成为传统自回归生成器的强大替代方案，实现了文本、代码和计划的快速、非顺序生成。然而，掩码标记被“解掩”（即解码路径）的顺序会显著影响最终输出质量。本文形式化了该问题，提出了一种新的不确定性度量——Denoising Entropy，并展示了如何引导解码过程以获得更高质量的结果。

关键贡献

• 解码路径敏感性的形式化 – 证明 MDM 中输出差异源于所选生成轨迹上累计的预测不确定性。
• 去噪熵度量 – 一种可计算的、模型内部的度量，用于量化每一步去噪的不确定性。
• 两种基于熵的算法：
  1. 事后路径选择 – 在生成后评估多个采样路径，挑选总熵最低的路径。
  2. 实时引导 – 根据当前熵分布动态选择下一个要填充的掩码，实时引导生成过程。
• 实证验证 – 在一系列具有挑战性的基准（推理、规划、代码合成）上持续取得提升，常常超越强大的自回归基线。
• 开源工具 – 作者发布了计算去噪熵的代码以及将引导策略集成到现有 MDM 流水线的实现。

方法论

1. 量化不确定性

   • 在每个扩散步骤，模型预测当前被掩码位置的可能 token 值的分布。
   • 去噪熵 简单地是该分布的 Shannon 熵，对所有被掩码的 token 求和。熵越低表示模型对下一步去噪更有信心。

2. 路径优化策略

   • 事后选择：使用不同的随机掩码顺序多次运行 MDM，计算每条完整轨迹的总熵，并保留总和最小的轨迹。这种方式易于并行，且不需要修改模型本身。
   • 实时引导：在生成过程中，评估如果下一个候选 token 被解掩后会产生的熵。选择产生最低即时熵的 token（或小组 token），然后继续。这将解码过程转化为一种贪婪的、不确定性驱动的搜索。

3. 评估协议

   • 基准：GSM8K（数学推理）、MiniWoB（交互式规划）、HumanEval（代码生成）。
   • 指标：代码的 Exact match / pass@k、规划的成功率以及推理的准确率。
   • 基线：使用随机解码顺序的标准 MDM，以及强大的自回归 Transformer（例如 GPT‑Neo、CodeGen）。

Results & Findings

• 两种基于熵的方式始终在推理与规划任务上比普通 MDM 高出 7–9 个百分点，在代码合成上高出 约 6 个百分点。
• 实时指导的效果与事后选择相当，但每个 token 只需一次前向传播，因而在生产环境中更具实用性。
• 消融实验表明，性能提升主要来源于熵驱动的排序，而非额外计算。

实际意义

• 更高质量的非自回归生成：开发者现在可以部署基于 MDM 的服务（例如代码自动完成、计划合成），在保持并行解码速度优势的同时，不牺牲输出的忠实度。
• 即插即用的改进：由于熵度量来源于模型自身的 logits，现有的 MDM 检查点只需添加熵计算和贪婪选择器，即可轻松升级，工程工作量极小。
• 资源高效的采样：后置方法利用并行硬件同时探索多种解码顺序，提供一种随 GPU 数量扩展的 “best‑of‑N” 策略。
• 不确定性感知的调试：可以可视化 Denoising Entropy，以定位模型不确定的步骤，帮助工程师诊断生成文本或代码的失败模式。
• 更广泛的 AI 安全视角：通过将生成引导离开高熵（即不确定）区域，该方法可能降低下游应用中的幻觉或不安全输出。

限制与未来工作

• 计算开销：实时引导会增加每一步的适度成本（对每个候选掩码进行熵评估）。在极高吞吐量的场景下，这可能抵消部分并行性的收益。
• 贪婪性：当前的引导是局部最优的；更复杂的搜索（例如基于熵的束搜索）可能捕获更好的全局轨迹，但会增加复杂度。
• 领域特定性：实验聚焦于推理、规划和代码；尚不清楚熵引导解码在开放式文本生成（如故事写作）上的表现如何。
• 理论保证：虽然熵在经验上与质量相关，但去噪熵与下游任务指标之间的正式关联仍是一个未解的研究问题。

未来的工作可以探索将熵引导与学习策略相结合的混合方法，将该度量扩展到多模态扩散模型，并将其整合到基于强化学习的微调流水线中。

作者

• Ziyu Chen
• Xinbei Jiang
• Peng Sun
• Tao Lin

论文信息

• arXiv ID: 2512.21336v1
• 类别: cs.CL, cs.AI, cs.LG
• 出版日期: 2025年12月24日
• PDF: 下载 PDF