论文
Spectral Feedback:选择重编辑位置以对齐蛋白扩散生成
Spectral Feedback for Test-Time Alignment of Protein Diffusion Models
摘要
离散扩散模型的奖励最大化对齐方法主要集中于通过影响token logits或在中间步骤选择有利的序列来引导反向过程。这些方法在很大程度上将推理视为单向过程,缺乏重新审视不良标记选择的机制。我们引入了Spectral Feedback,这是一种在反馈循环中选择编辑位置的算法,允许模型迭代地纠正自己的生成。这种方法通过重新屏蔽和重新采样标记来利用离散扩散模型的掩模结构,类似于重新引入含噪潜变量并重新运行相反过程的图像编辑方法。虽然先前的对齐方法侧重于分配哪些令牌标签以最大化目标奖励,但我们将重新审视哪些令牌视为中心对齐问题。选择编辑位置具有挑战性,因为编辑效果是相互依赖的:修改一个标记的影响取决于同时编辑的其他标记。我们将编辑集定义为一组要重新屏蔽和重新采样的标记位置。受先前关于生物系统中稀疏相互作用的研究的启发,我们凭经验发现蛋白质逆折叠的编辑集值函数允许稀疏傅立叶表示。这种结构使Spectral Feedback能够有效地学习和优化编辑位置选择的值函数。Spectral Feedback与模型无关,可应用于预训练、测试时间对齐和微调的扩散模型。对于所有这些模型,该算法在不修改底层生成过程的情况下提高了对齐性能。应用于具有蛋白质稳定性奖励预言机的反向折叠时,预训练模型的稳定蛋白质增加了 32.3%,Best-of-10 增加了 24.8%,最先进的 RL 微调扩散模型增加了 5.8%。
