FEROM：掩码扩散语言模型的边沿策略优化新方法

引言

掩码扩散语言模型（Masked Diffusion Language Models）通过迭代地解掩码位置来生成文本，近年来在文本生成领域受到广泛关注。与自回归模型逐词生成不同，这类模型在每一步需要从模型自身的 logits 中选择要揭示的位置，因此揭示顺序（reveal order）成为了一个内生的潜在变量，其离散空间极其庞大。如何针对这一特性进行有效的策略优化，是当前研究的一个关键挑战。

现有方法的局限

传统的策略优化方法通常依赖于预定义的调度（pre-defined schedules）或采样的轨迹（sampled trajectories）来近似策略梯度。然而，这些方法要么引入偏差，要么导致高方差。具体来说，预定义调度无法适应模型自身的学习动态，而采样轨迹虽然灵活，但方差较大，影响训练稳定性。

FEROM 的核心思想

针对上述问题，论文提出了 FEROM（Frontier Endogenous Reveal-Order Marginal Policy Optimization），一种针对掩码扩散语言模型 rollout 诱导的边沿响应策略（marginal response policy）的近似策略梯度方法。FEROM 的关键创新在于：

• 推导出关于潜在揭示顺序的 Rao-Blackwellized 策略梯度恒等式，将估计量表示为揭示状态 DAG（有向无环图）上的后验边占用（posterior edge occupancy）。
• 引入 Frontier Reveal Marginalization，一种预算受限的估计器，通过结合基于 CDF 的局部概率计算与高概率部分状态的前沿扩展（frontier expansion），来近似边沿边结构。

Rao-Blackwellized 策略梯度

Rao-Blackwellization 是一种经典的方差缩减技术，通过条件期望来减少估计的方差。FEROM 将其应用于揭示顺序的潜在变量，从而得到更稳定的梯度估计。具体而言，FEROM 将策略梯度表示为 DAG 上边的后验占用概率，这允许我们利用图结构进行高效计算。

Frontier Reveal Marginalization

为了实际计算边沿化，FEROM 设计了 Frontier Reveal Marginalization。该方法在给定预算下，通过 CDF 计算内生揭示过程的局部概率，并扩展高概率的部分状态，从而近似边沿边结构。这种预算受限的估计器在计算效率和近似精度之间取得了平衡。

与 GRPO 的结合

FEROM 被集成到 GRPO 风格的目標函数中，用基于后验相关揭示边的 detached token-score 代理替换了单路径 log-scores。这种替换使得目标函数能够更好地反映边沿策略的梯度，同时保持了实现的简洁性。

实验评估

论文在数学和编码任务上进行了实验。在匹配评分状态预算的条件下，FEROM 取得了与现有方法相当或更好的结果。此外，一项离线代理研究表明，随着评分预算的增加，边沿梯度近似得到改善。这些结果验证了 FEROM 的有效性和可扩展性。

结论

FEROM 为掩码扩散语言模型的策略优化提供了一种新颖的边沿化方法，通过 Rao-Blackwellized 策略梯度和前沿揭示边沿化，有效降低了方差并提升了性能。未来工作可以进一步探索该方法在其他生成任务上的应用。