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.
Figures & tables
Figure 1: Spectral Feedback loop using targeted edits to iteratively improve attributes of a protein sequence with discrete diffusion.
Figure 2: Spectral Feedback in different protein model inference settings, aligning for stability ( ΔΔG - kcal/mol). These are reward evaluation distributions from experiments in section 5 .
Figure 3: Spectral analysis of the value function favg with the pretrained diffusion model. The bar graph shows the proportion that each order of Fourier coefficients contributes to the total variance of the coefficients. The plots compare the 2KVV and R6-560 (r6_560_TrROS_Hall) backbones from the Megascale dataset. Experimental details and additional plots are in Appendix D.2 .
Figure 4: Spectral Feedback learns a sparse Fourier representation of the edit-set value function, then solves for the max-reward set via integer optimization over the sparse coefficients.
Figure 5
Figure 5: (left) Spectral Feedback is applied to several underlying protein models: pretrained, Best-of-10 with pretrained, and DRAKES. Spectral Feedback improves the alignment reward of each algorithm in only a few iterations. (middle) Percentage of stable proteins (evaluation ΔΔG>0 ) and the resulting success rate when scRMSD<2 are shown. The scRMSD constraint and metric averages are in Table 1 . ( right ) ESMFold protein structures comparing the wild-type ( red ) and Spectral Feedback ( green ) sequences demonstrate that Spectral Feedback preserves the underlying inverse protein folding model’s capabilities.
Method
Align ΔΔG↑
Eval ΔΔG (avg) ↑
scRMSD (avg) ↓
scRMSD <2 (%) ↑
Log-Likelihood ↑
Pretrained
0.10
-0.49
1.10
91.8
-147.8
+ Spectral
0.53 (+0.43)
0.33 (+0.82)
1.06 (-0.04)
94.2 (+2.4)
-146.5 (+1.3)
BoN-10
0.45
0.09
1.09
92.4
-148.1
+ Spectral
0.70 (+0.25)
0.86 (+0.77)
1.07 (-0.02)
96.7 (+4.3)
-147.5 (+0.6)
DRAKES
0.69
0.99
1.17
93.4
-157.5
+ Spectral
0.83 (+0.14)
1.38 (+0.39)
1.53 (+0.36)
85.8 (-7.6)
-157.7 (-0.2)
Table 1: Spectral Feedback improves alignment and evaluation oracle rewards across all three base methods, with the largest gains on the Pretrained and Best-of-10 baselines. DRAKES + Spectral achieves the highest evaluation ΔΔG overall (1.38) but trades off structural quality (scRMSD). Each pair shows baseline followed by its Spectral-enhanced counterpart.
Figure 6: Spectral Feedback shows strong performance in scaling of edit-set samples and latency.
Figure 9: Exponential tilt function Fourier analysis for 2KRU with the pretrained model.
Figure 10: Exponential tilt function Fourier analysis for r6_560_TrROS_Hall with pretrained model.
Figure 11: scRMSD vs ΔΔG distributions show a worse tradeoff with scRMSD for DRAKES when applying Spectral Feedback.
Figure 12: Spectral Feedback with k=20 , D=8192 , and 5 feedback iterations.
Figure 13: ProtGPT2 Log-Likelihood alignment trajectory with Spectral Feedback. Hyperparameters are k=10 , D=2048 , n=16 (reward oracle calls per value function call).
Figure 14: An example multi-objective oracle balances between ΔΔG and ProtGPT2 Log-Likelihood.
Figure 15: Spectral Feedback performance comparison with ProxySPEX vs first-order LASSO as the edit selection method. We use the pretrained model and constrain edit-selections to at most k=20 positions. (left) Visualization of one held-out feedback trajectory for the R6-560 protein backbone. (right) Average Spectral Feedback alignment across the test set, comparing ProxySPEX and LASSO as the edit-selection strategies.
We propose D-DOIT (Discrete Doob-Oriented Inference-time Transformation), a training-free and efficient adaptation method for discrete diffusion models with generic rewards. D-DOIT formulates adaptation as sampling from a reward-tilted target distribution and realizes this transport through Doob's h-transform of the discrete diffusion reverse kernel, using only reward values rather than reward gradients. Unlike continuous diffusion, masked discrete diffusion samples categorical token-reveal transitions rather than continuous state updates. D-DOIT derives the corresponding discrete Doob's h-transform, which guides sampling by reweighting reverse transition probabilities instead of adding a drift correction. To make this transformation practical, D-DOIT avoids expensive future rollouts. At each guided step, D-DOIT samples candidate next states, uses the model prediction head to complete each candidate into a clean sequence, evaluates each completion with the reward oracle, and resamples the next state with probabilities proportional to the rewards. An optional late-stage best-of-K refinement further improves sample quality by branching trajectories only near the end of denoising, avoiding the K-fold cost over the full trajectory. Empirically, across regulatory DNA design and protein inverse folding benchmarks, D-DOIT outperforms training-free guidance baselines. It improves enhancer activity and cell-type specificity while preserving sequence naturalness, and achieves the highest success rate in protein inverse folding.
Jieke Wu, Qijie Zhu, Weimin Wu +3
Department of Computer Science, King Abdullah University of Science and Technology · Department of Statistics and Data Science, Northwestern University · Center for Foundation Models and Generative AI, Northwestern University +1
Online reinforcement learning is becoming increasingly important for aligning diffusion models with non-differentiable objectives. However, existing methods still face limitations in assigning fine-grained credit along denoising trajectories and in realizing stable value-based optimization. We propose a state-aligned latent actor-critic framework for diffusion post-training, in which the diffusion model serves as its own timestep-conditioned value function and predicts values directly on noisy latent states. This enables trajectory-level PPO training, supports stable actor-critic optimization with simple conditioning and value pretraining strategies, and naturally allows the learned critic to be reused for inference-time steering. We further extend the framework to multi-reward optimization, where joint training with complementary rewards helps alleviate reward hacking. Across both UNet- and DiT-based backbones, our method consistently outperforms prior group-relative RL and actor-critic baselines on single-reward and multi-reward benchmarks, while test-time steering provides additional gains in generation quality.
Zhengyang Liang, Qihang Zhang, Ceyuan Yang
University of Toronto · Vector Institute · The Chinese University of Hong Kong
Reinforcement learning based post training of diffusion models, such as Denoising Diffusion Policy Optimization (DDPO), optimizes a reverse diffusion process under a reward function. However, current approaches to reward optimizations do so at the cost of diversity and quality. In this paper, we provide better tradeoffs through careful theoretical considerations and method design. We analyze the theoretical framework and mathematically demonstrate that only-latter timestep updates of diffusion model may be harmful for diversity contrary to the conclusions presented in a previous work. Additionally, we propose an incremental Feynman-Kac training based on strong theoretical foundations in order to achieve the best-yet alignment-diversity tradeoffs. We perform extensive experiments and compare our method against related diffusion policy optimization approaches in three different tasks and also provide strong ablations for each component, thus validating strong performance gains in both alignment and diversity.