谱反馈用于蛋白质扩散模型的测试时对齐
Spectral Feedback for Test-Time Alignment of Protein Diffusion Models
- University of California, Berkeley（加州大学伯克利分校）
机构由 AI 辅助整理，请以论文原文为准。
AI总结：
谱反馈通过反馈循环选择编辑位置，利用稀疏傅里叶结构优化蛋白质逆折叠扩散模型的对齐，无需修改生成过程，显著提升稳定性。
AI中文摘要：
针对离散扩散模型的奖励最大化对齐方法主要集中于引导逆向过程，要么通过影响令牌对数概率，要么在中间步骤选择有利的序列。这些方法在很大程度上将推理视为单向过程，缺乏重新审视不良令牌选择的机制。我们引入了谱反馈（Spectral Feedback），一种在反馈循环中选择编辑位置的算法，使模型能够迭代地纠正自身的生成。该方法利用离散扩散模型的掩码结构，通过重新掩码和重新采样令牌，类似于图像编辑方法中重新引入噪声潜变量并重新运行逆向过程。虽然先前的对齐方法侧重于分配哪些令牌标签以最大化目标奖励，我们反而将重新访问哪些令牌视为核心对齐问题。选择编辑位置具有挑战性，因为编辑效果是相互依赖的：修改一个令牌的影响取决于同时编辑的其他令牌。我们将编辑集定义为要重新掩码和重新采样的一组令牌位置。受生物系统中稀疏相互作用的先前工作的启发，我们经验性地发现，蛋白质逆折叠的编辑集价值函数具有稀疏傅里叶表示。这种结构使谱反馈能够高效地学习和优化编辑位置选择的价值函数。谱反馈是模型无关的，可应用于预训练、测试时对齐和微调的扩散模型。对于所有这些模型，该算法在不修改底层生成过程的情况下提高了对齐性能。应用于具有蛋白质稳定性奖励预言机的逆折叠任务时，对于预训练模型，稳定蛋白质增加了32.3%，对于Best-of-10增加了24.8%，对于最先进的强化学习微调扩散模型增加了5.8%。
英文摘要：
Reward maximization alignment methods for discrete diffusion models have primarily focused on steering the reverse process, either by influencing token logits or by selecting favorable sequences at intermediate steps. These approaches largely treat inference as a unidirectional process, lacking mechanisms for revisiting undesirable token selections. We introduce Spectral Feedback, an algorithm that selects edit-positions in a feedback loop, allowing the model to iteratively correct its own generations. This approach leverages the mask structure of discrete diffusion models by re-masking and re-sampling tokens, analogous to image editing methods that reintroduce noisy latents and re-run the reverse process. While prior alignment methods focus on what token labels to assign to maximize a target reward, we instead treat which tokens to revisit as the central alignment problem. Selecting edit-positions is challenging because edit effects are interdependent: the impact of modifying one token depends on which others are edited simultaneously. We define an edit-set as a set of token positions to re-mask and re-sample. Motivated by prior work on sparse interactions in biological systems, we find empirically that edit-set value functions for protein inverse folding admit sparse Fourier representations. This structure enables Spectral Feedback to efficiently learn and optimize the value functions for edit-position selection. Spectral Feedback is model-agnostic and can be applied to pretrained, test-time aligned, and fine-tuned diffusion models. For all of these models, the algorithm improves alignment performance without modifying the underlying generative process. Applied to inverse folding with a protein stability reward oracle, it achieves a 32.3% increase in stable proteins for a pretrained model, 24.8% for Best-of-10, and 5.8% for a state-of-the-art RL fine-tuned diffusion model.
