← 研踪 PaperTrailarXiv:2609.30456 · v2
Spectral Feedback for Test-Time Alignment of Protein Diffusion Models
首次提交 2026-09-25 02:49
版本更新 2026-11-07 17:37
雷达观测 2026-11-10 11:14
摘要
Discrete diffusion models have become leading architectures for inverse protein folding. Test-time alignment guides these models during inference to generate proteins with desirable attributes. Alignment methods for discrete diffusion models often treat inference as a unidirectional process, lacking effective 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 masking structure of discrete diffusion models by re-masking and re-sampling tokens. 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 positions to edit by re-masking and re-sampling the corresponding tokens in a sequence. 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 pre-trained, test-time aligned, and fine-tuned diffusion models without modifying the underlying generative processes. Applied to inverse folding with a protein stability reward oracle, Spectral Feedback increases the proportion of stable proteins by 32.3 percentage points over a pre-trained model, 24.8 over Best-of-10, and 5.8 over a state-of-the-art RL fine-tuned diffusion model.
作者备注：Neurips 2026
观测到的版本
相邻记录可能跳过中间版本；首次发现 v2 及以上也会标为版本更新。
v2 · 2026-11-07 17:37 · 最新观测版本
Spectral Feedback for Test-Time Alignment of Protein Diffusion Models
Discrete diffusion models have become leading architectures for inverse protein folding. Test-time alignment guides these models during inference to generate proteins with desirable attributes. Alignment methods for discrete diffusion models often treat inference as a unidirectional process, lacking effective 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 masking structure of discrete diffusion models by re-masking and re-sampling tokens. 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 positions to edit by re-masking and re-sampling the corresponding tokens in a sequence. 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 pre-trained, test-time aligned, and fine-tuned diffusion models without modifying the underlying generative processes. Applied to inverse folding with a protein stability reward oracle, Spectral Feedback increases the proportion of stable proteins by 32.3 percentage points over a pre-trained model, 24.8 over Best-of-10, and 5.8 over a state-of-the-art RL fine-tuned diffusion model.
查看此版本解读 →查看此版本原文 ↗