Title: Inference-Time Diffusion Model Alignment via Random Ordinary Equations
Primary area: generative models
Keywords: Diffusion Models, Inference-Time Scaling, Diffusion Model Alignment, Random Ordinary Equations

Abstract: Existing diffusion model (DM) alignment methods necessitate intensive training, differentiable rewards, and suffer from sparse rewards. We proposes a parameter-efficient inference-time scaling framework to align DMs to any reward function. We aim to fully exploit the potential of a fixed pre-trained DM and a fixed initial noise, achieving DM alignment solely by altering the denoising trajectories. We first extract a random ordinary differential equation (RODE) sampling from DDIM for controllable and low-variance randomness. Then, we model the inference process of DMs as a Markov decision process with a low-dimensional action space and dense rewards.  Finally, we formulate the inference-time DM alignment as a max-reward optimal control problem, and solve it using a value-based Monte Carlo tree search with augmented RODE-based scaling and Beta policies. Experiments show that our method successfully aligns pre-trained DMs with reward functions defined on clean image domains, such as aesthetics and semantics. Our method outperforms traditional inference-step scaling, and achieves significantly higher parameter efficiency than existing approaches adopting high-dimensional action spaces. Our approach does not necessitate differentiable rewards and can be plug-and-play integrated into any multi-step inference DMs.

TL;DR: We scale low-dimensional actions with random ordinary equations to achieve inference-time diffusion model alignment.