Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
Authors: Zhiwei Tang
Jiangweizhi Peng
Jiasheng Tang
Mingyi Hong
Fan Wang
Tsung-Hui Chang
Venue: ICML 2025 poster (2025)
Keywords: Diffusion Models
Inference-Time Alignment
Optimization
RLHF
Primary area: deep_learning->generative_models_and_autoencoders
TL;DR: This work studies the problem of inference-time alignment of diffusion generative models with downstream objectives

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.

Inference-Tinee Alignment of Diffusion Models with Direct Noise Optimization
Zhiwei Tang" Jiangweizhi Pengà Jiasheng Tang* 24 Mingyi Hong?
Fan Wang" 24
Tsung-Hui Chang
Abstract
1. Introduction
Diffusion models work by learning to reverse the process
of diffusing the data distribution p(z) into noise, which can be described by stochastic differential equation (SDE) (Song al., 2020b; Karras et al., 2022): dr, f()) + g(L)dw, where dw, is the standard Wiener process, and f(() and g(t) are the drift and diffusion coefficients, respectively.
In this work, focus the alignment prob-
lem of diffusion models with a continuous re- ward function, which represents specific objec- tives 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 mod- els such that the generated samples maximize the target reward function. We propose novel align- ment approach, named Direct Noise Optimization (DNO), that optimizes the injected noise during the sampling process of diffusion models. By de- sign, DNO operates at inference-lime, and thus tuning-free and prompt-agnostic, with the align- ment occurring in an online fashion during gener- ation. We rigorously study the theoretical proper- ties 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 sam- ples 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 regu- larization technique. We conduct extensive ex- periments on several important reward functions and demonstrate that the proposed DNO approach can achieve state-of-the-art reward scores within reasonable time budget for generation.
The reverse process relies on the score function e(2,, t) def. Vz log Pe(x)) where pt denotes the p.d.f of noisy data 2y, and its closed-fone form can be expressed either as an ODE or as SDE: (Song et al., 2020b):
ODE: dæ, = (000))))0000000))1 dt,
(1)
SDE: dr, = (f(t)xe- - -2(())((r t)) dt + g(t) (2)
With the capability to evaluate c( t), it becomes possible to generate samples from noise by numerically solving ei- ther the ODE (1) or the SDE (2). The training process, there- fore, involves learning a paramcterized surrogate eø cø(z,, t) to approximate cxx;, t). following a denoising score match- ing framework described in (Song et al., 2020b; Karras et al., 2022). Despite the effectiveness of diffusion models in modeling continuous distributions, when deploying these generative models for specific tasks, it is not suitable to sam- ple from the original learned distribution directly, because this distribution has not been aligned with the task-specific objective. For instance, in image generation, users may wish to produce images that are aesthetically pleasing rather than mediocre, or generate images with enhanced bright- ness, darkness, compressibility. Recently, the alignment problem has drawn considerable interest in the context of diffusion models, as evidenced by a series of studies such as (Yuan et al., 2024; Song et al., 2023a; Dong et al., 2023; Prabhudesai et al., 2023; Black et al., 2023; Fan al., 2023).
SShooo of Science and Engincering, The Chinese University of Hong Kong. Shenzhen, China 'DAMO Academy, Alibaba Group 'Department of Electrical and Computer Engineering. Univer- sity of Minnesota, USA 'Hupan Lab, Zhejiang Province, China "Shenzhen Research Institute of Big Data. Correspondence to: Zhiwei Tang <zhiweitangl @ @wink.cchk.ccck. edu.e...
Alignment Problem for Diffusson Models. Given a diffu- sion model characterized by parameters O and its associated distribution po(a)), as well as reward function r(z: that can assign real-valued scores to generated samples, the central goal of the alignment problem is to adjust the distribution (z( such that the generated samples maximize the reward from r(x).. In this work, we consider the reward functions
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
to be continuous but possibly non-differentiable. In the fol- lowing sections, we will provide a comprehensive review of some well-established methods for aligning diffusion models.
Diffusion Sampling as Noise-to-Sample Mapping. From the diffusion sampling process described above, we observe that the sampling process can be conceptualized as an end- to-end mapping (=), which translates noise vectors 2, sampled from the standard Gaussian distribution, into gen- crated samples. Here, the noise vectors & serve as a unified abstraction for both the ør in the ODE-based sampling process and the (2T, {21, )T in the SDE-based sam- pling process. As we can see, the noise vector uniquely determines the generated sample from the diffusion models,
1.1. Methods for Aligning Diffusion Models
1.1.1. ONLINE/OFFLINE RL METHODS
A common mathematical formulation in methods based Reinforcement Learning (RL) is to maximize the expected reward while ensuring the resulting distribution does not deviate excessively from the original distribution. This can be expressed as the following Kl-rregllrrized optimization problem: maxp (xpp((>) [r(=)] - (p (po). Current RL- based methods can be categorized into online and offline methods based on the data used. In the online method, the algorithm has the capacity to query the reward function throughout the entire optimization process. Two notable on- line RL methods are DDPO (Black et al., 2023) and DPOK (Fan et al., 2023), which have been shown to improve down- stream objectives such as aesthetics and compressibility. Alternatively, research has also delved into the offline RL optimization setting, where an explicit reward function is not accessible and only fixed preference dataset is utilized. Noteworthy works in this category include Diffusion-DDe -PO (Wallace et al., 2023a) and SPIN-Diffusion (Yuan et al., 2024).
Two recent studies, AlignProp (Prabhudesai et al., 2023) and DRaFT (Clark et al., 2023), have proposed to di- rectly fine-tune diffusion models using differentiable re- wards. Specifically, their optimization objective is formu- lated as: maxę ( (+)) Both the AlignProp (Prabhudesai et al., 2023) and DRaFT (Clark et al., 2023) methods utilize the ODE-type DDIM solver for the sampling process, specifically employing Algorithm I with vj 0.
1.1.3. REWARD GRADIENT GUIDANCE
A recent work focusing on loss-guided diffusion mod- els (LGD) (Song et al., 2023a) also examines the con- cept of aligning diffusion models with differentiable re- wards. Unlike the methods mentioned previously, LGD is an Inference-Tinee alignment method, meaning it does not necessitate modifications to the pretrained model @ and only works by modifying the inference process. In essence, the core idea of LGD is that, during the sampling process for the ODE (I), it considers a modified version of the ODE by introducing a new term that guides the generation to- ward areas of higher reward. Specifically, the new ODE is: dx, (f(t)r, - } (t))(((),, t) + ((zo(z(()) dt, where (x)) denotes the solution of this ODE starting from Te. However, the gradient term V (()o(())) is not readily available during generation. To address this, the authors suggest utilizing Monte Carlo estimation to approximate the gradient. Nevertheless, this estimation tends to be noisy and imprecise, leading suboptimal performance, particularly with complex reward functions.
1.1.2. DIRECT FINE-TUNING METHODS
Before delving into the formal description of this method, it is useful to revisit the sampling process of diffusion mod- els, which also serves as the foundation for the rest of this work. Solving the ODE (1) or the SDE (2) typically in- volves discretizing the time steps into T steps. By starting with the initial noise IT ~ (0,01), I), the solution process for the ODE (1) can be viewed as a mapping that transforms the initial noise FT into less noisy data through the fol- lowing process: Fz-1 = ODE solver(x:), for = T ..... 1. After T steps, the output will be the generated sample To. Similarly, solving the SDE (2) can be seen as a map- ping that gradually converts both the initial noise Dr and the entire extra random noise vectors 2T: 21 into the generated sample Ho, e.g., through the following process: ft-1- SDE_solver(z,, 21), for t T...., 1, where 2, is also drawn from standard Gaussian distribution.
1.2. Comparing Existing Methods
Existing works can be generally categorized based on two criteria: whether it requires fine-tuning and whether the reward function needs to be differentiable.
Inference- Time or Tuning- Methods. All RL-based methods and the direct fine-tuning method are aaming-base., meaning they necessitate adjustments to the network mod- els 0. There are two main disadvantages associated with tuning-based methods. The first is that they require fine- tuning for new reward functions, a process that can consume considerable resources, especially when faced with exten- sive choices for reward functions. The second drawback
Remark 1.I. Throughout this work, for simplicity, we adopt only the DDIM sampling algorithm (Song et al., 2020a) for our experiments, as it remains one of the most popu- lar choices for diffusion sampling and, more importantly, supports both ODE-style and SDE-style sampling. To be self-contained, we summarize the notations and procedure of the DDIM sampling method in Appendix A.
2

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
is that the fine-tuning process typically relies on a limited set of input prompts, which challenges the model gen- eralize to new and unseen prompts. In contrast, methods such as LGD (Song et al., 2023a) belong to the inference- time category. The main advantage of this approach is its elimination of the need for fine-tuning, allowing it to be flexibly applied in conjunction with any reward functions and imput prompts. Further, inference-time methods require significantly fewer computing resources than tuning-baser methods. However, the major drawback of inference-time methods is the substantial increase in the time required for the generation process.
Through the experiments on several important image reward functions, we demonstrate that our proposed method can achieve state-of-the-art scores in compar- ison to existing alignment methods, without any fine- tuning on the parameters of diffusion models.
2. Direct Noise Optimization
Given the noise-to- tosse-ll mapping Ma described in Section 1.1.2, DNO can be mathematically formulated as follows:
max r(Ma
(3)
Differentiable or Non-Differentiable Rewards. Current RL.-based methods can work by utilizing solely the value or preference information of the reward functions, therefore they can still work even when the reward function is non- differentiable. In contrast, the existing direct fine-tuning methods and LGD require the reward function to be dif- ferentiable. In practice, working with non-differentiable reward functions is important due to their prevalence. This non mon-differentiabee property can arise from the simulation- based procedures used to compute the reward, or the reward function itself may be black box.
with ~ I) as the initial solution. As we will dis- cuss in Section 3, the Gaussian distribution (0 7) serves as an important prior on the optimization variables z. By solving this optimization problem, we can obtain the opti- mized noise vectors, which are then used to generate high- reward samples. When the reward function r(.) is differen- tiable, gradient-based optimization methods can be applied to solve the problem efficiently. That is, the following step can be performed iteratively until convergence: Znew optimizer_stepi
Ver (Ma (zota))) where the optimizer can be either vanilla gradient ascent or adaptive optimization algorithms like Adam (Kingma & Ba, 2014). When the gra- dient of reward function r(.) is not available, we can lever- age techniques from zeroth order optimization (Nesterov & Spokoiny, 2017; Tang et al., 2024a) to estimate the ground- truth gradient Var (Ma(zala). denoted as ((Zota). and then apply similarly znew = optimizer stepp (Zota)). In Sec- tion 4, we provide a dedicated diseussion on how to obtain better estimator for ğ(=ouu when the reward function is mon-differentiable.
1.3. Our Contributions
In this work, we focus on inference-time alignment of diffu- sion models, as we believe that flexibility with the choices of the reward functions, generalization on unseen prompts, and low computing requirements are critical for a broad range of real-world applications. Our primary goal is to design an inference- enference-time alignment method that can match the performance of tuning-based based methods by incurring a reasonable additional time cost, and is capable of handling both differentiable and non-differentiable objective func- tions. To this end, we focus under-explored technique for achieving inference-time alignment of diffusion mod- els Direct Noise Optimization (DNO). Specifically, we make the following contributions:
Several recent studies have explored similar formulations to (3) across different applications. (Wallace et al., 2023b) investigates the optimization of latent vectors obtained through DDIM-inversion (Song et al., 2020a), aiming to enhance the CLIP score (Radford et al., 2021) and Aesthetic Score (Schuhmann et al., 2022b) of given images. (Ben- Hamu al., 2024) discusses optimizing the initial noise dT for the ODE process to address inverse problems, leveraging the diffusion model as prior. (Novack et al., 2024) and (Karunratanakul et al., 2023) consider the optimization of initial noise Fr for the ODE with the objective of improving downstream objectives in robotics and audio.
We conduct a self-contained and comprehensive study for DNO, and demonstrate that noise optimization can be theoretically interpreted as sampling from an im- proved distribution.
We identify -ntoof-distribution reward-hacking as critical issue in DNO. To address this issue, we intro- duce novel probaaility--eggularieedd noise optimization method designed to ensure the generated samples re- main within the support of pretrained distribution.
While similar methods of DNO has appeared in previous works, many of its technical details remain insufficiently explored. On one hand, there is lack of a comprehensive framework concerning the design choices, theoretical un- derstanding, and practical challenges associated with DNO for aligning diffusion models. On the other hand, the field has yet to systematically investigate whether DNO, as a
By developing a novel and highly efficient hybrid gra- dient approximation strategy, we extend the DNO ap- proach to handle non- non-differentiable reward functions effectively.
3

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
inference-time method for aligning diffusion models, can achieve competitive performance compared to tuning-based methods. In this work, we aim to conduct thorough study on DNO. In the following two sections, we dive deep to understand the theoretical foundation of DNO and discuss a critical design choice ignored in previous works.
2.2. Optimizing ODE vs. Optimizing SDE
As previously introduced, there are two primary methods for sampling from pretrained diffusion models: one based solving the ODE (1) and the other solving the SDE (2). A critical difference lies in the fact that ODE sampling depends exclusively on the initial noise I7, whereas SDE sampling is additionally influenced by the noise a added at every step of the generation process. It has been noted that existing works on noise optimization (Ben-Hami et al., 2024; Novack al., 2024; Karunratanakul et al., 2023) have mainly concentrated on optimizing the initial noise 2T for the ODE sampler.
2.1. Understanding Direct Noise Optimization
In Appendix I. we present a simple example to visualize the process of DNO, where we observed that the distribu- tion of the generated samples shifts toward a distribution on the local maxima of the reward function. Inspired by this example, we propose to view DNO as sampling from improved distribution. To rigorously describe this evoly- ing process, we define an operator function g to represent
2.0
Aesthetic Score
a single gradient step, i.e., g(2) def
= z + E. Ver= o Mo(z),
Velue
0.5
where o denotes the function composition operator and denotes the step size for gradient ascent. Additionally, we
SDE
define the operator It, which denotes applying the gradient
ODE
-1.0
ascent step for t steps, i.e., 9e(2) = (()) with go being the identity mapping. With these notations, we can now express the distribution after t gradient steps as pt (x), which is the distribution of M=() with 2 ~ (0, In the following theorem, we demonstrate that there is rig- orous improvement after every single gradient step, i.e., the distribution Pt+ (ar) is provably better than (xr in terms of expected reward.
20
40
10
80
100
20
60
60
80
100
steps
steps
Figure ODE vs. SDE for optimization
Different from preceding studies, we explore the utilization of the SDE sampler for noise optimization. Specifically, we employ the DDIM with n = 1 (Song et al., 2020a) as the SDE sampler and propose to optimize both the added noise 2 at every timestep and the initial noise ær. In this context, the dimensionality of the optimized noise significantly sur- passes that of ODE sampler, typically T v.s. d, where d is the dimension for the learned distribution and the number of discretization steps in the sampling process. With this higher dimension for optimization, we have empirically observed that optimizing the SDE-based generation process can be significantly faster than its ODE-based counterpart. This is illustrated in Figure 1, where we juxtapose the opti- mization speeds between the ODE (DDIM with rj = 0) and SDE samplers using both the simple depicted in Figure (Left) and optimizing for stable diffusion in alignment with the aesthetic reward (Right), major experiment detailed in the subsequent section 5.2. In Appendix D, we also provide theoretical justification for the the better performance of DNO with SDF-BBased Sampling.
Theorem 2.1. Assuming that r o MM is L-smooth, namely, Vr o Mz)x) - Vr o Mo('')l & L||> - 2'1 for any & + 2', holds true that
E Exepi (x)p(1)00)((x)= 2 E (e)p((s)(())+
E2)
( ( 2
E B2o Ver Mo Mo(3))
(4)
In Theorem 1.1, we rely the smoothness assumption of the composite mapping ro Mg. We note that this is a reason- able assumption in practice, as the voise-to-sample mapping in diffusion models has been observed to be smooth. For in- stance, see Figure 4 in (Tang et al., 2024a). As described in (4), provided that the step size f is less than Z. the distribu- tion (2)) strictly better than the previous distribution Pa(x) in terms of expected reward, as long the second term is non-7ero. Based on this result, it is natural to ask: When does the distribution stop improving? Namely, when does the second term in (4) become zero. We provide a de- tailed discussion to answer this question and also the proof for Theorem 2.1 in Appendix
3. OOD Reward-Hacking in DNO
It has been observed that when aligning generative models (e.g., including autoregressive language models or diffu- sion models) with reward functions, one can experience the so-called reward-hacking, i.c.. the optimized samples yield high rewards but do not possess the desirable prop- erties (Miao et al., 2024; Chen et al., 2024a). Generally, there are two different types of reward- hacking.. In the first
4

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
type, the reward function admits some shorteuts, so the optimized samples score high rewards but remain barely dis- tinguishable from the samples of the pretrained distribution. The second type is related to the generative model used- the optimized samples no longer fall within the support of the pretrained distribution after optimization. We denote this second type of reward-hacking as Out-of-Distributioo (OOD)
m maxx vI+MI- 11---11--
wk/m
Pz (2(M) def.
max
2e
In practice, to determine if an n-dimensional vector 2 lies within a low- cow-probaiility region, we can factorize n as n = m k, and divide = into m subvectors: z = [=].... mm where n m k and 2; = [2] 24 ~ (+),, Then, we compute M (x)=) - 2 2i| and M2(2) m E 2.22 - Jx [l Finally, we can determine that z lies in low-probability region if both pi ())) and P2 (2)M2(2) are low.
In this work, we will focus on this second type and reveal that OOD seward-hacking is common issue in DNO. In Appendix B.2, we provide two visualized examples of the phenomenon of OOD Reward- Hacking using both a 2-dimensional diffusion model and an image diffusion model.
Remark 3.2. According to (Wainwright, 2019), these two inequalities in Lemma 3.1 are tight when m/k is large. On the other hand, h = is not advisable, as it examines only the mean and variance of the noise vector, but not the covariance of different subvectors. In this work, we empirically found that l 2 serves good default choice. In Appendix .7, we provide a more detailed analysis for choosing an appropriate k.
One of our key contributions in this work to identify one critical cause of OOD aeward-hacking in noise optimiza- tion. That is, the optimized noise vectors stray towards the low-probability regions of the high-dimensional standard Gaussian distribution; in other words, there is an extremely low chance of such noise vectors being sampled from the Gaussian distribution. As diffusion models are originally trained with Gaussian noise, when the noise vectors orig- inate from these low-probability areas-such as vectors comprised entirely of zeros the neural network within the diffusion models may incur significant approximation er- rors for these particular inputs. This error, in turn, leads to the generation of out-of-distribution samples. In the subse- quent section, we introduce a novel method to measure the extent to which noise is part of the low-probability region by leveraging the classical concentration inequalities for high-dimensional Gaussian distributions.
An important point to note is that the standard Gaussian distribution is invariant to permutation, i.c., for any per- mutation matrix II, if follows a standard Gaussian distri- bution, the permuted vector IIz will have the same prob- ability behavior. With this insight, to increase the ro- bustness of the probability measure Pi and p2, a natural idea is to examine the probability of many permuted vec- tors. Specifically, given q permutation matrices II.,.., Ig we define the following indicator metric, P(x) def. min {pr (MT:), P2 (M2(TT;=) , E (1,-...))-
Interpretation of P(z). If the probability P(=) is low, it implies that there exists permutation matrix II; such that the noise vector 11,2 is in the low-probability region of the standard Gaussian distribution. Therefore, due to the permutation-invarian property, the noise vector z is also less likely to be sampled from the standard Gaussian distribution. In practice, utilize randomly generated permutation matrices and have found that setting q 100 results in empirically good performance. In Appendix 55, we provide some visualized empirical evidence to show that P(z) serves as a good indicator for determining if the generated samples are OOD.
3.1. Quantifying Low-Probability Region via Concentration Inequalities
High-dimensional Gaussian distributions possess several unique properties. For instance, it is known that the all- zero vector is the most probable in terms of the probability density function (p.d.f) of the standard Gaussian distribu- tion. However, in practice, it is nearly impossible to obtain samples near the all-zero vector from a Gaussian distri- bution, as it resides within low- wwwppobbability region. In high-dimensional statistics, concentration inequalities are usually employed describe these distinctive properties and delineate the low-probability regions of high-dimensional distributions. In the following lemma, we present two clas- sical inequalities for the standard Gaussian distribution.
3.2.
Noise Optimization
With the insights discussed above, a natural idea for pre- venting OOD reward hacking is to regularize noise vectors to remain within the probabibbybi region of the Gaus- sian distribution. To achieve this, we propose the following Noise Optimization problem:
Lemma 3.1 (WWainwrihtt, 2019)). Consider that 21, ... Zm follow a k-dimensional standard Gaussian distribution. We have the following concentration inequalities for the mean
and covariance: Pr
[I|
m
zi|
M)
< pi (M) def.
max {2e-
mx²² -M 1 Pr
x=
1
M <
max maxx((M(=)) + Bnn (pr (M, (II=)) + (a(M)( (5)
5

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
where y is the coefficient used control the regularization effect. In particular, for the regularization term, we use the expectation of the log probabilities over the permuta- tion matrices, rather than the minimum probability P(z). This is because the expectation is smoother for optimization purposes.
Gz def. Hi(x) Mz))
where the main idea is to replace the ground truth Vr(*)
in the chain-rule of differentiating r(Mo(=)). We refer this method as Hybrid-1 in the following sections.
4. Tackling Non-Differentiable Rewards
Remark 4.1. As one can observe, the computation of (7) involves the Jacobian V. VaMa(2), However, it is important to note that when we only require the vector-Jacobian product Hi (I) V 2 Ma(),, it is unnecessary compute the full Jaco- bian V Ma )). In Appendix .6, we describe an elegant and efficient way to implement (7) using an auto-differentiation technique.
In the previous section, DNO method has been applied to optimize differentiable reward functions. However, in many applications the ground-truth gradient of the reward function is unavailable. Such non-differentiable properties can arise from various scenarios; here we present two rep- resentative cases. Firstly, the reward may be computed through simulation-based procedures, such as the JPEG- compressibility employed in DDPO (Black et al., 2023), which calculates the size of an image in bits after running JPEG compression algorithm. Additionally, the reward func- tion itself may be a black box provided through online API providers, as in the setting considered in (Sun et al., 2022). This scenario is common when the reward function is large neural network model, like those in (Wang et al., 2024a; Lin et al., 2024), making it impossible to directly obtain the gradient for optimization. To address these scenarios, we explore adapting the noise optimization approach to handle the optimization of non-differentiable reward functions by estimating the gradient with function values. Specifically, we explore three methods under this setting.
Method 3. There is a crucial drawback in the gradient estimator (6), that one needs to query the reward function r(+) with noisy input + ju.. When the reward function is only defined on some manifold M, e.g., defined on the image manifold, rather than the whole space RR, this can lead to severe problems, because. for some æ E .M, the noisy sample x + je may no longer stay within M. To remedy this issue, we propose to perturb the sample through the latent noise, rather than directly in the sample space. Specifically, our proposed new gradient estimator for Vr(x) is H2(r)= =:
EE [M + HE)) - r(x)) (Mo(z + uE) - x)]. (8)
Method 1. Concerning optimization with only function value, major family of optimization approaches is zeroth- order optimization algorithms, including ZO-SGD (Nes- terov & Spokoiny, 2017). This method treats the entire mapping r M-)-) as a black-box function and seeks to es- timate the gradient of r a Mo(-) via function value queries.
Following a similar proof in (Nesterov & Spokoiny, 2017),
we can also show that Hz(x) z C2Vr(x) for some con- stants C2. As we can see, when computing the gradient (8), we ensure that we query the reward function r(.) with only samples that are within the manifold of the pretrained distribution. Similar Hybrid-1, we can adopt gradient estimator G2(=) with H2(x);:
Method 2. is worth noting that for the mapping r Me(-),
only the gradient of the reward function r(x) is not avail- able, and are still able compute the gradient of Mo Therefore, a straightforward idea is to adopt hybrid gra- dient approach-one using the ground truth gradient for Mz).). Specifically, we denote that the initial noise is x, and the generated sample is x Mz).). Firstly, we can estimate the gradient of Vr(x) in a similar fashion with the ZO-SGD (Nesterov & Spokoiny, 2017):
to estimate the gradient of r, while
G2(2) del. Hz(x) V=Mo(>).
(9)
We refer this method (9) as Hybrid-2. As we will see in Section 5.3, this Hybrid-2 method is significantly faster than the other two in terms of optimization speed.
5. Experiments
(x) = (e- + ((())
(6)
In this section, we aim to demonstrate the effectiveness of the method proposed above. For all subsequent experiments, we utilize Stable Diffusion v1.5 (Rombach et al., 2022) as the base model for noise optimization. For each figure, we perform the optimization using 1,000 different random seeds and report the average value along with the standard
where j is the coefficient for perturbation, and C₁ is some constant. With the estimated gradient H, (z) for the reward function r(z), we can use the following estimated gradient G1 (2) for optimization:
6

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
deviation (std) of the results. For comprehensive details regarding the implementation of our proposed methods, as well as information on hyperparameters, we refer readers to Appendices .6 and .7. Additionally, we provide examples to visualize the optimization process in Appendix B. For all the following experiments, unless explicitly stated otherwise, a single run DNO is performed on a single A800 GPU.
W/ reg w/o reg
5.1. Experiments on Image Brightness and Darkness Reward Functions
Figure 2. Comparison of DNO with and without prohability regu- larization. Upper row: Optimizing for the brightness reward, i.e., the average value of all pixels in the images. Lower row: Opti- mizing for the darkness reward, i.c., the negative of the brightness reward. The x-axis refers to the number of gradient ascent steps during optimization.
Experiment Design, In this section, we design an experi- ment to demonstrate the effectiveness of DNO as described in Section 2, and the probability regularization proposed in Section 3.2. We consider two settings: The first involves optimizing the brightness reward, which is the average value of all pixels-the higher this value, the brighter the image becomes-win the prompt "black [animal]"", where the to- ken [animal] is randomly selected from a list of animals. The second setting involves optimizing the darkness reward, defined the negative of the brightness reward, with the prompt "white [animal]". The primary rationale behind designing such experiments is the inherent contradiction between the prompt and the reward, which makes it casier to trigger the OOD reward-hacking phenomenon. More- over, it is straightforward to verify whether the generated -istriiiiittionn by simply examining the color of the generated animals. In these experiments, we compare the noise optimization process with and without probability regularization to assess the capability of proba- bility regularization in preventing the OOD reward-hacking phenomenon.
Results. In Figure 2, we first observe that adding the regu- larization term leads to a mildly slower optimization process. However, the generated samples are much more consistent with the prompt throughout the entire optimization process, as reflected by the CS and ITM curves. The trajectory of P(z) further corroborates that it is good indicator of the OOD phenomenon, it is positively associated with CS and ITM. We also provide visualized examples in Appendix B.3, which also confirm that the probability regularization pro- posed in Section 3.2 can effectively prevent the generated samples from becoming OOD.
samples are
Generating Images with Purely Bright/Dark Back- grounds. From the images in Figure 7 of Appendix B.3, we observe that DNO with regularization can lead to images with purely dark or bright backgrounds. It is worth noting that this is a remarkable result, as DNO requires no fine- tuning. Notably, as discussed in (CrossLabs, 2023), such an effect can only be achieved by fine-tuning the diffusion models using technique called Offset-Noise.
Importance of the Brightness and Darkness Reward Functions. While the primary purpose of using these two re- ward functions is to better examine the effectiveness of prob- ability regularization, it is also important highlight their practical utility. There is often a genuine need to generate images with extremely dark or bright backgrounds, which cannot be achieved by the base models through prompting alone, as reported in the notable research by (CrossLabs, 2023).
5.2. Benchmarking on Three Human-Aligned Reward Functions
Setting. In this section, we investigate the performance of the proposed method using three common reward functions trained from human feedback data, specifically Aesthetic Score (Schuhmann et al., 2022b), HPS-v2 score (Wu et al., 2023), and PickScore (Kirstain et al., 2023), respectively. In this experiment, we also compare noise optimization with and without probability regularization. However, compared to the reward function used in the previous section, using these reward functions presents lesser chance for the opti- mized image to be OOD. In this case, to measure the benefit of probability regularization in maintaining the quality of the generated sample, we consider using the other two re- ward functions as test metrics when optimizing one of them. For the prompt dataset, we follow (Clark et al., 2023; Black et al., 2023) to use a simple animals prompt dataset.
Measuring the Degree of OOD. As observed in this ex- periment, the reward function contradicts the input prompt, leading to an inconsistency between the generated sam- ples and the prompt. Therefore, we utilize the CLIP Score (CS) (Radford et al., 2021), a commonly used metric for measuring the semantic similarity between images and text descriptions, to gauge the degree of OOD for the gener- ated samples. A higher CS indicates that the sample is less likely to be we also use an MLLM-bassed score to measure the degree of OOD, specifically employing the Image- Text Matching (TTM) score from (Wang et al., 2024a) as the metric.
(OOD). In addition to the CS,
7

---

Inference-Tine Alignment of Diffusion Models with Direct Noise Optimization
Results. Firstly, we observe that the effect of probability regularization is less pronounced than that in Section 5.1.
fusion steps can significantly decrease optimization time. In Table 1, we fix the diffusion steps to 50, as this is the setting used by all other algorithms to ensure fair comparison. In practice, 50 diffusion steps is a fairly long choice, and using 15-25 steps can still result in sufficiently good sam- ples. Therefore, in Table 2, we present the performance of DNO with different numbers of denoising steps, fixing the time budget minute, as well as other hyperparameters like learning rate and reguralization coefficient. As Table 2 shows, with 15 diffusion steps, DNO can achieve high reward values within 1 minute.
This observation is also reflected by our proposed indicator P(z); if no regularization is applied, the value of P(2) in Figure decreases much slower than that in Figure 2. Nonetheless, by adding the regularization term to noise optimization, we can stabilize the value of P(z) throughout the optimization process and also improve the test metries. For instance, when optimizing the aesthetic reward, the regularization has no significant effect on the optimization speed, while it prevents the test metrics, i.e., the HPS score and Pick Score, from decreasing throughout the process.
Finally, we would like to highlight that this alignment pro- cess can run with memory usage of less than 15GB, and thus can easily fit into just one consumer-level GPU. In contrast, current tuning-based methods require significantly more computing resources, typically 4-8 advanced GPUs like A A100 (Black et al., 2023; Prabhudesai et al., 2023).
Hi+
Figure 3. Comparison of running DNO with three human- re- ward functions, with and without regularization. When optimizing one reward function, the other two are used as test metries. A, H. P are short for Aesthetic Score, HPS Score, and Pick Score, respectively. The name for each line comprises the used reward function and whether the regularization used. For example, A + w/ reg means optimizing aesthetic score with regularization.
5.3. Optimization for Functions
Reward
Setting. In this section, we aim to explore the three exten- sions proposed in Section 4 for handling non non-differentiable reward functions. For Method I, we use ZO-SGD (Nes- terov & Spokoiny, 2017), and compare it with our proposed Hybrid-1 and Hybrid-2. We consider using two reward func- tions: The Jpeg Compressibility score, which is the file size of the image after compressing it using the JPEG algorithm. This reward function was also used in (Black et al., 2023) and is intrinsically non-differentiable. The second reward function is the Aesthetic Score used in the previous section. Unlike previous experiments, in this section we treat the Aesthetic Score as a non-differentiable reward function for optimization. The goal is to simulate scenario where the neural network model of the Aesthetic Score can only be queried via API provider that returns the score rather than its gradient. We compare the three methods in terms of optimization steps and the final score of the reward function. For ease of demonstration, we do not add the probability regularization term to the optimization process. Moreover, it important to note that in these three methods, the major time expenditure comes from estimating the gradient with finite samples. For a fair comparison among these methods, we set the number of samples for estimating the gradient separately so that the total time spent estimating the gradient is the same for all three methods
Comparison to Existing A Alignment Methods. We sum- marize the performance of DNO with probability regular- ization from Figure 3 into Table and compare it to the major existing alignment methods discussed in the introduc- tion. As shown, the performance of DNO matches that of state-of-the-arr suning-based alignment methods without any fine-tuning on the network models, all within a reasonable time budget for generation. More importantly, we demon- strate that DNO provides an worthwhile trade-off between inference time and the reward of the generated samples. On one hand, another inference-time method, LGD (Song et al., 2023a), performs poorly with these complex reward functions, as it is impossible to estimate the gradient of the reward functions without a complete generation process. On the other hand, we also examined the most fundamen- tal inference-time alignment algorithm, Best-of-N (BoN) Sampling, which generates N samples and selects the one with the highest reward. In this experiment, we fix the time budget for BoN to 10 minutes, and we observe that DNO outperforms it by large margin, demonstrating that DNO presents a highly advantageous trade-off between inference time and reward.
Results. The main results are visualized in Figure 4. For better visualization purposes, when plotting the line for ZO- SGD, we compute and plot the current best reward instead of the current reward due to the extremely high variance. Initially, we observe that the Hybrid-2 method is signifi- cantly faster than the other two methods, and the final score is also higher for both reward functions used in this experi-
Using Fewer Steps to Reduce Inference Time. In Figure 3, we show that our proposed DNO can achieve high re- ward values within approximately 3-5 minutes. While this already represents a highly advantageous performance, 3-5 minutes of optimization may still be prohibitively long in practice. However, we note that reducing the number of dif-
8

---

Inference- Time Alignment of Diffusion Models with Direct Noise Optimization
SD v1.5
LGD
SPIN 20h
DDPO ~ 56h
AlignProp 12h
DNO
(This work) min
BoN
Method
==10
1=100
min
min
10 min
Aesthetic
5,367
5,340
5.224
6.248
7.180
8,940
5.754
7.202
8,587
6.531
HPS f
0.278
0.276
0.271
0.276
0.287
0.330
0.285
0.303
0.324
0.298
PickScore t
21.11
21.01
21.09
22.00
/
21.25
23.17
25.13
22.09
Table 1. Performance comparison. For SD v1.5 and DNO, annotate the generation time below the name. For LGD, we annotate the number of samples used for Monte Carlo approximation. For SPIN, DDPO, and AlignProp. we annotate the estimated time for fine- tne-tuning. All time costs in the table are measured with respect to the GPU time single A800 GPU. Baselines: LGD (Song et al., 2023a), SPIN (Yuan et al., 2024), DDPO (Black et al., 2023), AlignProp (Prabhudesai et al., 2023), Best-of-N (BoN).
T
10
15
20
25
50
less, we argue that the additional time cost, being within a reasonable and acceptable range, is a worthwhile trade- off for attaining high-rewand generated samples a wide range of real-world applications. We anticipate DNO gain- ing greater attention future research and applications due to its flexibility to accommodate any reward function-an even a combination of different reward functions- de- manding only modest computing resources, which positions it as an accessible tool for many practitioners.
Aesthetic t HPP
6.992
7.496
6.773
6,381
5.754
0.342
0.341
0.306
0.293
0.285
PickScore 23.98
24.82
23.69
23.02
21.25
Table 2. Running DNO with different numbers of diffusion steps T, while fixing the time budget minute.
Hybrid-2
Hybrid-1
Z0-5GD
speg Compressibiit
-40
2.5
Acknowledgements
-00
This work was supported by Alibaba Group through Alibaba Research Intern Program. The work also supported by Shenzhen Science and Technology Program under Grant No. ZDSYS20233066091 and by Guangdong Provincial Key Laboratory of Big Data Computing.
-100
130
60
122
200
200
500
300
1000
Steps
staps
Figure 4. Comparing three methods on reward functions.
Impact Statement
ment. Furthermore, in both scenarios, the ZO-SGD method exhibits the slowest optimization speed. Another interesting finding is the poor performance of Hybrid-1 in optimiz- ing the Aesthetic Reward. This mainly occurs because the aesthetic reward function can only work with image in- puts, which validates our initiative to propose the Hybrid-2 method in Section 4. Finally, when comparing the opti- mization of the aesthetic score in Figures 3 and 4, we can also note that using the true gradient results in substantially fewer steps than using the estimated gradient.
This work focuses alignment for diffusion models. As far can see, there is no foresecable negative impact on the society.
References
Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., Das- Sarma, N., Drain, D., Fort, S.. Ganguli, D., Henighan. T., et al. Training a helpful and harmless assistant with rein- forcement learning from human feedback. arXiv preprint arXiv:2204.05862 2022.
6. Conclusions
Ben-Hamu, H., Puny, O., Gat, I., Karrer, B., Singer, U., and Lipman, Y. D-flow: Differentiating through flows for controlled generation. arXiv preprint 240. 2024.
In this work, we present a comprehensive study on Direct Noise Optimization (DNO) for aligning diffusion genera- tive models at inference-time. We introduce variants of DNO designed to efficiently address challenges such as out-of-distribution ceward-hacking and the optimization of non-differentiable reward functions. More significantly, we demonstrate the exceptional efficacy of DNO, underscoring its capacity to rival tuning- based methods. The primary limitation of DNO lies in its integration with the sampling process of diffusion models, leading to a substantial increase in processing time compared to direct sampling. Nonethe-
Bertsekas, D. P. Nonlinear programming. Journal of the
Operational Research Society, 48(3):334-334, 1997.
Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S.
Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. /3301, 2023.
Chen, I.., Zhu, C., Soselia, D., Chen, J., Zhou, T., Goldstein,

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
T., Huang, H., Shoeybi, M., and Catanzaro, B. Odin: Disentangled reward mitigates hacking in rlhf. arXiv preprint arXiv:2402.07319 2024a.
Ho, and Salimans, T. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12533 2022.
Ho, J., Jain, A., and Abbeel, Denoising diffusion proba- bilistic models. Advances in Neural Information Process- ing Systems, 33:6840-6851, 2020.
Chen, W. Zhang, J., Wu, J., Wu, H., Xiao, X., and Lin, L. Id-aligner: Enhancing identity preserving text-to-ime ext-to-immae generation with reward feedback learning, 2024b.
Ho, J., Chan, W., Saharia, C., Whang, J., Gao, R., Gritsenko, A., Kingma, D. P., Poole, B., Norouzi, M., Fleet, D. J., al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210002303, 2022.
Clark, K., VIcol, P.. Swersky, K., and Fleet, D. J. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400 2023.
CrossLabs.
Diffusion with offset noise. https: // www.:wrrollaa.....oo org/blog/
Hu, J., Tao, L.., Yang. J., and Zhou, C. Aligning language models with offline learning from human feedback, 2023.
diffusion www.ffusoowwwtto-oiteeeeee with- - noise, 2023. January 30, 2023.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-basee generative models. Advances in Neural Information Processing Systems, 35: 26565-26577, 2022.
Deng, F., Wang, Q., Wei, W., Grundmann, M., and Hou, T. Prdp: Proximal reward difference prediction for large-scale reward finetuning of diffusion models. arXiv preprint urXiv:2402.08714, 2024.
Karunratanakul, K.,, Preechakul, K., Aksan, E., Beeler, T., Suwajanakom, S., and Tang, S. Optimizing diffusion noise can serve universal motion priors. arXiv preprint arXiv:2312.1199 11994 2023.
Dhariwal, and Nichol, A. Diffusion models beat gans on image synthesis. Advances in Neural Information 2021.
Processing Systems,
Kingma, D. and Ba, Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 2014.
Dong, H., Xiong, W., Goyal, D., Pan, R., Diao, S., Zhang, J., Shum, K., and Zhang. T. Raft: Reward ranked fine- tuning for generative foundation model alignment. arXiv preprint SSTI:22304067666 2023.
Kirstain, Y., Polyak, A., Singer, U., Matiana, S., Penna, J., and Levy, O. Pick-a-pic: An open dataset of user preferences for text-to-image generation, 2023.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Dpok: Reinforcement learning for fine-tuning text-to- image diffusion models, 2023.
Korbak, T., Shi, K., Chen, A., Bhalerao, R., Buckley, C. L., Phang, J., Bowman, S. R., and Perez, E. Pretraining language models with human preferences, 2023.
Geng, Z., Pokle, A., and Kolter, J. Z. One-step diffusion distillation via deep equilibrium models. arXiv preprint arXiv:2401.08633 2023.
Kynkäääniemi, T., Karras, T., Laine, S., Lehtinen, J.. and Aila, T. Improved precision and recall metric for assess- ing generative models, 2019.
Gulcehre, C., Paine, T. L.., Srinivasan, S., Konyushkova, K., Weerts, L., Sharma, A., Siddhant, A., Ahern, A., Wang, M., Gu, C., Macherey, W., Doucet, A., Firat, O., and de Freitas, N. Reinforced self-training (rest) for language modeling, 2023.
Li, J., Li, D., Savarese, S., and Hoi, S. Blip-2: Bootstrapping language-image pre-training with frozen image encoders and large language models, 2023.
Lin, Z., Pathak, D., Li, B., Li, J., Xia, X., Neubig, G., Zhang. P., and Ramanan, D. Evaluating text-to-visual generation with image-to-tex generation. arXiv preprint arXiv:2404.0129/ 2024.
Guo, Z., Wu, Y., Chen, Z. Chen, L., and He. Q. Pulid: Pure and lightning id eustomization via contrastive alignment, 2024.
Liu, H., Sferrazza, C., and Abbeel, P. Languages are rewards: Hindsight finetuning using human feedback, 2023a. URL https: org/ abs/ 2302. 02676.
Hao, Y., Chi, Z. Dong, L., and Wei, Optimizing prompts for text-to-image generation. arXiv preprint arXiv:2212.096111 2022.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Liu, X., Zhang, X., Ma, J., Peng, J., and Liu, Q. Instaflow: One step is enough for high-quality diffusion-basse text- to-image generation. arXiv preprint arXiv:2309.06::3 2023b.
10

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
Lu, C., Zhou, Y., Bao, F,, Chen, J., Li, C., and Zhu, Dpm-solver: A fast ode solver for diffusion probabilis- tic model sampling in around 10 steps. arXiv preprint arXiv:2206. 00927, 2022.
image synthesis. In The Twelfth International Confer- ence Learning Representations, 2024. URL https: / //openeeview nett wet.://oooomn.ddddd2222222-
Prabhudesai, M., Goyal, A., Pathak, D., and Fragkiadaki, K. Aligning rext-to-image diffusion models with reward backpropagation, 2023.
Meng, C., He, Y., Song, Y., Song. J., Wu, J., Zhu, J.-Y., and Ermon, S. SDEdit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations, 2022.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp.
Meng, C., Rombach, R., Gao, R., Kingma, D., Ermon, S., Ho, J., and Salimans, T. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 14297-14306, 2023.
PMLR, 2021.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Miao, Y., Zhang, S., Ding, L., Bao, R., Zhang, L., and Tao, D. Mitigating reward hacking via information-Hen theoretic reward modeling. arXiv preprint arXiv:240200934,, 2024.
-hotorttin-thhorttic
Salimans, T. and Ho, Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005122 2022.
Mudgal, S., Lee, J., Ganapathy, H., Li, Y., Wang. T., Huang, Y., Chen, Z., Cheng, H.-T., Collins, M., Strohman, T., Chen, J., Beutel, A., and Beirami, A. Controlled decoding from language models, 2023.
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Nesterov, Y. and Spokoiny, V. Random gradient-free mini- mization of convex functions. Foundations of Computa- tional Mathematics, 17 17(2):577-566, 2017.
Sauer, A., Lorenz, D., Blattmann, A., and Rombach, Adversarial diffusion distillation. arXiv preprint arXhv:2311.1742,.0 2023.
Novack, Z., McAuley, J., Berg- Kirk-Kiikpatrick, T., and Bryan, N. J. Ditto: Diffusion inference-time t-optimization for music generation. arXiv preprint arXiv:2401.12179, 2024.
Schuhmann, C., Beaumont, R., Vencu, R.. Gordon, C., Wightman, R., Cherti, M., Coombes, T., Katta, A., Mullis, C., Wortsman, M., Schramowski, P., Kundurthy, S., Crow- son, K., Schmidt, L., Kaczmarczyk, R., and Jitsev, J. Laion-5b: An open large-scale dataset for training next generation image-text models, 2022a.
OpenAl. Chagpp..ttpss//oppnai..como bltg//hatoo/oo//ob..booo 2022. URL https:tt//// https://cppeaade wwtpbotw./oooooo..co...oo....ooo http://llog/hhtttt.././o chatqpt/.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to fol- low instructions with human feedback. arXiv preprint arXiv:2203.0m 02555, 2022.
Schuhmann, C., Beaumont, R., Vencu, R., Gordon, C., Wightman, R., Cherti, M., Coombes, T., Katta, A., Mullis, C., Wortsman, M., et al. Laion-5b: An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 35: 22278-25294, 2022b.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et Pytorch: An imperative style, higf-peeformance deep learning library. Advances in neural information processing systems, 32, 2019.
Song. J., Meng, C., and Ermon, Denoising diffusion im- plicit models. arXiv preprint arXiv:2010.0251 2020a.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. Proceedings of the IEEF/CVF Interna- tional Conference on Computer Vision, pp. 4195-4205 2023.
Song, J., Zhang, Q., Yin, H., Mardani, M., Liu, M.-Y., Kautz, J., Chen, Y., and Vahdat, A. Loss-guided diffu- sion models for plug-and and--and-llay controllable generation. In International Conference on Machine Learning, pp. 04483-32498. PMLR, 2023a.
Podell, D., English, Z., Lacey, K., Blattmann, A., Dock- horn, T., Müller, J., Penna, J., and Rombach, R. SDXL: Improving latent diffusion models for high-resolution
Song, Y., Sohl- Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. baore-based generative modeling

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
through stochastic differential equations. arXiv preprini arXiv:2011.13456, 2020b.
Wallace, B., Dang, M., Rafailov, R., Zhou, L.,, Lou, A., Pu- rushwalkam, S., Ermon, S., Xiong, C., Joty, S., and Naik, N. Diffusion model alignment using direct preference optimization. arXiv preprint
Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. Consis- tency models, arXiv preprint arXiv:2303.01465 2023b.
2023a.
Wallace, B., Gokul, A., Ermon, S., and Naik, N. End- to-end diffusion latent optimization improves classifier guidance. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7280-7290, 2023b.
Stiennon, N., Ouyang, L.., Wu, J., Ziegler, D., Lowe, R., Voss, C., Radford, A., Amodei, D., and Christiano, P. F. Learning to summarize with human feedback. Ad- vances in Neural Information Processing Systems, 33: 3008 3008-3021,22 2020.
Wang, G., Cheng, S., Zhan, X., Li, X., Song, S., and Liu, Y. Openchat: Advancing open-source language models with mixed-quality data, 2023a.
Sun, J., Fu, D., Hu, Y., Wang. S., Rassin, R., Juan, D.-C., Alon, D., Herrmann, C., van Steenkiste, S., Krishna, R., and Rashtchian, C. Dreamsyne: Aligning text-to-image generation with image understanding feedback, 2023.
Wang, W., Mrini, K., Yang, L., Kumar, S., Tian, Y., Yan, X., and Wang, H. Finetuned multimodal language models are high-quality image-text data filters. 2024a.
Sun, T., Shao, Y., Qian, H., Huang. X., and Qiu, X. Black- box tuning for language-modelas-a-sseee national Conference on Machine Learning, pp. 20841- 20855, PMLR, 2022.
Wang, W., Mrini, K., Yang. L., Kumar, S., Tian, Y., Yan, X., and Wang, H. Finetuned multimodal language models are high-quality image-text data filters. arXiv preprint arXiv:2403.000 026777 2024b.
In Inter-
Tang. Z., Rybin, D., and Chang, T.-H. Zeroth-order opti- mization meets human feedback: Provable learning via ranking oracles. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL https: //openreview net www.eorrrrooooooooo.....o.rieeie TVVVVig.u...........mm
Wang. Z. J., Montoya, E., Munechika, D., Yang, H., Hoover, B., and Chau, D. H. Diffusiondb: large-scale prompt gallery dataset for text-to-image generative mod- els, 2023b.
Wu, X., Hao, Y., Sun, K., Chen, Y., Zhu, F., Zhao, R., and Li, H. Human preference score v2: A solid benchmark for evaluating human preferences of text-to-image synthesis, 2023.
Tang. Z., Tang, J., Luo, H., Wang, F., and Chang, T.-H. Accelerating parallel sampling of diffusion models. arXiv preprint 202V:2200..09970, 2024b.
Touvron, H., Martin, I.., Stone, K., Albert, P., Almahairi,
Xing. Y., He, Y. Tian, Z. Wang, X., and Chen, Q. Seeing and hearing: Open-domain visual-audio gen- eration with diffusion latent aligners. arXiv preprint arXh::2402.17733, 2024.
A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshom, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril, T., Lee, J., Liskovich, D., Lu, Y., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P., Molybog, I. Nie, Y., Poulton, A., Reizenstein, J., Rungta, R., Saladi, K., Schelten, A., Silva, R., Smith, E. M., Subramanian, R., Tan, X. E., Tang, B., Taylor, R., Williams, A., Kuan, J. X., Xu, P., Yan, Z., Zarov, I., Zhang, Y., Fan, A., Kambadur, M., Narang, S., Rodriguez, A., Stojnic, R., Edunov, S., and Scialom, T. Llama 2: Open foundation and fine-tuned chat models, 2023.
Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q., Ding, M., Tang, J., and Dong, Y. Imagereward: Learning and evaluating human preferences for text-to-image generation, 2023.
Yang, L., Zhang, Z., Song, Y., Hong, S., Xu, R., Zhao, Y., Zhang, W., Cui, B., and Yang, M.-H. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1-39, 2023.
Yuan, H., Huang, K., Ni, C., Chen, M., and Wang, M. Reward-directed conditional diffusion: Provable distribu- tion estimation and reward improvement, 2023.
Uehara, M., Zhao, Y., Black, K., Hajiramezanali, E., Scalia, G., Diamant, N. L.. Tseng, A. M.. Levine, S., and Bian- calani, T. Feedback efficient online fine-tuning of diffu- sion models. arXiv preprint arXiv:2402.16333 2024.
Yuan, H., Chen, Z., Ji, K., and Gu, Q. Self-play fine-tuning of diffusion models for text-to-image generation. arXiv preprint arXiv:2402. 10210, 2024.
Zhang, S., Yang, X., Feng, Y., Qin, C., Chen, C.-C., Yu, N., Chen, Z., Wang, H., Savarese, S., Ermon, S., et al. Hive: Harnessing human feedback for instructional visual editing. arXiv preprint arXiv:2303. 096/8, 2023.
Wainwright, M. J. High-dimensional statistics: A non- asymptotic viewpoint, volume 48. Cambridge university press, 2019.
12

---

Inference-Time Alignment of Diffusion Models with Direct Noisc Optimization
Zhao, W., Bai, L.., Rao, Y., Zhou, J., and Lu, J. Unipe: A unified predictor-corrector framework for fast sampling of diffusion models. arXiv preprint arXiv:2302.04867, 2023.
13

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
Appendix
A. DDIM Sampling Algorithm
In Algorithm described below, we summarize the sampling algorithm for diffusion models, DDIM (Song et al., 2020a), which is essentially the Euler method for solving ODEs/SDEs. The diffusion coefficients ai.....: orr Algorithm computed using the coefficient functions f(() and 9()), with the detailed computation found in (Song al., 2020b; Karras al., 2022). The coefficient y in DDIM determines whether solving the ODE the SDE, with r) 0 corresponding to ODE and corresponding SDE.
Algorithm DDIM Sampling Algorithm
Require: Discretization timesteps T, diffusion coefficient oi,,.., Ot, initial noise ær NV 1), noise vectors 21,.., 2T (0, coefficient 1) E [0, 1] for balancing ODE and SDE, learned score network ca(:,)).
1: for to = to do 2: Compute oe TT(11 at-1)/(1 ax) at/aa-i
3: Jt-1 - / Jt (Vot-1(1 as)/ar vI 0:11 - of co( 1) oc2t
end
B. Visualization
To assist the reader in understanding the optimization process of DNO and also provide qualitative evaluation, we offer several visualizations in this section.
B. A SIMPLE EXAMPLE
In this section, we present a simple example to visualize the process of noise optimization. Specifically, we trained a toy diffusion model for generating uniform distribution on ring with a radius between 0.8 and 1.2, and the initial distribution visualized in Figure 5a, where each red point denotes a single sample drawn from the trained diffusion model. We then solve the noise optimization problem (3) with the reward function r(z) sin(4 [1][ + sin(4
((=[1) - 1)2 - x(2]²) /5. solve the DNO problem (3) using
highly nonconvex function with many local maxima. To perform this optimization, gradient ascent with the learning rate set to 0.01. In Figures 5b, 5c, and 5d, we visualize the optimized samples after 10, 50, and 100 gradient steps, respectively. We can observe that the distribution of the samples shifts toward distribution on the local maxima of the reward function.
1.0
0.8
0.5
0.5
0.5
0.5
0.6
0.3
0.0
0.0
0.0
0.4
-05
-055
0.2
-1.0
-1.0
-1.0
-1.0
-1.0
0.5
1.0
-1. -0.5
0.5
1.0
-100 -0.5
1.0
-1.0
1.0
0.0
(a) Initial
(b) 10 steps
(c) 50 steps
100 steps
Figure 5. Example 1: Evolution of the sample distribution of a toy diffusion model while running DNO to maximize non-convex reward function.
B.2. SIMPLE EXAMPLES FOR OOD HOWWARD-HACKING
In this section, we present two examples using both simple diffusion model from the Example 1 described above and the open-source image diffusion model SD v1.5 (Rombach et al., 2022). For the first example, we revisit the pretrained distribution displayed in Figure 5, modifying the reward function to r(x) = - (x(11]1.4)) (x 1 Figure 6a exhibits the pretrained distribution of the diffusion models, where we note that every sample stays within the support. However, as illustrated in Figure 6b, after 1000 gradient steps of optimization for the reward function, all the generated samples become
14

---

Inference-Timee Alignment of Diffusion Models with Direct Noise Optimization
out-of-distribution. In second example, we examine optimization for the "brightness" reward- the average value of all pixels in image-un using SD v1.5 the image diffusion model. We start with the prompt "black duck", with the initial image depicted in Figure 6c. After 50 gradient steps of optimization for the brightness reward. it becomes apparent that the generated samples diverge from the original prompt "black duck and transform into a "white duck". as evidenced in Figure 6d. Ideally, in the absence of reward-hacking, the generated samples should always adhere to the "black duck" prompt while incorporating the overall brightness in the images.
1.0
1.
100
0.8
0.5
0.5
0.6
0.4
0.2
1.0
-1.0
-1.0
0.5
1.0
-1.0
-0.5
0.5
1.0
0.0
(a) Initial
(b) 1000 steps
(c) Black Duck
(d) 50 steps
Figure 6. Examples of OOD Reward Haward-
B.3. COMPARING WITH AND WITHOUT REGULARIZATION
In this section, our objective demonstrate the impact of the probability regularization term on the optimization process through visual examples. Specifically, present examples from three settings. The first two examples derived from the experiments in Section 5.1, while the last example is from the experiment optimizing the aesthetic score in Section 5.2.
The examples can be seen in Figure 7. regularization term consistently prevents the generated samples from falling into the category of being Out-OF-Distribution (OOD).
As observed across all examples, the optimization process that integrates the
B.4. QUALITATIVE EXAMPIEES
In this section, we aim provide additional visualized examples for our proposed method.
Firstly, in Figure 8, present examples from the optimization of all three popular human-level reward functions discussed in Section 5.2. As can be observed, the optimization process indeed results in an increase in human preference throughout.
Furthermore, we also include examples from the experiments in Section 5.3, that is, optimizing JPEG Compressibility Score and Aesthetic Score using the Hybrid-2 method for non con-differentiable optimization. These examples Figure effectively showcase the efficiency of Hybrid-2 in both estimating the gradient and optimizing.
We also present some non-cherry-picked examples of aligning Stable Diffusion XI. (Podell et al., 2024) with DNO across four reward functions and four popular prompts from Reddit, see Figure 10. Note that this effect is achieved without fine-tuning the diffusion models. The experiment was conducted on single A800 GPU. and also the setting for Figure 10. In these examples, we only use the base model of SDXL (Podell et al., 2024) as the image diffusion model. We adopt the DDIM sampler with 50 steps and 1 = 1 for generation, and optimize all the injected noise the generation process, the same as most experiments this work. The classifier-free guidance is set 5.0. For each reward function, we adopt the same hyperparameters for the optimizer and regularization terms as the experiments in Section 5.1 and 5.2. From top to bottom Figure 10, the used prompts are listed follows:
1. dark alley, night, moon Cinematic light, intricate detail, high detail, sharp focus, smooth, aesthetic, extremely detailed
2. 1970s baseball player, hyperdetailed, soft light, sharp, best quality, masterpiece, realistic, Canon EOS R3, 20 megapixels.
3. a rabbit, wildlife photography, photograph, high quality, wildlife. ff 1.8, soft focus, 8k, national geographic, award winning photograph by nick nichols.
15

---

Inference- Time Alignment of Diffusion Models with Direct Noise Optimization
Initial
5 steps
10 steps
20 steps
30 steps
50 steps
w/o
reg
w/
reg
(a) Optimizing brightness reward with and without probability regularization. Prompt: "black duck".
Initial
5 steps
10 steps
20 steps
30 steps
50 steps
w/o
reg
w/
reg
(b) Optimizing darkness reward with and without probability regularization. Prompt: "white duck".
Initial
5 steps
10 steps
20 steps
30 steps
50 steps
w/o
reg
w/ reg
(c) Optimizing Aesthetic Score with and without probability regularization. Prompt: "yellow squirrel"'
Figure Examples of optimized samples with and without regularization
16

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
Initial
5 steps
10 steps
20 steps
30 steps
50 steps
(a) Optimizing Aesthetic Score with prompt "gray lion"
Initial
5 steps
10 steps
20 steps
30 steps
50 steps
(b) Optimizing HPS v2 Score with prompt "black deer"
Initial
5 steps
10 steps
20 steps
30 steps
50 steps
(c) Optimizing PickScore with prompt "black lizard"
Figure 8. Representative examples of optimizing reward functions trained on human feedback data.
17

---

Inference-Tinee Time Alignment of Diffusion Models with Direct Noise Optimization
4. A beef steak, depth of field, bokeh, soft light, by Yasmin Albatoul, Harry Fayt, centered, extremely detailed, Nikon D850, award winning photography
Initial
10 steps
20 steps
30 steps
50 steps
100 steps
(a) Optimizing Jpeg Compressiblity with Hybrid-2 gradient approximation. Upper: prompt "blue pig". Lower: prompt "yellow
rabbit".. Initial
50 steps
100 steps
200 steps
300 steps
500 steps
(b) Optimizing Aesthetic Score with Hybrid-2 gradient approximation. Upper: prompt "silver butterfly". Lower: prompt "yellow hedgehog".
Figure 9. Representative examples for non-differentiable optimization
B.5. VISUALIZATION OF DNO
We provide visual illustration in Figure 11 to describe the main procedure for DNO using the DDIM sampling procedure detailed in Algorithm 1. As shown, DNO, similar to LGD (Song et al., 2023a). operates at inference-time and does not require tuning the network parameter 0. However, it requires more time for generation compared to direct sampling, as the optimization is integrated with the sampling process, meaning the optimization is performed for each new sample generated. Despite this, as we will demonstrate Section 5, the extra time needed for the DNO approach a worthwhile trade-off for obtaining high-reward samples.
C. Theoretical Results
Proof for Theorem 2.1. To leverage the L-smoothness assumption, we need to state a classical lemma for smooth optimiza- tion.
18

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
Image Brightness
Image Darkness
Aesthetic Score
Pick Score
0
2
Time (minutes)
Figure 10. Running DNO with SDXL
19

---

Inference-Timg Alignment of Diffusion Models with Direct Noise Optimization
Lemma .1 (Descent Lemma (Bertsekas, 1997)). For any 21 and 2g, have
L ro Ma(za) 2 ro Mo(z)) Vr Ma(zi) (22 =11 - -11122 - =1112-
(10)
Now for any and steps t 2 1, with the descent lemma, we have
Mo(g++1(2)) Mo(gr(z)) + Vr Mo(a((z)) (9++(22) 9x(2)) (g(()))))))))))))))
(2) (x)))))))))=)
Notice that by the definition of 9t- we have
g(+1(2) - g((z) go(z) ivr Mo(g(())) - ge(z)
.Vr Mo(g((=))
Therefore, we have
ro Mo(g+1(2)) 2 Mc(g((=)) (8- ( -
Vro (=)
Taking the expectation over ~ NN(0 I) we have
Es,,(0,1) (r Mo(g(+1(2))) Ez-Nio. (ro Mo(or()))]+
L
E BEnNionD
[[Vr Mo(gx<)))1
By using the change of variable formula for distribution, we can easily see that
E (x))(() E2- [r Me(0)+1(3))).
and
E [x=)((x)T(x>> = -E-Nio. [ro Mo(g:
Therefore, we conclude with
PI) 2
(x)p+)()) (x))) + (
B /| Mo(e)
(11)
2 (x)))
When Does the Distribution Stop Improving? As observed, the distribution ceases to improve when the sec- ond term in Equation (11) becomes zero. Initially, we note that the optimized distribution stops improving when E V/ Ma(z))==99(to) 122 = 0. In statistical terms, this implies that V er Ma(=)) is zero vector with probability one.
To discern the circumstances under which this zero vector occurs, assume that 2g some fixed noise vector and consider the scenario where
V Var V==(1))==b (= 0,
(12)
Here, we denote (2)))) = (ar(z) (e= as the gradient of the reward functions on the generated sample, and G2(20) = (=)a(=)) representing the Jacobian matrix the roise-to-sample mapping. We categorize the situation in Equation (12) into three cases:
20

---

Inference- Time Alignment of Diffusion Models with Direct Noise Optimization
2T
1
Forward
Compute Gradient
Estimate Gradient
TT
Reward
DDIM step
DDIM
r T(xO)
step
TT-1
0
Prompt
Optimization Variables
Mo
"Black Duck"
Frozen Modules
Figure 11. Overview of the DNO procedure with the DDIM sampling algorithm: DNO seeks optimize only those Gaussian noise vectors (2T, 21, 22..., 2r} to maximize the reward value a single generated sample zo. To facilitate the gradient backpropagation from za {xI, 00, #2..,, hr leverage the technique gradient checkpointing. is worth noting that when using r for DDIM sampling, there is no need to compute the gradient for 21, 2r, as the generated sample zo depends exclusively on ær. When computing the gradient from r(ra) zo. we can use either ground-trute gradient or estimated gradient Vr. depending on whether the reward function r(-) is differentiable.
Type-I: (2o)1l 0 and ||G2(zo)]| O. Here, the gradient of the reward function the generated sample is zero, indicating that the generated sample has reached stationary point (or local solution) of the reward function.
Type-II: ||)](z))|| = 0 and [[G((2o) > 0. This indicates that the Jacobian matrix of the taise-to-sample mapping is zero, which often suggests that the generated sample is at the boundary of the support of the distribution, as a zero Jacobian means that changes the noise will not affect the generated sample.
Type-III: |[ [2 0 and |[2((z)]] 0, but (zg) G2(2)) 0. In this scenario, the gradient of the reward function on the generated sample is orthogonal to the Jacobian matrix of the noise-to-sample mapping.
In summary, the distribution will halt its improvement after the (-th step if it almost surely holds that zu corresponds to Type-1, Type-II, or Type-Ill noise vector.
We provide examples for the three scenarios, respectively, in the following figures. First, in Figure 12a, we display examples of Type-I and Type-II by reutilizing the experiment from Figure 5. To determine the type of the noise vector, we empirically compute ||G (=.) G2(zz))l ||. and (zo) G2(zo)] for each noise vector.
To showcase an example of Type-III noise vectors, introduce new toy example illustrated in Figure 12b. Specifically, the ground-truth distribution learned by diffusion models is uniform across horizontal line spanning from (-1, to (1, 0). The reward function is defined as r(z,, y) = It can be readily confirmed that, for every point on this line, the gradient of the reward function is orthogonal to the Jacobian matrix of the noiselto-sample mapping. Consequently, all points along the line segment [( (111 0), (1] qualify as Type-Ill noise vectors.
D. Understanding the Faster Speed of DNO with SDE-Based Sampling.
Intuitively, the acceleration in optimization speed can be attributed to the finer-graines control over the generation process afforded by the SDE sampler compared to the ODE sampler. To formally state this intuition, we revisit the DDIM sampling algorithm in Algorithm 1. We consider the procedure of SDE-based sampling algorithm, DDIM with = as defining the voise-to-sample mapping Mo 21,..., 2T). An important observation that the ODE sampling
21

---

Inference- Time Alignment of Diffusion Models with Direct Noise Optimization
Type-11
1.0
Type-1ll
1.0
1.0
0.8
0.5
1.0
0.8
0.6
0.0
0.5
0.6
0.4
-0.5
0.0
0.4
0.2
-1.0
-0.5
0.2
0.0
-1.0
0.5
0.0
0.5
1.0
-1.0
0.0
Type-1
-1.0
-0.5
0.0
0.5
1.0
(a) Example
(b) Example 2
Figure 12, Examples of generated samples with Type-I, Type-II and Type-III noise vectors the toy examples.
0.8
0.6
0.2
-1.0
-0.5
0.0
0.5
1.0
(a) Toy Example
(b) P(=) 0.36
(c) P(z) 0.00
(d) P(z) 0.02
Figure 13. Examples of generated samples with corresponding values of P(z).
0.6
0.4
P(z)
0.3
0.2
0.1
0.0
0
10
20
30
50
90
steps
Figure 14. Trajectory of P(z) optimizing brightness reward.
22

---

Inference-Timee Alignment of Diffusion Models with Direct Noise Optimization
1.0
P(z)
0.8
0.6
0.4
0.2
0.0
0
10
20
30
50 steps
90
Figure 15. Trajectory of P(z) optimizing Aesthetic Score.
algorithm, DDIM with = O, can also be expressed using the same (2T, 21...., 2T), with the distinction that the noise 21, 21,...,27 2r becomes deterministic and dependent on 27, rather than sampled from a Gaussian distribution. Specifically, if we define the deterministie noise vectors as: =t =0DE def. where (1 11,-1/(1 ax) V -1////t-1, then the sampling process of DDIM with y 0 can be expressed as .DE Ma(xT, 21 27 2DBB). In this context, the advantage of SDE-based sampling becomes evident:
co ( (), for 1 1,..., T,
dt
r (M ..,,).)T)))... 2 max (Ma(ur; ODE
LODE
max
meaning that running DNO with SDE will yield better results, or at least as good as DNO with ODE for aligning diffusion models. Based on this conclusion, our work will focus on optimizing the SDE-based sampling (DDIM with 7) 1) for the remainder of the study. Additionally, fix the number of generation steps T 50 throughout this work for simplicity.
.5. Empirical Investigation of P())
In this section, we provide several empirical evidence demonstrate that P(z) acts as an effective indicator for the out-of-distribution phenomenon.
Finstly, in Figure 13a, revisit the examples from Figures 6a and 6b, coloring each sample based the value of P(z). As depicted in Figure 13a, P(z) proves to be an efficient metric to separate in-distribution samples from out- out-of of-distribution samples; those in-distribution have high values for P(=), whereas those out-of-distributia exhibit values of P(z) near zero.
Secondly, we manually construct several noise vectors that reside in the low-probability region of the standard Gaussian distribution. To establish baseline comparison, we first draw one sample from the standard Gaussian distribution and use generate image with Stable Diffusion v1.5 and the prompt "black duck". As be seen, this leads normal image with P(z) also within reasonably large value. We then construct the oow-prbbability vectors in two ways. The first one is to use all-zero vectors, which obviously reside in the low-probability zone of high-dimensional Gaussian distributions. The generated images with all-zero vectors are visualized in Figure 13c, showcasing that there is nothing discernible in the image while P==) approximates zero. The second method to repeat parts the noise vectors, such that the noise vectors exhibit high covariance in the elements. Specifically, we construct the repeated vectors by first generating an n/4 dimensional 2o from the standard Gaussian distribution, and then constructing the noise vectors as z = [zo: 20, 20; zo]. making = an n dimensional vector. The figure corresponding these repeated vectors, shown in Figure 13d, once again results in a poor image, with P(2) illustrating that the noise vectors also come from low-probability region.
We further visualize the entire optimization trajectory for the examples in Figures 6c and 6d, i.e., optimizing the brightness reward for Stable Diffusion v1.5 with the prompt "black duck" in Figure 14. Specifically, from Figure 14 we can clearly see that the value of P(z) gradually decreases, and the generated image also gradually diverges from the distribution associated with black duck. Notably, at around 20 steps, the value of P(2) becomes near- near-zero, and at the same time, the generated image more closely resembles a blue duck rather than the specified black duck.
23

---

Inference-Timg Alignment of Diffusion Models with Direct Noise Optimization
Similarly, we visualize the optimization trajectory for optimizing the Aesthetic Score for SD v1.5 with the prompt "black duck". The results are in Figure 13. A clear conclusion is that in this case, it is less likely for the optimized samples be out-of-distribution. This is mainly because the Aesthetic Score itself penalizes those OOD samples. It is noteworthy to observe that this insight is also captured by our proposed indicator P(z), because when comparing the trend of P(z) in Figure 14 and Figure 15, we can see that optimizing the Aesthetic Score leads to a much less significant decrease in the P(z) value.
.6. Implementation Details
In this section, we discuss some implementation details of our proposed method, as well as clarify some omissions in the experimental section.
1. ALGORITHM IMPLEMENTATION
clear that to solve the direct noise optimization problem stated in Problem 3, differentiation of the noise-to-sample mapping Mo is required. It is worth noting that this differentiation cannot be handled by standard dittediiferentiittinn in PyTorch (Paszke et al., 2019), as it can lead to memory explosion. A common technique to resolve this issue is gradient checkpointing, which has also been adopted by other related works on noise optimization (Wallace et al., 2023b; Novack et al., 2024; Karunratanakul et al., 2023).
Here, we describe an efficient method to implement our proposed hybrid gradient estimators detailed in Section 4, along with the optimization process, by utilizing the built-in difeeffferentiiionnn in Py PyToorch (Paszke al., 2019). Specifically, suppose wish to use q samples to estimate the gradient in Equation (8); that is, we draw q noise vectors for perturbation:
E1,...,Eu Ey. We then generate the corresponding samples x; MM + μεx) for i = 1,... q. At this point, we should compute the estimated gradient of the reward functions in non-differenniiabl mode as follows:
H2(s)= = (r(Mo(z)))
x).
(13)
Finally, we can execute gradient backpropagation with the loss function,
loss(2)
which produces the exact gradient estimator for 2.
6.22. EXPERIMENT DETAILS
In this section, our goal is to provide the experimental details that were omitted from Sections 5.1, 5.2, 5.3.
Details for Section 5.1. In this experiment, to solve the probability-regularized noise optimization problem as formulated in Equation (5), we employ the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 0.01. For optimization with regularization, we set the regularization coefficient 7 to 1. To compute the minibatch stochastic gradient for the regularization term in Equation (5), we set the batch size b the number of random permutations drawn at each step to 100. For each optimization run, we utilize single A800 GPU, with the total memory consumption being approximately 15 GB.
Details for Section 5.2. In this second set of experiments, we continue using the Adam optimizer with a learning rate of 0.01. For optimization with regularization, though, we reduce the regularization coefficient to 0.1 because optimizing these human-like reward functions is less susceptible to the OOD soward-hacking issue, while maintaining the batch size for the permutation matrix d at 100. Each optimization run also uses a single A800 GPU, but the total memory consumption around 20 GB. Regarding the baselines Table 1, we implemented LGD (Song et al., 2023a) ourselves, following the algorithm from their paper on these reward functions. For other basclines, we reuse the statisties presented in their corresponding papers.
Details for Section 5.3. In this section, the primary hyperparameters for the three tested algorithms are the perturbation coefficient/ μ and the number of samples q used to approximate the gradient (as formulated in Equation (13)). Clearly, plays a crucial role in determining the running time of each algorithm. For an equitable comparison, we tune q separately for each algorithm to achieve roughly the same time cost per gradient step. Specifically, we set values for ZO-SGD, Hybrid-1, and Hybrid-2 to 16, 8, and 4 respectively. For μs we also adjust them individually for each algorithm, as they have varying
24

---

Inference-Time Alignment of Diffusion Models with Direct Noise Optimization
sensitivity to p. Through trial and error, we select j values of 0.01 for ZO-SGD and Hybrid-1, and 0.02 for Hybrid-2. Finally, for optimizing JPEG Compressibility. Score experiment, we reduce the learning rate to 0.001, the Aesthetic Score. Each optimization run continues to use a single A800 GPU.
use the Adam optimizer with learning rate of 0.01, but for the Aesthetic found that 0.01 can lead to divergence during optimization for
.7. Hyperparameters Analysis
In this section, we conduct thorough amalysis of the hyperparameters for the proposed method. Our objective is to offer concise guideline for selecting the hyperparameters in the proposed method.
As discussed in Section 3.1, the concentration inequalities involve a hyperparameter k, which represents the dimension of subvectors from the noise vectors z that we aim to assess probabilistically. As noted in Remark 3.2, the dimension should be neither too large nor too small. Additionally, another critical hyperparameter the number permutation matrices b employed compute the stochastic gradient for the probability regularization in Equation (5). Furthermore, aim explore the impact of the regularization coefficient y in the probability regularization term.
To examine the effects of k, b, and mitigating the OOD (Out-Of- Of-Distribution) reward haauidnn problem, we revisit the experiment of optimizing darkness reward with the prompt "white janimals;" from Section 5.1. In Figure 16, we illustrate how these three hyperparameters influence both the reward and the consistency score (CS), across four different values.
Firstly, Figure 16a supports the notion that should be carefully chosen-not too large, yet not overly small. We observe that k underperforms compared noise vectors. Conversely, k 100 proves to be a poor choice because it entails smaller m, potentially rendering the concentration inequalities detailed in Lemma 3.1 less precise.
and & 10, as selecting k fails to account for the covariance among
Secondly, as demonstrated in Figure 16b, the number of permutation matrices seems to have minor impact on the optimization process, provided sufficiently large. Based on empirical evidence, b 100 emerges as an optimal selection for the proposed method.
Lastly, the effects of y are depicted in Figure 16c. Adjusting the value of 7 clearly presents trade-off' between convergence speed and the propensity for OOD seward-hacking problems. Given this observation, we recommend empirically tuning the value of for different reward functions and prompts using limited number of samples and few optimization steps.
25

---

Inference- Time Alignment of Diffusion Models with Direct Noise Optimization
260
240
220
Clip Score
reward
200
k=1
k=2
k = 10
k = 100
120
20
100
20
100
steps
steps
The effect of l..
240
220
Clip Score
reward
200
26
1=1
160
10
24
0=100
b 1000
120
0
20
40
60
100
20
40
60
80
100
steps
steps
(b) The effect of b.
220
Clip Score
rewaro
2 0.3
160
1-000
1.00
V 0.10
0.01
steps
steos
steps
(c) The effect ry.
Figure 16. Hyperparamet Analysis
26


Reviews:
[{"review_id": "n0nYOgzoP7", "reviewer": "Reviewer_Cv15", "summary": "This paper introduces a novel approach for optimizing diffusion input noise based on a specified reward function. The key advancements over prior work include:\n- A new regularization technique that ensures the optimized noise remains within the distribution of the diffusion model.\n- A method for handling non-differentiable reward functions.\n- Optimization of a sequence of noise variables throughout the diffusion SDE sampling process, rather than just a single initial noise in ODE sampling.\n\nThese improvements enhance performance and yield strong empirical results using SDv1.5 as the teacher model, all without requiring any network optimization.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.\n\nRegarding your comment:\n\"Optimizing noise to increase darkness with the prompt 'white [animal]'—does not fully reflect real-world use cases. Are there more practical scenarios where this regularization would be particularly crucial?\"\n\nFirstly, we would like to clarify that although optimizing blackness/whiteness was chosen to better reflect an adversarial setting, these optimizations also represent realistic applications. There is a genuine need to generate images with highly dark or highly light backgrounds, which cannot be achieved by base models through either prompting or best-of-k selection, and this application is actually inspired by a well-known trick for diffusion models called offset-noise (see CrossLabs Blog https://www.crosslabs.org/blog/diffusion-with-offset-noise.).  This is why we believe these are also very important applications.\n\nFrom an academic perspective, the main goal of this section is to demonstrate that OOD Reward Hacking is more likely to occur in a strong adversarial setting, and we propose a solution to mitigate it. Interestingly, we have found that our proposed technique has been adopted in a more practical setting in a recent work (https://arxiv.org/pdf/2412.03876). In this work, the reward model evaluates whether the content is safe while the prompts are malicious. This resembles the adversarial setting in our brightness and darkness example, and they demonstrate that our method can mitigate the reward-hacking phenomenon to some extent.\n\nLooking ahead, as more reward models and applications emerge, we believe our proposed techniques will find even more important and crucial applications.", "claims_and_evidence": "The claims are substantiated with supporting evidence.", "methods_and_evaluation_criteria": "The evaluation is generally well-reasoned. However, Section 5.1 may require additional work to further validate the effectiveness of the proposed regularization. Specifically, the current adversarial setting—optimizing noise to increase darkness with the prompt \"white [animal]\"—does not fully reflect real-world use cases. Are there more practical scenarios where this regularization would be particularly crucial?", "theoretical_claims": "I briefly reviewed the equations, and they appear to be correct.", "experimental_designs_or_analyses": "Yes, but some baselines and comparisons are missing. For instance, the performance of ODE-based and SDE-based noise optimization should be directly compared. Additionally, it would be valuable to explore simpler regularization techniques, such as KL regularization between the optimized noise and a standard Gaussian distribution, as a baseline.", "relation_to_broader_scientific_literature": "This paper addresses the problem of noise optimization for improved image generation. Optimizing noise presents a promising avenue for further performance enhancement beyond weight optimization in pretrained networks and holds significant potential in the emerging trend of test-time scaling.", "essential_references_not_discussed": "it covers the related works well.", "other_strengths_and_weaknesses": "A few additional strengths of this work include:\n\n- The use of SDE and a more flexible optimization target, which significantly outperforms prior approaches and achieves performance comparable to more expensive weight tuning.\n- The ability to handle non-differentiable rewards, broadening the applicability of the method.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "n/a", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AogTzMwmK6", "reviewer": "Reviewer_6H3K", "summary": "This paper investigates the alignment problem of diffusion models during inference and proposes a tuning-free, prompt-agnostic method named Direct Noise Optimization (DNO). The authors theoretically investigate the properties of DNO and propose variants of DNO, aiming to solve problems of out-of-distribution reward hacking and optimization of non-differentiable reward functions. Experiments demonstrate that DNO achieves state-of-the-art performance on several reward functions.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.\n\n### 1. On the smoothness assumption.\n\nConceptually, it is quite straightforward to argue that the reward function is smooth with respect to pixel changes. Because small changes to the image pixels would not result in large differences in the reward function's score. To provide a more concrete answer, we conducted a quantitative analysis:\n\nWe first sampled a noise vector $x_1 \\sim \\mathcal{N}(0, I)$ and then generated a second noise vector $x_2$ in the neighborhood of $x_1$, i.e., $x_2 \\sim \\mathcal{N}(\\sqrt{0.9} x_1, \\sqrt{0.1} I)$. Using the Aesthetic Score as the reward function $r(\\cdot)$, we computed the following quantities:\n\n$A = E_{x_1, x_2} \\frac{||r(M_\\theta(x_1)) - r(M_\\theta(x_2))||}{||x_1 - x_2||}$\n\nand\n\n$B = E_{x_1, x_2} \\frac{||\\nabla r(M_\\theta(x_1)) - \\nabla r(M_\\theta(x_2))||}{||x_1 - x_2||}.$\n\nUsing 100 samples with random prompts from Section 5.1, we estimated these values to be $A=0.19$ and $B=6.93$, respectively. These results rigorously demonstrate that the composite mapping $r \\circ M_\\theta$ is indeed smooth. We will include this justification in a more formal way in the revised manuscript.\n\n\n### 2. Results on more diverse prompts.\n\nWe quickly conducted some additional quantitative experiments using DNO with SD v1.5 and set $\\lambda = 0.1$, similar to the setting in Section 5.2. We tested 1000 randomly selected prompts from the Pick-a-Pic test dataset https://huggingface.co/datasets/yuvalkirstain/pickapic_v1. Below is a new table that reports the average performance of our DNO. As shown, DNO can still perform well for complex prompts. This is not a surprising result, because by design, DNO optimizes noise vectors specific to each prompt, ensuring robust performance across diverse scenarios.\n\n|               | SD v1.5   | DNO (1 min) | DNO (3 min) | DNO (5 min) |\n|---------------|--------|-------------|-------------|-------------|\n| Aesthetic     | 5.769  | 6.013       | 6.993       | 8.305       |\n| HPS           | 0.270  | 0.279       | 0.291       | 0.326       |\n| PickScore     | 21.20  | 21.85       | 23.61       | 24.89       |\n\n### 3. Discussion on the two related works.\n\nThanks for mentioning these two works to us! After carefully reading them, we agree that these two works are indeed closely related to our work, and we will add them to our revised manuscript accordingly.\n\nWhile these two works also focus on noise optimization for diffusion models, there are several distinctions between them and our work:\n\n- **[1]** ReNo considers a similar reward-based gradient optimization for noise optimization. However, they only consider one-step distilled models, rather than the full-step diffusion models explored in our work. This makes their approach a simplified version of our proposed DNO method. Using one-step distilled models can result in faster optimization speed, but this inevitably sacrifices sample quality. Moreover, for one-step distilled models, there is only one noise to optimize, so they do not need to distinguish between ODE and SDE samplers. Finally, their work lacks a deeper analysis of several critical aspects covered in our paper, such as the OOD reward-hacking problem, convergence issues, and extensions to non-differentiability scenarios.\n\n- **[2]** This work considers a fundamentally different setting compared to our work and ReNo [1]. Their approach aims at constructing a \"reward function\" using only the sampling trajectory information, and then applies an optimization idea similar to DNO. Since they do not use any external reward function to directly improve the noise, it is natural to see (as demonstrated in Table 1 of their work) that their improvements are relatively marginal.\n\n### 4. Impact on diversity.\n\nThis is indeed an important question! Unfortunately, it is difficult to rigorously discuss diversity since it heavily depends on the reward function's optimization landscape, which is typically unknown in practice. For instance, if the reward function is strongly concave with a single global maximum, running DNO to convergence would collapse the distribution into a Dirac delta, eliminating diversity entirely. Conversely, if the reward function is constant, DNO would leave diversity unaffected.\n\nThanks for highlighting this point; we will include this discussion in our revised manuscript.\n\n\n### 5. Explanation of Table 2.\n\nThis occurs because we fix the time budget to 1 minute; thus, running DNO with smaller $T$ allows more optimization steps. Although using smaller $T$ can improve performance on these benchmarks, in our main experiments (Table 1), we set $T=50$ to maintain consistency with existing baselines.", "claims_and_evidence": "Most claims in the paper are well-supported by theoretical analysis and empirical results.", "methods_and_evaluation_criteria": "The paper primarily uses reward scores (e.g., Aesthetic Score, HPS Score, PickScore) and Out-of-Distribution (OOD) indicators (e.g., CLIP Score, ITM Score). These metrics are appropriate for assessing alignment performance.", "theoretical_claims": "The key assumption in Theorem 2.1 that the noise-to-sample mapping is smooth is not entirely convincing, and it is unclear how the cited Figure 4 in (Tang et al., 2024a) directly supports the claim. A more rigorous justification is needed.", "experimental_designs_or_analyses": "A potential limitation is the use of a simple animal prompt dataset, which may restrict the evaluation’s generality. Testing on more diverse prompts (e.g., scenes, abstract concepts) could better assess DNO’s robustness.", "relation_to_broader_scientific_literature": "The proposed Direct Noise Optimization (DNO) method contributes to the growing body of work on aligning diffusion models with specific tasks or reward functions. Previous methods like DDPO, DPOK, AlignProp, and DRaFT have explored various approaches to this problem, including reinforcement learning and direct fine-tuning. DNO belongs to the category of inference-time alignment methods, which also includes LGD.", "essential_references_not_discussed": "I noticed that there are some works on inference-time initial noise optimization, such as [1] () and [ 2]. These works seem to be closely related to the topic of this paper. However, they are not cited or discussed in the current version.\n\n---\n[1] Eyring L, Karthik S, Roth K, et al. Reno: Enhancing one-step text-to-image models through reward-based noise optimization. NIPS, 2024.\n\n[2] Qi Z, Bai L, Xiong H, et al. Not all noises are created equally: Diffusion noise selection and optimization. ArXiv 2024.", "other_strengths_and_weaknesses": "**Strengths**\n- DNO is a test-time optimization method that distinguishes itself from conventional approaches like reinforcement learning and direct fine-tuning of diffusion models.\n- The proposed method addresses OOD reward-hacking and non-differentiable rewards.\n- The paper provides detailed mathematical derivations and theory.\n\n---\n**Weaknesses**\n- It is unclear whether the proposed method might negatively impact the original model's capabilities, such as diversity in generation. \n- lacks a comparison with other inference-time initial noise optimization methods. \n- The paper tests on simple animal prompts and does not evaluate DNO performance on commonly used t2i or complex prompts. Additionally, the paper claims that DNO is prompt-agnostic, but human preference metrics like HPS and PickScore actually require prompt consideration.", "other_comments_or_suggestions": "No further comments or suggestions.", "questions_for_authors": "Why does DNO achieve significantly better performance with very few diffusion steps (e.g., 10 or 15 steps) compared to the standard setting of 50 steps (Table 2)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C19vcAikBI", "reviewer": "Reviewer_F5TP", "summary": "This paper conducts a comprehensive investigation in optimizing the noise in the sampling process of diffusion models for alignment. The main contributions are: 1) giving a rigorous definition of noise optimization and extending it to SDE sampling, 2) explaining and quantifying the root of OOD issues in noise optimization and providing a regularization, and 3) extending noise optimization to non-differentiable rewards by estimating their gradients. Generally, this paper delves into some crucial technical issues of noise optimization, a new branch for diffusion alignment, and proposes basic solutions.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.\n\n\n### 1. \"What is the level of M1 and M2 value for diffusion models with hacked reward alignments?\"\n\nIn this work, we did provide a more direct illustration of this point. In lines 262–263, we define the metrics $P(z)$ using the values of $M_1$ and $M_2$. Then, in Figures 2 and 13–15, we show that the value of $P(z)$ diminishes to zero when reward-hacking occurs.\n\n\n\n### 2. \"Can DNO be combined with other tuning-based methods?\"\n\nYes, and we believe the most natural way to combine DNO with tuning-based methods is to directly apply DNO on fine-tuned models. In this way, our proposed DNO can continue to improve samples generated by aligned models at test time. To validate this point, we conducted a quick experiment. Following the setting in Section 5.2, we directly applied DNO to the model fine-tuned by DDPO. We observed that DNO can indeed continue to improve the sample quality to achieve higher rewards and can reach the level of reward achieved by running DNO for 5 minutes using the base model, but with only a 1-minute budget.\n\n| Method               | DDPO     | DNO (5 min) | DDPO + DNO (1 min) |\n|----------------------|----------|-------------|--------------------|\n| Aesthetic            | 7.180    | 8.587       | 8.761              |\n| HPS                  | 0.287    | 0.324       | 0.319              |\n\nIn this sense, applying DNO with fine-tuning-based methods can be viewed as a way to accelerate the DNO algorithm, making it more practical, especially for time-sensitive applications.\n\n\n### 3. \"Could you elaborate on the optimization process of SDE sampling as advised above?\"\n\nYou might be overthinking the complexity here. Using ODEs and SDEs for DNO has almost the same time cost and memory usage. Gradient backpropagation is conducted over the entire sampling process for both ODEs and SDEs using automatic differentiation.\n\nTo clarify how our implementation works, here is a core snippet of our code:\n\n```python\n# SDE sampling logic. To change it into ODE, we only need to modify the sampler and dimension of noise vectors accordingly\nnoise_vectors = torch.randn(args.num_steps + 1, 4, 64, 64, device=args.device)\nnoise_vectors.requires_grad_(True)\nsampler.initialize(noise_vectors, \"sde\")\n\n# Sampling process\nwhile not sampler.is_finished():\n    model_kwargs = sampler.prepare_model_kwargs(prompt_embeds=prompt_embeds)\n    model_output = checkpoint.checkpoint(unet, **model_kwargs)\n    sampler.step(model_output) \n\n# Gradient computation\nsample = sampler.get_last_sample()\nloss = - reward_function(sample)\nloss.backward()\n```\n\n### 4. Relationship to the work [arXiv:2501.09732]\nThank you for pointing out this work to us! We also noticed this work after the ICML submission deadline. For simplicity, we will refer to this work as ITS below. From our perspective, the ITS work represents a different approach to handling test-time scaling for diffusion models using combinatorial search techniques, while we focus on continuous search techniques using gradient-based optimization.\n\nTo illustrate the difference:\n\n* If the reward model is not continuous or differentiable, the method proposed by ITS is more useful, as it does not require continuity or differentiability by design.\n* On the other hand, when the reward model is continuous and smooth, our proposed DNO is more favorable because it leverages more information for optimization, resulting in much faster convergence compared to the ITS work. This is also evident if we compare Figure 9 in the ITS work and Figure 3 of our work, where we show that using the gradient of the Aesthetic reward leads to significantly faster optimization.\n\nAs illustrated above, we acknowledge that this is indeed a very important concurrent work, and we will include a discussion of this work in the revised manuscript.", "claims_and_evidence": "For all contributions mentioned above, there are all clearly convinced by either theoretical or empirical results. Additionally, for 2), I expect a more direct demonstration of the connection between the low probability of z and the OOD example, e.g. what is the level of M1 and M2 value for diffusion models with hacked reward alignments?", "methods_and_evaluation_criteria": "Yes. This paper discusses an underexplored branch of diffusion alignment: optimizing the noise, which makes sense for deeper understanding of diffusion models.", "theoretical_claims": "The correctness of Theorem 2.1, results about SDE optimization advantages, and Lemma 3.1 are all checked. To the best of my knowledge, I do not see any issues.", "experimental_designs_or_analyses": "The experiment is extensive and persuasive. I checked the improvement by introducing regularization and the comparison to existing alignment methods. As an underexplored method, DNO yields considerable performance. \n\nIssues: I wonder whether DNO could be combined with fine-tuning based methods for even more superior performance.", "relation_to_broader_scientific_literature": "this paper delves into some crucial technical issues of noise optimization, an underexplored branch for diffusion alignment, and proposes basic solutions. It provides a fundamental baseline for this field and guides necessary attention to this branch, which I believe is significant.", "essential_references_not_discussed": "One existing (or maybe concurrent) work should be noticed: Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps (https://arxiv.org/abs/2501.09732). I am curious about the difference between this paper and the mentioned work.", "other_strengths_and_weaknesses": "Weaknesses:\n\n[1] The difference between steps of ODE optimization and those of SDE optimization should be elaborated. For example, does optimizing one step for SDE sampling optimize all TxD noise? Since this should be a long serial optimization, how do all gradients computed?", "other_comments_or_suggestions": "I think this paper studies a significant problem. If my concerns are addressed, I would be pleased to raise my scores.", "questions_for_authors": "[1] Can DNO be combined with other tuning-based methods? I believe this could be extremely important for potential real-world applications, since training cost may not be a real concern for industry deployment of diffusion models.\n\n[2] Could you elaborate the optimization process of SDE sampling as advised above? How is the optimization conducted?\n\n[3] Please explain the relationship between the work Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps (https://arxiv.org/abs/2501.09732).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6y1NHgr6fc", "reviewer": "Reviewer_gBKq", "summary": "The paper proposes a novel approach called Direct Noise Optimization (DNO) for aligning diffusion models with continuous reward functions at inference time. DNO optimizes the injected noise during the sampling process to maximize the reward function, without requiring any fine-tuning of the model parameters. The key contributions are:\n\nOut-of-Distribution Reward Hacking‌: The paper identifies out-of-distribution reward hacking as a critical issue in DNO. It proposes a probability regularization technique to ensure the generated samples remain within the support of the pretrained distribution.\n\nNon-Differentiable Reward Functions‌: The paper extends DNO to handle non-differentiable reward functions by developing a hybrid gradient approximation strategy.\n\nExperimental Results‌: Extensive experiments demonstrate that DNO can achieve state-of-the-art reward scores for various image reward functions, within a reasonable time budget for generation. DNO is shown to outperform tuning-based methods in terms of reward scores while requiring significantly fewer computing resources.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.\n\n### 1. Comparing to prior works\nPlease allow us to emphasize and reiterate our main contributions here, which are distinct from all prior works on noise optimization.\n\n* **A more comprehensive formulation for noise optimization with theoretical understanding:**\nFirstly, we reveal a fundamental insight: every stochastic element in the diffusion sampling process can be harnessed and optimized. In contrast, previous works have focused solely on optimizing the initial noise. As another major contribution, we elucidate the underlying mechanism of optimizing noise vectors, showing that it can be viewed as sampling from a provably better distribution.\n* **Identification and mitigation of the OOD reward-hacking problem in noise optimization:**\nFor example, in our experiments on brightness and darkness enhancement, we show that without proper regularization, noise optimization cannot be applied effectively. Our work highlights the standard Gaussian distribution as a crucial prior for noise vectors, enabling successful regularization for such applications.\n* **Extension to non-differentiable reward functions:**\nThis innovation significantly broadens the applicability of DNO to a wider range of scenarios, as requiring reward functions to be differentiable can be a highly restrictive condition in real-world applications.\n\nWith these three important contributions developed in this work, we believe it lays a stronger foundation for more applications of noise optimization in the future.\n\n### 2. Optimization time\nWe would like to point out that the optimization time in DNO is actually controllable, which provides significant flexibility for different applications.\n\nFor time-sensitive applications, fewer optimization steps can be used, making the additional time less prohibitive, although the sample improvement may also be limited in such cases. For less time-sensitive applications, more optimization steps can be employed to achieve better results.\n\nMore importantly, as also discussed in our second response to Reviewer F5TP, directly combining DNO with fine-tuning-based methods will make the time cost required by DNO more acceptable.\n\n### 3. More modality\nFor audio, we believe the answer is yes. For text, however, we think the technique of DNO is not directly applicable. This is because, in discrete diffusion, it is not possible to compute the gradient from the sample back to the latent noise due to the combinatorial nature of the problem. The reasons we chose to experiment with image diffusion models are mainly threefold:\n\n1. There are many excellent open-sourced image diffusion models available for us to conduct experiments.\n2. There are numerous existing baselines for the image alignment problem, which allow for meaningful comparisons. In contrast, there are very few implementations addressing the alignment problem for audio.\n3. Most importantly, there are many powerful, open-sourced reward models for images, which are trained on high-quality human-ranked datasets. Such resources are currently lacking for other modalities like audio.\n\nIn the future, once the community for audio diffusion achieves the same maturity as the image diffusion community, we believe DNO will continue to demonstrate its value in this domain.\n\n### 4. Scaling behavior of computational resources\nThank you for bringing this important question to our attention. This is indeed a crucial aspect to discuss!\n\nDue to limited tokens in the rebuttal response, here we can only provide a quick response and leave the evidence for the revised manuscript or later discussion period of rebuttal. Here, we will elaborate on how memory usage scales with fine-tuning methods compared to our DNO method, as memory usage is typically the dominant factor in determining the number of GPUs required for these tasks. In summary, our conclusions can be drawn as follows:\n\n1. Assuming the memory usage of direct sampling as 1 unit, fine-tuning methods would require at least 10 times more units to get started, while our DNO method requires approximately 1.5 times more units. Generally, memory usage scales linearly as the model size grows.\n2. One gradient step in DNO roughly takes 2.5–2.7 times longer than direct sampling. Generally, the time cost of direct sampling scales sublinearly or linearly as the model size grows, depending on the level of parallelism in the architecture.\n\nThank you again for raising this excellent discussion. We believe this analysis will make our manuscript more informative, and we will include these insights and the rigorous measurements in the revised version.\n\n### 5. Presentation\n\nCould you please specify which parts you feel need the most improvement or reorganization? Your detailed suggestions would be greatly appreciated. Thank you for your valuable feedback!", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence. The authors provide theoretical analysis, extensive experiments, and visual examples to back up their claims about the effectiveness and efficiency of Direct Noise Optimization (DNO) for aligning diffusion models at inference-time. \n\nThe authors present a comprehensive theoretical study of DNO, including a theorem that demonstrates the improvement of the distribution after each gradient step. They also propose variants of DNO to handle non-differentiable reward functions and address the out-of-distribution (OOD) reward hacking problem. The theoretical foundation is solid, with proofs and justifications provided in the appendices.\n\nThe authors conduct extensive experiments on several important reward functions, including brightness, darkness, aesthetic score, HPS-v2 score, and PickScore. They compare DNO with existing alignment methods like LGD, SPIN, DDPO, and AlignProp, showing that DNO can achieve state-of-the-art reward scores within a reasonable time budget. The results are presented in tables and figures, demonstrating the superiority of DNO in various settings.\n\nThe authors explore three methods for optimizing non-differentiable reward functions, including the proposed Hybrid-2 method. They provide experimental results showing that Hybrid-2 is significantly faster and more effective than other methods like ZO-SGD and Hybrid-1. The experiments on JPEG Compressibility and Aesthetic Score reward functions validate the effectiveness of their approach.\n\nTo prevent OOD reward hacking, the authors introduce a novel probability regularization technique. They provide visual examples and quantitative metrics (like CLIP Score and ITM score) to show that this technique effectively keeps the generated samples within the support of the pretrained distribution. The regularization term is shown to stabilize the optimization process and maintain sample quality.\n\nThe authors argue that DNO is efficient and practical, requiring significantly fewer computing resources than tuning-based methods. They provide details about the memory usage and time budget for experiments, showing that DNO can run on a single consumer-level GPU with memory usage of less than 15GB. This is supported by the implementation details and experimental settings described in the paper.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in the paper make good sense for the problem of aligning diffusion models with reward functions at inference-time. Here's why:\n\nThis approach is well-suited for inference-time alignment as it doesn't require modifying the pretrained model parameters. By optimizing the injected noise during the sampling process, DNO can effectively align the generated samples with the target reward function while maintaining the pretrained distribution's support.\n\nThe introduction of probability regularization to prevent out-of-distribution (OOD) reward hacking addresses a critical issue in alignment methods. This technique ensures that optimized samples remain within the support of the pretrained distribution, which is essential for maintaining sample quality and relevance.\n\nFor non-differentiable reward functions, the proposed hybrid gradient methods (especially Hybrid-2) provide an efficient solution. These approaches combine estimated gradients of the reward function with true gradients of the noise-to-sample mapping, offering a practical way to handle real-world scenarios where reward functions may not be differentiable.\n\nThe authors use established benchmark datasets and reward functions relevant to diffusion model alignment, such as aesthetic score, HPS-v2 score, and PickScore. These are appropriate and widely recognized metrics in the field of generative models and alignment research.\n\nThe paper compares DNO against several existing alignment methods (LGD, SPIN, DDPO, AlignProp), providing a comprehensive evaluation of its performance relative to state-of-the-art approaches.\n\nVisual examples and optimization trajectories are provided to qualitatively assess the effectiveness of DNO and its variants. This complements the quantitative results and helps in understanding the behavior of the proposed methods.", "theoretical_claims": "I checked the proof for Theorem 2.1, which is the main theoretical claim in the paper. This theorem demonstrates that under the assumption of L-smoothness for the composite mapping r ◦ Mθ, the expected reward improves after each gradient step in the Direct Noise Optimization (DNO) process.\n\nThe proof follows these key steps:\n1. It leverages the Descent Lemma from optimization theory, which is a classical result for smooth functions.\n2. It applies this lemma to the specific context of the noise optimization problem.\n3. It substitutes the gradient step definition into the inequality from the Descent Lemma.\n4. It arrives at the final inequality showing that the expected reward improves with each gradient step.\n\nI didn't find any issues with this proof. The assumptions are clearly stated, and the logical flow from premises to conclusion appears sound. The application of the Descent Lemma is appropriate, and the algebraic manipulations seem correct.", "experimental_designs_or_analyses": "I checked the soundness and validity of several key experimental designs and analyses in the paper.", "relation_to_broader_scientific_literature": "The key contributions of this paper are closely related to several areas of prior research in diffusion models, reinforcement learning, and optimization. Previous approaches to aligning diffusion models with reward functions have primarily focused on fine-tuning the model parameters through reinforcement learning (RL) or direct fine-tuning. Notable examples include: DDPO. DNO represents a different approach to inference-time alignment by directly optimizing the noise vectors rather than modifying the model parameters or the sampling dynamics.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "While DNO is presented as a novel framework, some components build directly on existing ideas from noise optimization in diffusion models . The paper could benefit from a more explicit discussion of how it advances beyond these prior works. \n\n\nThe practical impact might be limited by the computational requirements of the optimization process, though the authors argue that the time costs are reasonable. For some applications, the additional optimization time might still be prohibitive.\n\nThe paper is difficult to understand and should be reorganized before being published.", "other_comments_or_suggestions": "Ensure consistent capitalization in figure captions (e.g., \"Figure 1. ODE vs. SDE for optimization\" should be \"Figure 1. ODE vs. SDE for Optimization\")\n\nEnsure all acronyms are defined upon first use (e.g., OOD, , ODE)", "questions_for_authors": "The experiments in the paper focus primarily on image generation tasks. Could the authors discuss whether DNO can be applied to other types of diffusion models (e.g., text, audio) and what modifications might be necessary? Are there fundamental limitations to DNO that would prevent its application in these domains?\n\n\nThe authors argue that DNO offers a favorable trade-off between inference time and reward, especially compared to fine-tuning methods. Could the authors provide more detailed comparisons of computational resources required for DNO versus fine-tuning methods, particularly for larger models like SDXL? How does the optimization time scale with model size?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper has received mixed reviews in the final recommendations, with three acceptances and one weak rejection. It presents a novel approach for optimizing diffusion input noise by leveraging a specified reward function. While the reviewers raised concerns regarding the insufficient clarification of the optimization process in SDE sampling and the lack of comparison with other noise optimization methods, the authors' rebuttal effectively addressed several of these issues. This ultimately led to a consensus among three reviewers in favor of acceptance. As a result, the Area Chair has decided to accept this paper.
Decision: Accept (poster)