让扩散模型自动学习最优掩码解码顺序，提升文本与蛋白质生成效果。
Adaptive Order Policies for Masked Diffusion
- 用轻量级策略网络学习掩码解码顺序，根据去噪器置信度动态调整。
- 在组合任务和蛋白质生成上优于随机或启发式顺序，性能显著提升。
- 适合需要精确序列生成的场景，如生物序列设计与复杂语言建模。
掩码扩散模型在文本、蛋白质等离散序列领域表现优异，通过从全掩码状态逐轮解码生成数据，传统方法通常采用随机或基于去噪器概率的启发式顺序。本文提出一种新方案：在扩散模型上增加轻量级策略网络，学习最优解码顺序。通过重加权掩码扩散损失，使策略倾向于选择去噪器更可能正确的位置。研究了两种设置：（i）仅训练策略网络，冻结预训练去噪器；（ii）联合训练策略与去噪器，实现相互适应。实验表明，该方法在对顺序敏感的任务（如组合问题和蛋白质生成）中显著优于常见启发式策略。
原文摘要 · Abstract (English)
Masked diffusion models have seen great success in capturing data distributions over discrete sequences in domains such as text and proteins. These models generate data by iteratively unmasking tokens starting from a fully masked sequence, with the unmasking order typically chosen at random or using a heuristic based on denoiser probabilities. In this work, we propose a scheme for learning the unmasking order using an additional lightweight policy network on top of a diffusion model. Our proposed loss reweights terms in the masked diffusion loss according to policy probabilities, and results in a policy that prefers positions where the denoiser is more likely to be correct. We study this loss in two settings: (i) training solely the policy while using a frozen pre-trained denoiser, and (ii) training the policy and denoiser jointly with the weighted loss to allow for mutual adaptation. We demonstrate that our approach outperforms common heuristics on problems that are sensitive to token ordering, such as combinatorial tasks and proteins.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。