# Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design

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

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 34 of its 58 references that are in the map.
- First posted on arXiv 20 Feb 2025.

## 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, 30 of 34)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Maximum Likelihood Training of Score-Based Diffusion Models](https://mlsanity.com/paper/2101.09258/) (2021)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [Universal Guidance for Diffusion Models](https://mlsanity.com/paper/2302.07121/) (2023)
- [Generating Sequences by Learning to Self-Correct](https://mlsanity.com/paper/2211.00053/) (2022)
- [Practical and Asymptotically Exact Conditional Sampling in Diffusion Models](https://mlsanity.com/paper/2306.17775/) (2023)
- [Diffusion Model-Based Image Editing: A Survey](https://mlsanity.com/paper/2402.17525/) (2024)
- [BERT has a Mouth, and It Must Speak: BERT as a Markov Random Field Language Model](https://mlsanity.com/paper/1902.04094/) (2019)
- [Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control](https://mlsanity.com/paper/2402.15194/) (2024)
- [Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding](https://mlsanity.com/paper/2408.08252/) (2024)
- [Unlocking Guidance for Discrete State-Space Diffusion and Flow Models](https://mlsanity.com/paper/2406.01572/) (2024)
- [Dirichlet Flow Matching with Applications to DNA Sequence Design](https://mlsanity.com/paper/2402.05841/) (2024)
- [Generating Sentences by Editing Prototypes](https://mlsanity.com/paper/1709.08878/) (2017)
- [Fine-Tuning Discrete Diffusion Models via Reward Optimization with Applications to DNA and Protein Design](https://mlsanity.com/paper/2410.13643/) (2024)
- [Monte Carlo guided Diffusion for Bayesian linear inverse problems](https://mlsanity.com/paper/2308.07983/) (2023)
- [Particle Denoising Diffusion Sampler](https://mlsanity.com/paper/2402.06320/) (2024)
- [Elements of Sequential Monte Carlo](https://mlsanity.com/paper/1903.04797/) (2019)
- [DPLM-2: A Multimodal Diffusion Protein Language Model](https://mlsanity.com/paper/2410.13782/) (2024)

## Built on by (most cited first, 25 of 25)

- [Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models](https://mlsanity.com/paper/2508.09968/) (2025)
- [SteerX: Creating Any Camera-Free 3D and 4D Scenes with Geometric Steering](https://mlsanity.com/paper/2503.12024/) (2025)
- [Inference-Time Scaling of Discrete Diffusion Models via Importance Weighting and Optimal Proposal Design](https://mlsanity.com/paper/2505.22524/) (2025)
- [Autonomous Agents for Scientific Discovery: Orchestrating Scientists, Language, Code, and Physics](https://mlsanity.com/paper/2510.09901/) (2025)
- [It's Never Too Late: Noise Optimization for Collapse Recovery in Trained Diffusion Models](https://mlsanity.com/paper/2601.00090/) (2025)
- [CORE: Context-Robust Remasking for Diffusion Language Models](https://mlsanity.com/paper/2602.04096/) (2026)
- [MIRA: Towards Mitigating Reward Hacking in Inference-Time Alignment of T2I Diffusion Models](https://mlsanity.com/paper/2510.01549/) (2025)
- [Steering Generative Models with Experimental Data for Protein Fitness Optimization](https://mlsanity.com/paper/2505.15093/) (2025)
- [Ctrl-DNA: Controllable Cell-Type-Specific Regulatory DNA Design via Constrained RL](https://mlsanity.com/paper/2505.20578/) (2025)
- [Posterior Inference in Latent Space for Scalable Constrained Black-box Optimization](https://mlsanity.com/paper/2507.00480/) (2025)
- [Inference-Time Alignment for Diffusion Models via Variationally Stable Doob's Matching](https://mlsanity.com/paper/2601.06514/) (2026)
- [VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL](https://mlsanity.com/paper/2505.15791/) (2025)
- [ProSpero: Active Learning for Robust Protein Design Beyond Wild-Type Neighborhoods](https://mlsanity.com/paper/2505.22494/) (2025)
- [Language Models for Controllable DNA Sequence Design](https://mlsanity.com/paper/2507.19523/) (2025)
- [LENS: Low-Frequency Eigen Noise Shaping for Efficient Diffusion Sampling](https://mlsanity.com/paper/2605.07253/) (2026)
- [Science sandboxes measure the scientific capability of AI agents](https://mlsanity.com/paper/2608.30165/) (2026)
- [From Snapshots to Symphonies: The Evolution of Protein Prediction from Static Structures to Generative Dynamics and Multimodal Interactions](https://mlsanity.com/paper/2603.18505/) (2026)
- [LPDP: Inference-Time Reward Control for Variable-Length DNA Generation with Edit Flows](https://mlsanity.com/paper/2605.11368/) (2026)
- [Generative Refinement for Low-Budget Black-Box Optimization](https://mlsanity.com/paper/2607.00691/) (2026)
- [On the Design Space of Discrete Diffusion Online Adaptation for Molecular Optimization](https://mlsanity.com/paper/2607.02834/) (2026)
- [When Genomic Masking Priors Fail to Transfer: Strong Variant Prediction, Weak Functional Generation](https://mlsanity.com/paper/2609.04861/) (2026)
- [Steer2Grasp: Inference-Time Embodiment-Aware Steering for Diverse Physically Feasible Grasp Diffusion](https://mlsanity.com/paper/2609.33546/) (2026)
- [ENCORE: Exact Non-equilibrium COntrol with Replica Exchange for Diffusion Generation](https://mlsanity.com/paper/2610.02538/) (2026)
- [D-DOIT: Training-free Adaptation of Discrete Diffusion via Doob's h-Transform](https://mlsanity.com/paper/2610.04938/) (2026)
- [DiMOS: Doob-Guided Inference-Time Multi-Objective Search for Scientific Design](https://mlsanity.com/paper/2610.05808/) (2026)

## Abstract (arXiv)

To fully leverage the capabilities of diffusion models, we are often interested in optimizing downstream reward functions during inference. While numerous algorithms for reward-guided generation have been recently proposed due to their significance, current approaches predominantly focus on single-shot generation, transitioning from fully noised to denoised states. We propose a novel framework for inference-time reward optimization with diffusion models inspired by evolutionary algorithms. Our approach employs an iterative refinement process consisting of two steps in each iteration: noising and reward-guided denoising. This sequential refinement allows for the gradual correction of errors introduced during reward optimization. Besides, we provide a theoretical guarantee for our framework. Finally, we demonstrate its superior empirical performance in protein and cell-type-specific regulatory DNA design. The code is available at hrefhttps://github.com/masa-ue/ProDifEvo-Refinementhttps://github.com/masa-ue/ProDifEvo-Refinement.