论文

掩蔽扩散的自适应顺序策略

Adaptive Order Policies for Masked Diffusion

模型推理解码与生成控制

摘要

掩蔽扩散模型在捕获文本和蛋白质等领域离散序列的数据分布方面取得了巨大成功。这些模型通过从完全屏蔽的序列开始迭代地揭开标记来生成数据,揭开顺序通常是随机选择的或使用基于降噪器概率的启发式选择。在这项工作中,我们提出了一种在扩散模型之上使用额外的轻量级策略网络来学习揭露顺序的方案。我们提出的损失根据策略概率对掩蔽扩散损失中的项进行重新加权,并产生更喜欢降噪器更可能正确的位置的策略。我们在两种情况下研究这种损失:(i)在使用冻结的预训练降噪器的同时仅训练策略,以及(ii)与加权损失联合训练策略和降噪器以允许相互适应。我们证明,我们的方法在对标记排序敏感的问题上优于常见的启发式方法,例如组合问题、蛋白质以及各种编码和语言任务。