Inference Alignment for Discrete Diffusion Models via Targeted Sequence Editing
Shai Dickman
EECS Department, University of California, Berkeley
Technical Report No. UCB/EECS-2026-140
May 14, 2026
http://www2.eecs.berkeley.edu/Pubs/TechRpts/2026/EECS-2026-140.pdf
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. In addition to introducing Spectral Feedback, we conclude by formalizing an alternative perspective of randomized edit-position selection that aligns an underlying protein model to a Gibbs-distribution. We verify convergence guarantees of this alternative method and pose it as a way to re-adapt Spectral Feedback for preventing over-fitting to a reward oracle.
Advisors: Kannan Ramchandran
BibTeX citation:
@mastersthesis{Dickman:EECS-2026-140,
Author= {Dickman, Shai},
Title= {Inference Alignment for Discrete Diffusion Models via Targeted Sequence Editing},
School= {EECS Department, University of California, Berkeley},
Year= {2026},
Month= {May},
Url= {http://www2.eecs.berkeley.edu/Pubs/TechRpts/2026/EECS-2026-140.html},
Number= {UCB/EECS-2026-140},
Abstract= {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. In addition to introducing Spectral Feedback, we conclude by formalizing an alternative perspective of randomized edit-position selection that aligns an underlying protein model to a Gibbs-distribution. We verify convergence guarantees of this alternative method and pose it as a way to re-adapt Spectral Feedback for preventing over-fitting to a reward oracle.},
}
EndNote citation:
%0 Thesis %A Dickman, Shai %T Inference Alignment for Discrete Diffusion Models via Targeted Sequence Editing %I EECS Department, University of California, Berkeley %D 2026 %8 May 14 %@ UCB/EECS-2026-140 %U http://www2.eecs.berkeley.edu/Pubs/TechRpts/2026/EECS-2026-140.html %F Dickman:EECS-2026-140