Spectral feedback improves protein design by iterative editing
Spectral Feedback for Test-Time Alignment of Protein Diffusion Models
Artificial Intelligence
Summary
Designing proteins that fold stably is a big challenge, and current AI methods often make mistakes they can't easily fix. The authors introduce Spectral Feedback, a way for protein-generating models to identify which parts of a protein sequence should be rethought and edited repeatedly to improve results. This method uses mathematical tools to efficiently pick the best spots to revise, boosting the number of stable proteins created. It works with existing models without needing to change their internal workings.
What this means in practice
- •For biotech engineers: Generate more stable proteins by iteratively refining sequences during generation using Spectral Feedback with existing diffusion models.
- •For software engineers building ai systems: Enhance discrete diffusion models by adding test-time feedback loops that select and re-sample tokens for improved output alignment without retraining.
Authors
Shai Dickman, Mert Cemri, Landon Butler, Kannan Ramchandran
Abstract
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.