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

TL;DR

This paper introduces a novel iterative refinement framework for reward-guided inference in diffusion models, inspired by evolutionary algorithms, demonstrating improved performance in protein and DNA design tasks.

Contribution

It presents a new iterative reward-guided inference method for diffusion models with theoretical guarantees, extending beyond single-shot approaches.

Findings

Superior empirical performance in protein design

Effective application to cell-type-specific DNA design

Theoretical guarantee for the proposed framework

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 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 \href{https://github.com/masa-ue/ProDifEvo-Refinement}{https://github.com/masa-ue/ProDifEvo-Refinement}.