Inference-Time Alignment of Diffusion Models with Direct Noise Optimization

TL;DR

This paper introduces Direct Noise Optimization (DNO), a novel inference-time method for aligning diffusion models with specific reward functions, improving sample quality according to desired objectives without retraining.

Contribution

The paper presents DNO, a tuning-free, inference-time alignment method for diffusion models, with theoretical analysis, variants for non-differentiable rewards, and solutions to reward hacking issues.

Findings

DNO achieves state-of-the-art reward scores in experiments.

DNO operates efficiently at inference-time without retraining.

Effective regularization mitigates reward hacking problems.

Abstract

In this work, we focus on the alignment problem of diffusion models with a continuous reward function, which represents specific objectives for downstream tasks, such as increasing darkness or improving the aesthetics of images. The central goal of the alignment problem is to adjust the distribution learned by diffusion models such that the generated samples maximize the target reward function. We propose a novel alignment approach, named Direct Noise Optimization (DNO), that optimizes the injected noise during the sampling process of diffusion models. By design, DNO operates at inference-time, and thus is tuning-free and prompt-agnostic, with the alignment occurring in an online fashion during generation. We rigorously study the theoretical properties of DNO and also propose variants to deal with non-differentiable reward functions. Furthermore, we identify that naive implementation of DNO occasionally suffers from the out-of-distribution reward hacking problem, where optimized samples have high rewards but are no longer in the support of the pretrained distribution. To remedy this issue, we leverage classical high-dimensional statistics theory to an effective probability regularization technique. We conduct extensive experiments on several important reward functions and demonstrate that the proposed DNO approach can achieve state-of-the-art reward scores within a reasonable time budget for generation.