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

Masatoshi Uehara, Xingyu Su, Yulai Zhao, Xiner Li, Aviv Regev, Shuiwang Ji, Sergey Levine, Tommaso Biancalani

TL;DR: Diffusion models

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.

Lay Summary: To get the most out of diffusion models, we often want to improve certain outcomes or "rewards" when generating results. Many recent methods have tried to do this, but they typically rely on a one-step process that goes directly from random noise to a final output. In contrast, we introduce a new approach that improves results step by step. Each step involves adding some noise and then refining the output based on the reward we want to achieve. This gradual process helps fix mistakes along the way. We also show that our method is both theoretically sound and performs better in practice, especially for designing proteins and DNA.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Link To Code: https://github.com/masa-ue/ProDifEvo-Refinement

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: Diffusion models

Submission Number: 12256