Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design
Abstract
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. 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. Finally, we provide a theoretical guarantee for our framework. Finally, we demonstrate its superior empirical performance in protein and DNA design.
Cite
Text
Uehara et al. "Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design." Proceedings of the 42nd International Conference on Machine Learning, 2025.Markdown
[Uehara et al. "Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design." Proceedings of the 42nd International Conference on Machine Learning, 2025.](https://mlanthology.org/icml/2025/uehara2025icml-rewardguided/)BibTeX
@inproceedings{uehara2025icml-rewardguided,
title = {{Reward-Guided Iterative Refinement in Diffusion Models at Test-Time with Applications to Protein and DNA Design}},
author = {Uehara, Masatoshi and Su, Xingyu and Zhao, Yulai and Li, Xiner and Regev, Aviv and Ji, Shuiwang and Levine, Sergey and Biancalani, Tommaso},
booktitle = {Proceedings of the 42nd International Conference on Machine Learning},
year = {2025},
pages = {60515-60529},
volume = {267},
url = {https://mlanthology.org/icml/2025/uehara2025icml-rewardguided/}
}