PaperlayerView exact revision on arXiv

Source-linked AI summary

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization

Zhiwei Tang, Jiangweizhi Peng, Jiasheng Tang, Mingyi Hong, Fan Wang, Tsung-Hui Chang

arXiv:2405.18881v3cs.LGcs.AI

TL;DR

Diffusion models may need task-specific alignment to make generated samples maximize continuous rewards. This paper introduces inference-time Direct Noise Optimization, which adjusts sampling noise; experiments report that it matches state-of-the-art tuning-based alignment methods within a reasonable generation-time budget.

  • Problem

    The alignment problem is to adjust diffusion models’ learned distribution so generated samples maximize task-specific rewards, such as brightness or darkness.

    Source: §1 · pp. 2–3

  • Method

    Direct Noise Optimization (DNO) performs inference-time gradient ascent on sampling noise to improve the reward of generated samples.

    Source: §2.1 · p. 7, §6 · p. 15

  • Results

    DNO with probability regularization matches state-of-the-art tuning-based alignment methods within a reasonable generation-time budget.

    Source: §5.2 · p. 13

  • Takeaways & Limitations

    DNO offers a trade-off between inference time and the reward of generated samples.

    Source: §5.2 · p. 13

    Tuning-based alternatives require fine-tuning for new reward functions and may struggle to generalize from limited prompts to unseen ones.

    Source: §1.2 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames diffusion alignment as adapting generation to continuous task rewards and motivates inference-time Direct Noise Optimization (DNO), which optimizes sampling noise without fine-tuning model parameters.

1 Introduction

  • Diffusion alignment adjusts the generated-sample distribution to maximize real-valued rewards, which may encode human feedback or custom metrics such as brightness and darkness.Source: §1 · pp. 2–3

1.1.1 Online/Offline Reinforcement Learning for Fine-tuning Diffusion Models

1.1.2 Direct Fine-tuning of Diffusion Models with Differentiable Rewards

  • AlignProp and DRaFT [Clark et al., 2023] fine-tune diffusion models with differentiable rewards using ODE-type DDIM sampling.The sampling process can also be represented as a noise-to-sample mapping in which noise uniquely determines the generated sample.Source: §1.1.2 · p. 4

1.1.3 Loss-Guided Diffusion

  • LGD guides inference-time sampling toward higher rewards, but its Monte Carlo gradient estimates are noisy and imprecise, especially for complex rewards.Source: §1.1.3 · p. 4

1.2 Comparing Existing Methods

  • Inference-time methods avoid fine-tuning and can flexibly handle reward functions and prompts with lower compute needs, but increase generation time; direct fine-tuning and LGD require differentiable rewards.Non-differentiable rewards are common when computed through simulations or exposed as black boxes.Source: §1.2 · pp. 4–5

1.3 Our Contributions

  • DNO interprets noise optimization as sampling from an improved distribution, regularizes against out-of-distribution reward hacking, and uses hybrid gradient approximation for non-differentiable rewards.The probability regularization is designed to keep generated samples within the support of the pretrained distribution.Source: §1.3 · p. 5
  • Experiments on several image rewards show the proposed method achieves state-of-the-art scores against existing alignment methods without fine-tuning diffusion-model parameters.Source: §1.3 · p. 5

2 Direct Noise Optimization for Aligning Diffusion Models

Direct Noise Optimization (DNO) aligns diffusion models by optimizing sampling noise at inference time without tuning model parameters. This section formalizes reward improvement under gradient ascent and explains why SDE-based optimization can be faster and at least as effective as ODE-based optimization.

2 Direct Noise Optimization for Aligning Diffusion Models

  • DNO optimizes the Gaussian noise vectors used to generate each sample, using reward gradients when available or zeroth-order estimates when they are not.The method runs during inference without tuning θ, but optimization is repeated for each new sample and takes longer than direct sampling.Source: §2 · pp. 5–6, Figure 2 · p. 6

2.1 Understanding Direct Noise Optimization

  • In a toy diffusion model, gradient ascent shifts the generated-sample distribution toward local maxima of a highly nonconvex reward.The example uses a ring-shaped target distribution and shows sample evolution after 10, 50, and 100 gradient steps.Source: §2.1 · pp. 6–7, Figure 3 · p. 7
  • If r ◦ Mθ is L-smooth and ℓ < 2/L, each gradient-ascent step strictly improves expected reward when the theorem’s second term is nonzero.The distribution after t steps is that of Mθ(g^t(z)) for z ∼ N(0, I), where g applies one gradient-ascent step to the noise.Source: §2.1 · p. 7

2.2 Optimizing ODE vs. Optimizing SDE

  • SDE-based DNO can optimize significantly faster than ODE-based DNO and is guaranteed to yield results at least as good for alignment.The authors attribute the speed advantage to the finer-grained control of SDE sampling and use DDIM with η = 1 for the remainder of the study, fixing T to 50.Source: §2.2 · p. 8, Figure 4 · p. 8

3 Out-Of-Distribution Reward-Hacking in Noise Optimization

The section identifies out-of-distribution (OOD) reward hacking in DNO as a consequence of optimized noise entering low-probability Gaussian regions, and presents concentration-based probability measurement and regularization to address it.

3 Out-Of-Distribution Reward-Hacking in Noise Optimization

  • Reward hacking can exploit reward shortcuts while remaining near the pretrained distribution, or push samples outside its support; the paper focuses on the latter, common in DNO.The paper calls the second type OOD reward hacking.Source: §3 · pp. 8–9
  • After 1,000 reward-gradient steps, all toy samples became OOD; after 50 brightness steps, SD v1.5 changed the prompted “black duck” into a “white duck”.The brightness reward is the average pixel value; ideally, the samples would remain consistent with the prompt while becoming brighter.Source: §3 · p. 9, Figure 5 · p. 9
  • Optimized noise can enter low-probability regions of the high-dimensional Gaussian, where diffusion-model approximation errors can produce samples outside the pretrained distribution’s support.The paper gives noise vectors composed entirely of zeros as an example of inputs in such regions.Source: §3 · p. 9

3.1 Quantifying Low-Probability Region via Concentration Inequalities

  • The probability measure uses concentration inequalities for Gaussian means and covariances and checks permuted noise vectors; low P(z) indicates that some permutation lies in a low-probability region.The paper reports that 100 randomly generated permutation matrices work well empirically.Source: §3.1 · p. 10

3.2 Probability-Regularized Noise Optimization

  • Probability-Regularized Noise Optimization constrains noise toward high-probability Gaussian regions, using a coefficient γ to control regularization and expected log probabilities over permutations for smoother optimization.The expected log probabilities are used instead of the minimum probability P(z) because they are smoother for optimization.Source: §3.2 · p. 10

4 Tackling Non-Differentiable Reward Functions

Section 4 adapts DNO to rewards with unavailable gradients by estimating gradients from function values and exploring three methods, including hybrid estimators that retain the differentiability of the noise-to-sample mapping.

Motivation

Method 1

  • ZO-SGD treats the full reward-after-sampling mapping as a black box and estimates its gradient through function-value queries.Source: §4 · p. 11

Method 2

  • Hybrid-1 estimates only the reward gradient and combines it with the true noise-to-sample gradient through the chain rule.The required vector-Jacobian product can be implemented with autodifferentiation without computing the full Jacobian.Source: §4 · p. 11

Method 3

  • Hybrid-2 estimates reward gradients by perturbing latent noise, keeping reward queries on pretrained-distribution samples; it is reported significantly faster than the other two methods.This avoids querying rewards on noisy samples that may leave the image manifold.Source: §4 · p. 11

5 Experiments

Experiments evaluate DNO with brightness and darkness rewards, three human-aligned rewards, and non-differentiable rewards. They examine probability regularization, comparisons with existing alignment methods, and faster inference settings.

5.1 Experiments on Image Brightness and Darkness Reward Functions

  • Probability regularization mildly slows brightness and darkness optimization but keeps samples more prompt-consistent at comparable rewards, as measured by CS and ITM.P(z) tracks CS and ITM positively; without regularization, optimized animals are more likely to conflict with the prompt.Source: §5.1 · p. 12, Figure 6 · p. 12, Figure 7 · p. 13
  • Regularized DNO produces images with purely bright or dark backgrounds without fine-tuning, an effect CrossLabs reports otherwise requiring Offset-Noise fine-tuning.The experiments use brightness reward with “black <animal>” prompts and darkness reward with “white <animal>” prompts to expose reward hacking.Source: §5.1 · pp. 12–13

5.2 Benchmarking on Three Human-Aligned Reward Functions

  • Regularization stabilizes P(z) and improves test metrics for human-aligned rewards; for aesthetic optimization, it prevents HPS and PickScore from decreasing without significantly changing optimization speed.Its effect is less pronounced here than for the contradictory brightness and darkness prompts.Source: §5.2 · p. 13, Figure 8 · p. 13
  • DNO with probability regularization matches state-of-the-art tuning-based methods, while outperforming Best-of-N by a large margin; LGD performs poorly on the complex rewards.The comparison is reported within a reasonable generation-time budget, and DNO offers a trade-off between inference time and reward.Source: §5.2 · pp. 13–14, Table 1 · p. 14
  • With 15 diffusion steps, DNO achieves high reward values within a 1-minute budget, and its memory use is under 15GB.The authors contrast this with tuning-based methods, which typically require 4–8 advanced GPUs; 50 steps were used for the Table 1 comparison.Source: §5.2 · p. 14, Table 2 · p. 14

5.3 Optimization for Non-Differentiable Reward Functions

  • The non-differentiable reward experiments compare ZO-SGD, Hybrid-1, and Hybrid-2 on JPEG Compressibility and Aesthetic Score, matching total gradient-estimation time across methods.Probability regularization is omitted, and Aesthetic Score is treated as an API-returned score without access to its gradient.Source: §5.3 · pp. 14–15
  • Hybrid-2 optimizes JPEG Compressibility and Aesthetic Score faster and reaches higher final scores than ZO-SGD and Hybrid-1; ZO-SGD is slowest.Hybrid-1 performs poorly on Aesthetic Score because that reward accepts only image inputs; true-gradient optimization also takes substantially fewer steps than estimated-gradient optimization.Source: §5.3 · p. 15, Figure 9 · p. 15

6 Conclusions · Appendix · A Visualization of Optimization with Images

The paper concludes that inference-time Direct Noise Optimization (DNO) can rival tuning-based methods, with variants for reward-hacking and non-differentiable rewards. Its flexibility and modest computing requirements come with increased sampling time, and the appendix provides qualitative optimization visualizations.

6 Conclusions

  • DNO demonstrates exceptional efficacy and can rival tuning-based methods for aligning diffusion models at inference time.The conclusion presents this as the central finding of its comprehensive study.Source: §6 · p. 15
  • DNO variants address out-of-distribution reward hacking and enable optimization with non-differentiable reward functions.The conclusion describes these variants as designed to address those challenges efficiently.Source: §6 · p. 15
  • Integrating DNO with diffusion sampling substantially increases processing time compared with direct sampling.The conclusion identifies this added processing time as DNO’s primary limitation.Source: §6 · p. 15
  • The authors consider the additional time cost reasonable and worthwhile for obtaining high-reward samples in real-world applications.They frame the added time as a trade-off for high-reward generation.Source: §6 · p. 15
  • DNO can accommodate any reward function or combinations of rewards while requiring only modest computing resources.The conclusion presents this flexibility and modest resource demand as making DNO accessible to practitioners.Source: §6 · p. 15

A Visualization of Optimization with Images

A.1 Comparing With and Without Regularization · A.2 More Examples

The examples show that probability regularization prevents optimized samples from becoming out-of-distribution, while additional results illustrate improved human preference and Hybrid-2 optimization.

A.1 Comparing With and Without Regularization

  • Across the examples, adding probability regularization consistently prevents generated samples from becoming out-of-distribution.The examples are shown in Figure 10.Source: §A.1 · p. 21

A.2 More Examples

  • Optimization increased human preference across all three human-level reward functions shown in Figure 11.Source: §A.2 · p. 21
  • Figure 12 examples show Hybrid-2 effectively estimating gradients and optimizing JPEG Compressibility Score and Aesthetic Score.These examples are from non-differentiable optimization experiments.Source: §A.2 · p. 21

B Theoretical Results

The optimized distribution stops improving when the reward gradient propagated through the noise-to-sample mapping is zero almost surely. This occurs when the generated sample is stationary, the mapping Jacobian vanishes, or the two are orthogonal.

B Theoretical Results

  • DNO’s distribution stops improving after step t when noise vectors are almost surely Type-I, Type-II, or Type-III.These cases correspond, respectively, to a zero reward gradient at the generated sample, a zero noise-to-sample Jacobian, or a nonzero reward gradient orthogonal to a nonzero Jacobian.Source: B Theoretical Results · p. 23
  • In the Type-III toy example, every point on the horizontal segment from (−1, 0) to (1, 0) qualifies when the reward is r(x, y) = y.The ground-truth distribution is uniform along that segment, and the reward gradient is orthogonal to the noise-to-sample Jacobian at every point.Source: B Theoretical Results · p. 23

C Empirical Investigation of P(z)

The section presents P(z) as an indicator of whether generated samples remain in-distribution. Its examples show P(z) falling as brightness optimization moves away from the prompt, while aesthetic-score optimization produces a smaller decline.

C Empirical Investigation of P(z)

  • All-zero and repeated noise vectors produce poor images with P(z) near zero, unlike a standard-Gaussian baseline that generates a normal image with a reasonably large P(z).Repeated vectors are constructed as z = [z0, z0, z0, z0], where z0 is sampled from an n/4-dimensional standard Gaussian.Source: C Empirical Investigation of P(z) · p. 24, Figure 14 · p. 25
  • In-distribution samples have high P(z), while out-of-distribution samples have values near zero, making P(z) an effective indicator of distribution membership.The samples are colored by P(z) in Figure 14a.Source: C Empirical Investigation of P(z) · p. 24, Figure 14 · p. 25
  • At around 20 brightness-optimization steps, P(z) is near zero and the generated image resembles a blue duck rather than the prompted black duck.Across the trajectory, P(z) decreases as the image diverges from the black-duck distribution.Source: C Empirical Investigation of P(z) · p. 24, Figure 15 · p. 25
  • Aesthetic-score optimization produces a much less significant decrease in P(z), as the reward itself penalizes out-of-distribution samples.The comparison is between trajectories for the aesthetic score and brightness reward on Stable Diffusion v1.5 with the prompt “black duck.”Source: C Empirical Investigation of P(z) · p. 24, Figure 16 · p. 25

D Implementation Details

Section D discusses implementation details of the proposed method and clarifies omissions in the experimental section.

D Implementation Details

D.1 Algorithm Implementation

The implementation avoids memory-intensive differentiation through the noise-to-sample mapping and uses a hybrid estimator with PyTorch autodifferentiation to optimize the noise.

D.1 Algorithm Implementation

  • Differentiating the noise-to-sample mapping is required, but standard PyTorch autodifferentiation can cause memory explosion; gradient checkpointing is a common remedy used in related work [Wallace et al., 2023b, Novack et al., 2024, Karunratanakul et al., 2023].The paper notes that standard PyTorch autodifferentiation [Paszke et al., 2019] cannot handle this differentiation without the memory issue.Source: §D.1 · p. 26
  • For q perturbation noises, the method generates samples Mθ(z + µξi), estimates reward gradients in non-differentiable mode, then backpropagates a loss to obtain the exact gradient estimator for z.The procedure uses built-in PyTorch autodifferentiation [Paszke et al., 2019].Source: §D.1 · p. 26

D.2 Experiment Details

This appendix specifies the generation, optimization, and compute settings for Figure 1 and the experiments in Sections 5.1–5.3. The settings include reward-dependent regularization and learning rates, plus algorithm-specific choices for fair runtime comparison.

Figure 1 details

  • Figure 1 uses SDXL base with 50-step DDIM sampling (η = 1), optimizes all injected noise, and sets classifier-free guidance to 5.0.The four prompts depict a dark alley, a 1970s baseball player, a rabbit, and a beef steak; optimizer and regularization settings match Sections 5.1 and 5.2.Source: §D.2 · p. 26

Section 5.1 details

  • Section 5.1 uses Adam at learning rate 0.01, regularization coefficient γ = 1, permutation batch size b = 100, and one A800 GPU with about 15 GB memory.The batch size is the number of random permutations drawn per step for the minibatch stochastic gradient of the regularization term.Source: §D.2 · p. 26

Section 5.2 details

  • Section 5.2 retains Adam at learning rate 0.01 and b = 100, but sets γ = 0.1 because human-like rewards are less susceptible to OOD reward hacking.Runs use one A800 GPU and about 20 GB memory; the authors implement LGD themselves and reuse other baselines’ published statistics for Table 1.Source: §D.2 · pp. 26–27

Section 5.3 details

  • Section 5.3 matches approximate time cost per gradient step with q = 16, 8, and 4 for ZO-SGD, Hybrid-1, and Hybrid-2, respectively.The respective µ values are 0.01, 0.01, and 0.02; Adam uses learning rate 0.01 for JPEG Compressibility and 0.001 for Aesthetic Score because 0.01 caused divergence, with one A800 GPU per run.Source: §D.2 · p. 27

E Hyperparameters Analysis

The analysis varies k, b, and γ to assess their effects on reward and consistency score when mitigating OOD reward hacking in darkness-reward optimization.

E Hyperparameters Analysis

  • In the Figure 17 experiment revisiting darkness-reward optimization with the prompt “white <animals>” from Section 5.1, k = 1 underperforms k = 2 and k = 10, while k = 100 is a poor choice.k = 1 fails to account for covariance among noise vectors; k = 100 entails a smaller m, potentially rendering the concentration inequalities detailed in Lemma 1 less precise.Source: E Hyperparameters Analysis · p. 28
  • The number of permutation matrices b has little impact once sufficiently large, and b = 100 emerges as an empirical choice for the proposed method.b is used to compute the stochastic gradient for probability regularization.Source: E Hyperparameters Analysis · p. 28
  • Changing γ trades off convergence speed against OOD reward-hacking propensity, so its value should be tuned empirically for each reward function and prompt.The authors recommend tuning γ with a limited number of samples and a few optimization steps.Source: E Hyperparameters Analysis · p. 28
Loading 2405.18881v3…