Meta：扩散模型轨迹概率高效优化

标题dTRPO: Trajectory Reduction in Policy Optimization of Diffusion Large Language Models来源arXiv, 2603.18806v1摘要扩散大语言模型diffusion Large language ModelsdLLM为语言生成引入了一种新的范式这反过来又为将它们与人类偏好相匹配带来了新的挑战。在这项工作中我们的目标是通过降低轨迹概率计算的成本来改进dLLM的策略优化从而实现规模化的离线策略训练。我们证明i在参考策略正则化下新的未屏蔽令牌的概率比是中间扩散状态的概率比的无偏估计以及ii整个轨迹的概率可以用重新计算的单个前向通过来有效地估计通过将这两个轨迹缩减策略集成到策略优化目标中我们提出了轨迹缩减策略优化dTRPO。我们在7B dLLM上跨指令评估dTRPO跟踪和推理基准。结果表明它大大提高了最先进的dLLM的核心性能在STEM任务上实现了高达9.6%的增益在编码任务上实现了高达4.3%的增益在推理跟踪任务上实现了高达3.0%的增益。此外dTRPO由于其离线、单向前性质而表现出强大的训练效率并通过高质量输出实现了提高的生成效率。️文章简介研究问题如何在扩散大语言模型dLLM的策略优化中以极低的计算成本准确估计复杂的生成轨迹概率从而实现可扩展的离线训练主要贡献论文提出了 dTRPO 方法通过理论证明将轨迹概率比简化为新解掩码 token 的乘积仅需单次前向传播即可完成高效的离线策略优化。重点思路状态缩减策略利用分块注意力机制证明在每个生成块中仅采样一个时间步即可无偏估计整个扩散轨迹的对数概率将计算量从多步降低为单步。比率缩减策略在参考策略正则化下推导出当前策略与参考策略的转移概率比中依赖于调度表的系数会相互抵消仅保留新解掩码 token 的分类概率比。构建 dTRPO 目标函数将上述两种缩减策略集成到直接偏好优化DPO框架中设计出新损失函数使其能够像自回归模型一样通过单次前向传播计算偏好损失。推理对齐调度采用与推理时解码策略一致的基于置信度的解掩码调度器确保训练时的概率估计分布与实际生成过程相匹配提升稳定性。分析总结性能显著提升在 7B 参数模型上的实验显示dTRPO 在 STEM 任务上提升高达 9.6%编码任务提升 4.3%指令遵循任务提升 3.0%大幅缩小了与自回归模型的差距。训练效率极高该方法完全基于离线数据每个训练样本仅需 4 次前向传播策略与参考模型各两次计算成本与传统自回归模型的 DPO 训练相当远低于在线强化学习方法。通用性与鲁棒性该方法不仅适用于分块扩散模型也适用于长块扩散模型如 LLaDA且对投影函数选择和超参数设置表现出较强的鲁棒性。生成质量优化相比其他基线方法dTRPO 生成的回答更简洁准确同时在保持高吞吐量的前提下显著减少了无效生成导致的推理延迟。个人观点论文打破了扩散模型策略优化必须依赖昂贵多步 rollout 的固有认知通过数学推导将复杂的轨迹概率估计坍缩为简单的 token 级概率比。附录