Algorithms for Optimal Adaptation ofDiffusion Models to Reward Functions

Krishnamurthy Dj Dvijotham, Shayegan Omidshafiei, Kimin Lee, Katherine M. Collins, Deepak Ramachandran, Adrian Weller, Mohammad Ghavamzadeh, Milad Nasr, Ying Fan, Jeremiah Zhe Liu

Keywords: diffusion models; optimal control

Abstract: We develop algorithms for adapting pretrained diffusion models to optimize reward functions while retaining fidelity to the pretrained model. We propose a general framework for this adaptation that trades off fidelity to a pretrained diffusion model and achieving high reward. Our algorithms take advantage of the continuous nature of diffusion processes to pose reward-based learning either as a trajectory optimization or continuous state reinforcement learning problem. We demonstrate the efficacy of our approach across several application domains, including the generation of time series of household power consumption and images satisfying specific constraints like the absence of memorized images or corruptions.

Submission Number: 112