扩散语言模型训练中的规划感知路径学习
机器学习
2026-03-09 v3
摘要
扩散语言模型已成为自回归模型的强有力替代品,能够通过更灵活且并行的生成路径实现快速推理。这种生成路径的灵活性是通过新的工程化采样策略——即规划器——实现的,这些规划器通过迭代规划来选择更有利的生成路径,而不是均匀随机地选择去噪位置。然而,通过规划修改逆向路径后,规划器会在训练时假设的均匀随机去噪路径与基于规划的推理之间产生不可撤销的不匹配。在本文中,我们系统性地调查了离散扩散训练与规划下推理之间的这种不匹配,并理论上证明了标准离散扩散训练证据下界 (ELBO) 并不准确地描述使用非均匀规划器的去噪器。为解决这一差距,我们推导了新的计划证据下界 (P-ELBO),该下界直接将规划器基的逆向动力学纳入训练目标。使用 P-ELBO,我们引入了规划感知路径学习 (PAPL),这是一种新颖的训练方案,在计划去噪器下对齐训练与推理。PAPL 实现为标准遮盖离散扩散损失的简单而有效的修改,具有广泛适用性和易于采纳性。经验上,我们展示 PAPL 在多个领域均实现了持续的改进,包括在蛋白质序列上实现 40% 的相对改进,提升文本生成的最高达 4 倍的相对 MAUVE 增益,以及在代码生成 HumanEval pass@10 上实现 23% 的相对改进。代码已在 github.com/pengzhangzhi/PAPL 上提供。
引用
@article{arxiv.2509.23405,
title = {Planner Aware Path Learning in Diffusion Language Models Training},
author = {Fred Zhangzhi Peng and Zachary Bezemek and Jarrid Rector-Brooks and Shuibai Zhang and Anru R. Zhang and Michael Bronstein and Alexander Tong and Avishek Joey Bose},
journal= {arXiv preprint arXiv:2509.23405},
year = {2026}
}
备注
Camera ready version for ICLR2026