Trust-Region Noise Search for Black-Box Alignment of Diffusion and Flow Models

Optimizing the noise samples of diffusion and flow models isan increasingly popular approach to align these models to target rewardsat inference time. However, we observe that these approaches are usuallyrestricted to differentiable or cheap reward models, the formulation ofthe underlying pre-trained generative model, or are memory/computeinefficient. We instead propose a simple trust-region based search algorithm(TRS) which treats the pre-trained generative and reward models as ablack-box and only optimizes the source noise. Our approach achieves agood balance between global exploration and local exploitation, and isversatile and easily adaptable to various generative settings and rewardmodels with minimal hyperparameter tuning. We evaluate TRS acrosstext-to-image, molecule and protein design tasks, and obtain significantlyimproved output samples over the base generative models and otherinference-time alignment approaches which optimize the source noisesample, or even the entire reverse-time sampling noise trajectories in thecase of diffusion models. Our source code is publicly available⋆ .