Spectral Feedback for Test-Time Alignment of Protein Diffusion Models

Alignment methods for discrete diffusion models have primarily focused on steering the denoising process, either by influencing token logits or by selecting favorable sequences at intermediate steps. However, these approaches largely treat inference as a unidirectional process, lacking effective mechanisms for revisiting undesirable token selections. We introduce Spectral Feedback, an algorithm that uses a feedback loop that selects sets of positions to edit, allowing the model to iteratively correct its own generations. This approach leverages the masking structure of discrete diffusion models where tokens can be re-masked and re-sampled, analogous to image editing methods that reintroduce noise and re-run a reverse diffusion process. While prior methods focus on what token values to assign, 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. Motivated by prior work on sparse interactions in biological systems, we show empirically that, for protein inverse folding, the expected reward over edit-position sets admits a sparse Fourier representation. This structure enables efficient learning and optimization of a value function over edit-position sets, which Spectral Feedback uses to select promising edits. The algorithm is model-agnostic and can be applied to pretrained, test-time aligned, or fine-tuned protein diffusion models, improving alignment performance without modifying the underlying generative process. Applied to inverse folding with a protein stability oracle, it achieves a 30.6% increase in stable proteins for a pretrained model, 24.8% for Best-of-10, and 8.5% for the SOTA RL-tuned diffusion models.