arXiv:2502.03540cs.LGcs.AI2025-02被引 77

提出路径规划策略,让掩码扩散模型能迭代修正已生成内容,提升生成质量。

Path Planning for Masked Diffusion Model Sampling

  • 分两阶段生成:先规划需更新的标记,再用去噪器采样
  • 在蛋白质、代码等任务上性能提升22%至68%
  • 支持自适应规划与训练规划,适用于多种离散数据生成

使用掩码扩散模型(MDMs)进行离散数据的任意顺序生成,为缺乏自然因果顺序的领域提供了优于传统自回归模型的替代方案。然而,现有主流MDMs在推理时简化了掩码机制,导致未掩码的令牌无法迭代优化——即使存在错误也无从修正。本文提出一种新的采样策略‘路径规划’(P2),将每一步生成分解为规划与去噪两个子阶段。在该策略下,规划器每步选择需更新的令牌,随后由去噪器采样。我们证明P2可统一现有所有MDM采样策略,并通过允许对已有未掩码令牌进行修正,显著提升生成质量。理论证明,P2建立了数据对数边缘似然的新扩展证据下界(ELBO)。我们设计了三类规划器:自规划、BERT规划和训练规划,其中训练规划器在多个任务上达到当前最优表现。仅采用P2推理,便在蛋白质折叠性上提升22%,RNA pLDDT提升8%,数学推理提升4%,故事生成(ROUGE)提升68%,代码生成(pass@1)提升33%。

原文摘要 · Abstract (English)

Any order generation of discrete data using masked diffusion models (MDMs) offers a compelling alternative to traditional autoregressive models, especially in domains that lack a natural causal ordering of data. However, current popular MDMs depart from their successful continuous diffusion model counterparts with simplified masked inference wherein unmasked tokens cannot be iteratively refined -- even if there is a mistake. In this paper, we extract the full power of MDMs by introducing a novel inference sampling strategy termed Path Planning (P2) that decomposes each generation step into two sub-stages: planning and denoising. Under P2, the planner at every step selects appropriate tokens that are marked to be updated, which can then be sampled using the denoiser. We demonstrate that P2 generalizes all existing sampling strategies for MDMs and critically enhances generative quality through the new capability of refining and updating existing unmasked tokens. We theoretically prove that P2 establishes a (new) expanded evidence lower bound (ELBO) on the log marginal likelihood of data. We instantiate P2 with a family of planners including: 1.) Self-Planning, 2.) BERT-Planning, and 3.) Trained-Planning with a learned planner leading to SOTA generative performance for MDMs on a suite of domains. Specifically, solely using P2 inference, we observe relative improvements of 22% in protein sequence foldability, 8% in RNA sequence pLDDT, 4% in math reasoning, 68% in story generation (ROUGE score), and 33% in code generation for the challenging pass@1 metric.

扩散模型路径规划生成质量序列建模

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。