掩码扩散模型采样的路径规划
机器学习
2026-03-06 v5 人工智能
摘要
使用掩码扩散模型（MDMs）进行离散数据的任意顺序生成，为传统自回归模型提供了一种引人注目的替代方案，尤其是在缺乏自然因果顺序的数据领域。然而，当前流行的MDM与其成功的连续扩散模型对应物不同，它们采用简化的掩码推理，其中未掩码的token无法被迭代优化——即使存在错误。在本文中，我们通过引入一种名为路径规划（P2）的新型推理采样策略，充分挖掘了MDM的潜力，该策略将每个生成步骤分解为两个子阶段：规划与去噪。在P2框架下，规划器在每一步选择适当的、被标记为待更新的token，然后使用去噪器对其进行采样。我们证明P2概括了MDM现有的所有采样策略，并通过优化和更新现有未掩码token的新能力，显著提升了生成质量。我们从理论上证明，P2在数据的对数边际似然上建立了一个（新的）扩展证据下界（ELBO）。我们通过一系列规划器实例化P2，包括：1.) 自规划，2.) BERT规划，以及3.) 使用学习型规划器的训练规划，在一系列领域上为MDM带来了最先进的生成性能。具体而言，仅使用P2推理，我们观察到蛋白质序列可折叠性相对提升22%，RNA序列pLDDT提升8%，数学推理提升4%，故事生成（ROUGE分数）提升68%，以及在具有挑战性的pass@1指标上代码生成提升33%。
引用
@article{arxiv.2502.03540,
title = {Path Planning for Masked Diffusion Model Sampling},
author = {Fred Zhangzhi Peng and Zachary Bezemek and Sawan Patel and Jarrid Rector-Brooks and Sherwood Yao and Avishek Joey Bose and Alexander Tong and Pranam Chatterjee},
journal= {arXiv preprint arXiv:2502.03540},
year = {2026}
}