Inference-Time Alignment of Diffusion Models with Direct Noise Optimization

Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.

1. Comparing to prior works

Please allow us to emphasize and reiterate our main contributions here, which are distinct from all prior works on noise optimization.

• A more comprehensive formulation for noise optimization with theoretical understanding: Firstly, 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.
• Identification and mitigation of the OOD reward-hacking problem in noise optimization: For 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.
• Extension to non-differentiable reward functions: This 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.

With these three important contributions developed in this work, we believe it lays a stronger foundation for more applications of noise optimization in the future.

2. Optimization time

We would like to point out that the optimization time in DNO is actually controllable, which provides significant flexibility for different applications.

For 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.

More 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.

3. More modality

For 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:

1. There are many excellent open-sourced image diffusion models available for us to conduct experiments.
2. 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.
3. 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.

In 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.

4. Scaling behavior of computational resources

Thank you for bringing this important question to our attention. This is indeed a crucial aspect to discuss!

Due 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:

1. 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.
2. 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.

Thank 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.

5. Presentation

Could 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!

Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.

1. "What is the level of M1 and M2 value for diffusion models with hacked reward alignments?"

In 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.

2. "Can DNO be combined with other tuning-based methods?"

Yes, 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.

In 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.

3. "Could you elaborate on the optimization process of SDE sampling as advised above?"

You 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.

To clarify how our implementation works, here is a core snippet of our code:

# SDE sampling logic. To change it into ODE, we only need to modify the sampler and dimension of noise vectors accordingly
noise_vectors = torch.randn(args.num_steps + 1, 4, 64, 64, device=args.device)
noise_vectors.requires_grad_(True)
sampler.initialize(noise_vectors, "sde")

# Sampling process
while not sampler.is_finished():
    model_kwargs = sampler.prepare_model_kwargs(prompt_embeds=prompt_embeds)
    model_output = checkpoint.checkpoint(unet, **model_kwargs)
    sampler.step(model_output) 

# Gradient computation
sample = sampler.get_last_sample()
loss = - reward_function(sample)
loss.backward()

4. Relationship to the work [arXiv:2501.09732]

Thank 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.

To illustrate the difference:

• 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.
• 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.

As 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.

Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.

1. On the smoothness assumption.

Conceptually, 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:

We 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:

$A = E_{x_1, x_2} \frac{||r(M_\theta(x_1)) - r(M_\theta(x_2))||}{||x_1 - x_2||}$

and

$B = E_{x_1, x_2} \frac{||\nabla r(M_\theta(x_1)) - \nabla r(M_\theta(x_2))||}{||x_1 - x_2||}.$

Using 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.

2. Results on more diverse prompts.

We 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.

3. Discussion on the two related works.

Thanks 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.

While these two works also focus on noise optimization for diffusion models, there are several distinctions between them and our work:

• [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.

• [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.

4. Impact on diversity.

This 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.

Thanks for highlighting this point; we will include this discussion in our revised manuscript.

5. Explanation of Table 2.

This 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.

Thank you for your time in reviewing our work. Here, please allow us to provide specific responses to your major comments.

Regarding your comment: "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?"

Firstly, 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.

From 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.

Looking ahead, as more reward models and applications emerge, we believe our proposed techniques will find even more important and crucial applications.