通过量化不确定性优化掩码扩散模型中的解码路径

论文信息

标题: Optimizing Decoding Paths in Masked Diffusion Models by Quantifying Uncertainty

作者: Ziyu Chen, Xinbei Jiang, Peng Sun, et al.

发布日期: 2025-12-24

arXiv ID: 2512.21336v1

PDF 链接: 下载 PDF

3 分钟速览

• 研究问题：掩码扩散模型（MDM）的生成质量高度依赖于解码路径的选取，但缺少从全局不确定性角度分析这一问题的理论工具和优化策略。
• 核心方法：提出可量化的 “去噪熵”（包括状态熵 $h_{\mathrm{DE}}$ 和路径熵 $H_{\mathrm{DE}}$），用以评估解码路径中累积的预测不确定性，并基于此设计两个路径搜索算法：后验选择的 E‑BoN 与实时引导的 E‑SMC。
• 关键结果：在 LLaDA‑Instruct‑8B 上，将熵引导算法与强基线 PC‑Sampler 结合，GSM8K 准确率从 79.3% 提升至 81.2%（+1.9%），Countdown 规划任务从 36.3% 提升至 40.4%（+4.1%）（见表 2）。
• 主要局限：方法需要额外计算开销（多粒子采样与重采样），熵最小化如果过度（如贪婪搜索）会导致多样性崩溃（见附录表 6）；算法对温度、重采样间隔等超参数敏感。
• 适合读者：关注扩散语言模型解码策略、不确定性量化、推理与规划任务优化的研究人员，以及希望提升 MDM 生成质量的工程师。

论文背景和研究动机

掩码扩散模型通过随机顺序的掩码与去噪训练，天然支持任意解码顺序，理论上可在庞大的解码路径空间中寻找到优于自回归固定顺序的生成结果。然而，实践中 MDM 在复杂任务上的表现往往不及自回归模型，解码策略的选择成为瓶颈。简单的随机顺序策略将搜索最优路径等同于抽奖，而基于局部贪婪信号（如最高置信度、最低熵、最大边距）的方法则缺乏对整条生成路径的全局视角，容易陷入 “局部最优” 而导致整体质量下降。

受自回归模型中利用熵衡量生成不确定性的启发，作者首次将 MDM 输出质量的波动归因为路径不确定性：即模型在整条解码路径上的累积预测不确定性。如果能够量化这种全局不确定性，就有可能在生成过程中主动规避高不确定性路径，从而稳定地提升输出质量。基于此，论文提出了 “去噪熵” 这一可内部计算的不确定性度量，并以此为核心开发出两套路径优化算法，将 MDM 的灵活性从劣势转化为优势。

核心方法和技术细节

路径不确定性与去噪熵

在 MDM 的逆向去噪过程中，一条解码路径 $\tau = (\mathbf{z}_{t_N}, \dots, \mathbf{z}_{t_0})$ 由 $N$ 个逐步退掩码的状态构成。对整个路径的全局不确定性，作者定义为路径不确定性，并引入去噪熵进行量化。

状态熵 $h_{\mathrm{DE}}(\mathbf{z}_t)$ 定义为在时间步 $t$、给定当前状态 $\mathbf{z}_t$ 时，对所有掩码位置预测分布香农熵的平均值：

$$h_{\mathrm{DE}}(\mathbf{z}_t) \triangleq \frac{1}{|\mathcal{M}_t|}\sum_{\ell\in\mathcal{M}_t} H\big(p_{\boldsymbol{\theta}}(X_0^\ell|\mathbf{z}_t,t)\big)$$

其中 $\mathcal{M}_t$ 是当前被掩码的位置集合，$p_{\boldsymbol{\theta}}$ 是模型输出的概率分布。状态熵是 “预言机状态不确定性” 的实用上界（命题 1），且在 $\epsilon$ 准确模型的假设下，与瞬时训练损失之间仅相差 $\mathcal{O}(\epsilon)$（命题 2），因此可以作为模型在每个解码步骤中内在困惑度的良好代理。

路径熵 $H_{\mathrm{DE}}(\tau)$ 将状态熵沿时间积分（或离散求和）得到整条路径的累积不确定性：

$$H_{\mathrm{DE}}(\tau) \triangleq \int_0^1 h_{\mathrm{DE}}(\mathbf{z}_t)\, dt \;\approx\; \frac{1}{N}\sum_{i=1}^N h_{\mathrm{DE}}(\mathbf{z}_{t_i})$$

理论分析进一步表明，模型生成路径分布与真实参考路径分布之间的 KL 散度，被由路径熵期望差所构成的下界所约束（命题 3），因此朝着减小路径熵的方向搜索，能有效推动生成分布靠近数据流形。

熵引导的解码算法

基于去噪熵，论文提出两种优化解码路径的方法。

熵基最优‑N (E‑BoN)：先并行生成 $M$ 条候选解码路径，然后从中选择路径熵 $H_{\mathrm{DE}}$ 最小的路径作为最终输出。这是一种后验选择方案，完全无需改变单次采样的内部机制，实现简单但计算预算均匀分配在所有路径上。

熵引导的顺序蒙特卡洛 (E‑SMC)：在实时生成过程中维护 $M$ 个粒子（即候选解），每隔 $\Delta i_r$ 步进行一次评估与重采样。每个粒子当前的状态熵越低，被赋予的权重就越高，从而更可能被复制；高熵粒子则被逐步淘汰。这一机制使得计算资源能够动态地向更可靠的路径倾斜，主动避开高不确定性区域。具体采用 Gibbs 势函数进行加权：

$$w^{(m)} \propto \exp\!\Big(\lambda\cdot\big(1 - \frac{h_{\mathrm{DE}}(\mathbf{z}_{t_{i-1}}^{(m)})}{\log K}\big)\Big)$$

其中 $\lambda$ 控制选择压力，$K$ 为词表大小。通过调节 $\lambda$ 和重采样间隔，E‑SMC 能够在避免过度优化导致的多样性崩溃（如表 6 的贪婪搜索所示）的同时，显著降低平均路径熵。

创新点和贡献

1. 首次形式化路径不确定性，将 MDM 生成质量的波动归因于整条解码路径上的累积预测不确定性，填补了这一理论空白。
2. 提出去噪熵 $h_{\mathrm{DE}}$ 与 $H_{\mathrm{DE}}$ 作为可内部计算的不确定性度量，并通过理论证明它们是对理想不确定性的紧上界、与训练损失高度相关，为在线路径评估提供了严格工具。
3. 设计了两种熵引导的解码算法——E‑BoN 和 E‑SMC，分别实现了后验选择和实时主动搜索，不依赖任何外部模型或标签，通用性强。
4. 在多类推理和规划任务上验证了方法的有效性，与强基线 PC‑Sampler 结合在 GSM8K、MATH500、Countdown 等数据集上稳定提升，同时保持了输出多样性。

实验结果分析

小规模验证：去噪熵与文本质量

在 OpenWebText 上使用 1.3 亿参数的 MDLM 进行无条件生成，用更大的 GPT2‑Large 计算困惑度（PPL），观察到路径熵 $H_{\mathrm{DE}}$ 与 $\ln(\text{PPL})$ 呈强烈的近线性正相关（相关系数在 0.85 左右，图 3）。随去噪步数 $S$ 从 16 增至 1024，$H_{\mathrm{DE}}$ 从 6.81 降至 5.12，同时 $\ln(\text{PPL})$ 从 5.73 降至 3.73，说明更精细的生成过程降低了模型内在累积不确定性，也带来了更高的输出质量。

路径搜索优化

在相同模型上比较普通均匀采样、E‑BoN 和 E‑SMC（表 1）。以 256 步、4 粒子、重采样间隔 32 为例，E‑SMC 将 PPL 从 68.5 大幅降至 40.4（LLaMA‑3‑8B 评估），同时多样性分数仅从 5.45 微降至 5.31，表明搜索到的不确定性更低的路径并未牺牲输出的丰富度。增加粒子数、提高重采样频率进一步提升了性能，尤其是 E‑SMC 在粒子数从 4 增至 12 时 PPL 从 47.7 降至 36.1，体现了更好的预算扩展性。

大规模推理和规划任务

在 LLaDA‑Instruct‑8B 和 LLaDA‑1.5‑8B 上，将 E‑BoN 或 E‑SMC 叠加于当前最强的解码策略 PC‑Sampler 之上，在所有五个推理/规划基准上均观察到了提升（见表 2）。其中，GSM8K 准确率在两款模型上分别提升了 1.9% 和 1.0%，Countdown 规划任务的提升尤为显著（+4.1% 和 +4.3%），这印证了面向全局路径的优化对于需要多步连贯推理的任务特别关键。

在 Open‑dCoder‑0.5B 的代码生成任务上，将熵引导方法同时应用于 P2、Uniform、Confidence、Entropy 和 Margin 五种不同策略，每种策略的平均准确率都获得了稳定的提升（见图 4）。这表明去噪熵是一种与底层解码策略解耦的通用增强器。

消融与预算效率

在 Sudoku 任务上，通过调整 PC‑Sampler 的顺序性参数，控制了不同的解码路径，并发现路径熵与最终准确率呈负相关（图 5），进一步支持了 “低 HDE 路径对应高质量生成” 这一核心假设。此外，在相同 5 粒子的预算下，熵引导方法明显优于多数投票（见表 2 对应行），说明基于不确定性的内部信号比简单集成更能高效分配采样计算。

实践建议

对于希望在 MDM 的实际部署中提升生成质量、尤其是提升复杂推理和规划任务成功率的开发者，可以考虑以下实践方案：

• 引入路径熵作为内部质量监控信号。在推理时计算 $H_{\mathrm{DE}}$，可以在不依赖外部评分器的情况下过滤低质量输出，适用于需要快速筛选候选的在线服务。
• E‑BoN 可作为无侵入的即插即用增强。当推理系统已支持并行多采样时，直接选择路径熵最小的样本即可获得显著的 PPL 和准确率提升，无需改动模型或解码器的内部逻辑。
• E‑SMC 适用于预算受限但需要更高质量的场景。通过将重采样间隔设在总步数的约 1/81/4（如 256 步中每 32 步评估一次），并设置适中的温度 $\lambda$（如 5.010.0），能在不过度增加计算耗时的情况下动态淘汰低质量路径。实验显示（表 7），E‑SMC 的并行版本延迟仅比 E‑BoN 略高，在相同粒子数下仍具备线性加速潜力。
• 注意防止过度熵优化。极端的最小化熵（如贪婪搜索）会使输出陷入重复循环，多样性剧降（见表 6）。建议在部署时监控输出的唯一 token 比例，当多样性低于阈值时适当提高温度或减少重采样频率。
• 针对不同任务调节超参数。对数学推理（如 GSM8K）可选择较小的 $\lambda$ 和稍大的重采样间隔，避免过早收敛；对强规划任务（如 Countdown、Sudoku）可增大 $\lambda$ 并适度提高重采样频率，更主动地偏离高不确定性区域。

总之，去噪熵为 MDM 的解码路径优化提供了可计算、可理论的全局视角，工程实现成本可控，适合作为提升 MDM 生成质量的基础组件。