Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Niklas Schweiger, Daniel Cremers, Karnik Ram

Abstract

Optimizing the noise samples of diffusion and flow models is an increasingly popular approach to align these models to target rewards at inference time. However, we observe that these approaches are usually restricted to differentiable or cheap reward models, the formulation of the underlying pretrained generative model, or are memory/compute inefficient. We instead propose a simple trust-region based search algorithm (TRS) which treats the pre-trained generative and reward models as a black-box and only optimizes the source noise. Our approach achieves a good balance between global exploration and local exploitation, and is versatile and easily adaptable to various generative settings and reward models with minimal hyperparameter tuning. We evaluate TRS across text-to-image, molecule and protein design tasks, and obtain significantly improved output samples over the base generative models and other inference-time alignment approaches which optimize the source noise sample, or even the entire reverse-time sampling noise trajectories in the case of diffusion models. Our source code is publicly available.

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

Abstract

Optimizing the noise samples of diffusion and flow models is an increasingly popular approach to align these models to target rewards at inference time. However, we observe that these approaches are usually restricted to differentiable or cheap reward models, the formulation of the underlying pretrained generative model, or are memory/compute inefficient. We instead propose a simple trust-region based search algorithm (TRS) which treats the pre-trained generative and reward models as a black-box and only optimizes the source noise. Our approach achieves a good balance between global exploration and local exploitation, and is versatile and easily adaptable to various generative settings and reward models with minimal hyperparameter tuning. We evaluate TRS across text-to-image, molecule and protein design tasks, and obtain significantly improved output samples over the base generative models and other inference-time alignment approaches which optimize the source noise sample, or even the entire reverse-time sampling noise trajectories in the case of diffusion models. Our source code is publicly available.
Paper Structure (23 sections, 8 equations, 7 figures, 3 tables, 1 algorithm)

This paper contains 23 sections, 8 equations, 7 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Diffusion and Flow-based Models
  4. Reward alignment of diffusion and flow models.
  5. Gradient-based guidance.
  6. Noise sequence search.
  7. Black-box search.
  8. Methodology
  9. Problem Statement
  10. Trust-Region Search (TRS)
  11. Warm-up.
  12. Trust-region iterations.
  13. Perturbations.
  14. Trust-region adaptation.
  15. Experiments
  16. ...and 8 more sections

Figures (7)

  • Figure 1: Progression of output samples produced by our algorithm (TRS) across iterations $i \in \{0, \dots, 15\}$ for text-to-image, molecules, and proteins. Top: Aesthetic reward alignment for the prompt "Animated movie poster...". Middle: Molecule property alignment. Bottom: Protein designability. The horizontal arrow indicates the direction of optimization, while the arrows $\uparrow / \downarrow$ indicate whether the objective is maximized or minimized.
  • Figure 2: Illustration of the black-box, defined by the generative model $\mathcal{F}$ and the reward function $R(\mathbf{x}_1)$, which connect noise samples $\mathbf{x}_0$ with scalar rewards $r$.
  • Figure 3: Illustration of the trust region search algorithm with a two-region example for the prompt "A panda making latte art" from DrawBench saharia2022photorealistic. (a) Samples from the noise space $\mathbb{R}^M$ are mapped to the data manifold in $\mathbb{R}^D$ via the generative model $\mathcal{F}$. Generated samples from the same region ($\mathcal{T}_j$) exhibit visual similarity; here, $\mathcal{T}_1$ shows markedly better prompt alignment than $\mathcal{T}_2$. (b) New candidates $\mathbf{x}_{0,j,b}$ are generated by adding masked ($\mathbf{m}_{j,b}$) relative perturbations ($\tilde{\mathbf{x}}_{0,j,b}$) to the current center $\mathbf{x}_{0,j}^{\mathrm{c}}$. (c) Comparison of Sobol and Gaussian perturbation schemes used to fill the trust-region hypercube. (d) The update and shift logic: $\mathcal{T}_1$ expands upon identifying a top-$k$ sample, while the underperforming $\mathcal{T}_2$ is re-centered (shifted) to a more promising region.
  • Figure 4: Examples of optimized samples from different algorithms for all methods in \ref{['sec:exp:text-to-image']}, with SDXL-Lightning lin2024sdxllightning. The first row is optimized with ImageReward xu2023imagereward and the lower two rows with HPSv2 wu2023human. All prompts are from DrawBench saharia2022photorealistic. Outputs from TRS adhere to the prompt more closely in terms of specified animal count, text, and relative positions. Further examples (incl. randomized) are given in the appendix.
  • Figure 5: Here we plot the mean best rewards for SD1.5/SDXL optimizing for HPSv2 and ImageReward across different NFE budgets. TRS shows the best scaling performance among all methods.
  • ...and 2 more figures