Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models
Authors: Luca Eyring
Shyamgopal Karthik
Alexey Dosovitskiy
Nataniel Ruiz
Zeynep Akata
Venue: NeurIPS 2025 poster ()
Keywords: generative models
diffusion models
hypernetworks
preference optimization
flow matching
reward fine-tuning
noise optimization
Primary area: deep_learning
TL;DR: We propose a post-training technique, that uses a hypernetwork to efficiently steer diffusion model initial noise towards distributions favored by reward models, enhancing generation quality at low computational cost.

Abstract: The new paradigm of test-time scaling has yielded remarkable breakthroughs in Large Language Models (LLMs) (e.g. reasoning models) and in generative vision models, allowing models to allocate additional computation during inference to effectively tackle increasingly complex problems. Despite the improvements of this approach, an important limitation emerges: the substantial increase in computation time makes the process slow and impractical for many applications. Given the success of this paradigm and its growing usage, we seek to preserve its benefits while eschewing the inference overhead. In this work we propose one solution to the critical problem of integrating test-time scaling knowledge into a model during post-training. Specifically, we replace reward guided test-time noise optimization in diffusion models with a Noise Hypernetwork that modulates initial input noise. We propose a theoretically grounded framework for learning this reward-tilted distribution for distilled generators, through a tractable noise-space objective that maintains fidelity to the base model while optimizing for desired characteristics. We show that our approach recovers a substantial portion of the quality gains from explicit test-time optimization at a fraction of the computational cost.

Noise Hypernetworks: Amortizing Compute in Diffusion Models
Luca Eyring 1,2,3 1.2.3
Shyamgopal Karthik 1223,4 Alexey Dosovitskiy
Nataniel Ruiz
Zeynep Akata! 1,2,3*
Technical University of Munich
Munich Center of Machine Learning
Helmholtz Munich
University of Tübingen "Iceeptive
"Google
luca. Juca.eyring@t.. eyring@tum.de
Abstract
The new paradigm of test-time scaling has yielded remarkable breakthroughs in Large Language Models (LLMs) (c.g. reasoning models) and in generative vision models, allowing models to allocate additional computation during inference to effectively tackle increasingly complex problems. Despite the improvements of this approach, an important limitation emerges: the substantial increase in computation time makes the process slow and impractical for many applications. Given the success of this paradigm and its growing usage, we seek to preserve its benefits while eschewing the inference overhead. In this work we propose one solution the critical problem of integrating test-time scaling knowledge into a model during Specifically, we replace reward guided test-time noise optimization in diffusion models with a Noise Hypernetoor that modulates initial input noise. We propose theoretically grounded framework for learning this reward-tilted distribution for distilled generators, through a tractable noise-space objective that maintains fidelity to the base model while optimizing for desired characteristics.
We show that our approach recovers substantial portion of the quality gains from explicit test-time optimization at a fraction the computational cost. Code available https::/ thub. com/ cwm/Expl httpllaamaaaHyerlllllooooo
1 Introduction
Recently, inference-time scaling has made remarkable breakthroughs in Large Language Models [28, 40, 84] and generative vision models, enabling models spend more computation during inference to solve complex problems effectively. Drawing from the success and growing usage of test-time compute LLMs, several methods have attempted to apply similar ideas the context of diffusion models for generation [6, 22, 59, 66, 67, 80, 88, 90, 92, 93, 97]. The goal of this process is to spend additional compute during inference to obtain generations that better reflect desired output properties.
Diffusion model test-time techniques that optimize the initial noise or intermediate steps of the diffusion process, often guided by feedback from pre- pre-trained reward models [48, 54, 101, 102, 103, 107], have demonstrated significant promise in improving critical attributes of the generated outputs, such as prompt following, aesthetics, quality and composition [11, 22, 44, 59, 66, 67, 97]. These methods generally fall into two broad categories: gradient-based based optimization, which typically requires substantial GPU memory for backpropagation through the full model [6, 22, 46, 67, 97], and gradient-free optimization, which often necessitates very large number of function evaluations (NFEs), sometimes thousands, of the computationally expensive denoising network [44,
"Equal Supervision
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

SANA-Spant
SANN Sprin Hyper
FLUX-Schmell
FLU SLUU--Sobooe
glant chare red
blue bridge green morK
emt 1:ng xiice of pizza
EASIDOat 2ed
"ountaing"
red Lantarne retlectione"
doong yoga aomttai top"
Lighthouse
RUEDEA"
Figure I: The same initial random noise is used for the base generation and the initialization of the noise hypernetwork. HyperNoise significantly improves upon the initially generated image with respect to both prompt faithfulness and aesthetic quality for both SANA-Sprint and PLUX-Schned
59, 93]. While both strategies can effectively boost output quality, they introduce considerable latency (exceeding 10 minutes for one generation), severely limiting their practical utility, particularly for real-time applications. This is an instantiation of global problem of test-time scaling methods that we seek to tackle this work. The core hypothesis of our work is whether it is possible capture portion of test-time scaling benefits by integrating this knowledge into neural network during training time?
To address this, one might consider directly fine-tuning the diffusion model using reward signals [12, 14, 18. 53, 74, 89, 91, 108] or with Direct Preference Optimization (DPO) [34, 45, 52, 77, 98]. The objective here can be formulated as learning tilted distribution (Equation 3). which upweights samples with high reward while maintaining fidelity to a base model's distribution. These methods usually expensive to train due to the need for backpropagation through the sampling process. Instead, one might consider directly fine-tuning a step-distilled generative model learn this target distribution. However, this approach typically involves a KI. regularization to the base model that intractable for distilled models. An imbalance or poor estimation of this can lead to the model "reward-haw -hackingk, superficially maximizing the reward metric while significantly deviating from the desired underlying data distribution, thus not achieving the genuine desired improvements.
In this work, we propose different path to realize the benefits of the target tilted distribution (Equation 3), particularly for step-distilled generative models. Our core hypothesis is that instead of modifying the parameters of the base generator, we can achieve the desired output distribution by learning to predict an optimal initial noise distribution.
first show that such an optimal lilted noise distribution P exists (characterized by Equation 5). When samples from this Ph are passed through the frozen generator, they naturally produce outputs that are distributed according to the target data-space space tilted distribution. To lear this tilted noise distribution, we introduce lightweight

---

network, fe. that transforms standard Gaussian noise into modulated, improved noise latent. The crucial advantage of this approach lies in its optimization objective. In particular. the regularization term, a KI. divergence between the modulated noise distribution and the standard Gaussian prior, defined entirely in the noise space. We show that this noise-space KI. divergence can be made tractable and effectively approximated by an L penalty on the magnitude of the noise modification.
This lightweight network forms the core of our approach, which we term Noise Hypernetworks. It functions akin to a hypernetwork [2, 30, 31, 39, 63, 72, 95, 105] as rather than generating the final image, it produces a specific, optimized starting latent for the main frozen generative model. This effectively guides the output of the base model without any changes to its parameters. Broadly, a hypernetwork is an auxiliary model trained to generate crucial inputs or parameters of primary model. Our f embodies this concept by learning to predict the optimized initial noise as input to the frozen generator. Consequently, our proposed approach is effectively training a Noise Hypernetwork to perform the task of test-fime noise optimization, by learning to directly output an optimized noise latent in single step sidestepping the need for expensive, iterative test-time optimization.
Our practical implementation of the noise hypernetwork utilizes Low-Rank Adaptation (LoRA), ensuring it remains parameter-efficient and adds negligible computational cost during inference. We apply our method to text-to-image generation, conducting evaluations with an illustrative "redness" reward task to demonstrate core mechanies, as well as complex aligmments using sophisticated human-preefe preference reward models. We demonstrate the efficacy of our approach by applying it to distilled diffusion models SD-Turbo [83], SANA-Sprint [13], and FLUX-Schnell Overall, our experiments show that we can recover a substantial portion of the quality gains from explicit test-time optimization at fraction of the computational inference cost. In summary, our contributions are:
1. We introduce HyperNoise, novel framework that learns to predict an optimized initial noise for a fixed distilled generator, effectively moving test-time noise optimization benefits and computa- tional costs into a one-time post-training stage.
2. We propose the first theoretically grounded framework for learning the reward tilted distribution of distilled generators, through a tractable noise-space objective that maintains fidelity to the base model while optimizing for desired characteristics.
3. We demonstrate through extensive experiments significant enhancements in generation quality for sate-of-the-art distilled models with minimal added inference latency, making high-quality, reward-aligned generation practical for fast generators.
2 Background
Preliminaries. Recent generative models are based on time-dependent formulation between standard Gaussian distribution Xy ~ Po = (O, I) and a data distribution X1 ~ Pdata- These models define an interpolation between the initial noise and the data distribution, such that
x 0x + oxx1,
where at is a decreasing and o, is an increasing function of E [0, 1]. Score-based diffusion [33, 43, 47, 85, 86] and flow matching [3. 55, 56] models share the observation that the process x; can be sampled dynamically using a stochastic or ordinary differential equation (SDE or ODE). The neural networks parameterizing these ODEs/SDEs are trained to learn the underlying dynamics, typically by predicting the score of the perturbed data distribution or the conditional vector field, Generating a sample then involves simulating this learned differential equation starting from Xo Po-
Step-Dist Distilled Models. The iterative simulation of such ODEs/SDEs often requires numerous steps, leading to slow sample generation. To address this latency, distillation techniques have emerged as a powerful approach. The objective is to train a "student" model that emulates the behavior of a pre-trained "teacher" model (multi-step ODE/SDE simulation) but achieves this with drastically fewer, or even a single, evaluation step(s). Prominent distillation methods such as Adversarial Diffusion Distillation [83], Consistency Models [58, 87], or Flow Maps [9, 10] have enabled the development of highly efficient few-step generative models, like SDubDXL-Turbo [83] and SANA-Sprint [13]. In this work, we denote such distilled generator by go- The significantly reduced number of sampling steps in these distilled models makes them more amenable to various optimization techniques and practical for real-time applications, which is why they are the focus of our work.
Test-Time Noise Optimization. Test-time optimization techniques aim to improve pre-trained generative models on per-sample basis at inference. One prominent gradient-based strategy is
3

---

test-time noise optimization [6, 29, 46, 67, 90, 97]. Given pre-trained generator ge (which could be a multi-step diffusion or flow matching model), this approach optimizes the initial noise Xp for each generation instance. The objective is to find an improved xi that maximizes a given reward ((go(xo)) often subject to regularization and can formulated as
x = arg max
- Reg(xo)),
(2)
X
where (xg( is a regularization term designed to keep xi within a high-density region of the prior noise distribution Po. thus ensuring the generated sample ge(x;) remains plausible. ReNO [22] adapted this framework for distilled generators ge. enabling more efficient test-time optimization compared full diffusion models. However, this per-sample optimization still incurs significant computational costs at inference, involving multiple forward and backward passes, and increased GPU memory. This inherent latency and computational burden motivate the exploration of methods that can imbue models with desired properties without per-instance test-time optimization.
Reward-based Fine-tuning and the Tilted Distribution. To circumvent the per-sample inference costs associated with test-time optimization, an alternative paradigm is to directly fine-tune the generative model go to align with reward function. We consider the pre-trained base distilled diffusion model ga. which transforms an initial noise sample Xy into an output sample x = go (xo)- The distribution of these generated output samples is the pushforward of po by ge. which we denote as huse preference of samples x, our objective is to learn the so called tilted distribution
(ge :Po. Given and a differentiable reward function r(x) : IR + R that quantifies the
p" (x) p pose (x) exp(r(x)).
This target distribution is defined to upweight samples with high rewards under r(x) while staying close to the original phase (x). We would like to Iearn p" (x) by minimizing the KL divergence DEL (p=))"" Here, pa is the distribution generated by modifying the base process using trainable parameters d. e.g. e could correspond to fine-tuned version of 0. This objective can be decomposed such that
min (p (I min DKL. (p"ppp - Ex [r(x)];
where we omit the normalization constant of p" (x) which is constant w.r.t. & (see Appendix A.2). This objective encourages the learned model p" to generate high-reward samples while regularizing its deviation from the original base distribution passe
Challenges in Direct Reward Fine-tuning of Distilled Models. Directly optimizing Equation 4 by fine-tuning the parameters of distilled, e.g. one-step, ge poses significant challenges. The term DxL (p*) b ) requires evaluating the densities of pe and prase, For typical neural network generators, these densities involve Jacobian determinants through the change-of-variabe formula, which are often intractable or computationally prohibitive compute for high-dimensional data [70]. Previously, a line of work has analyzed fine-tuning Diffusion [89, 91] and Flow matching [18] models based on Equation through the lens of Stochastic Optimal Control. However, this formulation relies on dynamical generative models (SDEs) and its application to distilled models is not straightforward, as these often lack the explicit continuous-time dynamical structure (e.g., an underlying SDE or ODE) that these fine-tuning techniques leverage.
3 Noise Hypernetworks
Given the challenges in directly fine-tuning go, we introduce Noise Hypernetworks (HyperNoise), a novel theoretically grounded approach to learn p² for distilled generative models. The core idea is to learn a new distribution for the initial noise, Po- such that samples Xn ~ P when passed through the fixed generator ge, produce outputs x go (Xo) that are effectively drawn from the target tilted distribution p" (x) (Equation 3). Instead of modifying the parameters 0 of the base generator, we keep go fixed. This requires P to approximate
optimal modulated noise distribution, P This tilted noise distribution, which precisely steers go to p", can be characterized by (Appendix A.3)
(x0) (x() (go()))))))))
(5)
To realize the modulated noise distribution ph we parameterize it using a learnable noise hyper- network fo (with parameters 6). This network defines transformation T3 that maps initial noise

---

Training Noise Hypernetworks
Inference
X
Initial
LoRA Noise Network
Noise-based Loss:
(()) :11Axoll* r(x1)
Base Distilled
Base LoRA(e)
Generative Model
Optimizeed
X
Axi
r(xi
Noise
LoRA
Weights
K1. Regularization
Reward
Base (0)
Residual
Base Distilled 0 Generative Model
X
X
Initial Noisc
Optimized Noise
Figure 2: Illustration our proposed HyperNoise approach. During training, the LoRA parameters are trained to predict improved noises and are optimized by reward maximization subject to KL regularization. During inference, the noise hypernetwork directly predicts the improved noise initialization which is used for the final generation.
samples Xg ~ Po to modulated samples &y via a residual formulation such that
(6) The distribution of these modulated samples, po is thus the pushforward of po by T i.e., Ph = (T>):Po- We propose to train the parameters 6 of the noise modulation network fa by mini- mizing the KL divergence DKL (p lpj)) This can be shown to be equivalent to minimizing the loss function
Xy = (x()) = Xo + fx(o).
Lnoise Dri (piloo)) Exrpp (x))())))]-
Analogously Equation 4, this objective encourages Pe (and thus fo) produce initial noise samples Xo that effectively steer the fixed generator go towards high-reward outputs The KL term DxL (po lIPo) regularizes this steering by ensuring pa remains close the original noise distribution Po. Next, we show that in contrast to Equation 4, Lnoise can be made tractable.
3.1 KL Divergence in Noise Space
The resulting KL divergence term DxL (po [Po) is derived in detail in Appendix A. The derivation involves the change of variables formula, simplification of Gaussian log- terms, and an application of Stein's 's Lemma. This leads to the following expression for the KL divergence:
DKL (p (pPo) = Ex ( + Tr(J (xx)) - log det(L + J/ (xo))]l,
(8)
where Jx,(x) (Xo) is the Jacobian of f4 with respect to xo. Let ((A) = Tr(A) - log | det (7+ A)]. Then Equation can be rewritten DRL (pg llpo) = Exo II (xo) EJ,,, (xo)).. To simplify this
expression, analyze the error term E EJ (xo)). The following Theorem provides bound on this term under a Lipschitz assumption on sá-
Theorem 1 (Bound on Nog-Determinant Approximation Error). Let A Jjs (Xx) be the d d Jacobian matrix of Jo (xo). Assume Je is L-Lipschitz continuous, such that its Lipschitz constant
L This implies that the spectral radius p(A) S L < T. Then, the error term E(A)
Tr(A) - -og det((T T + A)] is bounded by
(8(A)) < d(--log(1 L) L).
See Appendix A.4 for the full proof. Theorem shows that if the Lipschitz constant L of fo is sufficiently small (specifically, L < 1), the error term [E)) is bounded. For small log L) L x 1,2/2, /2, making the bound approximately d Thus, the expected error Ex [E[,, (Xo) ))] becomes negligible if is kept small. Under this condition, we can approximate the KL. divergence
with
DKL (pollll Exs po l (x)) [13).
(10)
5

---

This approximation simplifies the KL divergence term in our objective to computationally tractable L2 penalty on the magnitude of the noise modification fx))). Substituting into our initial noise modulation objective (Equation 7), we arrive at the final loss to minimize
Lnoise (0) Ex po (xo) II r(go (xo fx(x)))]]]
(11)
Connection to test-time noise optimization. Our proposed method addresses the same fundamental goal as Noise Optimization (Equation 2) of steering generation towards high-reward outputs while maintaining fidelity to the base distribution. However, instead of performing iterative optimization for each sample at inference time, we amortizes this optimization into one-time post-training process. By learning the noise modulation network fe. we effectively pre-computes a general policy for transforming any initial noise Xo. Consequently, steered generation with HyperNoise remains highly efficient at inference, requiring only single forward pass through fe and then ge.
Theoretical Justification via Data Processing Inequality. The KL divergence term DRi. (pg) (po)) in our objective (Equation 7) provides principled way to regularize the output distribution in data space. The Data Processing Inequality (DPI) [15] states that for any function, such as our fixed generator ge- the KL divergence between its output distributions is upper-bounded by the KL divergence between its input distributions. In our context, where Pu
& (To)aPo is the distribution of modulated noise
Xo = (xx) and hhse (go (go))Po is the base output distribution, the DPI implies
DxL (P (p0)) 2 DrL ((p))))) (ge))pp).
(12)
Thus, by minimizing DxL (Pê llPo) in the noise space, we effectively minimizes an upper bound on the KL divergence between the steered output distribution (90)spg and the original base distribution pluse This offers a theoretically grounded mechanism for controlling the deviation of the generated data distribution, complementing the empirical reward maximization, even when direct computation of data-space KL divergences (as in Equation 4) is intractable.
3.2 Effective Implementation
To implement Noise Hypernetworks efficiently and ensure stable training, we adopt several key strategies for the noise modulation network fo and the training process, summarized in Algorithm 1. Note that our training algorithm (Equation 11) does not require target data samples from p", pbase nor Pdata- It only requires: (1) base noise samples X Po. (2) the fixed generator go, and (3) the reward function r(.). For conditional go (-|c), it additionally requires the conditions c.
Lightweight Noise Hypernetwork with LoRA. The noise modulation network fa is instantiated by reusing the architecture of the pre-trained generator ge and mak- ing it trainable via Low-Rank Adaptation (LoRA) [35]. The original ge weights are frozen, and only the LoRA adapter param-
Algorithm HyperNoise
1: Input: go (distilled generative Model), r (reward fn), Optional C = { (condition dataset)
2: Initialize Noise Hypernetwork f.) = through
LoRA weights e applied on top of ga
3: while training do
4: Sample noise xo (0.
eters in fe are learned. This approach is -fffiiter-effcieet, reducing memory and
5:
if C then
computational overhead as we only need
6:
Sample condition c~C C
7: Predict modulated noise Axo= = fo(xx c)
to keep go in memory once. It also allows
So to inherit useful inductive biases from go 's architecture. For conditional mod-
8:
Generate x1 go (Xo Axp. c)
9:
Compute Loss Looise (Ø)
| - r(xi)
els (-(-),) fo (xole (c) can similarly leverage learned conditional representations by ap- plying LoRA to conditioning pathways, e.g. the learned text-conditioning of text-to-image model.
10:
Gradient step V Voonise (Ø)
11: return Noise Hypernetwork LoRA weights Φ
Initialization. We propose to initialize fa such that its output fa(-) 0. This is crucial for training stability and supports the validity of the L2 approximation for DxL ( PiPo) (Equation 10) from the start of training. Specifically, we modify the final layer of Jo to output only the LoRA- LoRA-ggnerated perturbation, which are initialized to output 0 (this is achieved by setting the second LoRA matrix, often denoted B, to zero), without using any frozen base weights. This ensures that at initialization fo(.) 0 such that effectively Xo xo + (X0) Xo. making Po Po-
@

---

HyperNoise (Ours)
Optimize Redness Reward
Training Steps
LoRA
Fine-tune
Figure 3: An illustrative example of optimizing for learning the tilted distribution with an image redness reward. We show direct LoRA fine-tuning of SANA-Sprint [13] in comparison to training a noise hypernetwork with our proposed objective. Notably, when training with our objective, model optimizes the desired reward while staying considerably closer to p showcased by the model not diverging from the image manifold, unlike in direct LoRA fine-tuning.
Experiments
Our experimental evaluation is designed assess the efficacy of our objective for the popular setting of cext-to-image (T2I) models. We benchmark the noise hypernetwork against established methods, primarily direct LoRA fine-tuning of the base generative model [74], and investigate its capacity to match recover the performance gains typically associated with test-time scaling techniques like ReNO [22], but through a post-training approach. To clearly delineate these comparisons, we structure our experiments as follows: We first present an illustrative experiment employing a "redness reward". This controlled setting designed demonstrate the advantages of our training objective, particularly its ability to optimize for target reward while mitigating divergence from the base model's learned data manifold basee Subsequently, we extend our evaluation to more complex and practical scenarios, focusing aligning generative models with human-prefereme reward models.
4.1 Redness Reward
We begin our evaluation with the goal of learning the tilted distribution (Equation 3) given a redness reward. This metric helps showcase the potential underlying issue of directly fine-tuning the generation model Io (a fine-tuned variant of the base model go). For this experiment, the redness re- ward r(x) is defined as the difference be- tween the red channel intensity and the av- erage of the green and blue channel inten- sities: r(x) = x [x' + x²), where x' denotes the i-th color channel of the gener- ated image x and is used to train the recent SANA-Sprint [13] model, for full details see Appendix B.I.
Reword Score
Redness Reward
a2s
LafA Fite-tan:: Rath Rovard
Latk Tiis tune: bragstimort Sew
lypeticie:: liadurss liwwort
Hype Inoundo Stare
75
15:
190
175
Training Steps
Figure 4: Trade-off between the redness reward objec- tive and image quality metric, ImageReward, for di- rect fine-tuning and Noise Hypernetworks. As opposed direct fine-tuning, our proposed method optimizes the redness objective while not significantly dropping image quality as indicated by the ImageReward score.
The primary concern with directly fine- tuning So to maximize a reward the risk of significant deviation from the original data distribution base This deviation can lead to high
7

---

Table 1: Quantitative Results GenEval. Our Noise Hypernetwork combined with (1) SD- Turbo [83], (2) SANA-Sprint 0.6B [13], and Flux-Schnell consistently improving results while maintaining few-step denoising, fast inference, and minimal memory overhead. Results from best-of- n sampling [44], ReNO [22], and prompt optimization [4, 61] are greyed out to provide a reference upper-bound in terms of applying optimization at inference. Prompt optimization i additionally requires a significant amount of calls to an LLM, either locally or through an API.
Model
Params (B)
Time (s)
Mean
Singlet
Two
Counting
Colorst
Position
Attribution
SD v2.1 [79]
0.8
1.9
0.50
0.98
0.51
0.44
0.85
0.07
0.17
SDXI,
2.6
0.55
0.98
0.74
0.39
0.85
0.15
0.23
DPO-SDXLI 1981
2.6
0.59
0.99
0.84
0.49
0.87
0.13
0.24
Hyper-SDX1.
2.6
0.56
1.00
0.76
0.43
0.87
0.10
0.21
Flux-dev
12.0
23.0
0.68
0.99
0.85
0.74
0.79
0.21
0.48
SDD-Mediium [2]]
2.0
0.70
1.00
0.90
0.72
0.87
0.31
0.66
Turbu [83]
0.8
0.2
0.49
0.99
0.51
0.38
0.85
0.07
0.14
HyperNoise
0.57
0.22
Prompt Optimization 61]
0.81
96.01
0.59
0.99
0.76
0.53
0.88
0.10
0.28
Best-of-N [44]
0.8
10.0
0.60
1.00
0.78
0.55
0.88
0.10
0.29
ReNO [22)
0.8
20.0
0.63
1,00
0.84
0.60
0.90
0.11
0.36
SANA-Sprint [13]
0.6
0.2
0,70
1.00
0.80
0.64
0.86
0.41
0.51
HyperNolse
0.75
0.71
0.85
0.51
0.55
Prompt Optimization [1]
0.61
96.01
0.75
0.99
0.91
0.82
0.89
0.36
0,56
Best-of-N [44]
0.6
15.0
0.79
0.99
0.92
0.72
0.91
0.53
0.65
RoNO [22]
0.6
30.0
0.81
0.99
0.93
0.74
0.92
0.60
0.67
Se..e-lShheell (4-step)
12.0
0.7
0.68
0.99
0.88
0.66
0.78
0.27
0.48
Hyper'Nolse
13.0
0.72
0.99
0.93
0.67
0.83
0.30
0.59
ReNO [22]
12.0
40.0
0.76
0.99
0.94
0.70
0.86
0.39
0.65
DKL. (pollp) where pe is the distribution induced by the fine-tuned model Sa. Such a divergence often manifests as degradation overall image quality or a loss of diversity, even if the target reward (e.g. redness) is achieved. Figure 4 quantitatively illustrates this trade-off by plotting the redness reward against a general image quality metric (ImageReward),, comparing our Noise Hypernetwork approach with LoRA fine-tuning, while Figure 3 visually corroborates these results.
4.2 Human-preference Reward Models
Implementation Details. We conduct our primary experiments aligning text-to-image models with human preferences using SD-Turbo [83], SANA-Sprint [13] and FLUX-Schnell. Notably, SANA-Sprint and FLUX-Schnell exhibit strong prompt-followee proprietary models, making them robust base models for our evaluations. For the reward signal r()) essential to our objective (Equation 11) and for the direct fine-tuning baseline, utilize the exact same composition of reward models proposed in ReNO [22] consisting of ImageReward [103], HPSv2.1 [101], Pickscore [48], and CcIP-score. For the noise hypernetwork, we use LoRA [35] module on the base distilled model with the proposed initialization as described in Section 3.2.
capabilities competitive with
Training for the noise hypernetwork is performed using -70k prompts from Pick-a-Picv2 [48], T2I- Compbench train set [37], and Attribute Binding (ABC-6K) [25] prompts. Our evaluations of the trained models are performed on GenEval [26], ensuring that the training and evaluation prompts do not have any overlap, measuring the generalization of the noise hypernetwork to unseen prompts. We mainly compare HyperNoiss with three different test-time techniques: Best-of-N sampling [44, 59], ReNO [22], and LLM-based prompt optimization [4, 61]. As detailed in Table 1, all of these incur significantly increased computational costs at test-time, ranging from 33x to 300x slower inference compared to HyperNoise, making them impractical for large-scale deployment where efficiency is paramount. Full experimental details are provided in Appendix B.2.
Quantitative Results. We present our main quantitative results on the GenEval benchmark in Table I. Our Noise Hypernetwork training scheme consistently yields significant performance gains across all model scales while maintaining near-baseling inference costs. When applied to SD-Turbo, our method nearly recovers most of the improvements from inference-time noise optimization, achieving an overall GenEval performance of 0.57 that even surpasses SDXL (which has 2x more parameters and 25x NFEs), clearly highlighting the benefits from our noise hypernetwork training. With SANA- Sprint, we observe consistent improvements (0.75 vs 0.70) over the base model, achieving the same performance as M----d prompt optimization while being 300x faster, and recovering about half of the performance gains achieved by ReNO with minimal GPU memory overhead. Notably, we observe
8

---

SD SD3.5- 5-Tubbo
Flux Schmell
FLAIX Schnell HyperSone
SANA Speine
SANA Sprzak HyperNoie
Japtos top Luddy beaz."
alaghant greg
rEStng
green gieafte blue pig.
Esoka EsLL
pagoda... Ganural
cheasy
Figure 5: Qualitative comparison our proposed noise hypernetwork with popular distilled models such as Flux-Schnell, Turbo, SANA-Sprint for 4-step generation. Both SANA-Sprint and FLUX-Schnell share the initial noise for the base and HyperNoise generation.
similar trends for the larger 12B parameter FLUX-Schnell, where we again recover substantial performance gains (0.71 vs 0.68) while maintaining the efficiency advantages that make our approach practical for real-world deployment. The consistent efficiency gains across model scales demonstrate that our approach successfully amortizes the optimization cost during training, enabling high-quality generation without the prohibitive test-time computational overhead of alternative methods.
Superiority over Direct Fine-tuning and Multi-Step Generalization. In Tab. 8, we show the generalization of our training on multi-step inference despite being trained only with one-step generation. We obtain consistent im- provements over SANA-Sprit for one, two, and four step generation. Notably, our model with one-step generation noticeably outperforms SANA-Sprint with 4 steps. We also illustrate how direct fine-tuning of the base model with the same reward objective can lead to significantly worse results, highlighting the necessity of preventing "reward- "reward-hacking" in a principled fashion. We visualize this in Appendix C.4, where observe similar patterns as previous works for reward-hacking [14, 53, 91].
Table 2: Mean GenEval results for SANA-Sprint highlighting generaliza- tion across inference timesteps of our Noise Hypernetwork and failure of di- rect LoRA fine-tuning.
SANA-Spriit [13]
NFEs
GenEval Meun
Oue-stop
fine-tune [14, 103J
HyperNoise
0.75
Twa-step
0.72
LoRA fine-tune 114, 103]
HyperNoise
0.76
Four-step
[14, 103]
HyperNoise
0.77
Qualitative Results. We illustrate examples of generated images in Fig. showing our method applied to both SANA- and FLUX FLUX-Schnell -chnnell, alongside comparisons SD3.5-Turbo, Our noise hypernetoork demonstrates consistent improvements across

---

both base models. For SANA-Sprint, the improvements are substantial: we observe both correction of generation artifacts and significantly enhanced prompt following for complex compositional requests. When applied to the already high-quality FL.U-SSchnell our method still provides noticeable im- provements in detail quality and prompt adherence, demonstrating that our approach can enhance even strong base models while maintaining the efficiency advantages essential for practical deployment.
5 Related Work
Test-Time Scaling. The paradigm of test-time scaling has yielded remarkable breakthroughs, with models allocating additional computation during inference to solve increasingly complex problems. In language models, this has manifested through process reward models [60, 84, 109] and reinforcement learning from verifiable rewards [49, 64], leading to systems like ol [40] and DeepSeek-RI [28]. Beyond scaling denoising steps in diffusion models, test-time techniques improve generation quality by finding better initial noise or refining intermediate states during inference, often guided by pre- trained reward models. These methods fall into two categories: search-based approaches [44, 59, 92, 93] that evaluate multiple candidates, and optimization-basee approaches [6, 19, 29, 46, 67, 90, 97] that iteratively refine noise or latents through gradient descent. Although both strategies achieve significant quality improvements, they introduce substantial computational overhead, with generation times frequently exceeding several minutes per image.
Aliguing Diffusion Models with Rewards. Reward models [48, 101, 102, 103, 107] have been effectively used to directly fine-tune diffusion models using reinforcement learning [8, 12, 17, 24, 108] or direct reward fine-tuning [14, 18, 41, 50, 53, 74, 75, 103]. Alternatively, Direct Preference Optimization (DPO) [34, 45, 52, 77, 98] learns from paired comparisons rather than absolute rewards. A particular instance of reward fine-tuning [18, 89, 91] analyzes learning the reward-tilted distribution through stochastic optimal control. Uehara et al. [91] fine-tune continuous-time diffusion models by jointly optimizing both the drift term and initial noise distribution, but their SDE-based formulation requires continuous-tine dynamics and backpropagation through the full sampling process, making computationally expensive and inapplicable to step-distilled models. For distilled models, concurrent work [42, 62, 68] has explored preference tuning, though without the theoretical foundation for sampling from the target-tilted distribution that our approach provides. Wagenmaker al. [96] apply similar noise-space optimization principles to diffusion policies in robotic control, demonstrating efficient adaptation while preserving pretrained capabilities across diverse domains.
Hypernetworks. Auxiliary models [30] that predict parameters of task specific models have been used for vision [2, 31] and language tasks [39, 63, 72]. For generative models, they have been used to
generate weights through diffusion [20, 99] and to speed up personalization [2, 81.. NoiseRefine [1] and Golden Noise [110] train hypernetworks to predict initial noise to replace classifier-free guidance or find reliable generations by selecting 'ground- trounndtruth'' noise pairs as supervision, as opposed to the
end-to-end training in our framework. Work on diffusion priors [5, 16, 23, 27, 65, 82] also adapts the noise distribution, but these approaches modify the training process rather than enabling post-hoc adaptation of pre-trained models. Concurrently, Venkatraman et al. [94] explore sampling from reward-tilted distributions for arbitrary generators, but our work demonstrates this approach at scale with comprehensive evaluation across multiple model architectures and unseen prompt distributions.
6 Conclusion
In this work we provide fresh perspective for post-training diffusiin models through the introduction of HyperNoise, noise modulation strategy. Our principled training objective coupled with the efficient training scheme is able to achieve a meaningful improvements in performance across multiple models while avoiding neward-hackingg for aligning diffusion models with downstream objectives finds use across a wide variety of domains and use cases, especially in cases where test-time optimization would be prohibitively expensive.
We hope that our efficient and effective solution
Limitations. Preference-tuning diffusion models heavily relies on strong pre-trained base models and meaningful reward signals. While constant improvements are made to develop better pre-trained base models, specific focus should be devoted to improving reward models that can give meaningful feedback variety of aspects that are important for high-qualit, generation.
10

---

Acknowledgements
This work was partially funded by the ERC (853489 DEXIM) and the Alfried Krupp von Bohlen und Halbach Foundation, which we thank for their generous support. Shyamgopal Karthik thanks the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. Luca Eyring would like to thank the European Laboratory for Learning and Intelligent Systems (ELLIS) PhD program for support.
References
[1] Donghoon Ahn, Jiwon Kang, Sanghyun Lee, Jaewon Min, Minjae Kim, Wooseok Jang, Hyoungwon Cho, Sayak Paul, SeonHwa Kim, Eunju Cha, al. noise is worth diffusion guidance. arXiv preprint arXiv:24/2.03889 2024.
[2] Yuval Alaluf, Omer Tov, Ron Mokady, Rinon Gal, and Amit Bermano. Hyperstyle: Stylegan inversion with hypernetworks for real image editing. In CVPR, 2022.
[3] Michael S Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In ICLR, 2023.
[4] Kumar Ashutosh, Yossi Gandelsman, Xinlei Chen, Ishan Misra, and Rohit Girdhar. Lims can see and hear without any training, 2025. URL atttps://rrie org/ httg/ab//200 /2s//bbb/500..... 18096.
[5] Grigory Bartosh, Dmitry Vetrov, and Christian A. Naesseth. Neural flow diffusion models: Learnable forward process for improved diffusion modelling, 2025. URL https://arxiit org/ abaabbs/4444 12940.
[6] Heli Differentiating through flows for controlled generation. ICMI., 2024.
Omri Puny. Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow:
[7] Samarth Bhatia and Felix Dangel. Lowering pytorch's memory consumption for selective differentiation. 2024.
[8] Kevin Black, Michael Janner, Yilun Du, llya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In ICLR, 2024.
[9] Nicholas M. Boff, Michael S. Albergo, and Eric Vanden-Eijnden. How to build a consistency model: Learning flow maps via self-distillation, 2025. URL https://arxiv org/ abs/ 2505 18825.
[10] Nicholas M. Boffi, Michael S. Albergo, and Eric Vanden-Eijnden. Flow map matching with stochastic interpolants: mathematical framework for consistency models, 2025. URL https://arxiv. org/ /a//a////24// 07507.
[11] Hila Chefer, Yuval Alaluf, Yael Vinker, Lior Wolf, and Daniel Cohen-Or. Attend-and-excice Attention-based semantic guidance for text-to-image diffusion models. SIGGRAPH, 2023.
[12] Chaofeng Chen, Annan Wang, Haoning Wu, Liang Liao, Wenxiu Sun, Qiong Yan, and Weisi Lin. Enhancing diffusion models with text-encoder reinforcement learning. In ECCV, 2024.
[13] Junsong Chen, Shuchen Xue, Yuyang Zhao, Jincheng Yu, Sayak Paul, Junyu Chen, Han Cai, Enze Xie, and Song Han. Sana-sprint: One-step diffusion with continuous-tim consistency distillation. arXiv preprint arXiv:2503.
2025.
[14] Kevin Clark, Paul Vicol, Kevin Swersky, and David J Fleet. Directly fine-tuning diffusion models on differentiable rewards. In ICLR, 2024.
[15] Thomas M. Cover and Joy A. Thomas. Elements of Information Theory 2nd Edition (Wiley Series in Telecommunications and Signal Processing). Wiley-Interscience, July 2006. ISBN 0471241954.
[16] Cecilia Curreli, Dominik Muhle, Abhishek Saroha, Zhenzhang Ye, Riccardo Marin, and Daniel Cremers. Nonisotropic gaussian diffusion for realistic 3d human motion prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2025.
[17] Fei Deng, Qifei Wang, Wei Wei, Matthias Grundmano, and Tingbo Hou. Prdp: Proximal reward difference prediction for large-scale reward finetuning of diffusion models. In CVPR, 2024.
11

---

[18] Carles Eamingo-Enrich. Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint (Xiv:2449.08866) 2024.
[19] Jesse Engel. Matthew Hoffman, and Adam Roberts. Latent constraints: Learning to generate conditionally from unconditional generative models, 2017. URL https: peattts://arxie http:///axvvvooooeeeee abs/
1711.05772.
[20] Ziya Erkoç, Fangchang Ma, Qi Shan, Matthias Nießner, and Angela Dai. Hyperdiffusion: Generating implicit neural fields with weight-space diffusion. In ICCV, 2023.
[21] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow trans- formers for high-resolution image synthesis. arXiv preprint arXiv:2403.03320 2024.
[22] Luca Eyring, Shyamgopal Karthik, Karsten Roth, Alexey Dosovitskiy, and Zeynep Akata, Reno: Enhancing one-step text-to-image models through reward-based noise optimization. In NeurIPS, 2024.
[23] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in neural monge maps improves unpaired domain translation. The Twelfth International Conference on Learning Representations, 2024. URL
met/forumiid-ea0 -ddoorrrrddn22nn3aaaa 2UnCj3jeao.
[24] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Reinforcement learning for fine-tuning eot-tt--mmae diffusion models. NeurIPS, 2023.
[25] Weixi Feng, Xuchai Hc, Tsu-Jui Fu, Varun Jampani, Arjun Reddy Akula, Pradyumna Narayana, Sugato Basu, Xin Eric Wang, and William Yang Wang. Training free structured diffusion guidance for compositional text-to-image synthesis. In ICLR, 2023.
[26] Dhruba Ghosh, Hanna Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused frame- work for evaluating text-to-image alignment. In NeurlPS, 2023.
[27] Sang gil Lee, Heeseung Kim, Chaehun Shin, Xu Tan, Chang Liu, Qi Meng. Tao Qin, Wei Chen, Sungroh Yoon, and Tie- Yan Liu. Priorgrad: Improving conditional denoising diffusion models with data-dependent adaptive prior, 2022. URL https: /tattps://rxivv /axxiv.ivvvv org/ aa//a/aa///10 /216. 06406.
[28] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.1299 2025.
[29] Xiefan Guo, Jinlin Liu, Miaomiao Cui, Jiankai Li, Hongyu Yang, and Di Huang. Initno: Boosting text-to-image diffusion models via initial noise optimization. In CVPR, 2024.
[30] David Ha, Andrew Dai, and Quoc V Le. Hypernetworks. arXiv preprint arXiv: arXiv:/609.09)) 1609.09106, 2016.
[31] Eric Hedlin, Munawar Hayat, Fatih Porikli, Kwang Moo Yi, and Shweta Mahajan. Hypernet fields: Efficiently training hypemetworks without ground truth by learning weight trajectories. arXiv preprint arXiv:2412.17044 2024.
[32] Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning, 2022.
[33] Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.
[34] Jiwoo Hong, Sayak Paul, Noah Lee, Kashif Rasul, James Thome,, and Jongheon Jeong. Margin- aware preference optimization for aligning diffusion models without reference. arXiv preprint arXiv:2406.0665 06-42,, 2024.
[35] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In ICLR, 2022.
[36] Xiwei Hu, Rui Wang, Yixiao Fang, Bin Fu, Pei Cheng, and Gang Yu. Ella: Equip diffusion models with Il for enhanced semantic alignment. arXiv preprint arX3v:2403.057335 2024.
12

---

[37] Kaiyi Huang, Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. c2i-compbench: A comprehensive benchmark for open-world compositional -ext-to-image generation. In NeurIPS, 2023.
[38]
Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannanch Hajishirzi. Ali Farhadi, and Ludwig Schmidt. Openclip, July 2021. URL https://doli org/10///0 10. 5281/
zenodo zenodo.5143733.
[39]
Hamish Ivison, Akshita Bhagia, Yizhong Wang, Hannaneh Hajishirzi, and Matthew Peters. Hint: hypemeework instruction tuning for efficient zero-& few-shot generalisation. arXiv preprint arXiv:2212.10311 2022.
[40] Aaron Jacch, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, al. Openai ol system card. arXiv preprint arXiv:2412.16722 2024.
[41] Rohit Jena, Ali Taghibakhshi, Sahil Jain, Gerald Shen, Nima Tajbakhsh, and Arash Vahdat. Elucidating optimal reward-diversity tradeoffss in text-to-image diffusion models. arXiv preprint arXiv:2409.0644 2024.
[42] Zhiwei Jia, Yuesong Nan, Huixi Zhao, and Gengdai Liu. Reward fine-tuning two-step diffusion models via learning differentiable latent-space surrogate reward. arXiv preprint arXiv:2411. 200::2211.15247, 2024.
[43] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In NeurlPS, 2022.
[44] Shyamgopal Karthik, Karsten Roth, Massimiliano Mancini, and Zeynep Akata. If at first you don't 't succeed, try, try again: Faithful diffusion-based text-to-image generation by selection. arXiv preprint arXiv:2305. 13308, 2023.
[45] Shyamgopal Karthik, Huseyin Coskun, Zeynep Akata, Sergey Tulyakov, Jian Ren, and Anil Kag. Scalable ranked preference optimization for text-to-image generation. arXiv preprint arXiv:2410.1800 2024.
[46] Korrawe Karunratanakul, Konpat Preechakul, Emre Aksan, Thabo Beeler, Supasorn Suwa- janakorn, and Siyu Tang. Optimizing diffusion noise can serve as universal motion priors, In CVPR. 2024.
[47] Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. NeurlPS, 2021.
[48] Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick-a-pic: An open dataset of user preferences for text-to-image generation. In NeurIPS, 2023.
[49] Nathan Lambert. Jacob Morrison, Valentina Pyatkin, Shengyi Huang, Hamish Ivison, Faeze Brahman, Lester James V Miranda, Alisa Liu, Nouha Dziri, Shane Lyu, et T ulu Pushing frontiers in open language model post-training. arXiv preprint arXiv::411. 2.221.115111 15124,, 2024.
[50] Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302. 12/92, 2023.
[51] Junnan Li, Dongxu Li, Caiming Xiong, and Steven Hoi. Blip: Bootstrapping language- image pre-training for unified vision- languaee understanding and generation. In International Conference on Machine Learning. 2022.
[52] Shufan Li, Konstantinos Kallidromitis, Akash Gokul, Yusuke Kato, and Kazuki Kozuka. Aligning diffusion models by optimizing human utility. arXiv preprint arXiv:2404.04 2024.
[53] Yanyu Li, Xian Liu, Anil Kag. Ju Hu, Yerlan Idelbayev, Dhritiman Sagar, Yanzhi Wang. Sergey Tulyakov, and Jian Ren. Textcraftor: Your text encoder can be image quality controller. In CVPR. 2024.
[54] Zhiqiu Lin, Deepak Pathak, Baiqi Li, Jiayao Li, Xide Xia, Graham Neubig, Pengchuan Zhang, and Deva Ramanan. Evaluating text-to-visur generation with image-to-text generation. arXiv preprint arXiv:2404.000 2024.
13

---

[55] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In ICLR, 2023.
[56] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209, 030003,0000000000000000000000333 2022.
[57] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
[58] Simian Luo, Yiqin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthesizing nigh-ressoluton images with stw-step inference. arXiv preprint arXiv:2310.043378 2023.
[59] Nanye Ma, Shangyuan Tong, Haolin Jia, Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang. Yandong Li, Tommi Jaakkola, Xuhui Jia, et al. Inference-time scaling for diffusion models beyond scaling denoising steps. arXiv preprint arXiv:2501.097. 2025.
[60] Qianli Ma, Haotian Zhou, Tingkai Liu, Jianbo Yuan, Pengfei Liu, Yang You, and Hongxia Yang. Let's reward step by step: Step-level reward model as the navigators for reasoning. arXiv preprint arXiv:2310. 10080, 2023.
[61] Oscar Mañas, Pietro Astolfi, Melissa Hall, Candace Ross, Jack Urbanek, Adina Williams, Aishwarya Agrawal, Adriana Romero- Soriano, and Michal Drozdzal. Improving text-to-image consistency via automatic prompt optimization, 2024. URL https: / /arxxvv.vvvv org/ abs/
2403.17804.
[62] Zichen Miao, Zhengyuan Yang. Kevin Lin, Ze Wang. Zicheng Liu, Lijuan Wang, and Qiang Qiu. Tuning timestep-distiledd diffusion model using pairwise sample optimization. arXiv preprint arXiv:2410.03190 2024.
[63] Jesse Mu, Xiang Li, and Noah Goodman. Learning to compress prompts with gist tokens. Advances in Neural Information Processing Systems, 36: 36:19327-19352 2023.
[64] Niklas Muennighoff, Zitong Yang. Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling. arXiv preprint arXiv:2501. 2025.
[65] Beatrix Miranda Ginn Nielsen, Anders Christensen, Andrea Dittadi, and Ole Winther. Diffene: Variational diffusion with learned encoder. In The Twelfth International Conference on Learn- ing Representations, 2024. URL https: /ttttpss:/opprreeeeeeee
eepperreeeee...m net/forum?id-Snry 1bQTTG
[66] Zachary Novack, Julian McAuley, Taylor Berg- Kirkpatrick, and Nicholas Bryan. Ditto-
2: Distilled diffusion inference-time t-optimization for music generation. arXiv preprint arXiv:2405.2022 2024.
[67] Zachary Novack, Julian McAuley, Taylor Berg- Kirk-iiika, and Nicholas J. Bryan. Ditto: Diffusion inference-time t-optimization for music generation, 2024. URL https://arxiii org/ abs//401...... 12179.
[68] Owen Oertell, Jonathan D Chang, Yiyi Zhang, Kianté Brantley, and Wen Sun. RI for consis- tency models: Faster reward guided text-to-image generation. arXiv preprint orXiv:2404.03673 2024.
[69] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo. Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.077 193, 2023.
[70] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference, 2021. URL
https://arzii: http:////xwvvo.r//bbbbe wwtg//a////ae/ /99.2 02762.
[71] Dong Huk Park, Samanch Azadi, Xihui Liu, Trevor Darrell, and Anna Rohrbach. Benchmark for compositional texi-to-image synthesis. In NeurIPS Datasets and Benchmarks Track, 2021.
[72] Jason Phang, Yi Mao, Pengcheng He, and Weizhu Chen. Hypertuning: Toward adapting large language models without back-propagation. In ICML., pages 27854-27875, 2023.
[73] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis, 2023.
14

---

[74] Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text- to-image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.037.9, 2023.
[75] Mihir Prabhudesai, Russell Mendonca, Zheyang Qin, Katerina Fragkiadaki, and Deepak
Pathak. Video diffusion alignment via reward gradients. arXiv preprint urXiv:2407.08737, 2024.
[76] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICMI., 2021.
[77] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model secretly reward model. NeurIPS, 2023.
[78] Yuxi Ren, Xin Xia, Yanzuo Lu, Jiacheng Zhang, Jie Wu, Pan Xie, Xing Wang, and Xuefeng Xiao. Hyper-sd: Trajectory segmented consistency model for efficient image synthesis, 2024.
[79] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In CVPR, 2022.
[80] Litu Rout, Yujia Chen, Nataniel Ruiz, Abhishek Kumar, Constantine Caramanis, Sanjay Shakkottai, and Wen-Sheng Chu. Rb-modulation: Training-free personalization of diffusion models using stochastic optimal control. arXiv preprini arXiv:2405. I 17401, 2024.
[81] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Wei Wei, Tingbo Hou, Yael Pritch, Neal Wad- hwa, Michael Rubinstein, and Kfir Aberman. Hyperdreambooth: Hypernetworks for fast personalization of text-to-image models. CVPR, 2024.
[82] Subham Sekhar Sahoo, Aaron Gokaslan, Christopher De Sa, and Volodymyr Kuleshov. Dif- fusion models with learned adaptive noise. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https://opent:i www.rrreiiie..ee/ornnnnn. id-1oMa9944p8.
[83] Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint XXiv:2311.11044 2023.
[84] Charlie Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute opti- mally can be more effective than scaling model parameters. arXiv preprint arXiv:2408.03334 2024.
[85] Jiaming Song, Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. In ICLR, 2021.
[86]
Yang Song. Jascha Sohl- Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021.
[87] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In ICMI., 2023.
[88] Aravindan Sundaram, Ujjayan Pal, Abhimanyu Chauhan, Aishwarya Agarwal, and Srikrishna Karanam. Cocono: Attention contrast-and-complee. for initial noise optimization in text-to- image synthesis. arXiv preprint arXiv:241 1. 16783, 2024.
[89] Wenpin Tang. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond, 2024. URL. https: //arxiaa.vvvv///////// org/abs/ abs//403, /2403.06279.
[90] Zhiwei Tang, Jiangweizhi Peng, Jiasheng Tang, Mingyi Hong, Fan Wang, and Tsung-Hun Chang. Inference-tin time alignment of diffusion models with direct noise optimization. arXiv preprint arXiv:2405. 1888/, 2024.
[91] Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine- tuning of continuous-time diffusion models as hatropy-regularized control, 2024. URL
bttps://arx:r http//rxiiv.g/aaaaaa///// herg/bb//4400 22020 15194.
[92] Masatoshi Uehara, Xingyu Su. Yulai Zhao, Xiner Li, Aviv Regev, Shuiwang Ji, Sergey Levine, and Tommaso Biancalani. soward-guided iterative refinement in diffusion models at test-time with applications to protein and dna design, 2025. URI. https: /ttp///xvvv/////////////////aev/ wtwpii..rggbbbs///beee 2502. 14944.
15

---

[93] Masatoshi Uchara, Yulai Zhao, Chenyu Wang, Xiner Li, Aviv Regev, Sergey Levine, and Tommaso Biancalani. inference-time aligmment in diffusion models with reward -uided generation: Tutorial and review, 2025. URL. hnttps://arxxt org/ aataaaas//255 .9685..
[94] Siddarth Venkatraman, Mohsin Hasan, Minsu Kim, Luca Scimeca, Marcin Sendera, Yoshua Bengio, Glen Berseth, and Nikolay Malkin. Outsourced diffusion sampling: Efficient posterior inference in latent spaces of generative models. arXiv preprint arXiv:2502.06990 2025.
[95]
Johannes Von Oswald, Christian Henning, Benjamin F Grewe, and João Sacramento. Continual learning with hypernetworks. arXiv preprint arXiv:1906.00695, 2019.
[96] Andrew Wagenmaker, Mitsuhiko Nakamoto, Yunchu Zhang, Seohong Park, Waleed Yagoub, Anusha Nagabandi, Abhishek Gupta, and Sergey Levine. Steering your diffusion policy with latent space reinforcement learning, 2025. URL. https http:///aavvvvv.eeeaeee.e.aaa.la org/ aa/abb//22/ /2006. 15799.
[97] Bram Wallace, Akash Gokul, Stefano Ermon, and Nikhil Naik. End-to-end diffusion latent optimization improves classifier guidance. In ICCV, 2023.
[98] Bram Wallace, Meihua Dang. Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In CVPR, 2024.
[99] Kai Wang, Zhaopan Xu, Yukun Zhou, 7elin Zang. Trevor Darrell, Zhuang Liu, and Yang You. Neural network diffusion. arXiv preprint 44X1v:2402.13144, 2024.
[100] Zijie J Wang. Evan Montoya, David Munechika, Haoyang Yang. Benjamin Hoover, and Duen Horng Chau. Diffusiondb: A large-scale prompt gallery dataset for text-to-image generative models. arXiv preprint arXiv:2210. 10..2296, 2022.
[101] Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of lext-to-imaae synthesis. arXiv preprint arXiv:2306 09341, 2023.
[102] Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, and Hongsheng Li. Better aligning text-to- image models with human preference. In ICCV, 2023.
[103] Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. In NeurlPS, 2023.
[104] Beichen Zhang, Pan Zhang, Xiaoyi Dong, Yuhang Zang, and Jiaqi Wang. Long-clip: Unlocking the long-text capability of clip. In European Conference on Computer Vision, pages 310-325. Springer, 2024.
[105] Chris Zhang, Mengye Ren, and Raquel Urtasun. Graph hypernetworks for neural architecture search. arXiv preprint arXiv:18/0.055455 2018.
[106] Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreason- able effectiveness of deep features a perceptual metric. In CVPR, 2018.
[107] Sixian Zhang, Bohan Wang, Junqiang Wu, Yan Li, Tingting Gao, Di Zhang, and Zhongyuan Wang. Learning multi-dimensional human preference for text-to-image generation. In CVPR, 2024.
[108] Yinan Zhang, Eric Tzeng, Yilun Du, and Dmitry Kislyuk. Large-scale reinforcement learning for diffusion models. In ECCV, 2024.
[109] Zhenru Zhang, Chujie Zheng, Yangzhen Wu, Beichen Zhang, Runji Lin, Bowen Yu, Dayiheng Liu, Jingren Zhou, and Junyang Lin. The lessons of developing process reward models in mathematical reasoning. arXiv preprini arXiv:2501.0733
2025.
[110] Zikai Zhou, Shitong Shao, Lichen Bai, Zhiqiang Xu, Bo Han, and Zeke Xie. Golden noise for diffusion models; A learning framework. arXiv preprint arXiv:2411.099 09502, 2024.
16

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract clearly demarcates the aspirational goals for the field along with the results targeted by the paper and the theory/results achieve the conrete claims made in the abstract
Justification:
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, there is clear discussion of limitations of the both the theory and experiments in the paper.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Yes, both the main paper and the appendix contains a clear discussion of the assumptions made to achieve the theoretical results shown in the paper.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Yes, the main paper already provides the most crucial details required reproduce the results in the paper. The supplementary material provides complete clarity. Additionally, we utilize only open-source models and data for our experiments, enhancing reproducing further.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: Yes, we released the entire codebase, training scripts and data to reproduce our results under https: /ttp::///ttttte hht.tttt.......g/ com/ httm//Ellaanablll//ppeeeee
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes. complete experimental details for the training and test are provided in the supplementary material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: For the most significant evaluations (e.g. Tab. 1), results are averaged over seeds as is standard practice for the GenEval benchmark.
17

---

8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: Yes, main experiments were performed on academic compute budget of 6H100 GPUs for days, which is clearly stated in the experiment section.
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http//eeurspp.ppeee httppp/ctt//tttictccciiiiiaoaaaa
Answer: [Yes]
Justification: All the authors have read the NeurlPS Code of Ethics and agree that all aspects our research conforms with the NeurIPS Code of Ethics.
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Image generation models have become widespread over the past 3 years and we believe that their overall societal impact is relatively well-understood at this point. We do not foresee major additional societal impact specifically from our work. We discuss these details the appendix.
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our method aims to improve the aligmment of existing open-soures methods. To the extent that they are publicly available without major safeguards, we beleieve that there are no other major additional safeguards that are required for our work.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have appropriately cited existing assets.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We shall provide proper documentation for our released models and code and the time of the public release.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: No experiments with crowdsourcing experiments with human subjects done in the main paper.
Institutional review board (IRB) approvals or equivalent for research with human subjects
18

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No IRB approvals equivalent was required for the experiments in paper.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: No non-standard use of LLMs beyond typographical and grammatical edits of the paper.
19

---

Appendix
The Appendix is organized as follows:
Section provides all of our theoretical derivations.
Section outlines the implementation details.
Section presents further quantitative and qualitative analysis.
A
Theoretical Derivations
This section provides rigorous derivations for reward-tilted noise distribution and our tractable training objective. We include temperature parameter > 0 for completeness, though the main paper uses a 1.
A.1 Setup and Standing Assumptions
Let (xo) denote the standard Gaussian density I
1 - 2 1x001)1))
Px(x)) = (2 exp
(13)
Let ge Rd + Rd be the pre-trained distilled generator and Id - IR be the reward function.
Standing Assumptions.
Throughout this section, assume:
1. The generator go Rdd + Rd is measurable.
2. The reward function R + IR is measurable and Lex (e for our chosen temperature a >
3. For any x Range(go), the preimage set 90 ((x)) has a well-defined measure structure.
These assumptions are mild and realistic for neural network generators.
Pushforward Measure and Base Distribution. The base generator density basse is the density of the pushforward measure (ga);Po. where Po is the probability measure corresponding to po (xx). Formally, (ga) ) is defined such that for any Borel set A C Rdd
/
(())))P))(4)) = Po (A))
Po(Xo)dxx
(14)
Under our standing assumptions, the density (xuse()) can be written using the Dirae delta as:
6(x - go buse (x) = / Rd
(15)
Note that in the main text, with slight abuse of notation, we write (ge):Po instead of (g
KL Divergence. The Kullhack-Leibler (KL) divergence between two probability densities g(v) and p(v) is defined as:
DRL(gll) / q(v))og log q(v) dv, p(v)
(16)
provided the integral exists and is finite.
20

---

A.2 The Reward-Tilter Output Distribution
The primary goal is align the generator with the reward function r(x) by targeting output distribution p* (x) that upweights high-reward samples while maintaining similarity to the base distribution.
Definition (Reward- RReterd-Tilted Output Distribution). The target reward-tilted output density p" (x) is defined by upweighting samples from the base generator density p phise (x) according to the reward r(x):
1 hase (x) exp p" p'(x): = Z
(r(x)
(17)
where Z is the normalization constant ensuring p" (x) integrates to one:
z z=:= hase (x) (
r(x)
dx.
(18)
Under our standing assumptions, we have z- < 00. We denote p as the probability measure corresponding to p
Interpretation.
The temperature parameter a > 0 controls the strength of the reward signal:
When a + oo, we have (x) + p base (x) (no reward influence)
When -> 0, the distribution concentrates on high-reward regions
a provides a natural balance between reward optimization and staying close to the base distribution
Objective for Thne-Tunig Generator Parameters. If we aim to fine-tune the generator parameters from 0 to ø, leading to a new output density (x) (when input is from (xo))), a principled approach is to minimize the KL divergence DKL (p"lp""))
Proposition 2 (KL Objective for Generator Fine-tuning). Minimizing DKL (p*Ilp") with respect to the generator parameters & is equivalent to minimizing:
Jgon Dx... 11p Exmpp [r(x)]
(19)
Proof. Using the definition of p² (x) from Equation (17):
DKL. DKI. (2*11=")) = [ (x) log (x) dx
p (x(x)
f 10 (x))oog log
1 (x)
r(x) dx
pbme (x) exp
[ (xx) [log (x())
r(x)
+ log Z' dx
pruse (x)
DKL (101pp base)
Exupo (r(x)] + log z
(20)
Since log Z* is constant with respect to o, minimizing DKLL Pxl(p"1pp") is equivalent to minimizing Vge(())
Challenges with Direct Generator Fine-tuning. While Proposition provides a theoretically sound objective, directly optimizing it for distilled models poses significant challenges:
1. Intractable KI. term: Computing DKL blpse requires evaluating densities of high- dimensional neural network generators, which involves intractable Jacobian determinants
2. No continuous- time structure: Unlike full diffusion models, distilled generators often lack explicit SDE/ODE structure that would enable techniques from stochastic optimal control
21

---

3. Reward hacking: Without proper regularization, optimization can lead to adversarial exploitation of the reward model, generating unrealistic samples that achieve high reward scores
These challenges motivate our alternative approach of modifying the input noise distribution while keeping the generator fixed, which we develop in the next section.
A.3 The Reward- Tilted Noise Distribution
An alternative to modifying the generator ge is modify the input noise density po (Xo) while keeping ge fixed. We seek an optimal tilted noise density po (xo) such that its pushforward through 9o results in the target output density p" (x).
Normalization Constant in Noise Space. First, we show that the normalization constant Z from Equation (18) can be expressed as an integral over the noise space. Using Equation (15) for plsse (x) in the definition of z
r(x)
Z
=
exp
8(x
(x)))xx))) dx
Rd
Rd
r(x)
exp
8(x go (x(x)))xx po po(x))dx (Fubini's theorem)
Rd
Rd
r(g0(x)))
exp
Po (x(x)) (sifting property of Dirac delta)
JRd
a
(r(90(x)))
exp
Po (20xx)xx ddx..
(21)
JRd
Definition 2 (Tilted Noise Distribution). The rilted noise density ph (Xb) is defined as:
1 po (xx) exp (r(9)(X)))
(xo)
(22)
7.
where Z* is the normalization constant from Equation (18), which by Equation (21) can be computed noise space.
Theorem 3 (Properties of the Tilted Noise Distribution). Let pū (xx) be the tilted noise density defined in Definition and P be the corresponding probability measure. Under our standing assumptions:
1. Pushforward Identity: The density of the pushforward measure (go);PT is p (x)).
2. KL Projection: The density pā(xx) uniquely minimizes DRL l(po) among all noise den- sities (x(x)) such that the density of (go);Qo (where Qo is the measure for qo) equals pf (x).
Proof. Part 1: Pushforward Identity. We need to show that (go); Po has density p² p'(x). For any bounded measurable set A C Rdd, we have:
(())) Po (A)) f ug (A)
Pû (x(xx))ee (xo)
/ 97 ((A)(A) Z po exp
1
((oo(x)o) dxo.
(23)
To evaluate this integral, we use the fundamental property of pushforward measures. measurable function h: R2 + R:
For any
[ (o))
f hhgo Po
(24)
Applying this with h(x)= = 1 1x(x) exp ((x)))
1 exp
(((00(z)))
Po Po(x)) / exp
r(x) a
p base base (x)dxx.
(25)
a
22

---

Substituting back into Equation (23):
J exp
r((x)
(A))))
buse (x)
p
Z
a
/
1
r(xx
D
base
(x)exp exp
dx
7.
p" (x)dx.
(26)
Since this holds for all measurable sets A, the pushforward Po has density p" (xx).
Part KL Projection Characterization Consider the constrained optimization problem: min (oo.(oll1)) subject to (g has density p'
(27)
We use the method of Lagrange multipliers. Introduce a multiplier function À IRd - R and consider the functional:
L(go, A) = f (x(x)) log (x(x)) (xo) =
dxo + / A(x) (x)) Poo (x)) dx,
(28)
where Poo (x) the density of (go):Qu- For the constraint term, we can write:
X(x) (x)xx= J (00((x)))
(29)
using the change of variables formula for pushforward measures. Therefore:
L(go A) - / 40 (xx)
[log
(x(xx) (x()
((oo(x)))) dxo + / >(x)D² (x)dxx
(30)
Taking the functional derivative with respect to qo (xxo) and setting zero:
(x) = log go (Xo)
8C
+ A(go(xo)) 0.
(31)
po(xo)
This yields:
(o(x)) (x)) exp (xo)) 1]-
(32)
To satisfy the constraint, we need density Qo to equal p (x). Using Part in reverse, this happens when:
(x(xx) = z po (xo) exp (((g((x))
(33)
Comparing with the optimality condition, need:
(xgo(x))) r(go(xo)) log Z"
(34)
a
Setting A(x) = r(x) log Z* + 1, we obtain Qo = Pô- Uniqueness follows from the strict convexity of the KI. divergence in its first argument.
L
Objective for Learning the Tilted Noise Distribution. learn a parameterized noise density Ph (Xo) that approximates p (xo). we minimize DKL l.t)).
Proposition 4 (KL Objective for Learning Tilted Noise Density). Minimizing DKL (pg ((Ø)) with respect to & equivalent to minimizing:
J Jnnise (e) DKI. ORI((pllaa (p8l1p))
a Xompp (x()(())))]
(35)
23

---

Proof. Using the definition of (xx) from Equation (22):
/
(xo)
DRI.(p
Po o
(xo) log
dxo
pú (xo)
/ po (xx) log
Po (x))
dxo
(x())) exp
r(au(xo))
| 1 (xx)
Ph de log
(((x))
r(go(xu)) + log =" dxo
(x()
(po
E Exomp (r(go + log Z"
(36)
a
Since log Z- is constant with respect to o, minimizing (RI(PO (ln)) is equivalent to minimizing Jnoise (o).
A.3.1
Connection to Stochastic Optimal Control
We now show how our result connects to the sophisticated stochastic optimal control framework of Uchara al. [91] for fine-tuning continuous-time diffusion models, demonstrating that their approach naturally reduces our simpler result for one-step generators.
Continuous- Time Framework. Uchara et [91] consider the entropy cetropy-regularized control prob- lem:
max Ep.. (r(p)).(((((r))))))))))))))))))) aEpe. [ 20
T [|u(t, x;) dt + log v(xo)
(37)
pu(xo)
where pu, is the path measure induced by the SDE with drift f(t, x) + u(t, x) and v the initial distribution optimize.
Reduction One-Step Generators. For a one-step generator x = go (xo), the stochastic process degenerates:
The evolution deterministic: Xy go(X))
No drift control is needed: optimal
Only the initial distribution requires optimization
The objective reduces to:
max Exp (r(go (xx))) a DKL(U|lpo)
(38)
Optimal Initial Distribution.
According to their Corollary 2, the optimal initial distribution is:
v (x.) exp(vg (xo)) (xo)/x) (x()))
(39)
Z
where uf (Xn) is the value function at time t = O.
Value Function for Deterministic Generators. From their Lemma the value function satisfies:
(Feynman- Kac formulation),
exp (xo)/aa = Epo.. exp (((xT)
(40)
For the deterministic generator go:
Elexp(r(xx)/a))xol) = Elexp(r(gg (xo))/a)
= exp(r(gn(xo))/a) (deterministic given xo)
(41)
Therefore: (xo) r(go (xo)).
24

---

Final Result and Validation.
Substituting back into the optimal distribution formula: v" v'(xn) exp(r(go(xx)) (x)))/) po(x))
(42)
zz
where 2 = exp(((90(x0))) (xo) )xxo
This is precisely our p in Definition 2. This alignment between the two frameworks is significant, as it confirms that:
1. Our direct variational approach and the general stochastic control theory yield the same optimal noise distribution.
2. This equivalence arises because for one-step generators, the continuous time framework naturally collapses to our setting, with their value function vn simplifying the composed reward r go.
This connection not only validates our result but also situates it as an important special case within the broader theory of stochastic control.
A.4 Tractable KL Divergence for Noise Modification
We derive a tractable expression for DKL (pollp)) [[o)) where P6 is the density of modified noise &o- Ta (Xo) with To (x,) = X + ff (xo). This derivation involves the change of variables formula, simplification of Gaussian log- IDg-PDF terms, and application of Stein's Lemma.
Setup and Minimal Assumptions.
Let Te Rd + Fdd be the residual transformation:
To (xo) = X + fx()))
(43)
where Lo Rd R& is a learned perturbation function with Jacobian Jj, (xo)
Ofe(Ko)
8x.
Assumption (Regularity Conditions). We assume:
1. fe is continuously differentiable
2. T is a global diffeomorphism (invertible with continuous derivatives)
3. fo satisfies the regularity conditions for Stein's lemma: B[]lf((xo]11] oo for X ~ I),
and
[[|[x]| fl
Sufficient Condition for Global Diffeomorphism. While Assumption requires To to be a global diffeomorphism, we provide practical sufficient condition:
Lemma 5 (Lipschitz Condition for Invertibility). If fo is L-Lipschitz continuous with L < 1, then To is global diffeomorphism.
Proof. Bi-Lipsccitz bounds: for any Xo. x/i
(xx) To (x6) < I - x/1 + (xf - (x(x)) $ (1 4) ||x xíl.
(44)
(xo) -
> | - x61 - (xx) - (x(x) 2 (1 - L) ||oo x
(45)
Hence T, is injective. For surjectivity, fix any target and define Gy y y-fzzz) contraction with constant L < 1. By Banach's fixed-point theorem there exists a unique z" with z (z*). i.., To (z() y. Finally, Jr, (Xu) I Jo (Xo) is invertible for all xy (its smallest singular value least 1 T, > 0), and the inverse C* by the inverse function theorem. Thus Tę is global C diffeomorphism.
KL Divergence via Change of Variables. Under Assumption 1, we apply the change of variables formula. The KL divergence is:
log Po (Xx))
DKL (pollpo) ExOrpO
(46)
po (Xo)
Po 2 (xo))
Exorpo
log
(47)
po (T (xn))
25

---

By the change of variables formula:
Po (xx()))) = po(x)) det (xo))
(48)
Since Jr. (xo) 1 + Jzs (xn), substituting into Equation (47):
(P 11pp))) Exo=po [log (Xo) - loggoo (xo)) log det + Jr (xo))]
(49)
Specialization Gaussian Base Distribution. For po(o)) = -N(0, 1). the log-density difference simplifies:
1
logPo(Xo) (ogPo(Ts(xo)) -- |xoo|
[x
(50)
2
1
|x0||
+ 2 xxo fo(xx)11
(51)
= (xo) + 2
(52)
Substituting Equation (52) into Equation (49):
(pL (po) = Exx~N(o.)) x xxxo(x))= + [1x((x0)]11 log det(/ + Js, (xo))
(53)
Application of Stein's Lemma. applies:
Under the regularity conditions in Assumption 1, Stein's 's lemma
Lemma (Stein's Lemma for Vector Fields). Let x NN(,,1) and h Rd Rd satisfy
[]][((x)] co and [B [|()]]| < co. Then:
E [[x] = ([Tr(](
(54)
Applying Lemma 6 to Equation (53), we obtain:
- [1(x((x))1 + (x)),, - log [og]des det(! + (xx))]
1
DKL (Ipo) = Exoo~po
(55)
This is exactly the expression referenced in the main text.
Nog-Determinant Approximation Analysis. Let E(A) Tr(A) log /det(I + A)]. Then Equation (55) can be rewritten as:
DKLlplllp Exo~po [ (1 &E(JJ (x01))
(56)
To simplify this expression, we analyze the error term E(J (xo)). The following theorem provides a bound on this term under Lipschitz assumption on for
Theorem 7 (Bound on Log- Determinant Approximation Error). Let A Jjs (xx) be the d d Jacobian matrix of fo (xo). Assume fe is L--iissciizz continuous, such that its Lipschitz constant This implies that the spectral radius p(A) I L < 1. Then, the error term E(A)
Tr(A) - log det(1 A)] is bounded by:
((AA) S log(1 L) L)
(57)
Proof. Since fo is T-Lipschitz, the spectral norm of its Jacobian satisfies ||A||2 < T. This implies the spectral radius p(A) & ||Al|2 I I. < 1, ensuring all eigenvalues (A( satisfy (>((A) < 1.
Since 1 + >x(A) > 0 for all i, we have det(1 + A) > 0, so log det(1 + A)I log det + A).
For p(A) < 1, the matrix logarithm series converges;
1-1)
log do + A)
Tr((a)
(58)
k=1
26

---

Therefore:
(-1)6-1
E(A) -r(A)-
Σ
Tr(Ak)
(59)
k
k=1
(-1)
Σ
Tr(
(60)
k
2=
Taking absolute values and using [Tr(A*) - d. p(A)* <d L.:
d L
(((A)
(61)
2=
I
d
Σ h-1
h - L
)
(62)
d( log[1 - L) -
(63)
Practical Approximation and Final Objective. Theorem 7 shows that if the Lipschitz constant of fo is sufficiently small (specifically, L 1), the error term |E(A) bounded. For small I., log( L) L T,2 /2, making the bound approximately d.T.2 Thus, the expected error Expepoo [ (xo))] becomes negligible if is kept small. Under this condition, we can approximate the KL divergence with:
DKL (pūl so Extrpo
(64)
This approximation simplifies the KL divergence term in our objective to computationally tractable L2 penalty on the magnitude of the noise modification fo (Xo).
Integration with Main Objective. Combining our approximation with Proposition 4, and sub- stituting Equation (64) into our initial noise modulation objective, we arrive at the final loss to minimize:
1 ((0 f6(xo)))
nnois (Q) Exor~po 2/
(xx))1
(65)
This objective balances reward maximization against the KL regularization term, providing princi- pled and computationally tractable approach to learning the reward-tilted noise distribution.
Practical Implementation Considerations. The validity of our approximation depends on main- taining small Lipschitz constants. In practice, this is supported by:
Initialization: Setting fa(-) ensures {(A) = initially
2. Regularization: The term (x)))l naturally penalizes large perturbations, helping maintain small eigenvalues of J
While we do not explicitly enforce L < 1 during training, these practical measures help maintain fo in regime where our approximation remains accurate throughout the optimization process.
B Experimental and Implementation Details
In this Section we report the details for all of our experimental results. We mainly use the SANA- Sprint 0.6B [13] model, and train it using one-step generation. Additionally, we use the default guidance scale of 4.5 for all experiments. After training, we evaluate our models using different amounts of NFEs with one forward pass of Noise Hypernetwork beforehand.
27

---

LoRA parameterization We parameterize our noise hypernetwrkk fo with LoRA weights on top of the base distilled generative model. We found this to be important mainly to reuse the conditional pathways learned by the base model. This is especially important for complex conditioning, like text. Without this paramertization, which we also explored initially, we found it difficult for the noise hypernetwork to learn an effective conditioning with limit data. While larger-scale training could be a solution to this, we found this LoRA parameterization to be an efficient solution. For a condition independent reward, c.g. the redness one, it is less important to choose such parametrization.
Initialization As described in Section 3.2, we initialize the noise network to output fa(+) 0) at the start of training. We implement this by setting the output of the last base layer to O and initializing the LoRA weights of the second LoRA weight matrix (also reffered to B) to 0. This effectively initializes f+(+) 0). For stable training, this initialization is important as the model go go(So (x) + xo)) generates meaningful images at the start of training. In that way La only needs to leamn how to refine Xo.
Memory efficient implementation. Section 3.2, we train our noise hypernetwork Sa as special LoRA version of our base model go, which ignores the last layer of the base model. As visualized in Figure 2, we only need to keep the base model in memory once. Thus, the GPU memory overhead is just the added LoRA weights Φ Additionally, we employ Pytorch Memsave [7] to all models, which further reduces the needed GPU memory during training enabling us to use larger batch sizes. We run all experiments in bfloat16. Additionally, we can leverage gradient checkpointing on the first call of the model with activated LoRA parameters to further reduce memory. We use this for our FLUX-Schnell training.
B.1 Redness Reward
For the Redness Reward, we use SANA-Sprint 0.6B [13] as the base model. We train the model with the redness reward
1 (x' + x²)),
r(x)
100 (x"
2
where x" denotes the i-th color channel of x. We use the same amount of LoRA parameters for fine-tuning and noise hypernetwork training. In general, we keep the hyperparameters for our comparison between fine-tuning and noise hypernetwork training exactly the same. Due to the sake of illustration, we lower the learning rate for fine-tuning in this case as otherwise the model collapses to generating pure red images after a few training steps. We train on 30 prompts from the GenEval [26] promptset and evaluate a dolphin", "A photo of a train", "A photo of a car"]. After each epoch on the 30 prompts, we compute the redness reward as well as "imageness score" for each of the 4 evaluation prompts and average. For the imageness score, we use the Reward [103] numan-preferees reward model it was shown to correctly quantify prompt-following capabilities. We provide the full hyperparameters in Table This experiment was conducted on H100 GPU.
the four unseen prompts ["A photo of parrot", "A photo of
Table 3: Hyperparameters for the Redness Reward setting
Fine-tuning
Noise Hypernetwork
Model
SANA-Sprint [13]
SANA-Sprint [13]
Learning rate
le 4
le 3
GradNorm Clipping
1.0
1.0
LoRA rank
128
128
LoRA alpha
256
256
Optimizer
SGD
SGD
Batch size
3
3
Training epochs
200
200
Number of training prompts
30
30
Image size
1024 x 1024
1024
1024
28

---

B.2 Human Preference Reward Models
For our large-scale experiments, we consider SD-Turbo [83] and SANA-Sprint [13] as our two base models. For SD- Durbo we generate images in 512 x 512 while for SANA-Sprint we generate them of size 1024 x 1024. The training for the noise hypernetwork is done using -70k prompts from Pick-a-Picv2 [48]. T21-Compbench train set [37], and Attribute Binding (ABC-6K) [25] prompts. As the reward we follow ReNO [22] and use combination of uuman-preference trained reward models consisting of ImageReward [103], HPSv2.1 [101], PickScore [48]. and CLIP-Score [38]. To balance
these, we weigh each reward model with the same weightings as proposed ReNO [22] and employ
them with the following implementation details. All training runs were conducted on 6 H100 GPUs.
Human Preference Score v2.1 (HPSv2.1) HPSv2.1 [101] is an improved version of the HPS [102] model, which uses OpenCLIP ViT-H/14 model and is trained on prompts collected from Diffu- sionDB [100] and other sources.
PickScore PickScore also uses the same ViT-H/14 model, however trained on the Pick-a-Pic dataset which consists of 500k+ preferences that are collected through crowd-sourced prompts and comparisons.
ImageReward ImageReward [103] trains MLP over the features extracted from a BLIP model [51]. This is trained on dataset of images collected from the DiffusionDB [100] prompts.
CLIPScore Lastly, we CLIPScore [32, 76], which not designed specifically as a human preference reward model. However, measures the text-image alignment with a score between and 1. Thus, it offers way of evaluating the prompt faithfulness of the generated image that can be optimized. We use the model provided by OpenCLIP [38] with a ViT-H/14 backbone.
Table 4: Hyperparameters for the Human-preference Reward setting
Noise Hypernetwork
Fine-tuning
Noise Ilypemetwork
Noise llypemetwork
Model
SD- Turbo [83]
SANA-Sprint [13]
SANA-Sprint [13]
FLUX Schnell
Learning rate
lc
3
GradNorm Clipping
10.0
1.0
1.0
10.0
LoRA rank
128
128
128
128
LoRA alpha
128
256
256
5 * 128
Optimizer
SGD
SGD
SGD
SGD
Batch size
18
18
Accumulation Steps
3
Training Epochs
Rs 25
25
25
Ro 25
Number of training prompts
nd 70k
2d 70k
R2 70k
RJ 70k
Image size
512 x 512
1024 x 1024
1024 x 1024
512 512
GenEval Our main evaluation metric is GenEval, object-focused framework introduced by Ghosh et al. [26] for evaluating the alignment between text prompts and generated images from Text-to-Image (T2I) models. GenEval leverages existing object detection methods to perform fine-grained, instance-level analysis of compositional capabilities. The framework assesses various aspects of image generation, including object co-occurrence, position, count, and color. By linking the object detection pipeline with other discriminative vision models, GenEval can further verify properties like object color. All the metries on the GenEval benchmarks are evaluated using MaskFormer object detection model with Swin Transformer [57] backbone. Lastly, GenEval is evaluated over four seeds and reports the mean for each metric, which we follow. Note that our BLUX-Schhnel differ from the ones in Eyring et al. [22] as we use bfloat16 instead of float16.
B.3 Test-time techniques
For ReNO [22], we use the default parameters as described in their paper with 50 forward passes for one image generation. For Best-of-N [44] we use N = 50 with the same reward ensemble for a fair comparison. For L.LM-based prompt optimization [4, 61], we use the default setup from the MILS [4] repository (https:/ /////gtttuuu http///cceeboeeeeeerrr//....... bltb//llbbb//loloo//o/ main/main
29

---

py) with local Llama 3.1 8B Instruct as the LLM. The time reflected Table I reflects these local LLM calls. Note that we left the GPU memory to just the base image generation model. We modify the hyperparameters to 5 prompt proposals for each LLM call and 10 iterations, such that we also end up with 50 image evaluations for fair comparison.
C Additional results
In this section we report additional quantiative ablation results and further qualitative results.
C.1 Additional Benchmarks
Here, we report further results two more benchmarks commonly employed in the evaluation of T2I generation. Note that again, none of the prompts in the used benchmarks are part the training data, showcasing the generalizability of the Noise Hypernetwork to unseen prompts and also that our optimization objective through human- cemaan-reference reward mdoels disentangled from these benchmarks.
T2I-CompBench. T2I- CompBench is comprehensive benchmark proposed by Park et al. [71] for evaluating the compositional capabilities of text-to-image generation models. We evaluate on the Attribute binding tasks, which includes color, shape, and texture sub- categories. where the model should bind the attributes with the correct objects generate the complex scene. The attribute binding subtasks are evaluated using BLIP-VQA (i.e., generating questions based the prompt and applying VQA the generated image). We perform these evaluations on the validation set of prompts and results are shown in Tab. and observe consistent improvements across steps and categories.
Table 5: Quantitative Results on T2I- D2I-CompBench. The Noise Hypemetwork consistently improves performance.
SANA-Sprint 0.6B [13]
NFEs
Color
Shapet
Texturet
One-step
0.49
0.63
Noise Hypernetwork
2
0.75
0.53
0.64
Two-step
2
0.73
0.50
0.64
Noise Hypernetwork
3
0.76
0.53
0.64
Four-step
4
0.73
0.50
0.64
+ Noise Hypernetwork
5
0.76
0.54
0.65
DPG-Bench. We provide results on DPG-Bench [36] in Tab. 6. Broadly, while performance increases for all mod- note that
Table 6: DPG-Bench results for SANA-Sprint high- lighting generalization across inference timesteps of our Noise Hypernetwork.
els with increasing timesteps, the results for the four step SANA-Sprint model is nearly matched by the one-step model with our noise hypernetwork. We also note that the DPG-Bench score of 80.82 surpasses powerful models such as SDXL [73], Pix and only surpassed by much larger models such as SD3 [21], and Flux. Finally, also note that the preaan-reeeernce reward models that we utilize all have a CLIP/BLIP encoder that limits the length of the captions to < 77 tokens, which offers minimal scope of improvements for benchmarks involving much longer prompts that exceed this context window. Future reward models that either utilize different CLIP models (e.g. Long-CLIP [104]) or LLM-based decoders (e.g. VQAScore [54]) would enable improving prompt following of these models more dramatically the case of long prompts.
SANA-Sprint 0.6B [13] NFEs
DPG-Bench Scoret
One-step
77.59
Noise Hypernetwork
79.20
Two-step
79.07
+ Noise Hypernetwork
79.74
Four-step
79.54
+ Noise Hypernetwork
5
80.82
30

---

+ LoRA Fine-tune
SANA-Sprint
HyperNoise
"A red dog and green cat"
train and orange car"
"A majestic eagle soaring over snow- capped mountains at sunset, dramatic clouds in the background
cozy coffee shop interius with lighting, vintage books on shelves, steam rising from cappuccino"
"A cozy seading nock with velvet armchair, soft blanket, stack of books, and against the window"
Figure 6: Examples of artifacts introduced by directly Direct Fine-tuning diffusion models on rewards [14, 53, 74] for the same reward objective in comparison to Noise Hypemetwork training with same initial noise.
C.2 Diversity Analysis
We also investigate the impact of the diversity of the generated out- puts as the result of our hypernetwork. For this purpose, we generate 50 im- ages by varying the seed from the 553 prompts of the GenEval bench- mark. The average similarity of dif- ferent images for the same prompt measured using similarities from
Table 7: We measure the average LPIPS and DINO similarity scores over images generated for 50 different seeds for the 553 prompts from GenEval.
LPIPS
DINO 1
SANA-Sprint
0.608 +0.074
0.780 #0.103
+ Noise HyperNetwork
0.592 +0.059
0.825 +0.090
31

---

LPIPS [106] and DINOv2 [69] embeddings. The results in Tab. indicate that the noise hypernetwork hacking" and broadly, the diversity of the generated
does cause any collapse due images is in the same ballpark as the base model.
C.3 Multi-step analysis
Here, in addition the main Table 2, we analyze the behavior of Noise Hyper-
Table 8: Mean GenEval results for SANA-Sprin high- lighting generalization across inference timesteps of Noise Hypernetwork.
networks when moving beyond the few- step regime of
steps. Remarkably, even when going up to 32 inference steps, we find that Noise Hypernetworks trained with the one-step generator, improve per- formance. We find that as we increase the NFEs, the added performance boost of the Noise Hypemetwork reduces. How- ever, note that underlying model SANA- sprint [13] was not trained to be used in the multi-step regime, but specifically for few-step generation.
SANA-Sprint [13]
NFEs
GenEval Mean
One-step
0.70
+ Direct fine-tune [74]
1
0.67
+ Noise Hypernetwork
2
0.75
Two-step
2
0.72
+ Direct fine-tune [74]
2
0.66
+ Noise Hypernetwork
3
0.76
Four-step
0.73
+ Direct fine-tune [74]
4
0.62
+ Noise Hypernetwork
5
0.77
C.4 Challenges with Direct Fine-tuning
Eight-step
8
0.74
+ Noise Hypernetwork
0.76
We also qualitatively illustrate the prob- lems with directly fine-tuning diffusion models on differentiable rewards in Fig- ure As visualized, there are drastic arti- facts introduced on the image which signif- icantly contribute to improving the reward scores. These artifacts the development of several regularization strategies to address these issues. However as explained in Section 2, in the few-step regime the KL regularization term to the base model is difficult to be made tractable and thus, to the best of our knowledge there exists no theoretical grounded approach learn the reward tilted distribution (Equation 3) with a one-step generator. The Noise Hypernework strategy on the other hand, ensures that the images remain the original data distribution with its principled regularization.
Sixteen-step
0.73
Noise Hypernetwork
17
0.75
Thirty-two-ster
32
0.71
+ Noise Hypernetwork
33
0.72
very similar the ones noticed in several works [14, 53] and require
C.5 LoRA Rank analysis
Here, we ablate the LoRA rank for both HyperNoise and direct fine-tuning on SprinSp We find that rank of 64 also seems to be sufficient to achieve almost the same improvements rank 128, while a lower rank seems not to be expressive enough. On the other hand, fine-tuning seems to be suffering from increased overfitting the reward.
C.6 Qualitative Results
We provide additional qualitative samples for the base SANA-Sprint result along with the generation with proposed noise hypernetwork in Figures 6 and 8. We broadly observe improved prompt following as well as superior visual quality in the generated images.
32

---

Table 9: GenEval results for HyperNoise on SANA-Sprint, showing generalization across timesteps.
Method
NFEs
GenEval Meant
SANA-Sprint (One-step)
0.70
LoRA-Rank 128 HyperNoise
0.75
LoRA-Rank 64 + HyperNoise
0.75
LoRA-Rank 16 + HyperNoise
0.71
LoRA-Rank 8 HyperNoise
0.70
SANA-Sprint (Two-step)
0.72
HyperNoise
0.76
SANA-Sprint (Four-step)
0.73
HyperNoise
0.77
HyperNoise (LoRA-Rank=64))
0.76
Table 10: GenEval results for direct LoRA fine-tuning on SANA-Sprint.
Method
NFEs
GenEval Meant
SANA-Sprint (One-step)
0.70
LoRA-Rank 128 LoRA fine-tune
1
0.67
LoRA-Rank 64 LoRA fine-tune
0.68
LoRA-Rank 16 LoRA fine-tune
0.65
LoRA-Rank + LoRA fine-tune
0.59
SANA-Sprint (Two-step) LoRA fine-tune
2
0.72
2
0.66
SANA-Sprint (Four-step)
4
0.73
LoRA fine-tune
4
0.62
33

---

+ Noise Hypernetwork
SANA -Sprint
"A waterfall crashing ento stones, belaw rainbou, with colorful nist'
"Guests watching the Northern Lights dancing above an ice hotel"
epic oil painting: silver knights charging across nud, red banhers weving, fiery eclipse averhaad, stone castle on cliffs'
Figure 7: More qualitative results on the human- prefeeeneee reward setting. Base SANA-Sprint compared to HyperNoise with same initial noise.
SANA-Sprint HyperNoise
SANA-Sprint
Figure 8: Non-cherry picked results on the human-preference reward setting. Base SANA-Sprint compared to HyperNoise with same initial noise.
34


Reviews:
[{"review_id": "5mCqO2JAHh", "reviewer": "Reviewer_RnrZ", "summary": "This paper addresses the problem of test-time scaling for distilled diffusion models, with the goal of quickly adapting pretrained image generators to align with user-specific preferences. Instead of relying on direct fine-tuning of the base model or performing costly per-sample test-time noise optimization, the authors propose a “pre-trained” LoRA-based Noise Hypernetwork that learns to modulate the input noise in accordance with a given reward function. By keeping the base model parameters fixed and learning only a lightweight noise transformation network, the method significantly reduces inference-time overhead while preserving or improving output quality. The key contribution lies in a theoretically grounded training objective that includes a tractable KL regularization in the noise space to ensure fidelity to the original model's distribution.", "questions": "1. The method requires training a separate hypernetwork for each reward function. While this is understandable, it limits generality—especially since simple reward functions (e.g., color biases) are not meaningful, and complex ones (e.g., composed human preference models like those in ReNO) tend to yield narrowly specialized models with limited generalization. It would be helpful if the authors could discuss the intended scope of applicability and limitations of the proposed framework.\n2. While ReNO requires expensive test-time optimization, it is end-to-end and can be applied on a per-sample basis. The proposed method, by contrast, incurs a one-time training cost. Could the authors elaborate on the relative advantages in low-resource vs. large-scale generation settings, and provide a cost-benefit comparison (e.g., amortized training time vs. repeated test-time inference)?\n3. In Figure 3, although the method preserves the global image structure better than LoRA fine-tuning, it appears that the main object in the image becomes progressively smaller during training. Does this suggest that the reward function (e.g., “redness”) is insufficiently informative, encouraging the model to optimize a narrow visual cue rather than a meaningful preference? This seems to suggest a form of reward hacking; could the authors comment on how robust the method is under poorly specified or degenerate rewards?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The theoretical motivation and derivation of the method are solid and rigorous, particularly the approximation of the KL divergence in the noise space and the use of the data processing inequality for distributional control.\n2. The core idea—to amortize test-time noise optimization via a pretrained hypernetwork—is novel, elegant, and highly practical.\n3. The KL divergence regularization is applied in a clever and tractable form, leading to better preservation of image quality while optimizing for reward-aligned outputs.\n4. The paper is well-written and well-structured, making the contributions easy to understand and follow.\n\nWeaknesses:\n1. The main comparison is against ReNO. While this is a reasonable choice given the shared setup, the paper could benefit from broader qualitative comparisons, e.g., prompting-based baselines or naive fine-tuning.\n2. The paper emphasizes inference-time efficiency, but provides limited discussion of the cost and scalability of training the noise hypernetwork itself, especially with large or complex reward models.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their thorough and constructive feedback. We're pleased that the reviewer found our theoretical contributions \"solid and rigorous\" and our core idea \"novel, elegant, and highly practical.\" Below, we address all the points raised.\n\n> \"... ***the paper could benefit from broader qualitative comparisons, e.g., prompting-based baselines or naive fine-tuning.***\n\nWe thank the reviewer for this valuable suggestion. To clarify, our method is the first to learn reward-optimized distributions for distilled diffusion models, which makes defining direct baselines challenging. However, we agree that a broader comparison against different test-time methods is beneficial here. Thus, we have now benchmarked our approach against two additional test-time techniques in the human-preference reward setting:\n1. **LLM-based prompt optimization** following [1, 2]. The prompt optimization is handled by an LLM (Llama 3.1 8B), which is tasked to optimize a score, where we employ the same reward ensemble as used for ReNO and Noise Hypernetworks. For a fair comparison, we benchmark it with a total of 50 evaluated images (same as for ReNO), where in one iteration, the LLM proposes 5 new prompts, and we run for 10 iterations. We used the parameters from the official MILS repository [2], modifying only the reward ensemble and number of generated images per prompt.\n2. **Best-of-N sampling** as e.g. used in [3]. Here we sample N=50 noises for each prompt and choose the best image based on, again, the used reward ensemble.\n#### Results on SANA-Sprint\n\n| Method                           | Inference Time (s) | Relative Speed (vs. Base) | GenEval Score |\n| :------------------------------- | :----------------: | :-----------------------: | :-----------: |\n| SANA-Sprint                      |        0.2         |          **1x**           |     0.70      |\n| **+ Noise Hypernetworks (Ours)** |      **0.3**       |         **1.5x**          |   **0.75**    |\n| *+ Best-of-N*                    |       ~15.0        |           ~75x            |     0.79      |\n| *+ Prompt Optimization*          |       ~115.0       |           ~575x           |     0.75      |\n| *+ ReNO*                         |       ~30.0        |           ~150x           |     0.81      |\n\n#### Results on SD-Turbo\n\n| Method                           | Inference Time (s) | Relative Speed (vs. Base) | GenEval Score |\n| :------------------------------- | :----------------: | :-----------------------: | :-----------: |\n| SD-Turbo                         |        0.2         |          **1x**           |     0.49      |\n| **+ Noise Hypernetworks (Ours)** |      **0.3**       |         **1.5x**          |   **0.57**    |\n| *+ Best-of-N*                    |       ~10.0        |           ~50x            |     0.60      |\n| *+ Prompt Optimization*          |       ~110.0       |           ~550x           |     0.59      |\n| *+ ReNO*                         |       ~20.0        |           ~100x           |     0.63      |\n\n- These results highlight the practical efficiency advantage of our method. Noise Hypernetworks achieves significant quality improvements with only a 1.5x latency increase, while alternative test-time methods require 50x to 575x more computation time. Notably, our approach achieves similar results to LLM-based prompt optimization while being ~380x faster, making it practical for real-time applications where other methods would be prohibitive.\n\n- Beyond direct comparison, our approach offers a unique advantage: it's fully compatible with existing test-time methods. Test-time approaches like ReNO could initialize from our network's predicted noise rather than random noise, potentially achieving better results or faster convergence by starting from a more promising region.\n\n**Fine-tuning comparison**: Finally, to demonstrate the importance of our method's formulation, we also benchmarked against direct LoRA fine-tuning (detailed in Section 4.1). This approach consistently fails to learn the reward-tilted distribution, as the absence of our proposed KL-regularization term leads to a divergence from the natural image manifold (Figure 3). We observe this same failure mode in the human-preference setting (Table 2), with visual examples of image degradation from reward-hacking provided in the Appendix (Figure 6). This underscores that simply fine-tuning with LoRA is insufficient and highlights the critical role of our proposed KL-regularization term in preventing reward-hacking and preserving image quality.\n\n---\n\n_**On Training Cost, Generality, and Amortization:**_\n\nWe thank the reviewer for raising these important points about our method's practical applicability. We'll address them together since they're closely related.\n\n- **Training Cost & Scalability:** First, some context on the computational requirements. Training our LoRA-based Noise Hypernetwork takes about 10 minutes on 1 H100 GPU for simple rewards (like redness) and 34 hours on 6 H100 GPUs for human preference rewards. While 34 hours sounds substantial, this is a one-time cost that's small compared to training the base diffusion models themselves. More importantly, this upfront investment \"compiles\" a complex iterative optimization process into a single efficient network that adds almost no latency at inference.\n\n- **Amortized Cost-Benefit:** The economics become compelling at scale. Our training uses 204 H100-GPU hours total. At inference, ReNO takes ~30 seconds per image on an A100 while inference with our Hypernetwork takes only ~0.3 seconds. Using a conservative 1:2 price ratio for A100 vs H100, our training investment equals 408 A100-equivalent hours. Since we save 29.7 seconds of A100 time per image, the training cost is fully recovered after generating about **50k images**. For any application generating hundreds of thousands of images, this represents massive long-term savings in both cost and latency.\n\n- **Scope and Generality:** This analysis highlights where our method shines: scenarios with a well-defined, valuable reward function that needs to be deployed at scale (e.g., human-preference reward models, a complex brand style, or a photorealism enhancer). While test-time optimization offers more flexibility for exploration, our approach is designed for the regime where a certain reward has been identified and now needs to be deployed efficiently at a larger scale.\n\nWe'll add this detailed cost-benefit analysis and scope discussion to the appendix of our revised manuscript.\n\n---\n\n> ***\"... it appears that the main object in the image becomes progressively smaller during training. Does this suggest that the reward function (e.g., “redness”) is insufficiently informative, encouraging the model to optimize a narrow visual cue rather than a meaningful preference? This seems to suggest a form of reward hacking; could the authors comment on how robust the method is under poorly specified or degenerate rewards?\"***\n\nWe thank the reviewer for pointing out the object shrinking behavior. This is actually a feature, not a bug, and shows why our approach is more robust than alternatives.\n\nThe \"redness\" reward we use here is intentionally narrow. A model trying to maximize redness could easily turn everything red and destroy the image completely. What's noteworthy is that our method doesn't do this. Instead, it finds a clever workaround: shrink the object so more red background shows through. This satisfies the reward while keeping the image intact.\n\nCompare this to direct LoRA fine-tuning, which breaks down and diverges from the natural image manifold and thus, from the base model (Figure 3). Our KL-regularized version finds a reasonable solution that games the reward in a harmless way. The ImageReward scores in Figure 4 confirm that image quality stays stable throughout training.\n\nThis behavior is exactly what we want from a robust system. Real-world reward functions are often imperfect or incomplete. When that happens, we need methods that bend rather than break. Our KL regularization provides exactly this by keeping the model close to its original distribution, preventing the kind of catastrophic failures that make direct fine-tuning impractical.\n\n---\n\n[1] Mañas et al. \"Improving Text-to-Image Consistency via Automatic Prompt Optimization\". TMLR 2024.\n\n[2] Ashutosh et al. \"LLMs can see and hear without any training\". ICML 2025.\n\n[3] Ma et al. \"Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps\". CVPR 2025."}, {"review_id": "wGor5Kl9Np", "reviewer": "Reviewer_YHFj", "summary": "We would often like to tweak a trained diffusion model so that it is more likely to produce certain kinds of outputs; for example, we might prefer if its outputs are somewhat more 'red' than usual. While this can in principle be done via prompting, the current excitement over test-time scaling has suggested it might also be done by spending more compute at inference time. One idea is to, at inference time, push a diffusion model towards certain kinds of samples by training the noise seed to reflect a reward model. While this works, solving an optimization problem like this at inference time can be somewhat time-consuming.\n\nThe authors propose to amortize this computation by training a \"noise hypernetwork\", i.e., to learn a reward-sensitive noise seed through additional training time. Importantly, this involves training an additional function, but not changing the weights of the trained diffusion model. They show that this approach works well in practice and has a reasonable theoretical basis.", "questions": "Can the authors link their core idea to work on the 'priors' of diffusion models? It seems like changing the prior (to reflect the reward function of interest) is precisely what the method does, since it doesn't actually change the downstream sampling process at all. The idea made me think of the (now kind of 'old') PriorGrad paper, although there are some key technical differences, like the fact that in that paper the changed prior also changed the training process.", "limitations": "This is obvious, but it may be worth mentioning that the proposed approach reduces flexibility in order to improve efficiency. If for some reason the user wants to use a variety of reward functions, one has to train a model for each. The inference time approach is much more costly, but (as a result) doesn't have this issue.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths.**\n- The paper is clear and well-written. The idea makes a lot of sense and has an appealing link to modifying the 'prior' of diffusion models, as well as previous work on spending test-time compute.\n- The authors do a lot of work to test that their idea performs well, and provide both numbers and qualitative examples to show that it does. I was pretty impressed by some of these demonstrations (e.g., Fig 1, Figs 3 through 5).\n- To validate their method, they do experiments to check for 'dumb' possibilities. For example, maybe noise hypernetworks help, but only about as much as direct fine-tuning. Table 2 shows that this isn't true. Table 1 makes a variety of important points that make sense; for example, spending lots of inference time compute a la ReNO still performs better than the noise hypernetwork approach, although not by an extreme amount. All of these experiments together paint a convincing picture that the method works well, even compared to many existing alternatives.\n- Some of the precise numerical claims are impressive, e.g., the NFEs statement on line 281.\n\n\n**Weaknesses.** I don't have much to say here. One thing that kind of bugged me was the presented motivation for this approach. The authors discuss excitement over improving performance by increasing inference time compute, which works but has the obvious drawback that inference time and cost go up. Then they say that they can mitigate this issue by spending that compute during training instead. But this logic is weird to me, since getting a model to do things efficiently by spending a lot of resources during training has been business as usual for a long time! It might be better to argue that, given how efficient the standard approach is, it's worth seeing if it can also be used to make the noise seed reward-sensitive.\n\nThe paper may contain some small typos. One I noticed is an apostrophe issue on line 336.", "was_revised": true, "final_justification": "The authors addressed the (admittedly minor) concerns I brought up related to novelty/motivation, the relationship between this work and diffusion model priors, and the flexibility vs efficiency tradeoff. I still feel like the paper is well-written and that the authors' proposed method has a solid quantitative performance.", "rebuttal": "We sincerely thank the reviewer for their thoughtful and encouraging feedback. We're delighted that the reviewer found our paper \"clear and well-written\" with an idea that \"makes a lot of sense\" and demonstrations that were \"pretty impressive.\" We appreciate the detailed evaluation and address the points raised below.\n\n**On Motivation and Framing:**\n\nWe thank the reviewer for this valuable perspective on our motivation. The reviewer is absolutely right that spending training resources to improve efficiency isn't new. What makes our approach distinctive is that we're specifically translating insights from successful test-time scaling methods into a post-training framework. Rather than designing optimization strategies from scratch, we're taking proven test-time techniques (like ReNO) and \"compiling\" them into efficient networks that preserve their benefits without the inference cost.\n\nThis represents a fundamentally different paradigm from standard training: we're adapting and amortizing successful test-time innovations rather than traditional end-to-end optimization. The key insight is that test-time scaling successes can be translated into post-training adaptations that maintain the core benefits while dramatically improving practical deployment.\n\n**On Connection to Diffusion Priors:**\n\nWe appreciate this insightful connection to diffusion priors. Our method does fundamentally modify the effective prior distribution of the diffusion process without requiring changes to the base model training or sampling procedure. This connects our work to prior methods that learn adaptive prior distributions for diffusion models. Our approach achieves prior modification through post-training adaptation while this line of work typically modifies the training process itself or requires architectural changes to the base model, e.g. [1, 2]. We'll incorporate this connection to adaptive prior methods in our related work section to better position our contribution within this line of work.\n\n_**On Flexibility vs. Efficiency Trade-off:**_\n\nThe reviewer raises an important consideration about our approach requiring separate models for different reward functions, which trades flexibility for efficiency compared to test-time methods. This trade-off becomes highly favorable for production scenarios with stable reward functions. We discuss this in detail in our response to Reviewer RnrZ, including a comprehensive cost-benefit analysis, which we will add to the appendix. In summary, that analysis shows:\n- **Scope:** Our method is ideal for production scenarios with a well-defined, valuable reward function that needs to be deployed at scale (e.g., human-preference models, brand styles, or photorealism enhancers).\n- **Amortization:** For complex rewards, the one-time training cost is fully recovered after generating approximately 50k images, leading to very significant long-term savings in both cost and latency.\n\nLastly, we thank the reviewer for catching the apostrophe issue on line 336 and will fix this along with a careful proofreading pass.\n\n[1] Lee et al. \"PriorGrad: Improving Conditional Denoising Diffusion Models with Data-Dependent Adaptive Prior\". ICLR 2022.\n\n[2] Bartosh et al. \"Neural Flow Diffusion Models: Learnable Forward Process for Improved Diffusion Modelling\". NeurIPS 2024."}, {"review_id": "cTHqqyRgy5", "reviewer": "Reviewer_sWqF", "summary": "The authors solve the critical problem of integrating test-time scaling knowledge into a model during post-training. The authors replace reward guided test-time noise optimization in diffusion models with a Noise Hypernetwork that modulates initial input noise. It demonstrates through extensive experiments significant enhancements in generation quality for state-of-the-art distilled models", "questions": "Refer to weakness", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "strength\n1. The authors predict an optimized initial noise for a fixed distilled generator with Lora, and theoretically learn through a tractable noise-space objective\n2. The authors benchmark the noise hypernetwork against established methods, primarily direct LoRA fine-tuning of the base generative model and demonstrate the effectiveness of the proposed noise network.\n\n\nweakness\n1. How do you avoid the risk of significant deviation from the original data distribution? I saw the ImageReward figure in the experiment section, can the authors elaborate more on it?\n2. Did the authors have the ablations on Hypernetworks's architecture? The authors now use Lora, and how does the number of rank or learnable prompts affect the results?", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their positive assessment and for their constructive technical questions. We are glad they recognized that our work \"solves the critical problem of integrating test-time scaling knowledge into a model during post-training\" with \"significant enhancements in generation quality.\"\n\nWe address their specific concerns below, including results from a new ablation study prompted directly by their valuable feedback.\n\n---\n\n>\"***How do you avoid the risk of significant deviation from the original data distribution? I saw the ImageReward figure in the experiment section, can the authors elaborate more on it?\"***\n\nWe thank the reviewer for this critical question about maintaining fidelity to the original data distribution. This is indeed a fundamental challenge when optimizing generative models with external rewards, and our framework addresses it through principled regularization mechanisms.\n\n• **The Core Challenge.** When fine-tuning generative models to optimize specific rewards, there is an inherent risk that the model will find solutions outside the natural data distribution by achieving high reward scores through unrealistic or artifact-heavy samples that deviate significantly from the training manifold. This distributional drift can lead to outputs that, while scoring well on the target metric, lack the quality and realism of the original model.\n\n• **Our Solution: Noise-Space KL Regularization.** We address this challenge through explicit regularization in our objective function (Equation 7). The KL divergence term $D_{KL}(p_0^{\\phi} \\|\\| p_0)$ directly constrains the learned noise distribution to remain close to the original Gaussian prior. By preventing the optimized noise from deviating too far from the initial noise, we mathematically constrain the model's outputs $g_\\theta(\\hat{x}_0)$ to remain close to the base data manifold.\n\n• **Quantitative and Qualitative Validation:**  The \"redness\" experiment provides compelling empirical proof of this mechanism. As the reviewer noted in **Figure 4**, direct fine-tuning (which lacks our regularization) leads to a catastrophic drop in the general ImageReward score as it over-optimizes for redness. This is visually confirmed in **Figure 3**, where direct fine-tuning deteriorates into unrealistic, monochromatic images, demonstrating a complete departure from the learned image manifold. In stark contrast, our method achieves the target reward while maintaining a high and stable ImageReward score, with generated images that preserve natural texture and semantic structure. This confirms our approach successfully optimizes _within_ the learned data distribution, rather than deviating from it. This principle of controlled optimization extends to our main results, as evidenced by the consistent quality improvements on the GenEval benchmark.\n\n• **Theoretical Guarantees.** This approach is grounded in the Data Processing Inequality (Equation 12), which provides a formal guarantee: minimizing KL divergence in noise space upper-bounds the KL divergence in the final image space. Crucially, unlike data-space KL terms that are intractable for distilled models, our noise-space formulation can be efficiently computed as an L2 penalty (Equation 10).\n\n• **Broader Validation.** This controlled optimization behavior extends beyond the redness experiment. Our GenEval results across multiple evaluation dimensions similarly show consistent quality improvements without the artifacts or unrealistic outputs typically associated with distributional deviation in reward-optimized models as visualized in Figure 6 (in the Appendix).\n\n---\n\n> ***LoRA Architecture Choice and Rank Ablation***\n\nWe thank the reviewer for this excellent question on our architectural choices. The selection of a LoRA-based architecture is a deliberate and critical decision, driven by the fundamental need to produce **prompt-dependent** noise modifications that can generalize to unseen prompts.\n\n**Prompt-Dependent Noise and Generalization**:\n- The central task of our Noise Hypernetwork, $f_\\phi$, is to predict an optimal noise perturbation $f_\\phi(x_0|c)$ that is conditioned on the semantic content of a given text prompt, $c$. The optimal noise required to improve an image of \"a pink elephant and a grey cow\" is fundamentally different from the noise needed to enhance \"a brain riding a rocketship towards the moon.\"\n\n- Training a standalone network to learn this complex mapping from scratch would be exceptionally difficult. To achieve robust generalization to the vast space of unseen, open-world prompts, such a network would require a prohibitively large and diverse training dataset of prompts and a corresponding long training time.\n\n**LoRA as the Principled Solution**: \n- Our choice of a LoRA-based architecture directly solves this challenge. Instead of attempting to learn text-conditioning from scratch, LoRA allows our Noise Hypernetwork to inherit and leverage the powerful, pre-existing text-conditioning pathways of the frozen base model. By adapting key components like the cross-attention layers, our network efficiently taps into the rich semantic understanding that the base model has already acquired. \n\n- This approach ensures that our predicted noise modifications are semantically aware and generalize effectively to novel prompts, without the need for an large-scale amount of training data. It is the most parameter-efficient and direct way to make the noise prediction $f_\\phi(x_0|c)$  truly prompt-dependent. During our initial development, we explored different text-conditioned architectures trained from scratch but found these struggled to learn any meaningful text conditioning on small-scale prompt sets.\n\n### **LoRA Rank Ablation**\n\nTo directly answer the reviewer's question on architectural hyperparameters, we conducted a new ablation study on the LoRA rank using the SANA-Sprint model.\n\n| Method                             | NFEs | GenEval Mean↑ |\n| ---------------------------------- | ---- | ------------- |\n| SANA-Sprint Baseline               | 1    | 0.70          |\n| Noise Hypernetwork (LoRA-Rank 128) | 2    | **0.75**      |\n| Noise Hypernetwork (LoRA-Rank 64)  | 2    | **0.75**      |\n| Noise Hypernetwork (LoRA-Rank 16)  | 2    | 0.71          |\n| Noise Hypernetwork (LoRA-Rank 8)   | 2    | 0.70          |\n\n**Analysis of Ablation Results:** The new ablation results validate that while our chosen rank of 128 offers robust, high-quality performance, a rank of 64 achieves identical results (0.75 GenEval score). This is a valuable finding, prompted by the reviewer's question, that demonstrates the potential for greater parameter efficiency. Meanwhile, performance degrades at lower ranks (8-16), which likely lack the capacity to model the complex prompt-conditional noise modifications required. We thank the reviewer for pushing us to explore this dimension. We will add this full analysis to the Appendix."}, {"review_id": "FizDd1DfCX", "reviewer": "Reviewer_sVrs", "summary": "The paper introduces a new method called Noise Hypernetworks, aiming to incorporate the advantages of test-time scaling into the training phase of diffusion models. Noise Hypernetworks learns a new distribution for the initial noise, producing modulated initial noise through a residual formulation. The loss function is based on theoretical results that aim to minimize the KL divergence between the initial noises while simultaneously maximizing the reward for the diffusion outputs. Experimental results demonstrate that Noise Hypernetworks can generate high-reward samples without compromising overall image quality.", "questions": "See weaknesses.", "limitations": "The authors do discuss the limitations related to the need for improved reward models.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths\n\nThe paper is well written and easy to follow. The authors offer comprehensive theoretical justifications for their approach to tuning the initial noise distribution, establishing a solid theoretical foundation. The illustrative experiments are interesting, providing an intuitive insight into the method's performance.\n\n### Weaknesses\n\n1. Since the reward model is applied to the output samples of diffusion models, training the Noise Hypernetwork still requires costly backpropagation through the sampling process. This may restrict the training of the Noise Hypernetwork to one-step generation models, as indicated in Table 2. What would happen if the Noise Hypernetwork were trained on multi-step generation process? Would it still achieve high performance?\n\n2. The performance may heavily depend on the quality of the distilled generative model $g_\\theta$.\n\n3. Is it unclear that whether altering the initial noise distribution could decrease the diversity of the generated samples.", "was_revised": true, "final_justification": "The majority of my concerns have been addressed through this rebuttal and responses to other reviewers, so I have decided to raise the score from 4 to 5.", "rebuttal": "We thank the reviewer for their thoughtful assessment and for recognizing that our work is \"well written\" with \"comprehensive theoretical justifications\" and \"interesting illustrative experiments.\" We appreciate their constructive feedback and address their specific concerns below.\n\n>***\"... What would happen if the Noise Hypernetwork were trained on multi-step generation process? Would it still achieve high performance?\"***\n\nThe reviewer raises an excellent question about multi-step training feasibility. While multi-step training is theoretically possible through direct multi-step backpropagation or techniques like Adjoint Matching [1], which can compute gradients through multi-step ODE/SDE solvers efficiently, it remains unclear whether direct fine-tuning approaches or our hypernetwork method would be more effective in such regimes. While Adjoint Matching can directly learn reward-tilted distributions for multi-step generators through tractable KL regularization, their SDE-based formulation requires continuous-time dynamics making it computationally expensive and inapplicable to step-distilled models.\n\n- **Addressing a Critical Gap for Distilled Models**: Our work specifically targets distilled models, which represent an increasingly important class of generators due to their efficiency. This addresses a fundamental limitation where existing reward optimization methods require multi-step sampling processes that don't exist in distilled models.\n\n- **Practical Advantages of Our Approach**: Beyond theoretical considerations, our method offers compelling practical benefits. While our training is performed on one-step distilled generators, **it generalizes remarkably well to multi-step inference** as demonstrated in Tables 2 & 8, achieving significant improvements when deployed with 2-16 sampling steps. This approach maximizes training efficiency while maintaining deployment flexibility across different quality-speed trade-offs.\n\n- **Computational Efficiency**: Our approach also offers computational advantages during training. Rather than requiring expensive multi-step backpropagation through sampling processes, we can train efficiently on distilled models and still achieve strong performance across different inference regimes. This makes reward optimization accessible for the fast, practical generators that are increasingly dominant in real-world applications.\n\n---\n\n> ***\"The performance may heavily depend on the quality of the distilled generative model\"***\n\nThe reviewer correctly identifies that our method's performance depends on the quality of the underlying base model. We acknowledge this fundamental limitation of any post-training alignment approach. However, our experiments demonstrate consistent relative improvements across different base model architectures and scales, suggesting that our approach successfully identifies and exploits suboptimal regions in the learned noise-to-image mapping regardless of the base model's absolute capability level.\n\n---\n\n> ***\"Is it unclear that whether altering the initial noise distribution could decrease the diversity of the generated samples.\"***\n\nWe thank the reviewer for this important consideration about diversity preservation. As detailed in Appendix C.2 (Table 7), we conduct comprehensive diversity analysis using LPIPS and DINO similarity metrics across 500+ prompts with 50 images per prompt. Our results show a nuanced diversity profile: while DINO similarity shows a modest decrease, LPIPS actually indicates a slight increase. Crucially, there is **no diversity collapse**. This controlled diversity trade-off is accompanied by meaningful quality improvements, representing principled optimization toward higher-reward regions rather than problematic mode collapse. This behavior is attributable to our KL regularization term that prevents the learned noise distribution from deviating too far from the original Gaussian prior.\n\n[1] Domingo-Enrich et al. \"Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control\". NeurIPS 2024."}]

Author rebuttal:

Meta-review: The paper proposes a hypernetwork to predict residual noise from the original noise in the diffusion model to optimise certain reward functions. Compared to previous noise optimisation methods, the main contribution is to train the hypernetwork to avoid test-time optimisation to improve inference efficiency. The authors provide a very good story for this in the context of test-time scaling, which is a currently important topic, and the proposed hypernetwork moves the computation from test-time to post-training. Other implementation techniques include using LoRA and step-distilled generators in the hypernetwork and KL divergence regulation during training. 

Reviewers consistently provided positive comments regarding the method, motivation, and writing quality. Reviewers queried about more experiments, applicability to multi-step generator, diversity of the examples, etc., which are mostly addressed in the rebuttal. The main reason for acceptance is that there is no major weakness in all aspects, and the method is motivated in the context of the trending test-time scaling.

Reviewers' suggestions should be incorporated into the next version of this paper. In addition, it will also be beneficial to add discussions on the generalisation ability of the hypernetwork to different reward functions once trained, and also the robustness of the hypernetwork to reward hacking compared to other existing techniques.
Decision: Accept (poster)