# Spectral Feedback for Test-Time Alignment of Protein Diffusion Models

arXiv: https://arxiv.org/abs/2609.30456 · HTML: https://mlsanity.com/paper/2609.30456/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 16 of its 41 references that are in the map.
- First posted on arXiv 24 Sep 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 16 of 16)

- [Null-text Inversion for Editing Real Images using Guided Diffusion Models](https://mlsanity.com/paper/2211.09794/) (2022)
- [Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design](https://mlsanity.com/paper/2402.04997/) (2024)
- [Remasking Discrete Diffusion Models with Inference-Time Scaling](https://mlsanity.com/paper/2503.00307/) (2025)
- [Score-based Continuous-time Discrete Diffusion Models](https://mlsanity.com/paper/2211.16750/) (2022)
- [Protein Design with Guided Discrete Diffusion](https://mlsanity.com/paper/2305.20009/) (2023)
- [Unlocking Guidance for Discrete State-Space Diffusion and Flow Models](https://mlsanity.com/paper/2406.01572/) (2024)
- [Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment](https://mlsanity.com/paper/2503.21878/) (2025)
- [Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review](https://mlsanity.com/paper/2501.09685/) (2025)
- [Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design](https://mlsanity.com/paper/2410.13643/) (2024)
- [Graph Denoising Diffusion for Inverse Protein Folding](https://mlsanity.com/paper/2306.16819/) (2023)
- [DiffusER: Discrete Diffusion via Edit-based Reconstruction](https://mlsanity.com/paper/2210.16886/) (2022)
- [Discrete Diffusion Models: Novel Analysis and New Sampler Guarantees](https://mlsanity.com/paper/2509.16756/) (2025)
- [ProxySPEX: Inference-Efficient Interpretability via Sparse Feature Interactions in LLMs](https://mlsanity.com/paper/2505.17495/) (2025)
- [A Fast Hadamard Transform for Signals with Sub-linear Sparsity in the Transform Domain](https://mlsanity.com/paper/1310.1803/) (2013)
- [Effective Test-Time Scaling of Discrete Diffusion through Iterative Refinement](https://mlsanity.com/paper/2511.05562/) (2025)
- [ProteinGuide: On-the-fly property guidance for protein sequence generative models](https://mlsanity.com/paper/2505.04823/) (2025)

## Abstract (arXiv)

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.