Inference-Time Alignment of Diffusion Models with Direct Noise Optimization

## Abstract

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.

## 1 Introduction

In recent years, diffusion models have become the state-of-the-art model for generating high-quality images, with exceptional resolution, fidelity, and diversity (Ho et al., ; Dhariwal and Nichol, ; Song et al., ). These models are also remarkably easy to train and can be effectively extended to conditional generation (Ho and Salimans, ).
Specifically, diffusion models work by learning to reverse the process of diffusing the data distribution $p(x)$ into noise, which can be described by a stochastic differential equation (SDE) (Song et al., ; Karras et al., ):

$$ $\displaystyle d{x}_{t}=f(t){x}_{t}dt+g(t)dw_{t},$ (1) $$

where $dw_{t}$ is the standard Wiener process, and $f(t)$ and $g(t)$ are the drift and diffusion coefficients, respectively.
The reverse process relies on the score function $\epsilon({x}_{t},t)\stackrel{{\scriptstyle\text{def.}}}{{=}}\nabla_{x}\log p_{
t}({x})$ where $p_{t}$ denotes the p.d.f of noisy data $x_{t}$, and its closed-form can be expressed either as an ODE or as an SDE: (Song et al., ):

$$ $\displaystyle\text{ODE: }\ d{x}_{t}=\left(f(t){x}_{t}-\frac{1}{2}g^{2}(t) \epsilon({x}_{t},t)\right)dt,$ (2) $\displaystyle\text{SDE: }\ d{x}_{t}=\left(f(t){x}_{t}-g^{2}(t)\epsilon({x}_{t} ,t)\right)dt+g(t)d{w_{t}}.$ (3) $$

With the capability to evaluate $\epsilon({x}_{t},t)$, it becomes possible to generate samples from noise by numerically solving either the ODE (2) or the SDE (3). The training process, therefore, involves learning a parameterized surrogate $\epsilon_{\theta}({x}_{t},t)$ to approximate $\epsilon({x}_{t},t)$, following a denoising score matching framework as described in (Song et al., ; Karras et al., ).

Despite the effectiveness of diffusion models in modeling continuous distributions, when deploying these generative models for specific tasks, it is not suitable to sample 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 brightness, darkness, or compressibility. The topic of aligning generative models has gained significant attention following the success of aligning LLMs with Reinforcement Learning from Human Feedback (RLHF) techniques (Ouyang et al., ; Liu et al., ; OpenAI, ; Bai et al., ). Recently, the alignment problem has also drawn considerable interest in the context of diffusion models, as evidenced by a series of studies such as (Deng et al., ; Uehara et al., ; Yuan et al., ; Song et al., ; Dong et al., ; Sun et al., ; Zhang et al., ; Prabhudesai et al., ; Black et al., ; Fan et al., ).

Alignment Problem for Diffusion Models. Given a diffusion model characterized by parameters $\theta$ and its associated distribution $p_{\theta}(x)$, as well as a reward function $r(x)$ that can assign real-valued scores to generated samples, the central goal of the alignment problem is to adjust the distribution $p_{\theta}(x)$ such that the generated samples maximize the reward from $r(x)$. Such reward function $r(x)$ can be constructed by training a neural network on human feedback data (Kirstain et al., ) or by utilizing custom-designed metrics like brightness and darkness. In this work, we consider the reward functions to be continuous but possibly non-differentiable. In the following sections, we will provide a comprehensive review of some well-established methods for aligning diffusion models.

### 1.1 Methods for Aligning Diffusion Models

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

Aligning generative models with downstream objectives has attracted extensive research efforts, following a series of successful attempts that align LLMs with Reinforcement Learning from Human Feedback (RLHF). For diffusion models, developing alignment algorithms has also gained sufficient
attention recently, where a majority of
recent research has focused on leveraging Reinforcement Learning (RL)-based approaches. A common mathematical formulation in RL-based methods 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-regularized optimization problem:

$$ $\displaystyle\max_{p}\mathbb{E}_{x\sim p(x)}[r(x)]-\lambda\text{KL}(p||p_{ \theta}).$ (4) $$

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 online RL methods are DDPO (Black et al., ) and DPOK (Fan et al., ), which have been shown to improve downstream 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 a fixed preference dataset is utilized. Noteworthy works in this category include Diffusion-DPO (Wallace et al., ) and SPIN-Diffusion (Yuan et al., ).

#### 1.1.2 Direct Fine-tuning of Diffusion Models with Differentiable Rewards

Before delving into the formal description of this method, it is useful to revisit the sampling process of diffusion models, which also serves as the foundation for the rest of this work. Solving the ODE (2) or the SDE (3) typically involves discretizing the time steps into $T$ steps. By starting with the initial noise $x_{T}\sim\mathcal{N}(0,I)$, the solution process for the ODE (2) can be viewed as a mapping that transforms the initial noise $x_{T}$ into less noisy data through the following process:

$$ $\displaystyle x_{t-1}=\text{ODE_solver}(x_{t}),\text{ for }t=T,...,1.$ (5) $$

After $T$ steps, the output will be the generated sample $x_{0}$. Similarly, solving the SDE (3) can be seen as a mapping that gradually converts both the initial noise $x_{T}$ and the entire extra random noise vectors $z_{T},...,z_{1}$ into the generated sample $x_{0}$, e.g., through the following process:

$$ $\displaystyle x_{t-1}=\text{SDE_solver}(x_{t},z_{t}),\text{ for }t=T,...,1,$ (6) $$

where $z_{t}$ is also drawn from standard Gaussian distribution. In the example below, we provide a specific example of how the above process works in practice.

Example. In Algorithm 1 described below, we provide an example of the sampling algorithm for diffusion models, DDIM (Song et al., ), which is essentially the Euler method for solving ODEs/SDEs. The diffusion coefficients $\alpha_{1},\ldots,\alpha_{T}$ in Algorithm 1 are computed using the coefficient functions $f(t)$ and $g(t)$ in (1), with the detailed computation found in (Song et al., ; Karras et al., ). The coefficient $\eta$ in DDIM determines whether we are solving the ODE or the SDE, with $\eta=0$ corresponding to ODE and $\eta>0$ corresponding to SDE.

Figure: Algorithm 1 DDIM Sampling Algorithm

###### Remark 1 .

Throughout this work, for simplicity, we only adopt the DDIM sampling algorithm for our experiments, as it remains one of the most popular choices for diffusion sampling and, more importantly, supports both ODE-style and SDE-style sampling.

Diffusion Sampling as a Noise-to-Sample Mapping. From the diffusion sampling process described in (5) and (6), we observe that the sampling process can be conceptualized as an end-to-end mapping $M_{\theta}(z)$, which translates noise vectors $z$, sampled from the standard Gaussian distribution, into generated samples. Here, the noise vectors $z$ serve as a unified abstraction for both the $x_{T}$ in the ODE-based sampling process and the $(x_{T},\{z_{1},...,z_{T}\})$ in the SDE-based sampling process.
As we can see, the noise vector $z$ uniquely determines the generated sample from the diffusion models.

Two recent studies, AlignProp (Prabhudesai et al., ) and DRaFT (Clark et al., ), have proposed to directly fine-tune diffusion models using differentiable rewards. Specifically, their optimization objective is formulated as:

$$ $\max_{\theta}\mathbb{E}_{z\sim\mathcal{N}(0,I)}[r(M_{\theta}(z))].$ $$

Both the AlignProp (Prabhudesai et al., ) and DRaFT (Clark et al., ) methods utilize the ODE-type DDIM solver for the sampling process, specifically employing Algorithm 1 with $\eta=0$.

###### Remark 1 .

#### 1.1.3 Loss-Guided Diffusion

A recent work focusing on loss-guided diffusion models (LGD) (Song et al., ) also examines the concept of aligning diffusion models with differentiable rewards. Unlike the methods mentioned previously, LGD is an Inference-Time alignment method, meaning it does not necessitate modifications to the pretrained model $\theta$ and only works by modifying the inference process. In essence, the core idea of LGD is that, during the sampling process for the ODE (2), it considers a modified version of the ODE by introducing a new term that guides the generation toward areas of higher reward. Specifically, the new ODE is:

$$ $\displaystyle d{x}_{t}=\left(f(t){x}_{t}-\frac{1}{2}g^{2}(t)\epsilon({x}_{t},t )+\nabla_{x_{t}}r(x_{0}(x_{t}))\right)dt,$ (7) $$

where $x_{0}(x_{t})$ denotes the solution of this ODE starting from $x_{t}$. However, the gradient term $\nabla_{x_{t}}r(x_{0}(x_{t}))$ 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 to suboptimal performance, particularly with complex reward functions.

### 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-based Methods. All RL-based methods and the direct fine-tuning method are tuning-based, meaning they necessitate adjustments to the network models $\theta$. 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 extensive choices for reward functions. The second drawback is that the fine-tuning process typically relies on a limited set of input prompts, which challenges the model to generalize to new and unseen prompts. In contrast, methods such as LGD (Song et al., ) 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 input prompts. Further, inference-time methods require significantly fewer computing resources than tuning-based methods. However, the major drawback of inference-time methods is the substantial increase in the time required for the generation process.

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 differentiable. In practice, working with non-differentiable reward functions is important due to their prevalence. This non-differentiable property can arise from the simulation-based procedures used to compute the reward, or the reward function itself may be a black box.

### 1.3 Our Contributions

In this work, we focus on inference-time alignment of diffusion 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-time alignment method that can match the performance of tuning-based methods by incurring a reasonable additional time cost, and is capable of handling both differentiable and non-differentiable objective functions. To this end, we focus on an under-explored technique for achieving inference-time alignment of diffusion models—Direct Noise Optimization (DNO). Specifically, we make the following contributions:

- •
We conduct a self-contained and comprehensive study for DNO, and demonstrate that noise optimization can be theoretically interpreted as sampling from an improved distribution.
- •
We identify out-of-distribution reward-hacking as a critical issue in DNO. To address this issue, we introduce a novel probability-regularized noise optimization method designed to ensure the generated samples remain within the support of pretrained distribution.
- •
By developing a novel and highly efficient hybrid gradient approximation strategy, we extend the DNO approach to handle non-differentiable reward functions effectively.
- •
Through the experiments on several important image reward functions, we demonstrate that our proposed method can achieve state-of-the-art scores in comparison to existing alignment methods, without any fine-tuning on the parameters of diffusion models.

## 2 Direct Noise Optimization for Aligning Diffusion Models

Given the noise-to-sample mapping $M_{\theta}$ described in Section 1.1.2, Direct Noise Optimization (DNO) can be mathematically formulated as follows:

$$ $\displaystyle\max_{z}\;r(M_{\theta}(z)),$ (8) $$

with $z\sim{\mathcal{N}}(0,I)$ as the initial solution. As we will discuss in Section 3, the Gaussian distribution ${\mathcal{N}}(0,I)$ serves as an important prior on the optimization variables $z$.
By solving this optimization problem, we can obtain the optimized noise vectors, which are then used to generate high-reward samples.

When the reward function $r(\cdot)$ is differentiable, gradient-based optimization methods can be applied to solve the problem efficiently. That is, the following step can be performed iteratively until convergence:

$$ $\displaystyle z_{\text{new}}=\text{optimizer_step}(z_{\text{old}},\nabla_{z}r \left(M_{\theta}(z_{\text{old}})\right)),$ (9) $$

where the optimizer can be either vanilla gradient ascent or adaptive optimization algorithms like Adam (Kingma and Ba, ). When the gradient of reward function $r(\cdot)$ is not available, we can leverage techniques from zeroth order optimization (Nesterov and Spokoiny, ; Tang et al., ) to estimate the ground-truth gradient $\nabla_{z}r\left(M_{\theta}(z_{\text{old}})\right)$, denoted as $\hat{g}(z_{\text{old}})$, and then apply similarly

$$ $\displaystyle z_{\text{new}}=\text{optimizer_step}(z_{\text{old}},\hat{g}(z_{ \text{old}})).$ (10) $$

In Section 4, we provide a dedicated discussion on how to obtain a better estimator for $\hat{g}(z_{\text{old}})$ when the reward function is non-differentiable.

We provide a visual illustration in Figure 2 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., ), operates at inference-time and does not require tuning the network parameter
$\theta$. 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 in Section 5, the extra time needed for the DNO approach is a worthwhile trade-off for obtaining high-reward samples.

Figure: Figure 2: Overview of the DNO procedure with the DDIM sampling algorithm: DNO seeks to optimize only those Gaussian noise vectors $\{x_{T},z_{1},z_{2}...,z_{T}\}$ to maximize the reward value of a single generated sample $x_{0}$. To facilitate the gradient backpropagation from $x_{0}$ to $\{x_{T},z_{1},z_{2}...,z_{T}\}$, we leverage the technique of gradient checkpointing. It is worth noting that when using $\eta=0$ for DDIM sampling, there is no need to compute the gradient for $z_{1},...,z_{T}$, as the generated sample $x_{0}$ depends exclusively on $x_{T}$. When computing the gradient from $r(x_{0})$ to $x_{0}$, we can use either ground-truth gradient $\nabla r$ or an estimated gradient $\widehat{\nabla r}$, depending on whether the reward function $r(\cdot)$ is differentiable.
Refer to caption: extracted/5894777/figures/main_DNO.png

Several recent studies have explored similar formulations to (8) across different applications. (Wallace et al., ) investigates the optimization of latent vectors obtained through DDIM-inversion (Song et al., ), aiming to enhance the CLIP score (Radford et al., ) and Aesthetic Score (Schuhmann et al., ) of given images. (Ben-Hamu et al., ) discusses optimizing the initial noise $x_{T}$ for the ODE process to address inverse problems, leveraging the diffusion model as a prior. (Novack et al., ) and (Karunratanakul et al., ) consider the optimization of initial noise $x_{T}$ for the ODE with the objective of improving specific downstream objectives in robotics and audio.

While similar methods of DNO has appeared in previous works, many of its technical details remain insufficiently explored. On one hand, there is a lack of a comprehensive framework concerning the design choices, theoretical understanding, 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 inference-time method for aligning diffusion models, can achieve competitive performance compared to tuning-based methods. In this work, we aim to conduct a 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.1 Understanding Direct Noise Optimization

Example 1. First, we present a simple example to visualize the process of noise optimization. Specifically, we trained a toy diffusion model for generating uniform distribution on a ring with a radius between $0.8$ and $1.2$, and the initial distribution is visualized in Figure 3(a), where each red point denotes a single sample drawn from the trained diffusion model. We then solve the noise optimization problem (8) with the reward function $r(x)=\sin(4\pi x[1])+\sin(4\pi x[2])-\left((x[1]-1)^{2}-x[2]^{2}\right)/5$, a highly nonconvex function with many local maxima. To perform this optimization, we solve the DNO problem (8) using gradient ascent with the learning rate set to 0.01. In Figures 3(b), 3(c), and 3(d), we visualize the optimized samples after 10, 50, and 100 gradient steps, respectively. We can observe that the distribution of the samples shifts toward a distribution on the local maxima of the reward function.

Figure: (a) Initial
Refer to caption: extracted/5894777/figures/sec2_0.png

Noise Optimization as Sampling from an Improved Distribution. As is evident from Figure 3, optimizing the noise vectors for diffusion models indeed leads to an improved distribution in terms of the reward scores. To rigorously describe this evolving process, we define an operator function $g$ to represent a single gradient step, i.e., $g(z)\stackrel{{\scriptstyle\text{def}}}{{=}}z+\ell\cdot\nabla_{z}r\circ M_{
\theta}(z)$, where $\circ$ denotes the function composition operator and $\ell$ denotes the step size for gradient ascent. Additionally, we define the operator $g_{t}$, which denotes applying the gradient ascent step for $t$ steps, i.e., $g_{t}(z)=g(g_{t-1}(z))$ with $g_{0}$ being the identity mapping. With these notations, we can now express the distribution after $t$ gradient steps as $p_{t}(x)$, which is the distribution of $M_{\theta}(g_{t}(z))$ with $z\sim\mathcal{N}(0,I)$. In the following theorem, we demonstrate that there is a rigorous improvement after every single gradient step, i.e., the distribution $p_{t+1}(x)$ is provably better than $p_{t}(x)$ in terms of expected reward.

###### Theorem 1 .

Assuming that $r\circ M_{\theta}$ is $L$-smooth, namely, $\|\nabla r\circ M_{\theta}(z)-\nabla r\circ M_{\theta}(z^{\prime})\|\leq L\|z-
z^{\prime}\|$ for any $z\neq z^{\prime}$, it holds true that

$$ $\displaystyle\mathbb{E}_{x\sim p_{t+1}(x)}r(x)$ $\displaystyle\geq\mathbb{E}_{x\sim p_{t}(x)}r(x)+\left(\ell-\frac{\ell^{2}L}{2 }\right)\mathbb{E}_{z_{0}\sim N(0,I)}\left\|\nabla_{z}r\circ M_{\theta}(z)_{|z =g_{t}(z_{0})}\right\|^{2}_{2}.$ (11) $$

As described in (11), provided that the step size $\ell$ is less than $\frac{2}{L}$, the distribution $p_{t+1}(x)$ is strictly better than the previous distribution $p_{t}(x)$ in terms of expected reward, as long as the second term is non-zero.

Based on this result, it is natural to ask: When does the distribution stop improving? Namely, when does the second term in (11) become zero. We provide a detailed discussion to answer this question and also the proof for Theorem 1 in Appendix B.

Smoothness Assumption. In Theorem 1, we rely on the smoothness assumption of the composite mapping $r\circ M_{\theta}$. We note that this is a reasonable assumption in practice, as the noise-to-sample mapping in diffusion models has been observed to be smooth. For instance, see Figure 4 in (Tang et al., ).

###### Theorem 1 .

### 2.2 Optimizing ODE vs. Optimizing SDE

As previously introduced, there are two primary methods for sampling from pretrained diffusion models: one based on solving the ODE (2) and the other on solving the SDE (3). A critical difference lies in the fact that ODE sampling depends exclusively on the initial noise $x_{T}$, whereas SDE sampling is additionally influenced by the noise ${z_{t}}$ added at every step of the generation process. It has been noted that existing works on noise optimization (Ben-Hamu et al., ; Novack et al., ; Karunratanakul et al., ) have mainly concentrated on optimizing the initial noise $x_{T}$ for the ODE sampler.

Figure: Figure 4: ODE vs. SDE for optimization. The x-axis refers to the number of gradient ascent steps during optimization.
Refer to caption: extracted/5894777/figures/sec21.png

Different from preceding studies, we explore the utilization of the SDE sampler for noise optimization. Specifically, we employ the DDIM with $\eta=1$ (Song et al., ) as the SDE sampler and propose to optimize both the added noise $z_{t}$ at every timestep and the initial noise $x_{T}$. In this context, the dimensionality of the optimized noise significantly surpasses that of ODE sampler, typically $T\cdot d$ v.s. $d$, where $d$ is the dimension for the learned distribution and $T$ is 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 4, where we juxtapose the optimization speeds between the ODE (DDIM with $\eta=0$) and SDE samplers using both the Example 1 depicted in Figure 3 (Left) and optimizing for stable diffusion in alignment with the aesthetic reward (Right), a major experiment detailed in the subsequent section 5.2.

Understanding the Faster Speed of DNO with SDE-Based Sampling. Intuitively, the acceleration in optimization speed can be attributed to the finer-grained 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 $\eta=1$, as defining the noise-to-sample mapping $M_{\theta}(x_{T},z_{1},\ldots,z_{T})$. An important observation is that the ODE sampling algorithm, DDIM with $\eta=0$, can also be expressed using the same $M_{\theta}(x_{T},z_{1},\ldots,z_{T})$, with the distinction that the noise $z_{1},\ldots,z_{T}$ becomes deterministic and dependent on $x_{T}$, rather than sampled from a Gaussian distribution. Specifically, if we define the deterministic noise vectors as:

$$ $\displaystyle z_{t}^{\text{ODE}}\stackrel{{\scriptstyle\text{def.}}}{{=}}\frac {\sqrt{1-\alpha_{t-1}}-\sqrt{1-\alpha_{t-1}-\sigma_{t}^{2}}}{\sigma_{t}} \epsilon_{\theta}(x_{t},t),\quad\text{for }t=1,\ldots,T,$ (12) $$

where $\sigma_{t}=\sqrt{(1-\alpha_{t-1}/(1-\alpha_{t})}\sqrt{1-\alpha_{t}/\alpha_{t-1}}$, then the sampling process of DDIM with $\eta=0$ can be expressed as $M_{\theta}(x_{T},z_{1}^{\text{ODE}},\ldots,z_{T}^{\text{ODE}})$. In this context, the advantage of SDE-based sampling becomes evident:

$$ $\displaystyle\underbrace{\max_{x_{T},z_{1},\ldots,z_{T}}r\left(M_{\theta}(x_{T },z_{1},\ldots,z_{T})\right)}_{\text{DNO with SDE}}\geq\underbrace{\max_{x_{T} }r\left(M_{\theta}(x_{T},z_{1}^{\text{ODE}},\ldots,z_{T}^{\text{ODE}})\right)} _{\text{DNO with ODE}},$ (13) $$

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 $\eta=1$) for the remainder of the study. Additionally, we fix the number of generation steps $T$ to 50 throughout this work for simplicity.

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

It has been observed that when aligning generative models (e.g., including autoregressive language models or diffusion models) with reward functions, one can experience the so-called reward-hacking, i.e., the optimized samples yield high rewards but do not possess the desirable properties (Miao et al., ; Chen et al., ). Generally, there are two different types of reward-hacking. In the first type, the reward function admits some shortcuts, so the optimized samples score high rewards but remain barely distinguishable 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-Distribution (OOD) Reward-Hacking. In this work, we will focus on this second type and reveal that OOD reward-hacking is a common issue in DNO.

We begin by presenting two examples using both a simple diffusion model from the Example 1 described above and the open-source image diffusion model SD v1.5 (Rombach et al., ). For the first example, we revisit the pretrained distribution displayed in Figure 3, modifying the reward function to $r(x)=-(x[1]-1.4)^{2}-(x[2]-1.4)^{2}$. Figure 5(a) exhibits the pretrained distribution of the diffusion models, where we note that every sample stays within the support. However, as illustrated in Figure 5(b), after 1000 gradient steps of optimization for the reward function, all the generated samples become out-of-distribution. In our second example, we examine optimization for the "brightness" reward—specifically, the average value of all pixels in an image—using SD v1.5 as the image diffusion model. We start with the prompt "black duck", with the initial image depicted in Figure 5(c). 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 5(d). 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.

Figure: (a) Initial
Refer to caption: extracted/5894777/figures/sec3_0.png

The Reasons Behind OOD Samples. In this work, one of our key contributions is to identify one critical cause of OOD reward-hacking in noise optimization. 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 originate from these low-probability areas—such as vectors comprised entirely of zeros—the neural network within the diffusion models may incur significant approximation errors for these particular inputs. This error, in turn, leads to the generation of out-of-distribution samples. In the subsequent 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.

### 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 distribution. However, in practice, it is nearly impossible to obtain samples near the all-zero vector from a Gaussian distribution, as it resides within a low-probability region. In high-dimensional statistics, concentration inequalities are usually employed to describe these distinctive properties and delineate the low-probability regions of high-dimensional distributions. In the following lemma, we present two classical inequalities for the standard Gaussian distribution.

###### Lemma 1 ( (Wainwright, 2019 ) ) .

Consider that $z_{1},...,z_{m}$ follow a $k$-dimensional standard Gaussian distribution. We have the following concentration inequalities for the mean and covariance:

$$ $\displaystyle{\rm Pr}\left[\left\|\frac{1}{m}\sum_{i=1}^{m}z_{i}\right\|>M \right]<p_{1}(M)\stackrel{{\scriptstyle\rm{def.}}}{{=}}\max\left\{2e^{-\frac{ mM^{2}}{2k}},1\right\},$ (14) $\displaystyle{\rm Pr}\left[\left\|\frac{1}{m}\sum_{i=1}^{m}z_{i}z_{i}^{\top}-I _{k}\right\|>M\right]<p_{2}(M)\stackrel{{\scriptstyle\rm{def.}}}{{=}}\max\left \{2e^{-\frac{m\left(\max\left\{\sqrt{1+M}-1-\sqrt{{k}/{m}},0\right\}\right)^{2 }}{2}},1\right\}.$ (15) $$

In practice, to determine if an $n$-dimensional vector $z$ lies within a low-probability region, we can factorize $n$ as $n=m\cdot k$, and divide $z$ into $m$ subvectors: $z=[z_{1}^{1},...,z_{m}^{k}]$, where $n=m\cdot k$ and $z_{i}=[z_{i}^{1},...,z_{i}^{k}]\sim\mathcal{N}(0,I_{k}))$. Then, we compute $M_{1}(z)=\left\|\frac{1}{m}\sum_{i=1}^{m}z_{i}\right\|$ and $M_{2}(z)=\left\|\frac{1}{m}\sum_{i=1}^{m}z_{i}z_{i}^{\top}-I_{k}\right\|$. Finally, we can determine that $z$ lies in a low-probability region if both $p_{1}(M_{1}(z))$ and $p_{2}(M_{2}(z))$ are low.

###### Remark 2 .

As one can see, the computed probability is related to the factors $m$ and $k$ of $n$. According to (Wainwright, ), the two inequalities (14) and (15) are tight when $m/k$ is large. On the other hand, $k=1$ 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 $k=2$ serves as a good default choice. In Appendix E, we provide a more detailed analysis for choosing an appropriate $k$.

An important point to note is that the standard Gaussian distribution is invariant to permutation, i.e., for any permutation matrix $\Pi$, if $z$ follows a standard Gaussian distribution, the permuted vector $\Pi z$ will have the same probability behavior. With this insight, to increase the robustness of the probability measure $p_{1}$ and $p_{2}$, a natural idea is to examine the probability of many permuted vectors. Specifically, given $q$ permutation matrices $\Pi_{1},...,\Pi_{q}$, we define the following indicator metric,

$$ $\displaystyle P(z)\stackrel{{\scriptstyle\text{def.}}}{{=}}\min\left\{\min_{i \in\{1,...,q\}}p_{1}\left(M_{1}(\Pi_{i}z)\right),\min_{i\in\{1,...,q\}}p_{2} \left(M_{2}(\Pi_{i}z)\right)\right\}.$ (16) $$

Interpretation of $P(z)$. If the probability $P(z)$ is low, it implies that there exists a permutation matrix $\Pi_{i}$ such that the noise vector $\Pi_{i}z$ is in the low-probability region of the standard Gaussian distribution. Therefore, due to the permutation-invariant property, the noise vector $z$ is also less likely to be sampled from the standard Gaussian distribution. In practice, we utilize randomly generated permutation matrices and have found that setting
$q=100$ results in empirically good performance. In Appendix C, we provide some visualized empirical evidence to show that $P(z)$ serves as a good indicator for determining if the generated samples are OOD.

###### Lemma 1 ( (Wainwright, 2019 ) ) .

###### Remark 2 .

### 3.2 Probability-Regularized Noise Optimization

With the insights discussed above, a natural idea for preventing OOD reward hacking is to regularize noise vectors to remain within the high-probability region of the Gaussian distribution. To achieve this, we propose the following Probability-Regularized Noise Optimization problem:

$$ $\displaystyle\max_{z}r(M_{\theta}(z))+\gamma\mathbb{E}_{\Pi}\left[\log p_{1}(M _{1}(\Pi z))+\log p_{2}(M_{2}(\Pi z))\right],$ (17) $$

where $\gamma$ is the coefficient used to control the regularization effect.
In particular, for the regularization term, we use the expectation of the log probabilities over the permutation matrices, rather than the minimum probability $P(z)$. This is because the expectation is smoother for optimization purposes.

## 4 Tackling Non-Differentiable Reward Functions

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 representative cases. Firstly, the reward may be computed through simulation-based procedures, such as the JPEG-compressibility employed in DDPO (Black et al., ), which calculates the size of an image in bits after running a JPEG compression algorithm. Additionally, the reward function itself may be a black box provided through online API providers, as in the setting considered in (Sun et al., ). This scenario is common when the reward function is a large neural network model, like those in (Wang et al., ; Lin et al., ), making it impossible to directly obtain the gradient for optimization.

To address the scenarios mentioned above, in this section, 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 1. Concerning optimization with only function value, a major family of optimization approaches is zeroth-order optimization algorithms, including ZO-SGD (Nesterov and Spokoiny, ). This method treats the entire mapping $r\circ M_{\theta}(\cdot)$ as a black-box function and seeks to estimate the gradient of $r\circ M_{\theta}(\cdot)$ via function value queries.

Method 2. It is worth noting that for the mapping $r\circ M_{\theta}(\cdot)$, only the gradient of the reward function $r(x)$ is not available, and we are still able to compute the gradient of $M_{\theta}(z)$, since the noise-to-sample mapping itself is intrinsically differentiable. Therefore, a straightforward idea is to adopt a hybrid gradient approach—only to estimate the gradient of $r$, while using the ground truth gradient for $M_{\theta}(z)$. Specifically, we denote that the initial noise is $z$, and the generated sample is $x=M_{\theta}(z)$. Firstly, we can estimate the gradient of $\nabla r(x)$ in a similar fashion with the ZO-SGD (Nesterov and Spokoiny, ):

$$ $\displaystyle H_{1}(x)=\mathbb{E}_{\xi\sim\mathcal{N}(0,I)}\left[\left(r(x+\mu \xi)-r(x)\right)\xi\right]\approx C_{1}\nabla r(x),$ (18) $$

where $\mu$ is the coefficient for perturbation, and $C_{1}$ is some constant. With the estimated gradient $H_{1}(x)$ for the reward function $r(x)$, we can use the following estimated gradient $G_{1}(z)$ for optimization:

$$ $\displaystyle G_{1}(z)\stackrel{{\scriptstyle\text{def.}}}{{=}}H_{1}(x)\cdot \nabla_{z}M_{\theta}(z)\approx C_{1}\nabla_{z}r\circ M_{\theta}(z),$ (19) $$

where the main idea is to replace the ground truth $\nabla r(x)$ in the chain-rule of differentiating $r(M_{\theta}(z))$. We refer this method as Hybrid-1 in the following sections.

###### Remark 3 .

As one can observe, the computation of (19) involves the Jacobian $\nabla_{z}M_{\theta}(z)$. However, it is important to note that when we only require the vector-Jacobian product $H_{1}(x)\cdot\nabla_{z}M_{\theta}(z)$, it is unnecessary to compute the full Jacobian $\nabla_{z}M_{\theta}(z)$. In Appendix D, we describe an elegant and efficient way to implement (19) using an auto-differentiation technique.

Method 3.
We found that there is a crucial drawback in the gradient estimator (18), that one needs to query the reward function $r(\cdot)$ with noisy input $x+\mu\xi$. When the reward function is only defined on some manifold $\mathcal{M}$, e.g., defined on the image manifold, rather than the whole space ${\mathbb{R}}^{n}$, this can lead to severe problems, because, for some $x\in\mathcal{M}$, the noisy sample $x+\mu\xi$ may no longer stay within $\mathcal{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 $\nabla r(x)$ is

$$ $\displaystyle H_{2}(x)=\mathbb{E}_{\xi\sim\mathcal{N}(0,I)}\left[\left(r(M_{ \theta}(z+\mu\xi))-r(x)\right)\left(M_{\theta}(z+\mu\xi)-x\right)\right].$ (20) $$

Following a similar proof in (Nesterov and Spokoiny, ), we can also show that $H_{2}(x)\approx C_{2}\nabla r(x)$ for some constants $C_{2}$. As we can see, when computing the gradient (20), we ensure that we query the reward function $r(\cdot)$ with only samples that are within the manifold of the pretrained distribution. Similar to Hybrid-1, we can adopt a gradient estimator $G_{2}(z)$ with $H_{2}(x)$:

$$ $\displaystyle G_{2}(z)\stackrel{{\scriptstyle\text{def.}}}{{=}}H_{2}(x)\cdot \nabla_{z}M_{\theta}(z).$ (21) $$

We refer this method (21) 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.

###### Remark 3 .

## 5 Experiments

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., ) 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 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 D and E. Additionally, we provide examples to visualize the optimization process in Appendix A. For all the following experiments, unless explicitly stated otherwise, a single run of DNO is performed on a single A800 GPU.

### 5.1 Experiments on Image Brightness and Darkness Reward Functions

Experiment Design. In this section, we design an experiment 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—with the prompt "black <animal>", where the token <animal> is randomly selected from a list of animals. The second setting involves optimizing the darkness reward, defined as 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 easier to trigger the OOD reward-hacking phenomenon. Moreover, it is straightforward to verify whether the generated samples are out-of-distribution 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 probability regularization in preventing the OOD reward-hacking phenomenon.

Importance of the Brightness and Darkness Reward Functions. While the primary purpose of using these two reward functions is to better examine the effectiveness of probability regularization, it is also important to 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, ).

Figure: Figure 6: Comparison of DNO with and without probability regularization. Upper row: Optimizing for the brightness reward, i.e., the average value of all pixels in the images. Lower row: Optimizing for the darkness reward, i.e., the negative of the brightness reward. The x-axis refers to the number of gradient ascent steps during optimization.
Refer to caption: extracted/5894777/figures/sec51-new.png

Measuring the Degree of OOD. As observed in this experiment, the reward function contradicts the input prompt, leading to an inconsistency between the generated samples and the prompt. Therefore, we utilize the CLIP Score (CS) (Radford et al., ), a commonly used metric for measuring the semantic similarity between images and text descriptions, to gauge the degree of OOD for the generated samples. A higher CS indicates that the sample is less likely to be out-of-distribution (OOD). In addition to the CS, we also use an MLLM-based score to measure the degree of OOD, specifically employing the Image-Text Matching (ITM) score from (Wang et al., ) as the metric.

Figure: (a) w/o reg
Refer to caption: extracted/5894777/figures/37_511.png

Results. In Figure 6, we first observe that adding the regularization 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 a good indicator of the OOD phenomenon, as it is positively associated with CS and ITM. We offer qualitative examples in Figure 7 to highlight the differences between optimized samples with and without regularization. Figure 7(a) and 7(b) are optimized images with brightness reward and prompt "black rabbit". Figure 7(c) and 7(d) are optimized images with darkness reward and prompt "white chicken". In the two pairs of images, we select the optimized images that achieved almost the same reward value from the optimizations with and without regularization, respectively. It is evident that at an equivalent level of reward, the non-regularized samples are more prone to be out-of-distribution, as the color of the generated animals is less likely to match the prompt.

Generating Images with Purely Bright/Dark Backgrounds. From the images demonstrated in Figure 7, 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, ), such an effect can only be achieved by fine-tuning the diffusion models using a technique called Offset-Noise.

### 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., ), HPS-v2 score (Wu et al., ), and PickScore (Kirstain et al., ), 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 a lesser chance for the optimized 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 reward functions as test metrics when optimizing one of them.

Results. Firstly, we observe that the effect of probability regularization is less pronounced than that in Section 5.1. This observation is also reflected by our proposed indicator $P(z)$; if no regularization is applied, the value of $P(z)$ in Figure 8 decreases much slower than that in Figure 6. 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 metrics. 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.

Figure: Figure 8: Comparison of running DNO with three human-like reward functions, with and without regularization. When optimizing one reward function, the other two are used as test metrics. 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 is used. For example, A + w/ reg means optimizing aesthetic score with regularization.
Refer to caption: extracted/5894777/figures/sec52.png

Comparison to Existing Alignment Methods.
We summarize the performance of DNO with probability regularization from Figure 8 into Table 1 and compare it to the major existing alignment methods discussed in the introduction. As shown, the performance of DNO matches that of state-of-the-art tuning-based alignment methods without any fine-tuning on the network models, all within a reasonable time budget for generation. More importantly, we demonstrate 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., ), 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 fundamental 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 a large margin, demonstrating that DNO presents a highly advantageous trade-off between inference time and reward.

Using Fewer Steps to Reduce Inference Time. In Figure 8, we show that our proposed DNO can achieve high reward 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 diffusion 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 a fair comparison. In practice, 50 diffusion steps is a fairly long choice, and using 15–25 steps can still result in sufficiently good samples. Therefore, in Table 2, we present the performance of DNO with different numbers of denoising steps, fixing the time budget to 1 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.

**Table 1: Performance comparison. For SD v1.5 and DNO, we 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-tuning. All time costs in the table are measured with respect to the GPU time on a single A800 GPU. Baselines: LGD (Song et al., ), SPIN (Yuan et al., ), DDPO (Black et al., ), AlignProp (Prabhudesai et al., ).**
| Method | SD v1.5 | LGD | SPIN | DDPO | AlignProp | DNO (This work) | BoN |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 2s | $n$=10 | $n$=100 | $\sim$ 20h | $\sim$ 56h | $\sim$ 12h | 1 min | 3 min | 5 min | 10 min |  |
| Aesthetic $\uparrow$ | 5.367 | 5.340 | 5.224 | 6.248 | 7.180 | 8.940 | 5.754 | 7.202 | 8.587 | 6.531 |
| HPS $\uparrow$ | 0.278 | 0.276 | 0.271 | 0.276 | 0.287 | 0.330 | 0.285 | 0.303 | 0.324 | 0.298 |
| PickScore $\uparrow$ | 21.11 | 21.01 | 21.09 | 22.00 | / | / | 21.25 | 23.17 | 25.13 | 22.09 |

**Table 2: Running DNO with different numbers of diffusion steps $T$, while fixing the time budget to 1 minute.**
| $T$ | 10 | 15 | 20 | 25 | 50 |
| --- | --- | --- | --- | --- | --- |
| Aesthetic $\uparrow$ | 6.992 | 7.496 | 6.773 | 6.381 | 5.754 |
| HPS $\uparrow$ | 0.342 | 0.341 | 0.306 | 0.293 | 0.285 |
| PickScore $\uparrow$ | 23.98 | 24.82 | 23.69 | 23.02 | 21.25 |

Finally, we would like to highlight that this alignment process 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 A100 (Black et al., ; Prabhudesai et al., ).

### 5.3 Optimization for Non-Differentiable Reward Functions

Setting.
In this section, we aim to explore the three extensions proposed in Section 4 for handling non-differentiable reward functions. For Method 1, we use ZO-SGD (Nesterov and Spokoiny, ), and compare it with our proposed Hybrid-1 and Hybrid-2. We consider using two reward functions: 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., ) 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 a scenario where the neural network model of the Aesthetic Score can only be queried via an 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 is 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.

Figure: Figure 9: Comparing three methods on two reward functions. 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.
Refer to caption: extracted/5894777/figures/sec53.png

Results.
The main results are visualized in Figure 9. Initially, we observe that the Hybrid-2 method is significantly faster than the other two methods, and the final score is also higher for both reward functions used in this experiment. Furthermore, in both scenarios, the ZO-SGD method exhibits the slowest optimization speed. Another interesting finding is the poor performance of Hybrid-1 in optimizing the Aesthetic Reward. This mainly occurs because the aesthetic reward function can only work with image inputs, which validates our initiative to propose the Hybrid-2 method in Section 4. Finally, when comparing the optimization of the aesthetic score in Figures 8 and 9, we can also note that using the true gradient results in substantially fewer steps than using the estimated gradient.

## 6 Conclusions

In this work, we present a comprehensive study on Direct Noise Optimization (DNO) for aligning diffusion generative models at inference-time. We introduce variants of DNO designed to efficiently address challenges such as out-of-distribution reward-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. Nonetheless, we argue that the additional time cost, being within a reasonable and acceptable range, is a worthwhile trade-off for attaining high-reward generated samples in a wide range of real-world applications. We anticipate DNO gaining greater attention in future research and applications due to its flexibility to accommodate any reward function—or even a combination of different reward functions—while demanding only modest computing resources, which positions it as an accessible tool for many practitioners.

## Appendix

### A Visualization of Optimization with Images

To assist the reader in understanding the optimization process, we offer visualizations of this process through images. The aim is to provide a qualitative evaluation of the proposed method.

Figure: (a) Optimizing brightness reward with and without probability regularization. Prompt: "black duck".
Refer to caption: extracted/5894777/figures/black_Exp.jpg

Figure: (a) Optimizing Aesthetic Score with prompt "gray lion"
Refer to caption: extracted/5894777/figures/aes_graylion.jpg

Figure: (a) Optimizing Jpeg Compressiblity with Hybrid-2 gradient approximation. Upper: prompt "blue pig". Lower: prompt "yellow rabbit".
Refer to caption: extracted/5894777/figures/non_diff_jpeg.jpg

#### A.1 Comparing With and Without Regularization

In this section, our objective is to demonstrate the impact of the probability regularization term on the optimization process through visual examples. Specifically, we present examples from three settings. The first two examples are derived from the experiments in Section 5.1, while the last example is from the experiment on optimizing the aesthetic score in Section 5.2.

The examples can be seen in Figure 10. As observed across all examples, the optimization process that integrates the regularization term consistently prevents the generated samples from falling into the category of being Out-Of-Distribution (OOD).

#### A.2 More Examples

In this section, we aim to provide additional visualized examples for our proposed method.

Firstly, in Figure 11, we 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-differentiable optimization. These examples in Figure 12 effectively showcase the efficiency of Hybrid-2 in both estimating the gradient and optimizing.

### B Theoretical Results

###### Proof for Theorem 1 .

To leverage the $L$-smoothness assumption, we need to state a classical lemma for smooth optimization.

###### Lemma 2 (Descent Lemma [Bertsekas, 1997 ] ) .

For any $z_{1}$ and $z_{2}$, we have

$$ $\displaystyle r\circ M_{\theta}(z_{2})\geq r\circ M_{\theta}(z_{1})+\nabla r \circ M_{\theta}(z_{1})\cdot(z_{2}-z_{1})-\frac{L}{2}\|z_{2}-z_{1}\|^{2}_{2}.$ (22) $$

Now for any $z$ and steps $t\geq 1$, with the descent lemma, we have

$$ $\displaystyle r\circ M_{\theta}(g_{t+1}(z))$ $\displaystyle\geq r\circ M_{\theta}(g_{t}(z))+\nabla r\circ M_{\theta}(g_{t}(z ))\cdot\left(g_{t+1}(z)-g_{t}(z)\right)-\frac{L}{2}\|g_{t+1}(z)-g_{t}(z)\|^{2} _{2}.$ $$

Notice that by the definition of $g_{t}{\cdot}$ we have

$$ $\displaystyle g_{t+1}(z)-g_{t}(z)$ $\displaystyle=g_{t}(z)+\ell\nabla r\circ M_{\theta}(g_{t}(z))-g_{t}{(z)}$ $\displaystyle=\ell\nabla r\circ M_{\theta}(g_{t}(z)).$ $$

Therefore, we have

$$ $\displaystyle r\circ M_{\theta}(g_{t+1}(z))$ $\displaystyle\geq r\circ M_{\theta}(g_{t}(z))+\left(\ell-\frac{\ell^{2}L}{2} \right)\|\nabla r\circ M_{\theta}(g_{t}(z))\|_{2}^{2}.$ $$

Taking the expectation over $z\sim{\mathcal{N}}(0,I)$ we have

$$ $\displaystyle\mathbb{E}_{z\sim{\mathcal{N}}(0,I)}\left[r\circ M_{\theta}(g_{t+ 1}(z))\right]\geq\mathbb{E}_{z\sim{\mathcal{N}}(0,I)}$ $\displaystyle\left[r\circ M_{\theta}(g_{t}(z))\right]+$ $\displaystyle\left(\ell-\frac{\ell^{2}L}{2}\right)\mathbb{E}_{z\sim{\mathcal{N }}(0,I)}\left[\|\nabla r\circ M_{\theta}(g_{t}(z))\|_{2}^{2}\right].$ $$

By using the change of variable formula for distribution, we can easily see that

$$ $\displaystyle\mathbb{E}_{x\sim p_{t+1}(x)}r(x)=\mathbb{E}_{z\sim{\mathcal{N}}( 0,I)}\left[r\circ M_{\theta}(g_{t+1}(z))\right],$ $$

and

$$ $\displaystyle\mathbb{E}_{x\sim p_{t}(x)}r(x)=\mathbb{E}_{z\sim{\mathcal{N}}(0, I)}\left[r\circ M_{\theta}(g_{t}(z))\right],$ $$

Therefore, we conclude with

$$ $\displaystyle\mathbb{E}_{x\sim p_{t+1}(x)}r(x)$ $\displaystyle\geq\mathbb{E}_{x\sim p_{t}(x)}r(x)+\left(\ell-\frac{\ell^{2}L}{2 }\right)\mathbb{E}_{z_{0}\sim N(0,I)}\left\|\nabla_{z}r\circ M_{\theta}(z)_{|z =g_{t}(z_{0})}\right\|^{2}_{2}$ (23) $\displaystyle\geq\mathbb{E}_{x\sim p_{t}(x)}r(x).$ $$

∎

When Does the Distribution Stop Improving?
As observed, the distribution ceases to improve when the second term in Equation (23) becomes zero. Initially, we note that the optimized distribution stops improving when $\mathbb{E}_{z_{0}\sim N(0,I)}\left\|\nabla_{z}r\circ M_{\theta}(z)_{|z=g_{t}(z
_{0})}\right\|^{2}_{2}=0$. In statistical terms, this implies that $\nabla_{z}r\circ M_{\theta}(z)_{|z=g_{t}(z_{0})}$ is a zero vector with probability one.

To discern the circumstances under which this zero vector occurs, let us assume that $z_{0}$ is some fixed noise vector and consider the scenario where

$$ $\displaystyle\nabla_{x}r(x)_{|x=M_{\theta}(g_{t}(z_{0}))}\cdot\nabla_{z}M_{ \theta}(z)_{|z=g_{t}(z_{0})}=\vec{0},$ (24) $$

Here, we denote $G_{1}(z_{0})=\nabla_{x}r(x)_{|x=M_{\theta}(g_{t}(z_{0}))}$ as the gradient of the reward functions on the generated sample, and $G_{2}(z_{0})=\nabla_{z}M_{\theta}(z)_{|z=g_{t}(z_{0})}$ representing the Jacobian matrix of the noise-to-sample mapping. We categorize the situation in Equation (24) into three cases:

Type-I: $\|G_{1}(z_{0})\|=0$ and $\|G_{2}(z_{0})\|>0$. Here, the gradient of the reward function on the generated sample is zero, indicating that the generated sample has reached a stationary point (or local solution) of the reward function.

Type-II: $\|G_{2}(z_{0})\|=0$ and $\|G_{1}(z_{0})\|>0$. This indicates that the Jacobian matrix of the noise-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 in the noise will not affect the generated sample.

Type-III: $\|G_{1}(z_{0})\|>0$ and $\|G_{2}(z_{0})\|>0$, but $\|G_{1}(z_{0})\cdot G_{2}(z_{0})|\ =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 $t$-th step if it almost surely holds that $z_{0}$ corresponds to a Type-I, Type-II, or Type-III noise vector.

We provide examples for the three scenarios, respectively, in the following figures. First, in Figure 13(a), we display examples of Type-I and Type-II by reutilizing the experiment from Figure 3. To determine the type of the noise vector, we empirically compute $\|G_{1}(z_{0})\|$, $\|G_{2}(z_{0})\|$, and $\|G_{1}(z_{0})\cdot G_{2}(z_{0})\|$ for each noise vector.

Figure: (a) Example 1
Refer to caption: extracted/5894777/figures/type12_exp.png

To showcase an example of Type-III noise vectors, we introduce a new toy example illustrated in Figure 13(b). Specifically, the ground-truth distribution learned by diffusion models is uniform across a horizontal line spanning from $(-1,0)$ to $(1,0)$. The reward function is defined as $r(x,y)=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 noise-to-sample mapping. Consequently, all points along the line segment $[(-1,0),(1,0)]$ qualify as Type-III noise vectors.

###### Proof for Theorem 1 .

###### Lemma 2 (Descent Lemma [Bertsekas, 1997 ] ) .

### C Empirical Investigation of P ⁢ ( z ) 𝑃 𝑧 P(z) italic_P ( italic_z )

In this section, we provide several empirical evidence to demonstrate that $P(z)$ acts as an effective indicator for the out-of-distribution phenomenon.

Firstly, in Figure 14(a), we revisit the examples from Figures 5(a) and 5(b), coloring each sample based on the value of $P(z)$. As depicted in Figure 14(a), $P(z)$ proves to be an efficient metric to separate in-distribution samples from out-of-distribution samples; those in-distribution have high values for $P(z)$, whereas those out-of-distribution 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 a baseline comparison, we first draw one sample from the standard Gaussian distribution and use it to generate an image with Stable Diffusion v1.5 and the prompt "black duck". As can be seen, this leads to a normal image with $P(z)$ also within a reasonably large value. We then construct the low-probability 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 14(c), showcasing that there is nothing discernible in the image while $P(z)$ approximates zero. The second method is to repeat parts of 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 $z_{0}$ from the standard Gaussian distribution, and then constructing the noise vectors as $z=[z_{0},z_{0},z_{0},z_{0}]$, making $z$ an $n$ dimensional vector. The figure corresponding to these repeated vectors, shown in Figure 14(d), once again results in a poor image, with $P(z)$ illustrating that the noise vectors also come from a low-probability region.

We further visualize the entire optimization trajectory for the examples in Figures 5(c) and 5(d), i.e., optimizing the brightness reward for Stable Diffusion v1.5 with the prompt "black duck" in Figure 15. Specifically, from Figure 15 we can clearly see that the value of $P(z)$ gradually decreases, and the generated image also gradually diverges from the distribution associated with a black duck. Notably, at around $20$ steps, the value of $P(z)$ becomes near-zero, and at the same time, the generated image more closely resembles a blue duck rather than the specified black duck.

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 16. A clear conclusion is that in this case, it is less likely for the optimized samples to 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 15 and Figure 16, we can see that optimizing the Aesthetic Score leads to a much less significant decrease in the $P(z)$ value.

Figure: (a) Toy Example
Refer to caption: extracted/5894777/figures/sec31_toy.png

Figure: Figure 15: Trajectory of $P(z)$ on optimizing brightness reward.
Refer to caption: extracted/5894777/figures/sec31_white_traj.png

Figure: Figure 16: Trajectory of $P(z)$ on optimizing Aesthetic Score.
Refer to caption: extracted/5894777/figures/sec31_aesthetic_traj.png

### D Implementation Details

In this section, we discuss some implementation details of our proposed method, as well as clarify some omissions in the experimental section.

#### D.1 Algorithm Implementation

It is clear that to solve the direct noise optimization problem stated in Problem 8, differentiation of the noise-to-sample mapping $M_{\theta}$ is required. It is worth noting that this differentiation cannot be handled by standard auto-differentiation in PyTorch [Paszke et al., ], as it can lead to a 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., , Novack et al., , Karunratanakul et al., ].

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 auto-differentiation in PyTorch [Paszke et al., ]. Specifically, suppose we wish to use $q$ samples to estimate the gradient in Equation (20); that is, we draw $q$ noise vectors for perturbation: $\xi_{1},...,\xi_{q}$. We then generate the corresponding samples $x_{i}=M_{\theta}(z+\mu\xi_{i})$ for $i=1,...,q$. At this point, we should compute the estimated gradient of the reward functions in a non-differentiable mode as follows:

$$ $\displaystyle\hat{H}_{2}(x)=\frac{1}{q}\sum_{i=1}^{q}(r(M_{\theta}(x_{i}))-r(x ))(x_{i}-x).$ (25) $$

Finally, we can execute gradient backpropagation with the loss function,

$$ $loss(z)=\langle\hat{H}_{2}(x),M_{\theta}(z)\rangle,$ $$

which produces the exact gradient estimator for $z$.

#### D.2 Experiment Details

In this section, our goal is to provide the experimental details that were omitted from Sections 5.1, 5.2, 5.3, and also the setting for Figure 1.

Details for Figure 1. In these examples, we only use the base model of SDXL [Podell et al., ] as the image diffusion model. We adopt the DDIM sampler with 50 steps and $\eta=1$ for generation, and optimize all the injected noise in the generation process, the same as most experiments in this work. The classifier-free guidance is set to 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 in Figure 1, the used prompts are listed as 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, f 1.8, soft focus, 8k, national geographic, award - winning photograph by nick nichols.
- 4.
A beef steak, depth of field, bokeh, soft light, by Yasmin Albatoul, Harry Fayt, centered, extremely detailed, Nikon D850, award winning photography

Details for Section 5.1.
In this experiment, to solve the probability-regularized noise optimization problem as formulated in Equation (17), we employ the Adam optimizer [Kingma and Ba, ] with a learning rate of $0.01$. For optimization with regularization, we set the regularization coefficient $\gamma$ to $1$. To compute the minibatch stochastic gradient for the regularization term in Equation (17), we set the batch size $b$—the number of random permutations drawn at each step—to $100$. For each optimization run, we utilize a 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 $\gamma=0.1$ because optimizing these human-like reward functions is less susceptible to the OOD reward-hacking issue, while maintaining the batch size for the permutation matrix $b$ at 100. Each optimization run also uses a single A800 GPU, but the total memory consumption is around 20 GB. Regarding the baselines in Table 1, we implemented LGD [Song et al., ] ourselves, following the algorithm from their paper on these reward functions. For other baselines, we reuse the statistics 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 $\mu$ and the number of samples $q$ used to approximate the gradient (as formulated in Equation (25)). Clearly, $q$ 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 $q$ values for ZO-SGD, Hybrid-1, and Hybrid-2 to 16, 8, and 4 respectively. For $\mu$, we also adjust them individually for each algorithm, as they have varying sensitivity to $\mu$. Through trial and error, we select $\mu$ values of 0.01 for ZO-SGD and Hybrid-1, and 0.02 for Hybrid-2. Finally, for optimizing JPEG Compressibility, we use the Adam optimizer with a learning rate of $0.01$, but for the Aesthetic Score experiment, we reduce the learning rate to $0.001$, as we found that $0.01$ can lead to divergence during optimization for the Aesthetic Score. Each optimization run continues to use a single A800 GPU.

Figure: (a) The effect of $k$.
Refer to caption: extracted/5894777/figures/secb1.png

### E Hyperparameters Analysis

In this section, we conduct a thorough analysis of the hyperparameters for the proposed method. Our objective is to offer a 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 2, the dimension $k$ should be neither too large nor too small.
Additionally, another critical hyperparameter is the number of permutation matrices $b$ employed to compute the stochastic gradient for the probability regularization in Equation (17).
Furthermore, we aim to explore the impact of the regularization coefficient $\gamma$ in the probability regularization term.

To examine the effects of $k$, $b$, and $\gamma$ on mitigating the OOD (Out-Of-Distribution) reward-hacking problem, we revisit the experiment of optimizing darkness reward with the prompt "white <animals>" from Section 5.1. In Figure 17, we illustrate how these three hyperparameters influence both the reward and the consistency score (CS), across four different values.

Firstly, Figure 17(a) supports the notion that $k$ should be carefully chosen—not too large, yet not overly small. We observe that $k=1$ underperforms compared to $k=2$ and $k=10$, as selecting $k=1$ fails to account for the covariance among noise vectors. Conversely, $k=100$ proves to be a poor choice because it entails a smaller $m$, potentially rendering the concentration inequalities detailed in Lemma 1 less precise.

Secondly, as demonstrated in Figure 17(b), the number of permutation matrices $b$ seems to have a minor impact on the optimization process, provided $b$ is sufficiently large. Based on empirical evidence, $b=100$ emerges as an optimal selection for the proposed method.

Lastly, the effects of $\gamma$ are depicted in Figure 17(c). Adjusting the value of $\gamma$ clearly presents a trade-off between convergence speed and the propensity for OOD reward-hacking problems. Given this observation, we recommend empirically tuning the value of $\gamma$ for different reward functions and prompts using a limited number of samples and a few optimization steps.

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: See the weakness section.
soundness: 3
strengths: - The idea is simple and neat. Directly optimizing the noise towards samples of maximum reward value is technically sound. The paper presents the idea in a clear way.

- The paper considers both differentiable and non-differentiable reward settings, making direct noise optimization a quite practical method to use.

- The paper provides some interesting theoretical insights to justify the drawback of direct noise optimization. Moreover, it provides a solution to avoid this drawback with a probability regularization.
summary: This paper introduces a method called direct noise optimization (DNO) for maximizing some reward function for the generated samples. Different from finetuning and RL methods, DNO operates as a form of testing-time optimization. DNO is also extended to the non-differentiable reward setting by leveraging zero-order optimization techniques. Some theoretical analyses are provided to account for its empirical effectiveness.
weaknesses: - As the paper mentioned, directly optimizing the input noise can result in degenerated solutions, which the authors describe as OOD reward hacking. I think the problem, despite being alleviated by some additional regularization, will fundamentally limit the generation quality. This is conceptually similar to the problem of adversarial examples. Moreover, regularizing the noise to lie in a high-probability region may alleviate the reward hacking problem, but it also limits the ability to maximize the reward function. How can the effectiveness of this probility regularization be measured in practice? How to set the hyperparameter $\gamma$ in practice?

- Zero-order optimization is applied to the non-differentiable reward setting, but I am highly skeptical about its empirical effectiveness, as the zero-order optimization is known to be difficult to converge, not to mention that the noise optimization could be very challenging and non-convex/non-smooth.

- The performance gain of DNO over the other baselines also seems be quite marginal from Table 1. For the other baseline methods, there are many hyperparameters that can be tuned. I think to better illustrate the performance gain, it will be good to compare the performance with different finetuning steps (or different inference-optimization steps).

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Please refer to Weakness part.
soundness: 2
strengths: + This study considers aligning diffusion models in the inference time, alleviating the computational cost of aligning the model.
+ The authors identify and address the problem with reward hacking in optimization.
+ The proposed method is applicable for both differentiable and non-differentiable reward functions.
summary: This paper focuses on the alignment of diffusion models by optimizing noises in the sampling process. The authors propose to optimize noises at each sampling step to improve the reward value, rather than optimize the model parameters. To this end, they discuss the effectiveness of the proposed DNO method, and address the reward hacking problem in optimization via a regularization term. Experiments show that DNO improves the reward value of generations.
weaknesses: - It seems that the idea of optimizing noises has been proposed by (Wallace et al., 2023b; Ben-Hamu et al., 2024; Novack et al., 2024; Karunratanakul et al., 2023). What is the essential difference between DNO and these prior works, except that DNO optimizes noises at each step?
- The assumption of Theorem 1 is not convincing. I am not sure how to understand the smoothness from Figure 4 (Tang et al., 2024a). Although the noise-to-sample mapping is smooth, the reward function is usually complicated, especially for reward functions related to human preference. Is there any direct justification for the smoothness of $r\odot M$?
- The faster convergence speed of DNO is not well supported by their theoretical analysis. Equation (5) only demonstrates the final performance of DNO with SDE may be better than DNO with ODE, but does not justify the optimization speed.
- In equation (10), there should be a $1/\mu$ term.
- There are several concerns about the experiment in Table 1. First, it is unclear what prompts and how many generations are used for evaluation. Second, for DNO, the annotated time is confusing. Does ``1 min’’ mean that generating one image using DNO costs 1 minute? Third, how about the performance of DNO with ZO-SGD, hybrid-1, and hybrid-2?
- The evaluation for reward hacking based on $P(z)$ is limited. $P(z)$ only indicates the distribution of initial noise $z$ but cannot describe the distribution of the final generated image. I’d like to see more comparisons between images generated by different methods, either quantitatively or qualitatively.
- Which experiment supports the claim in Line 452 `` while it prevents the test metrics, i.e., the HPS score and Pick Score, from decreasing throughout the process'’?
- Is the proposed DNO also effective for more complex prompts? Have the authors compared images generated by different methods on both simple and complex text prompts?

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. In fact, the model is repeatedly applied to intermediate noises too. Why there is no regularization of out-of-distribution intermediate noise inputs to the model?
2. Why there is no comparison to LGD and other baseline in Section 5.1? The performance of LGD is surprisingly bad in Section 5.2, even worse than SD 1.5. The authors argue that the reason is the complex reward used in 5.2, and I wonder whether LGD has a better performance in similar reward functions like 5.1.

======================

Post rebuttal:
The author rebuttal solved my concern about the motivation of the regularization technique. However, I feel like the paper still needs lots of improvement on writing and more comprehensive experiments. So I will raise my score to 6 but with lower confidence.
soundness: 3
strengths: 1. The paper provides a thorough treatment of DNO, including theoretical analysis, practical implementation details, and solutions to key challenges (OOD reward hacking, non-differentiable rewards). The probability regularization technique using concentration inequalities is particularly novel and well-motivated since it addresses important practical concerns without no fine-tuning, works with arbitrary reward functions, can handle non-differentiable rewards, etc.
2. The experimental evaluation is extensive and convincing, showing that DNO can match or exceed the performance of fine-tuning methods while requiring only inference-time optimization. The authors test on multiple reward functions and provide detailed ablation studies.
summary: This paper proposes Direct Noise Optimization (DNO), a novel approach for aligning diffusion models with continuous reward functions at inference time. The key idea is to optimize the injected noise during the sampling process to maximize desired reward functions, without needing to fine-tune the model parameters. The authors develop theoretical foundations for DNO, introduce probability regularization to prevent out-of-distribution samples, and extend the method to handle non-differentiable rewards. Extensive experiments demonstrate that DNO can achieve state-of-the-art reward scores within reasonable inference time budgets.
weaknesses: 1. Although the paper proposes a regularization technique to provide overfit, the resulting images still seem to overfit. Besides, the regularization only focuses on constraining the initial noise distribution to the high probability area. Also see question 1
2. Lack of baseline results: There is no inference time cost comparison with LGD since it's an important inference-time method as the baseline. Besides, a more relevant training method is DRAFT [1], which also uses direct backpropagation to optimize the reward function. It also faces a similar issue of reward hacking as the proposed method.


[1] Directly Fine-Tuning Diffusion Models on Differentiable Rewards
Kevin Clark, Paul Vicol, Kevin Swersky, David J Fleet

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How does DNO perform when used with other diffusion sampling methods besides DDIM? From my perspective of view, I think DDIM is the easiest and most direct way to model the mapping from noise to sample, i.e. $M_\theta(z)$. If using other schedulers other than DDIM, I wonder how DNP will work to simulate $M_\theta$. Moreover, if the algorithm leverages a differentiable reward function for optimization, the gradients are passed backward through all diffusion steps. I assume that would bring significant computational costs comparable to direct reward propagation. Can the authors comment on this?

2. I don't see how eq (2) brings a "prompt-agnostic" method. For example, when aligning stable diffusion, is the optimized initial noise a "universal" optimal value for all prompts?

3. Currently, there are already many works aimed at aligning diffusion models. Roughly, they can be classified into two categories,
      (1) directly fine-tune the diffusion models. See:  https://arxiv.org/abs/2407.13734, https://arxiv.org/abs/2410.08315.   While DNO includes comparisons with alignprop and ddpo in Table 1, comparisons with many later RL-based finetuning methods that alleviate reward hacking issues are missing. I wonder how DNO would work if compared to them. 
      (2) purely inference time techniques. While DNO does not fine-tune diffusion models, DNO still needs to optimize through the diffusion pipeline to get optimized noise. The computational costs are still heavy. Here are some works that are purely inference-time and prompt-agnostic: https://arxiv.org/abs/2410.08134, https://arxiv.org/abs/2408.08252.

Discussing and comparing these works quantitatively (or qualitatively) is important for positioning DNO. However, it might also be true that these works are concurrent with this submission. Finally, at least qualitative discussions are needed.

4. I assume the numbers in Table 1 do not consider reward hacking which is problematic for a method that claims to alleviate reward hacking. This is because alignprop would significantly suffer from mode collapse when aesthetic score > 8.5 from my experience. But Table 1 reports 8.9 for alignprop. Can the authors comment on this?
soundness: 3
strengths: 1. DNO's main feature is training-free. Current researchers have been more interested in inference-time techniques that do not require costly fine-tuning. Therefore, this work stands for an important and attentive direction of aligning generative models.
2. The integration of probability regularization to mitigate out-of-distribution reward hacking is interesting.
3. The authors applied DNO across various reward functions, highlighting its adaptability.
4. DNO is compared rather comprehensively with some baselines on a variety of reward functions, which is good.
summary: The paper addresses the alignment of diffusion models to downstream tasks with continuous reward functions. It proposes Direct Noise Optimization (DNO) as a tuning-free, inference-time method for adjusting generated samples to maximize target rewards. Moreover, in the face of reward hacking (aka over-optimization), the paper introduces a probability regularization for alleviation.
weaknesses: 1. While theoretical parts of the probability regularization are comprehensive, I find those parts hard to parse. Could the authors highlight the key message of such regularization, and why would it work from a high-level perspective?
2. Reward hacking is known to be a key concern in aligning LLM/diffusion models. The authors attempt to alleviate hacking through a regularize, which is good. However, the evaluation way remains problematic. In Fig 2, the authors presented that adding regularize sacrifices little performance for retraining fidelity. However the authors did not compare regularized DPO with other methods. In the end, the observation of Fig 2 is not a surprise to me because adding regularization almost surely leads to a tradeoff of performance and fidelity. There exist some works showing that simply adding small regularization to align-prop would greatly alleviate reward hacking. Therefore, Fig. 2 does not provide any further impressive information.
3. One primary concern is that this work lacks many comparisons with existing fine-tuning-based and inference-based techniques that appear earlier or concurrent to this work. See more details below.

Meta-review: This paper presents a method called *Direct Noise Optimization* (DNO) for maximizing a reward function associated with generated samples. Unlike finetuning or reinforcement learning approaches, DNO functions as a test-time optimization technique. The method is further extended to handle non-differentiable reward settings using zero-order optimization techniques. By addressing both differentiable and non-differentiable reward scenarios, DNO demonstrates practical applicability.  

The concept of optimizing noise has been previously explored in works by Wallace et al. (2023b), Ben-Hamu et al. (2024), Novack et al. (2024), and Karunratanakul et al. (2023). Additionally, related ideas, such as using diffusion models for black-box optimization (ICML 2023, NeurIPS 2024) and multi-objective optimization (ECML 2024), exhibit similarities to this approach. Consequently, the technical novelty of this work appears limited.  

Reviewers have raised concerns about the rationale behind regularizing noise vectors alone. For instance, maintaining all intermediate additive noise distributions while varying the initial noise distribution leaves the effect on the final output distribution unclear. Another critical issue is the lack of comparisons with prior fine-tuning-based and inference-based methods, whether preceding or concurrent to this work.  

The paper would benefit from significant improvements in writing clarity and more comprehensive experimental evaluations to establish its contributions more robustly.

Decision: Reject