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.