Test-Time Scaling of Diffusion Models via Noise Trajectory Search
Authors: Vignav Ramesh
Morteza Mardani
Venue: NeurIPS 2025 poster ()
Keywords: diffusion model
test-time scaling
Markov decision process
contextual bandit
optimization
search
MCTS
epsilon-greedy
Primary area: deep_learning
TL;DR: We present an algorithm for test-time scaling of SDE-based diffusion models by searching for noise trajectories which optimize arbitrary rewards, empirically matching/exceeding MCTS performance.

Abstract: The iterative and stochastic nature of diffusion models enables *test-time scaling*, whereby spending additional compute during denoising generates higher-fidelity samples. Increasing the number of denoising steps is the primary scaling axis, but this yields quickly diminishing returns. Instead optimizing the *noise trajectory*—the sequence of injected noise vectors—is promising, as the specific noise realizations critically affect sample quality; but this is challenging due to a high-dimensional search space, complex noise-outcome interactions, and costly trajectory evaluations. We address this by first casting diffusion as a Markov Decision Process (MDP) with a terminal reward, showing tree-search methods such as Monte Carlo tree search (MCTS) to be meaningful but impractical. To balance performance and efficiency, we then resort to a relaxation of MDP, where we view denoising as a sequence of independent *contextual bandits*. This allows us to introduce an $\epsilon$-greedy search algorithm that *globally explores* at extreme timesteps and *locally exploits* during the intermediate steps where de-mixing occurs. Experiments on EDM and Stable Diffusion reveal state-of-the-art scores for class-conditioned/text-to-image generation, exceeding baselines by up to $164$% and matching/exceeding MCTS performance. To our knowledge, this is the first practical method for test-time noise *trajectory* optimization of *arbitrary (non-differentiable)* rewards.

Test- Scaling of Diffusion Models via Noise Trajectory Search
Vignav Ramesh
Morteza Mardani
Harvard University
NVIDIA mmardani @wridiaii@iiiiia...m.. com
vignavraamesh@clllggo harvard. edu
Abstract
The iterative and stochastic nature of diffusion models enables test-time scaling, whereby spending additional compute during denoising generates higher- fidelity samples. Increasing the number of denoising steps is the primary scaling axis, but this yields quickly diminishing returns. Instead, optimizing the noise tru- jectory-the sequence of injected noise vectors- promising, as the specific noise realizations critically affect sample quality; but this challenging due to a high-dimensional search space, complex noise-outcome interactions, and costly trajectory evaluations. We address this by first casting diffusion as a Markov Decision Process (MDP) with terminal reward, showing tree-search methods such as Monte Carlo tree search (MCTS) to be meaningful but impractical. To balance performance and efficiency, we then resort relaxation of MDP, where we view denoising as a sequence of independent contextual bandits. This allows to introduce e-greedy search algorithm that globally explores at extreme timesteps and locally exploits during the intermediate steps where de-mixing oc- curs. Experiments EDM and Stable Diffusion reveal state-of-the-art scores for muss-cooddtiioned/ee--ooo--o -ox-mmmmmttttttteeeeeeeeeee generation, exceeding baselines by up to 164% and matching/exceeding MCTS performance. To our knowledge, this is the first practical method for test-time noise trajectory optimization of arbitrary (non- differentiable) rewards.
Introduction
Diffusion models have emerged as the de-facto standard for generative modeling of visual and, more recently, language data [Ho et al., 2020, Sohl-Dickstein et al., 2015, Nie et al., 2025]. A notable property of these models their test-time scaling: the iterative and stochastic nature of the generation process allows for enhanced sample quality by allocating additional computational resources during inference. primary strategy for leveraging this has been increase the number denoising steps. However, empirical observations indicate that performance gains from more denoising steps typically plateau after a few dozen iterations [Karras et al., 2022, Song et al., 2022, 2021].
The diminishing returns from merely increasing denoising steps motivate exploring alternative test- time optimizations. Optimizing the noise trajectory itself is promising direction. Diffusion models iteratively transform an initial noise distribution towards a target data distribution via a stochastic process (often stochastic differential equation (SDE)), where sequentially introduced fresh noise instances define random path. The characteristics of these injected noises, forming the noise trajectory, are crucial sample quality.
Recent studies show that injected noises are not equally effective; specific noise realizations can produce higher fidelity samples, with certain noises better suited for different objectives [Qi et al., 2024, Ahn et al., 2024]. However, optimizing this noise trajectory is challenging: the search
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(2)
1
0.47
R d
(2)
(N)
1
Xx
1
(1)
()
1
2
No of
X
0.98
Iatent X,
predicted cleon sample (1 Tweedio's formula
candidate x;:
reward function
Figure 1: (Left) Implicit denoising tree traversed by search algorithms. (Right) Visualization of local search in noise space to maximize reward at single timestep l.
space vast, the noise affects the generation intricately, and evaluating entire trajectories often computationally prohibitive.
Recent works have explored this noise search; notably, Ma al. [2025] explores zero-order local search of candidates for the initial noise in ODE-based diffusion sampling. However, no work as of yer has explored higher-order search over the entire noise trajectory for SDE-based diffusion models.
To develop a framework for finding the optimal trajectory of noise over denoising steps, we T-step Markov decision process (MDP) with
first frame the diffusion sampling process terminal reward that generally non-differentiable. This allows us to explore sophisticated tree search algorithms including Monte Carlo tree search (MCTS) over the denoising tree; however, -cile a pscudo-upper bound on performance in its approximation an exhaustive search over the denoising tree-is limited pre-selected set of noise candidates per timestep, and is quite computationally expensive (Table 3).
To address this, we introduce relaxation of the aforementioned MDP by formulating each of the T denoising steps as an independent contextual bandit problem. This approximation, which disregards inter-step dependencies inherent in the full MDP, enables more computationally tractable search strategies. Within this bandit framework, we propose a e-greedy variant of zero-order search. This method is designed to strike an effective balance between globul exploration of the noise space, particularly at extreme (initial and final) timesteps, and more focused local exploration during the critical intermediate denoising steps where demixing and mode selection occur predominantly [Karras et al., 2022, Hang et al., 2024, Wang et al., 2025b].
We evaluate our proposed test-time scaling approach on the Elucidated Diffusion Model (EDM) [Karras al., 2022] for cans-condiiional image generation ImageNet) and Stable Diffusion [Rombach et al., 2022] for text-to text-to-image generation. Our experiments employ task-specific, generally non-differentiable rewards, e.g. classifier-based reward for EDM to measure class adherence, and CLIP score to assess prompt-image alignment for Stable Diffusion. In both cases, our e-greedy noise trajectory search method achieves state-of-the-an generation quality. significantly outperforms established baselines, including zero-order search [Ma et al., 2025]. rejection sampling [Chatterjee

---

and Diaconis, 2017], and beam search [Fernandes et al., 2025]. Notably, our computationally efficient approach matches and exceeds MCTS, demonstrating its efficacy in navigating the noise trajectory search space.
All in all, our key contributions summarized as follows:
We propose effective test-time noise-trajectory search for diffusion models. Inspired by contextual bandits, it employs an e-greedy search to optimize arbitrary (non-differentiabbl) rewards and empirically outperforms MCTS.
We provide evidence and theoretical intuition that our c-greedy search algorithm explores globally and exploits locally. Notably, the search tends to be more local during middle diffusion steps where de-mixing occurs, and more global at large and small noise levels.
We perform extensive comparisons with alternative search algorithms,, including zero- order search, best-of-N sampling, beam search, and MCTS. These comparisons reveal the importance of adapting the search strategy (local vs. global) based on the diffusion step.
We demonstrate the efficacy of our method for noise optimization on EDM and Stable Dif- fusion. Under various rewards (e.g. classifier probability, VLM prompt-image alignment), our approach achieves up to 164% % improvement in sample quality over vanilla generation.
2 Related Work
Reward optimization with diffusion models. Several works have explored direct fine-tuning of diffusion models with reward functions human preferences [Black al., 2024, Domingo-Enre Enrich et al., 2025, Wallace et al., 2023, Wang et al., 2025a]. However, these approaches require expensive training and tie the model to single reward. Alternatively, Song et al. [2021] and Bansal al. [2023] have explored reward gradient-basee based guidance; but this is limited to differentiable reward functions and continuous state diffusion models. Therefore, steering arbitrary diffusion models at inference time with arbitrary rewards remains challenge; our work addresses this with the implementation of more sophisticated step-level, gradient-free search algorithms used solely at inference time.
Test-time scaling. While test-time scaling has been predominantly explored for language models [Yao et al., 2023, Muennighoff et al., 2025], recent work has begun to investigate test-time scaling behavior for diffusion. One means of investing more compute in diffusion model sampling is increas- ing the number of denoising steps; this generally leads to better generations, but with diminishing benefits, due to accumulation of approximation and discretization errors [Xu et al., 2023b]. In this work, fix the denoising steps and explore scaling along orthogonal axis: compute devoted to searching for more favorable noise candidates during the sampling process.
Few past works have considered this search problem. Ma et al. [2025] propose framework for searching for optimal noise in diffusion model sampling, characterized by the choice of algorithm (e.g. best-of-N search, zero-order search, and "search over paths") as well as reward function (e.g. FID [Heusel et al., 2018) 8], prompt-image embedding similarity [Radford et al., 2021, Caron et al., 2021]). However, all experiments are limited to the ODE formulation of diffusion models, which means that the mapping from the initial noise XT to the clean sample Xy is entirely detcrministie; that is, Ma et al. [2025] only have to search over the choices for xT. In this work, we extend this to the SDE formulation and consider an exponentially larger search space: noise candidates at every denoising step, {z, \T 1- Hence this opens up an investigation of more sophisticated step-level search algorithms (e.g. MCTS), which to our knowledge has never before been explored.
Zhang et al. [2025] propose T-SCEND, which sequentially performs best-of-N random search and MCTS as denoising proceeds; however this paper does not truly implement standard MCTS as (1) at timestep t they estimate reward by computing the predicted clean sample x0" via Tweedie's formula [Efron, 201 1] (which estimates the posterior mean E[x, x2, cl from noisy sample xx, see Appendix B) rather than doing a complete rollout and then backpropagating reward on the true clean sample, making their method still greedy search; and (2) their method performs best-of-N search for the majority of denoising steps, only doing their "MCTS" for small number (~3-5) of denoising steps near the very end of the sampling process. This version of MCTS can no longer be considered a true approximation of exhaustive tree search.
3

---

Background
3.1 Diffusion Models
We consider conditional diffusion models that learn distribution p(xo c) given data samples and contexts c. The continuous-timg forward process gradually perturbs clean data Xo towards noise distribution. This can be described by the stochastic differential equation (SDE):
dx; =
where wi a standard Wiener process and (t)) is a monotonically increasing noise schedule with (0) 0. This process yields x, (xx, or (t) (1). For sufficiently large (T") (c.g. o(T) for normalized data), the final distribution p(xr c) x (O (T))
Sampling begins with drawing XT ~ NO0 (T) and then simulating the corresponding reverse time SDE to obtain a clean sample Xg. This reverse process is given by:
dx, (o)o logp(x; (t)), c)dt + (t)
(2)
where dw, is a reverse-time Wiener process and (() do /dl. The score function Vx. p(x; (L), c) approximated by neural network ( The choice of sampler, which dis- cretizes Eq. (2), is also crucial. In our experiments, we use either the EDM sampler (Alg. 2) [Karras al., 2022] for class-conditional generation or the DDIM sampler (Alg. 3) [Song al., 2022] for text-to-image generation, enabling high-quality synthesis with few steps T (e.g. 18 for EDM, 50 for DDIM).
3.2 From Markov Decision Processes to Contextual Bandits
A Markov Decision Process (MDP) formalizes sequential decision-making problems. An MDP defined by a tuple (S, A, Po: P, R). where S is the state space, A is the action space, Po is the initial state distribution, P(S St, a²) is the state transition probability kernel, and R(s,, a,) is the reward function. At each timestep t, an agent observes a state S, E S, takes an action a, E A according its policy π ((a; St),, receives reward r, R(s,, ar), and transitions 8 P(- St, ar).
new state
A continuous contextual bandit problem can be viewed specialized single-state, one-step MDP with continuous action space. It is defined by tuple ( A, R). where & is the space of contexts, A is the continuous action space (shared with the definition of MDP), and R is the reward function specific to the bandit setting. In each round i, the agent observes a context Φ; E & (analogous initial state Sn from Po in MDP, but where contexts for different rounds typically drawn independently). The agent then selects an action a, E A based on its policy Tra; (x) and immediately receives a reward F; = R((o, a.).
4 Searching Over Noise Trajectories
Standard diffusion samplers use random noise per reverse step. To optimize generation for a downstream reward r(Xo, c), we instead search for the optimal T-step noise trajectory Z
Zi Given an initial Xy ~ NO, a II) a pre-trained model Dg, context c, and chosen sampler, the map from Z to the final sample xo (Z, c, x,) is deterministic.
[ØT ZT-1
The search space Z IR (where d is the noise dimension, i.e. CHW where C' x H X W is the image resolution) is sigh-dimensional and unbounded; optimization typically explores specific regions or targeted deviations from standard Gaussian noise (Figure 1).
4.1 Diffusion Finite-Horizon MDP
To formalize this search, we cast the T-step reverse diffusion process (denoising steps indexed I == T,..., 1) as a Markov Decision Process (MDP) with terminal reward:
States: (c, t). with t (T.... 0} indexing the reverse diffusion step.

---

Actions: Zt, the noise vector for t E 1I.
Initial State: ST (c, XT, T), where p(c) and xT ~ (N(,0,,,TTT)
Transitions: Deterministic. State S; and action a, z, yield next state St-1 (c, fx,,,, c, .t 1). The function f is one step of the chosen sampler (e.g. lines 2-6 of Alg. 2, lines 2-3 of Alg. 3) using 2 the injected noise.
Reward: Terminal reward R r(xo. c) at state so (after action a1): zero otherwise. r(.-) evaluates the generated sample Xo. potentially with respect to the conditioning c.
Finding the optimal Z can be written as solving:
r(xo
Z c, c,x)),))
(3)
ZER XX
where Xo (Z, x.) is the final sample generated using the noise sequence Z.
4.2 Monte-Carlo Tree Search (MCTS)
In this general formulation, the space of noise candidates Z. E IR 7xd is extremely large; hence, we first consider the simplification where we predecide a fixed of N noise candidates {z, (x())) Ji=1 for each timestep f. This gives rise to a "denoising tree" with depth T and branching factor N as shown in Figure 1, where each node x; is latent in the sampling process and each edge is noise candidate z; Finding the best noise trajectory now corresponds to finding the best root-to-leaf path this tree; hence, we can directly apply tree-search algorithms this setting, including Monte Carlo tree search (MCTS) as detailed in Algorithm 4.
MCTS approximates exhaustive search over the N possible noise trajectories, with the approxi- mation error going to as the number of simulations - co. Full MCTS though faces two key drawbacks: (1) it extremely computationally expensive, with its number of function evaluations (NFEs)-icc ic. number of forward passes through scaling quadratically in the number denois- ing steps T (see Table for each sampling method's NFE formula); and (2) MCTS by definition is restricted to a fixed denoising tree, i.e. at each timestep we cannot explore "noise space" beyond the N initially determined candidates. This is critical issue, the initial N candidates
sampled uniformly from our prior (i.e., (0, of I))) as opposed to any higher-reward regions of noise space that might exist. In other words, have notion of importance sampling: we are allocating the entire computation cost of running MCTS to samples that in general, likely yield lower-than-optimal rewards [Chatterjee and Diaconis, 2017].
To address these challenges, next we posit an alternative formulation of the denoising process that relaxes the previous MDP.
4.3 Noise Search via Contextual Bandits
Recall that the SDE diffusion formulation requires the injected noise samples at each timestep to be independent (and Normally distributed). Thus, we assume independence actions across timesteps, and ignore the impact of the action at timestep on the state at t 1, reformulate denoising as a sequence of I independent continuous contextual bandit problems; at each timestep f, we have the bandit characterized by:
Φ - (c. x2,
a;
Ze
R(o,, ar) - (x())
where x5" = E[xu X4, Z,, c) is computed via Tweedie's ' formula.
The "sequence-of-baadiss formulation for denoising introduces key tradeoff. By ignoring interac- tions across steps (i.e., how 2; affects x,-1 which in turn informs the choice of 7ı-1), and computing reward on our time-t predicted clean samples (t) rather than our actual clean samples Xp. we have worse approximation our original optimization problem; but we turn have vastly more NFEs to spend on exploration at each timestep (i.e., those that were allocated to simulation/backpoopgaaiii MCTS), and particular can search over more of noise space than the fixed N candidates we were previously limited to.
5

---

Algorithm 1 e-greedy noise search
Require:
Discretization timesteps T, context (e.g. class label) learned denoising network Dg. max. step size scaling factor A, number of local search iterations K, branching factor (number of noise candidates) N, sampling step function mixture proportion reward function initial noise sample XT do
1: for = T
Sample p ~ (0, oT // pivot
2: for 1 to K do
V N sample noise candidate z, as follows: with probability sample z
NO, I),, clse let 2 uv2dz p where (O,, Unif(0, X)
Set zi (4)
for s.t. one-step &o prediction using z attains highest score under r
3: end for
Set x = f(xt,,,c.
4: end for
5: return Xo
4.4 e-greedy
Formally, with sequence-of-bandit denoising, can explore "greedy algorithms" that attempt to independently solve at each timestep
(t) max r(x c)2
t= I..,,,
2e Edd
Inspired by the semi-uniform strategies originally proposed to approximately solve bandit problems [Slivkins, 2024], including the original c-greedy bandit algorithm [Sutton and Barto, 1998], we intro- duce a novel sampling algorithm that is e-contaminated version of zero-order search (henceforth simply "e-greedy" (Alg. 1), where at each timestep t: (1) We start with a random Gaussian noise p ,~(00, oT our pivot. (2) We then find N candidates {z, fi=1- via the following procedure: with probability c. 2 [O, III); else, we draw a random z, from the pivot's neighborhood (Av'dd-radius ball around p). (3) We compute the clean sample predicted x using Tweedie's 's
formula using each candidate 2 (() (4) Finally, we score each clean sample, set the pivot to the noise candidate leading to the clean sample with the highest scoring, and repeat steps 1-3 K times, finally setting ž, to the last value of p.
This simple change from the neighborhood in zero-order search to the e-contaminated mixture distribution
eN + (1=))
yields remarkable performance improvements, which we discuss in section 5. We also provide theoretical formalization of our algorithm's ''s effectiveness in Appendix D.
5
Experiments
5.1
ClassCConditioaal Image Generation
To further elucidate our design choices, we start by presenting a design walk-through for the class- conditional ImageNet generation task. We adopt EDM [Karras et al., 2022] pre-trained on ImageNet with resolution 64 x 64 and perform sampling with the stochastic EDM sampler (Alg. 2). The number of denoising steps is fixed to 18. Unless otherwise specified, we use the classifier-free guidance 1.0, focusing on the simple conditional generation task without guidance [Ho and Salimans, 2022].
5.1.1 Baselines
As baselines, we consider the following algorithms (see details in Appendix C):
Rejection (best-of-N) sampling [Chatterjee and Diaconis, 2017]. We randomly sample N root-to- leaf paths of the denoising tree, and choose the path that produces Xo with the highest reward.

---

Table 1: EDM results by sampling method. Each value in columns 2-4 is obtained by generating 36 images given random ImageNet class labels. The column header denotes the reward used both and during sampling if applicable. We set À = 0.15 and = 0.4. Note that we measure distance as ||2 pllr- hence let À a scaling factor applicd to (00.y=A(00..)d [||x yl|=] | v 2d (where || ||r is the Frobenius norm. Re. notation, N is branching factor, B is beam width, S is number of MCTS simulations, and K is the number of local search iterations. Generating each sample took <1 second for naive sampling (the lowest-compute method) and <1 minute for MCTS (the highest-compute method) on a single A100 (40GB). provide error bars for each reward and sampling method, computed as the standard deviation of scores for given prompt over 20 generations with variability from randomness of the noise draws.
the final images
Method
Brightness
Compressibility
Classifier
NFEs
Naive Sampling
0.4965+0.00
0.3563+0.07
0.3778=0.04
18
Best of Sampling
0.5767
0.4220=0.02
0.546110.00
72
Beam Search (N 4. B = 2)
0.6334+0.02
0.4679#0.05
0.5536=0.02
144
MCTS (N 4, 8)
0.7575:10.02
0.5395:10.04
0.9666:10.03
3888
Zero-Order Search (N =4, K 20)
0.6083::0.01
0.3751+0.02
0.6261#0.04
1440
6-greedy (N 4, K -20)
0.001
0.7208=0.08
0.9885-t0.44
1440
Beam search [Fernandes et al., 2025]. We start at the root and each node xr, traverse the edge to (6) the child x whose corresponding x (t-1)) (obtained via Tweedie's formula [Efron, 201 I]) has the highest reward.
Zero-order search [Ma et al., 2025]. At each timestep, we start with a random pivot and hill-climb with step size scaling factor for K iterations, greedily selecting the best candidate at each iteration. This can be seen as e-greedy with = O, and thus no global exploration.
5.1.2 Rewards
We consider the following reward functions:
Brightness. We compute the [0, normmalizeed luminance formula for perceived brightness, defined for an RGB image as 0.2126R + 0.7 D.7152G + 0.0722B.
Compressibility. Following Black al. [2024], we calculate the number of bytes b of the maximally- compressed JPEG encoding of the generated images and then bound this in [0, 1]-range via the formula 1 - -m where bmax 3000 is manually set based on empirical observations of 50 compressed real ImageNet images. Note that since we fix the resolution of diffusion model samples at 64 x 64. the file size is determined solely by the compressibility of the image.
Class probability under ImageNet classifier. We adopt the ImageNet- Ik 64 x 64 classifier of Dhariwal and Nichol [2021] I] for maassifier-guided diffusion. For each generated image-label pair, the reward is the classifier's probability assigned to the ground-truth class. The classifier is the downsampling trunk of a U-Net, with an attention pool at the 8 x 8 layer producing the final logits.
5.1.3 Results
Table 1 displays class-conditional image generation results. We discuss key observations below:
MCTS as approximate upper bound. Recall that rejection sampling, beam search, and MCTS are all methods to choose a root-to-leaf path in fixed denoising tree. Beam search, by virtue of greedily selecting the best noise candidate each timestep independently, fails consider interactions across trajectories; and rejection sampling, while considering trajectories as a whole, explores a very small number of them. (Indeed, rejection sampling faces the same issue as MCTS whereby have no notion of importance sampling-an we are allocating all computation to N trajectories that, in general, likely yield lower-than-optime rewards [Chatterjee and Diaconis, 2017]; but this issue is far more pronounced for rejection sampling due to the limited number trajectories seen during sampling).
MCTS solves these drawbacks by approximating exhaustive search over the NT possible root-to- leaf paths in the tree, with the approximation error going to 0 as the number of simulations S co.
7

---

vs. 0,
Lipschitz constant of reward landscape tt.
Estimoted Lipechitz consonn
10*
00
4-
10
10.
10'
16'
10'
10'
Figure 2: (Left) Average local-scarch iteration at which a random Normal candidate is chosen (k) as function of c. (Right) Estimated Lipschitz constant of the ImageNet reward as a function of x;, across d; values.
We indeed see this in Table 1: rejection sampling and beam search beat naive sampling across all reward functions; but excluding zero-order search, MCTS is best across the board, expected.
Leaving the denoising tree. Importantly, recall that rejection sampling, beam search, and MCTS all optimize a fixed denoising tree, i.e. at each timestep t we cannot explore "noise space" beyond the N initially determined candidates. Zero-order search and e-greedy address this, with the amount of exploration in noise space dictated by the number of local search iterations K, the maximum step size scaling factor X, and the global exploration factor e. Table 1 displays results for both these methods; somewhat surprisingly, the vanilla zero-order search lags behind the MCTS performance, while e-greedy far exceeds it.
Explore globally and exploit locally. The above results lead us to ask: why does simply drawing a Normal sample w.p. while otherwise sampling from the pivot's 's neighborhood (w.p. 1 e) yield such remarkable improvemen??
We hypothesize the following: at time t, conditioned on latent Xt, there is a specific "region" in the noise space (call it Z) containing noise realizations that best fit the downstream objective (e.g. the green region in Figure 1). In particular, 2 is not necessarily a subset of (O, (f, although it can.
restricted to the Äv Av 2d-radius ball around the original pivot, which is simply a random draw from (O, of More likely than not, the pivot p & 2; and over K local search iterations, we move an expected distance of K KEu-UUmif(o..) [u] /2d = K KA/dl away from p. But with e-greedy, we can explore the entire Normal with e probability; few such "random jumps" can move into Z, at which point we see large returns from hill-climbing. Indeed, we see this empirically: for each local search iteration l (1,..., K) of e-greedy, let 4 = I[Ee-greedy selects one of the random Normal draws instead of all the neighborhood draws in iteration k).. Define k = Ek= (k - 1) Note that e-greedy chooses one of the random draws at every iteration, k- = 10; but if < 10, e-greedy only chooses the random draw in carly iterations & and only hill-climbs (chooses candidates from the neighborhood) at later iterations. Figure 2 (left) plots & as a function a,. For initial denoising steps, the random noise is always chosen, hence & = 10. But in the intermediate denoising steps, k < 10 (0-supporting our hypothesis that we initially need a couple of random draws to get to 2, but for all subsequent iterations, hill climbing is required for us potentially leave the Normal distribution and find noises that maximize the reward.
During vanilla zero-order search, each local search iteration of each timestep
The aforementioned findings raise the question: in e-greedy, why do we only benefit from hill- climbing at intermediate denoising steps? We posit that the reward landscape as a function of x, is much smoother, perhaps near flat, at extreme time steps rather than intermediate ones, hence there is limited performance gain from hill climbing at those steps. This hypothesis aligns with the well-known result that denoising is harder at intermediate timesteps this is where demixing occurs (the result is already well-refined at later denoising steps, and latents are too noisy to mcaningfully understand the features of the image at earlier timesteps) [Karras et al., 2022, Hang et al., 2024, Wang et al., 2025b]. To test this, we measure the sensitivity of the reward landscape as a function of x2 for each t; specifically, we proxy the sensitivity with the Lipschitz constant estimator T (0)
c)ll2.
8

---

Table 2: Stable Diffusson results by sampling method. Following Ma et al. [2025]. we use natural language
versions of the same class labels used in EDM generation (e.g. "a toucan"). Generating each sample took second naive sampling (the lowest-compute method) and <2 minutes method) on a single A100 (40GB). We again provide as error bars the standard deviation of scores for given prompt over 20 generations with variability due to randomness of the noise draws.
MCTS (the highest-compute
Method
Brightness
Compressibility
CLIP
NFEs
Naive Sampling
0.4800=0.02
0.631440.03
0.2457=0.00
18
Best of Sampling
0.516010.04
0.684710.02
0.2673
72
Beam Search (N = 4, B = 2)
0.498110.00
0.701410.01
0.2682100.0
144
MCTS (N =
8)
0.6077#0.03
0.7322+0.02
0.2885+0.01
3888
PerorOrder Search (N =4, K 20)
0.5764=0.02
0.7281+0.06
0.2689+00.0
1440
c-greedy (N K 20)
0.7110#0.03
0.8052#0.05
0.2973#0.02
1440
Lee the local gradient norm of the reward function w.r.t. X₁. Figure 2 (right) shows a plot of the estimated Lipschitz constant by ds in which indeed see that the reward is most sensitive to X₁ perturbations at intermediate steps. To our knowledge, this is the first work to demonstrate the importance of adapting the search stralegy local vs. global based on the diffusion timestep.
Additional experiments, as well as qualitative results, are included in Appendices F, G and H. Notably, conduct hyperparameter sweeps and measure how each sampling method scales with increasing NFEs in Appendix G, finding that e-greedy achieves the best scaling laws empirically amongst all sampling methods.
5.2 Extending to Text-to-Imags Generation
Next, extend our experiments the more difficult task of text-to-image generation. We use Stable Diffusion v v1.5 [Rombach et al., 2022], setting all algorithm hyperparameters to their optimal values from EDM. We enable classifier-free guidance (default scale 7.5) and set T 50 here as is typical with DDIM sampling. The resolution of the generated images 512 x 512.
In lieu of the ImageNet classifier reward, consider very general purpose reward function for training a text-to-image model: prompt-image alignment. However, specifying a reward that captures generic prompt alignment is difficult, conventionally requiring large- scale human labeling efforts. Following Ma et al. [2025], we proxy human ratings of prompt-image alignment with a direct measure of such alignment in joint image-text embedding space; namely, compute the cosine similarity between the CLIP representations of both the image and original prompt [Radford et al., 2021]. Samples that more faithfully adhere to the semantics of the prompt should receive higher rewards under this framework, to the extent that those visual details are represented distinctly in CLIP embedding space.
The Stable diffusion results (Table 2) indeed affirm the superiority of our e-greedy search, which consistently outperforms other sampling methods across all reward functions. Specifically, e-greedy attains improvements 48/26/20% over vanilla sampling for the brightnss/coomressiittty/CC reward functions, even exceeding MCTS by up to 18%. This underscores the robustness and generality of our method across diffusion model architectures (image vs. latent models), number of denoising steps (18-50), and and context modalities (class label vs. text).
6 Conclusion
We present the first practical framework for test-time noise trajectory optimization in SDE-based diffusion models. We first cast denoising as an MDP, enabling us to apply sophisticated step-level search algorithms including MCTS; then, via relaxing the MDP to a sequence of independent contextual bandits, we design novel e-greedy search algorithm that adapts the search strategy (local exploitation vs. global exploration) based on the diffusion step, significantly outperforming vanilla sampling and even MCTS. We extensively discuss broader impacts, limitations, and future work in Appendices I and J.

---

References
D. Ahn, J. Kang, Lee, J. Min, M. Kim, W. Jang, H. Cho. S. Paul, Kim, E. Cha, K. Jin, and S. Kim. A noise is worth diffusion guidance, 2024. httth:///axiivi wwwpii..or//b///e////e abs/222b 03895.
A. Bansal, H.-M. Chu, A. Schwarzschild, S. Sengupta, M. Goldblum, J. Geiping, and T. Goldstein. Universal guidance for diffusion models, 2023.
www.v.org/oaaoo http//bbs2300.. 07121.
K. Black, M. Janner, Y. Du, I. Kostrikov, and Levine. Training diffusion models with reinforcement learning, 2024. URL patttss//aariin hwtp:///g///aa/aaee//eeoess/ hatg/abs/2230 abs/23/5...... 13301.
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers, 2021. URL https:///aaxi http///riii.orraa//////000/0/ 14294.
S. Chatterjee and P. Diaconis. The sample size required in importance sampling, 2017. URL https://arxiv. .rg/ atss///11.. 01437.
Dhariwal and A. Nichol. Diffusion models beat gans on image synthesis, 2021. URL https: //arxiv. org/ abss/105...e 05233.
C. Domingo-Enrich, M. Drozdzal, B. Karrer, and R. T. Q Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control, 2025. URL 08861.
B. Efron. Tweedie's formula and selection bias. Journal of the American Statistical Association, 106 (496):160--1614, 2011. doi: 10.1 198/jasa.201111 1tt11181. URL htttt:tt/pppbmmeeeee ncbi. mnbi.nlm. milh.oo//00 htlb.gov/22505788// Epub 2012-01-24.
G. Fernandes, V. Ramos, R. Cohen, I. Szpektor, and J. Magalhaes. Latent beam diffusion models for decoding image sequences, 2025. URL httph:tt//////tt /tatts://arxvv... http://awwww....abbs.aas//s// org/abs/2500 20429.
T. Hang, S. Gu, X. Geng, and B. Guo. Improved noise schedule for diffusion training, 2024. URL attps://arxie org/ /twp 032977.
M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium, 2018. URL https: org/ abs abs/1706.08t 08500.
and T. Salimans. Classifier-free diffusion guidance, 2022. URL UtL.itttss:/arxxw.rrrgg 2207. 12598.
J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models, 2020. URL https: //arxiv. org/ abs//006.. 11239.
Karras, M. Aittala, T. Aila, and Laine. Elucidating the design space of diffusion- based generative models, 2022. URL http:: //aa/xi/////ii/ arxiv aorg/bbs/2225 200pb//22200006....
D. Kim, M. Bae, K. Shim. and B. Shim. Visually guided decoding: Gradient-free hard prompt inversion with language models, 2025a. URI. https://t///arttt http://awvvvvvvaaaaaae argg/bbs///0// 08622.
J. Kim, T. Yoon, J. Hwang, and M. Sung. Inference-time scaling for flow models via stochastic generation and rollover budget forcing, 2025b. URL /tttttps://axiiv... org/ haa//bba/20000 19385.
X. Liu, X. Zhang, J. Ma. J. Peng, and Q Liu. Instaflow: One step enough for high-quality diffusion-based text-to-image generation, 2024. URL https::t//// //arxiv..vvvv/// org/ abs /2009. 2309.06380. 06380.
N. Ma. S. Tong, H. Jia, H. Hu, Y.-. Su, M. Zhang, X. Yang. Y. Li. T. Jaakkola, X. Jia, and S. Xie. Inference-time scaling for diffusion models beyond scaling denoising steps, 2025. URL org/ abs/ 22000 09732.
N. Muennighoff, Z. Yang, W. Shi, X. L. Li, L. Fei-Fei, H. Hajishirzi, L. Zettlemoyer, P. Liang, E. Candès, and Hashimoto. sl: Simple test-time scaling, 2025. URI. https///s// org/ abs/2501. 19393.
10

---

S. Nie, F. Zhu, Z. You, X. Zhang. J. Ou, J. Hu, J. Zhou, Y. Lin, J.-R. Wen, and C. Li. Large language diffusion models, 2025. URL https: //arxiv...iv wwwpi.oggaaaaaaa...a.oo/ta.a..ae wttg/ba////aaeeeeee 09992.
D. Podell, Z. English, K. Lacey, Blattmann, Dockhorn, Müller, Penna, and R. Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis, 2023. URL https: //arxiv. org/ ab2/2307 01952.
Z. Qi. L. Bai, H. Xiong. and Z. Xie. Not all noises are created equally.diffusion noise selection and optimization, 2024. URL https: / arxiv. org/abs/ aaaaabbs//00// 2207 14041.
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, Clark, G. Krueger, and Sutskever. Learning transferable visual models from natural language supervision, 2021. URL https: /ttttp:///rxiv.../ wtt./bs////aaa//ao/ 2102000000.00....
R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High High-resolution image synthesis with latent diffusion models, 2022. URL https::/t////... htapps:///axiv..///// wtt.//ow.vvoooo....a/././. /tt/////...... 10752.
Salimans, I. Goodfellow, W. Zaremba, Cheung, Radford, and X. Chen. Improved techniques for training gans, 2016. URL attts:://axtii wwgp/iivoor//oooooooooooo org/ org/abs// 11/1/666.... 03498.
A. Slivkins. Introduction to multi-armed bandits, 2024. URL UtL. 07272.
Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics, 2015. URL https: org/ abs 1503. 03585.
J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models, 2022. URL org/abs//20 02202
Y. Song, J. Gohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations, 2021. URL 13456.
R. Sutton and A. G. Barto. Reinforcement Learning: An Introduction. MIT Press, Cambridge, MA, 1998.
Z. Tang, J. Peng, J. Tang, M. Hong. F. Wang, and T.-H. Chang. Inference- time alignment of diffusion models with direct noise optimization, 2024. URL https: org/ abs//405.. 18881.
van der Maaten and G. Hinton. Visualizing data using t-sne. Journal of Machine Learning Research, (186):2579-205) 2008. URL /tttt://jmlll org/ wwrggpapers/v... /tt////anaaaaeen....... //tpperr////yydddraaaaeeeaaaaa
html.
B. Wallace, M. Dang, R. Rafailov, L. Zhou, A. Lou, S. Purushwalkam, S. Ermon, C. Xiong, S. Joty. and N. Naik. Diffusion model alignment using direct preference optimization, 2023. URL. https: //arxivv//xxvv org/ /bss//11.....aa 12908.
F.YY Wang, Y. Shui, J. Piao, K. Sun, and H. Li. Diffusion -po: Negative preference optimization for better preference aligned generation of diffusion models. In Proceedings of the International Conference on Learning Representations (ICLR), 2025a. URL https: http////wpeeeeeeeeeeeeeeeeeeeeee net/ forum?id=1J17nn50cce Poster.
K. Wang. M. Shi, Zhou, Z. Li, Z. Yuan, Y. Shang, X. Peng, H. Zhang, and Y. You. closer look at time steps is worthy of triple speed-up for diffusion model training, 2025b. URL https: //arxiv. , org/ abs/2405 17403.
J. Xu, X. Liu, Y. Wu, Tong, Q. Li, M. Ding, J. Tang, and Y. Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation, 2023a. URL https: org/ abs/2304 05977.
Y. Xu, M. Deng, X. Cheng, Y. Tian, Z. Liu, and T. Jaakkola. Restart sampling for improving generative processes, 2023b. URL
S. Yao, D. Yu, J. Zhao, I. Shafran, T. L. Griffiths, Y. Cao, and K. Narasimhan. Tree of thoughts: Deliberate problem solving with large language models, 2023. URL https: //arxiv.... org/ abs/ 2305. 10601.
11

---

H.-.. Yeh, K.-H. Lee, and J.-C. Chen. Training-free diffusion model alignment with sampling demons, 05760.
2025. URL https:://aaaat /awpiv.r/vvv///////////a///aaa org/
Zhang, J.-S. Pan, R. Feng, and T. Wu, T-scend: Test-time scalable mets-enhanced diffusion model, 2025. URL httpt: //axxiv org/ abs./b/2...00000 01989.
Z. Zhang. S. Zhang, Y. Zhan, Y. Luo, Y. Wen, and D. Tao. Confronting reward overoptimization for diffusion models: A perspective of inductive and primacy biases, 2024. URL https: https://arxii: //arxiv. org/abs /////b///// 2aabs/244 08552. 0855..
12

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All methods described in the abstract/intrroduction are explained in section 4, and all experimental results are detailed in section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
See Appendix
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
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: [NA]
13

---

Justification: No new theorems/proofs in the paper; all theoretical formulations discussed at length in section 4.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
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 affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All experimental setup details and described section 5, e.g. in captions of Table and Table 2. Code will also be uploaded by supplementary material deadline, providing full reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
14

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
Code and data are made publicly available at this GitHub link.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ttpp.cc/publll http:ccccppccccccceeeeeeecc tww//guiieee...e. wtw.eeeeeeooeebmmiicmce ssi aesionPooo onnolicy))) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All experimental setup details and described section 5, e.g. in captions of Table and Table 2. Code will also be uploaded by supplementary material deadline, providing full reproducibility.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental 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: See Table 1, Table 2.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
15

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: See section 5, specifically Table Table
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: : //nuuripss... cc/ httpppullcctttteee http:/hhhhcccceeeeeeeeeeee.
Answer: [Yes]
Justification: Have reviewed and ensured compliance with the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See Appendix I.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
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, scraped datasets)?
Question:
Answer: [NA]
Justification: No new assets released; safeguards to existing assets (publicly available diffusion models) with which our inference-tine search algorithms are used should apply.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
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 are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Citations of all past work is referenced appropriately, sections and 5,
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (e.g. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification: No new assets released, just inference-time algorithms for leveraging existing diffusion models.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
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 equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
18

---

16. Declaration of LLM usage
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.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: /ttprr//nnuu.//eeeeeee for what should or should be described.
19

---

Contents
Samplers
21
B Stepwise Noise Evaluation via Tweedie's Formula
21
Search Algorithms
22
e-greedy Regret Analysis
23
Visualizing e-greedy Noise Search
24
Additional Experiments
24
E.
Scaling up dataset size
24
F2 Increasing task complexity
25
F3 Increasing reward complexity.
25
F44 Varying model family, reward function
25
F.5 Tweedie estimates vs. learned reward models, few-step models
25
G Hyperparameter Settings & Scaling
26
G.I Scaling with Increasing N, K
26
G.2 Varying K across timesteps
G.3 Sweeps over X,
27
G.4 Varying classifier-free guidance scale, number of timesteps T
28
Qualitative Results
28
I Broader Impacts
33
J Discussion and Next Steps
33
J.1 Compute Efficiency
33
J.2 Population-Based Metrics
33
J.3 Reward Overoptimizztionn
J.4 Future Work
34
20

---

A Samplers
In all experiments we adopt either the EDM sampler (Alg. 2) [Karras et al., 2022] (for class- conditional generation) DDIM sampler (Alg. 3) [Song al., 2022] (for text-to-image generation) as it enables high-quality generation with low number of denoising steps T. Note that the choice of sampler fundamentally affects the formulation of denoising as an MDDP/squence--f-bannee problem, and hence our search over noise trajectories, since we define the transition function P(st-1 St, a+) _ (xc.f(x..f.c. f... where b is the Dirac delta function (places all probability mass single point y) and is step of the chosen sampler (e.g. lines 2-6 of Alg. 2, lines 2-3 of Alg. 3) using ž, as the injected noise.
Algorithm 2 EDM Sampling Algorithm
Require: Number of timesteps T, noise levels {t; with to = mmin and tyy = Emax expansion coefficients {7a )=a)== NW-1 initial sample XT ~ NO0,, Ey denoising network Da L, c), or, noise
samples {z, | ~ (O o7
1: for - to do
2: f = 711+7) t; & = x, + V 82 - t,² 2,
3:
d, = =xx: Do E,. c)) /fL; X;-1 = 8; + (t:-1 - f
4:
if if * then
5:
d' = (x-11 Da (x,-1, / /t: x x, (ti-1 - (5d: +
6: end
7:
end for
8: return X
Algorithm 3 DDIM Sampling Algorithm
Require: Discretization timesteps T, diffusion coefficients 01,..., OT, initial noise
(0, I).
context (e.g. text prompt) c, noise vectors 21...., zr (O,, coefficient r) 00 1] for balancing ODE and SDE, denoising network Do(x, c) for T do
2:
0 = // a 11////111///// 01/a1-1
3: X₁-1 v a-11/a (v (Var-1l1 - ai)/a VI a,-1 - -ot) co(x,,t,c) 0,2
4: end for
5: return X0
B Stepwise Noise Evaluation via Tweedie's Formula
The greedy algorithms treat as baselines (beam search, zero-order search), well e-greedy algorithm itself, require evaluating multiple noise candidates (z, ( J timestep independent
of future denoising steps / - I,....I. I. We operationalize this evaluation by letting the best noise candidate z!" be that which maximizes the reward of the expected clean sample conditioned on that noise candidate, i.e.,
arg max r (E[(x) f(x,,t,,,, (0) )),
(5)
We leverage Tweedie's formula [Efron, 2011] to compute the expected clean sample (!) as an estimate of the posterior mean E[xo Xr-1, c from the noisy sample x produced by doing denoising step from X, with noise candidate z, Formally, given Xt-1, 0z-1, we have
E[xo X;-1, cj X of V Vx logp(x:-1 0--1,
where the denoiser Da the diffusion model (e.g., EDM or Stable Diffusion v1.5) trained to approximate the score function Vx-, log,)ogg((xx111 c).
21

---

C Search Algorithms
Below. we provide detailed algorithms for the rejection sampling, beam search, zero-order search, and MCTS methods described in the main body the paper. Table displays the NFEs required per sample generation using each method; notably, all sampling methods but MCTS- PePiSS-incuiinggg e-greedy, which attains the highest performance across reward functions require NFEs only linear in the number of denoising steps T.
Algorithm MCTS noise search
Require: Discretization timesteps T, context (e.g. class label) c, learned denoising network Da.
branching factor N, sampling step function [0, 1]-bounded reward function r, fixed per-step
noise candidates {zl" (e..) (1..)) ((..... number of simulations S, exploration constant C 1.414, initial noise sample Xy
Define UCB score(x,, 1, 2 ( reward from (f(x..l.c. 2 ).--1) + +
log(# to (xx,)))
# visits to (s ( 1)
visits to (f (.)r) 1,t-1)
C(x,) - set of child nodes of x,, s.t. [C(x,)] E {0, N), (x,)) -
f(x,, 2, c, (1))
1: for t- = to do
2:
for to do
Let L' - while [C(xx)] + 0 do
3:
selection
(arg max, UCB scor( 22,)))))
Let xe
Xe,t, c, 2
+ end while
4:
5:
if then
// expansion
C(x²) (x)(x))) 6: end
For random h E {l...., denoise from ca (xe) to get Xn
|| simulation || backpropagation
7:
for 1"
to T
visits to (Xtr, (²) +=1
reward from (Xer, tn) += +=x c)
8: end for
reward from (e: (x: 1))
X1-1 +- arg (x.)
# vvisits to (cy (X+),(-1)
9: end for
end for
11: return Xo
Algorithm 5 Rejection Sampling
Require: Discretization timesteps T, context (e.g. class label) c, learned denoising network Da. number of samples N, sampling step function f. reward function r, initial noise sample Xy
1: for = to N do
2: for T' do
Sample noise candidate 2t NO,, o (n)
x (x,,,,,,,20 (n))
3:
4: end for
5: end for
(argmax, (E() (E(".....)))) ( (x("(x())))) (x))),))
6: return Xo
22

---

Algorithm 6 Beam Search
Require: Discretization timesteps T, context (e.g. class label) learned denoising network Do.
branching factor N, beam width B, sampling step function f, reward function r, initial noise
samples {x7- j=1
1: for = do
Sample N noise candidates {2) J NV
C=
2: for j to B do
3: for = to N do
Add f (0) ,t,, I0C
4:
end for 5: end for
Let xi-1 (1.-B) be the B latents in C whose one-step denoised predictions % via Tweedie's formula [Efron, 201 have the highest scores under r
(t-1)
(computed
6: end for
7: return x
r 5
ag m
Algorithm 7 Zero-Order Search
Require: Discretization timesteps T. context (e.g. class label) c, learned denoising network Dg. max. deviation scaling factor A, number of pivot updates K, number of samples N, sampling step function f, reward function r, initial noise sample x7
1: Sample p ~ (0,,,, // pivot
2: = do
Sample .V(0, o2 3: for Sample N noise candidates (
K do
~ NO ofI), for each let ( = X/2CHWz,
p
Set p zi for s.t. one-step & prediction using 2 attains highest score under r
4: end for
Set X,-1 = (xx,,,, p)
5: end for
6: return Xo
D e-greedy Regret Analysis
To strengthen our claims on e-greedy's superiority, we derive its regret bound, showing that even with only black-box access to a Lipschitz surrogate simple regret decays sublinearly in the total number global samples.
Theorem (Regret of e-greedy search) Let the noise at diffusion step be 24 NN,00 Ia). Fix confidence level 1) and truncate noise-spae space the high-probability ball I, 2 E Rdd |z| S orvddn((1n)) Assume the surrogate reward R. Т, - [0,1] is I-Lipschitz and satisfies
r(x))(x(2),c)) 58 AZET,
Run Algorithum I (e- (e-greedy local search) for K iterations with N candidate points per iteration and probability and let the final pivot be pt. Then the expected simple regret step
t obeys
max Rz(z) R.(pp)) < Ca(ENNK-1//
zEr.
where the constant C. > 0 depends only on the dimension d. Summing over T diffusion steps yields
T
R
max Re(z) - [c()
ST 8 C2dd
2E1'
23

---

Table 3: NFE formulas by sampling method using '' "imeettep. Note that NFEs denote "number of function evaluations" (i.e., number of calls Da) for sampling single image.
Method
NFE formula
Naive Sampling
T NT (N + BYT (N + SyT²
Best of Sampling
(N, B Search
(N, S)-MCTS
(N, K) ---ero--rrda Order Search
NKT NKT
(N, -gggeedy
Proof By Gaussian concentration, lies in I with probability at least 1 - so we may restrict attention to A at negligible cost. Each global-exploration draw is uniform on T and total the algorithm collects M = =NN such draws. It is well-known result that the best of M uniform samples over d-dimensional, Lipschitz-reward domain achieves expected simple regret on the order of M 11/d Since the e-greedy pivot never decreases in surrogate reward, its final surrogate value is at least that of the best global sample, inheriting the same ( 11/d rate. Finally, replacing the surrogate by the true reward adds at most 8, yielding the per-step bound; summing over T steps gives the stated total regret.
Note the following key points:
Polynomial decay in global samples. With M eNK global draws per step, regret scales ((M-1/d), so increasing the total sample budget drives simple regret to zero at the classic random-search rate.
Bias floor 8. In practice we can take 5 to be quite low as it is simply the approximation error induced by application of Tweedie's formula, [r(xo, c) r(x)) c)).
E Visualizing e-greedy Noise Search
We display a animation visualizing the e-greedy search process at this link. For each timestep t, the K = 20 local search iterations, plot the 2D t-SNE [van der Maaten and Hinton, 2008] projection of the pivot and track over time. The animation displays visually the emergent properties of e-greedy describe in subsubsection 5.1.3 of the main paper: at early and late denoising steps, the search is mostly random global exploration; but at intermediate steps, after a couple initial random draws, all subsequent iterations local hill-climbing.
F Additional Experiments
F.1 Scaling up dataset size
Below, we report EDM results using the ImageNet reward on 300 images, a 10x increase in dataset size. This confirms the superiority of our e-greedy search.
Table 4: EDM results by sampling method using classifier reward. We use the same algorithm hyperparame- ters as Table of the main paper.
Reward Naive
Best-of-4 Sampling
Beam Search
MCTS
Zero-Order Search
e-greedy
Classifier 0.3738
0.6413
0.6928
0.7311
0.6818
0.9790

---

F.2 Increasing task complexity
We provide results evaluating methods on far harder set of prompts containing compositional information Black et al. [2024], specifically those of the form "a [SUBJECT] [PHRASE]]" e.g. "a lion washing the dishes." Table 5 displays these results; e-greedy is again superior across the board.
Table 5: SD results by sampling method, across reward functions, on compositional prompts. We use the same hyperparameters as in Table 2 of the main paper.
Method
Brightness
Compressibility
CLIP
Naive Sampling
0.4579
0.6866
0.2578
Best-of-4 Sampling
0.5257
0.7639
0.2618
Beam Search
0.5179
0.7690
0.2646
MCTS
0.5921
0.7717
0.2954
Zero-Order Search
0.4573
0.7128
0.2802
e-greedy
0.6312
0.7869
0.3053
F.3 Increasing reward complexity
To test method on harder rewards capturing more fine-grained details of images, provide results evaluating each method using the counting reward from [Kim et al., 2025b], normalized [0, 1]. Each prompt is comma-separated list of 1-6 clauses, each of the form "[NUMBER] [OBJECT]," e.g. "five horses, three cars, one train, five airplanes." The counting reward function leverages open-source object detection models compute the MSE between the actual and ground truth number of objects in the generated image, averaged across all objects described in the prompt. Even with this hard, non-differentiable counting reward function, we again see that e-greedy is superior across the board.
Table 6: SD results by sampling method, using counting reward.
Reward
Naive
Best-of-4
Beam Search
MCTS
Zero-Order Search
e-greedy
Counting Reward
0.3076
0.3156
0.4318
0.4377
0.3846
0.5384
F.4
Varying model family, reward function
To support our claim of e-greedy's superiority, we display results using a larger text-to-image model (SDXL [Podell al., 2023]) and the text-to-image human preference reward model ImageReward [Xu et al., 2023a]. We use the same prompts and algorithm hyperparameters as in Table 2 of the main paper, and normalize the ImageReward score [0, 1]. Results are shown in Table 7.
c-greedy remains the highest-performing sampling method in this experimental setting as well, highlighting its superiority across reward functions and diffusion model families/sizes.
Table 7: SDXL results by sampling method, using ImageReward [Xu et al., 2023a] as the reward function.
Reward
Naive
MCTS
Zero-Order Search
c-greedy
Amage-Rewwar 0.4668
0.5152
0.5036
0.5830
F.5 Tweedie estimates vs. learned reward models, few-step models
A possible concern with our proposed approach is that predicting x solely using Tweedie's formula with a base model remains too crude, potentially yielding relatively low benefits, especially during the initial sampling stages. Alternatives include using learned soft value function latents few-step model distilled from the corresponding base model, which might lead to better results.
We test this empirically, showing that using Tweedie estimate is extremely competitive with, and often actually outperforms both these alternatives, across model architectures, sampling methods, and reward functions.
25

---

Brightnes Reward vs Noise Candidates
Compressibility Reward Noise Candidates
ImagoNet Reward e Noise Candidates
1.0
1.0
1.0
0.8
0.8
115
Aogguge Rewand
0.-
0.4
6.4
0.1
0.2
0.-
010
16
0.0
16
32
000
10
32
4 Noise Candidnnes per Timestep
Noris Candidates Timestop
Noise Cundiclones per Timestep
Sumpling Metheds
Order Search
Reward Functions
Rejection Sumpling
Sarch
siggttees
Compeesbbliyy
Imageheet
Figure 3: Performance scaling with number of noise candidates. Performance of sampling methods as the number of noise candidates N per timestep increases, measured three reward functions (brightness, compressibility. ImageNet). e-greedy achieves the best scaling law, attaining the highest rewards at the majority of N values; however, it experiences non-monotonic gains due to its greedy local selection, whereas MCTS by approximating exhaustive scarch- stcady, monotonic improvement.
First, in the class-conditional image generation setting, compare running the ImageNet classifier on the Tweedie estimate a versus running a classifier on the noisy latent 2, itself at each timestep 1. For this "noisy classifier" we adopt the ImageNet- Ik 64 x 64 classifier from Dhariwal and Nichol [2021], trained to classify latents across timesteps and noise scales. As with the regular elassifier reward, the reward here is the noisy classifier's 's probability assigned to the ground-truth class. Table 8 displays these results.
Table 8: EDM results by sampling method using either classifier reward on the Tweedie estimate, or learned classifier on the noisy latents.
Reward
Beam Search
Zero-Order Search
e-greedy
ImageNet classifier on Learned ImageNet classifier on latents
0.9666
0.6261
0.9885
0.9792
0.6099
0.9998
Second, in the harder text-to-image generation setting, we compare computing the CLIP reward on the Tweedie estimate, with instead computing the CLIP reward on the clean sample generated using one-step distilled model from the current latent and noise candidate, at each timestep t. As in the main paper, we use Stable Diffusion v1.5, and for the one-step model use InstaFlow-0.9B [Liu et al., 2024] which was distilled from SD1.5. Results are shown in Table 9.
Table 9: SD results by sampling method using CL.PP reward either the Tweedie estimate or the predicted clean sample using one-step distilled model.
Reward
Beam Search
Zero-Order Search
e-greedy
Tweedie-estinan
0. 06.6
0.2689
0.2973
One-step model distilled from SD1.5
0.2671
0.2825
0.2820
G Hyperparameter Settings & Scaling
G.1 Scaling with Increasing N, K
is important to recognize that e-greedy attains the best "scaling law" among different search methods empirically, although with drawbacks compared to MCTS. Figure 3 displays the perfor- mance by sampling method as the number of noise candidates per timestep is scaled up, with all other hyperparameters fixed. Although c-greedy consistently achieves the best performance for sufficiently high N as expected, does face the limitation of not improving monotonically when (!))
scaling up N. This is because e-greedy (per its name) greedily selects the best noise based on x)
26

---

Average Reward vs Average K
Average Reward
0.6
0.4
0.2
12
18
# Noise Candidates per Timestep (E((K))
Reward Functions
Brightness
Compressibility
ImageNet
Figure 4: EDM results by sampling method, varying E, [Ke). We use the same class labels and (other than K) algorithm hyperparameters in Table I of the main paper.
timestep t; hence, as we increase the number of noise candidates per timestep, it may greedily select better noise at the current step which ends up being worse one globally (i.e., locally worse noise would have led to a better overall noise trajectory). MCTS is the only method that should (and indeed does) monotonic improvements it approximates exhaustive search over the N possible noise trajectories.
For this same reason, shown in Figure 5 we similar lack of monotonicity in e-greedy and zero-order search's scaling as function of the number of search iterations per timestep; though again e-greedy is empirically extremely performant, attaining higher rewards than vanilla zero-order search across K E [1, 64] (leftmost figure).
G.2 Varying K across timesteps
Recall that e-greedy requires NKT NFEs for K local search iterations. Note, however, that this assumes K fixed across timesteps. Letting K, be the number of local search iterations at timestep t, can significantly bring down E, [K,] by noting that K, only need be high at intermediate denoising steps where de-mixing occurs, and can be low at extreme (beginning and end) timesteps. We report e-greedy performance
function of E, [K,) below in Figure 4.
Indeed, e-greedy attains strong performance across all reward functions, even with extremely low average K, (can retain performance within 0.002 of the highest value with [x[K,] as low as 4), highlighting our method's computational efficiency compared to MCTS.
G.3 Sweeps over X,
We additionally conduct hyperparameter sweeps for À E [0, 1] for zero-order and e-greedy) search, and the c parameter in c-greedy (bottom right plots of Figure 5). We find optimal c = 0.4 and (0.1..22 across sampling methods and reward functions. Although the drop-offs in performance with higher e and À (especially step for À) seem surprising, we purport it is actually expected behavior again due the greediness of zero-order and e-greedy search; the larger the and/or X, the more we can move in "noise space" at each timestep, meaning the more likely it is for us to find a locally better but globally worse noise candidate. For large À especially, we may observe noise candidates far outside (00, afI) for which x0" is no longer meaningful sample under the reward function, which explains the very low scores in the lower right plot. (This also explains why e-greedy suffers less from this problem than vanilla zero-order search; it always has nonzero probability of selecting a random Normal sample and escaping these degenerate noise regions.)
Overall, the optima we find seem to be a "sweet spot" this regime-there these values are large enough to permit exploration at each timestep, but not too large to succumb the aforementioned behavior.
27

---

Reward Search lierations per Timestep
Reward
Reward
1.0
1.0
1.0
118
0.8
115
Aorguae Rewand
0.4
0.4
0.4
0.2
0.-
0.0
012
94
0.6
10.
1.0
0.0
012
11.4
20
0.5
1.0
4 Soarch liettion per Timestep
Sampling Malladss Zeo Search
Rewand Functions Comgrescility
eggeedy
Brightnos
BrageNer
Figure 5: (Left) Average reward vs. number of search iterations per timestep. (Center) Sweep over for c-greedy. Highlights an optimal 0.4 that balances exploration and exploitation. (Right) Reward vs. maximum step size scaling factor À for zero-order and e-greedy search. Demonstrates optimal E [0.1, 0.2] and drastic performance drops at large due candidate set being overwhelmed by degenerate noise samples far outside the standard Normal distribution.
Table 10: SD results by sampling method, varying the classifier-free guidance scale. We evaluate only CLIP reward the brightness and compressibility rewards generally insensitive classifier-free guidance.
Method
cfg 1.5
cfg = 3.0
cfg = 7.5
cfg 9.0
cfg = 12.0
Naive Sampling
0.2446
0.2635
0.2457
0.2650
0.2671
Best-of-4 Sampling
0.2612
0.2743
0.2776
0.2783
0.2759
Beam Search
0.2673
0.2314
0.2682
0,2801
0.2780
MCTS
0,2828
0,2712
0,2885
0,2971
0.2827
Zero-Order Search
0.2539
0.2612
0.2689
0.2443
0.2801
6-greedy
0.2958
0.3092
0.2973
0.3055
0.2994
G.4 Varying classifier-free guidance scale, number of timesteps T
See Table 10 and Table 11 above. These results confirm the superiority of e-greedy irrespective of guidance scale and
H Qualitative Results
To assist the reader in understanding the optimization process, we offer visualizations the generated 512 x 512 resolution images from Stable Diffusion using our different sampling methods. The aim is to provide a qualitative evaluation of the proposed e-greedy method.
As can be seen in the images below, our inference- time algorithms and most noticeably e-greedy indeed yield samples with significantly higher rewards, while avoiding mode collapse and ensuring distributional guarantees by virtue of being solely inference-time noise search methods.
Not only do not cherry-pick examples, "hard/adversarial" with respect to the reward, i.e. "a black lizard for brightness or "a potted fern with
in fact display results using prompts that
Table 11: EDM results by sampling method, varying the value of Each cell contains two values, corresponding to scores under the Brightness/ Compressibility reward functions.
Method
T 18
= 32
T 64
128
T 256
Naive Sampling
0.4965/00.355
09912/0.3921
0.4621 0.4621 /0.4549
0.4849/0.4052
0.4541 /0.4250
Best-of-4 Sampling
0.5767 00/667/0,,4 00.7/20.2220 02.4220
0.5666/0. 4992
0.6475 /0.004
0.5801/0.5101
0.5923/0. 40.971
Beam Search
0.6334//0.465
0.6013/0.432.
0.5742
0.6162/0.4517
0.6618/0.4874
MCTS
0.7575/00 50.5395
0.6910/0.55 00.5133
0.6831/0.5624
0.6957/0.5015
0.7494/0.5789
Zero-Order Search
0. 0.6083/0..0 /0.3551
0.9999/ 0.. 7417
0.9999 00. 7.77
19999/0.760
0.9999/0.770
0.9999/0.766
e-greedy
0.9813/0.7200
0.9999 0.9999/0.744
0.7424
0.9999/0.7588
19999/0.7688
0.9999/0.7672
28

---

hundreds of leaves" for compressibility. The success of our method even on these samples highlights two desiderata our algorithms satisfy: faithfulness to prompt and lack of mode collapse.
Quite remarkably, we see that when using our sampling algorithms, the denoising process consistently uses the additional NEEs to employ shared across prompts and methods. Some examples of these common strategies, prefixed by the reward functions they apply to, include:
of limited number of reward optimization "strategies"
(compressibility) blurring content;
(brightness, compressibility) occluding or minimizing the subject of the image;
(compressibilltty changing style from photorealistic to "animateed/cartooniss"""
and (CLIP) emphasizing distinctive features of the image subject. Examples (see Figure 8) include the curved beak and wide brown region around the eyes of the käkāpö and the unique red nose of the red wolf. It's particularly remarkable that just noise optimization alone for the CLIP reward increases faithfulness to elements of the prompt like "three rockhoppers in a line:" following detailed prompts like these is notorious and long-standing problem in diffusion [Black et al., 2024, Kim et al., 2025a], and none of these features together were captured by the other sampling methods besides e-greedy.
29

---

"a panda wearing bowlle"
"a potted fern with hundreds leaves" "
"a black lizard"
Naive
Best-ef-4
Rejection
Zero-order
MCTS
E-greddy
Figure 6: Visualizations of images generated with Stable Diffusion v1.5, with varying sampling strategies. using the brightness reward function.
30

---

panda wearing bowtie"
"a potted fern with hundreds of leaves"
"a black lizard"
Naive
Best-of-4
Rejection
Zero-order
MCTS
E-greedy
Figure 7: Visualizations of images generated with Stable Diffusion v1.5, with varying sampling strategies, using the compressibility reward function.
31

---

"three rockhoppers in line
"a red wolf"
"a kākāpō
Naive
Best-of-4
Rejection
Zero-order
MCTS
E-greedy
Figure 8: Visualizations of images generated with Stable Diffusion v1.5, with varying sampling strategies, using the CLIP reward function.
32

---

I
Broader Impacts
Our inference-time optimization framework for diffuusion models empowers users steer generative outputs toward arbitrary, user-defined rewards. In positive scenarios, artists and designers can interactively refine visual concepts in real time, accelerating creative workflows; educators and accessibilitt advocates can tune outputs for enhanced contrast or simplified imagery to aid visually impaired learners; and scientists can optimize for perceptual clarity domain- speiiin--pifiic criteria (e.g. highlighting anatomical structures in medical images), improving interpretability of complex data visualizations.
However, the same mechanism that makes our method so flexible also amplifies risks. By optimizing for malicious objectives- as photorealistic deepfakes tailored to evade detection or synthetic imagery crafted to reinforce biased or harmful visual media, facilitate disinformation campaigns, violate individual privacy. Moreover, if reward functions inadvertently encode cultural or demographic biases, optimized outputs may perpetuate unfair or exclusionary representations.
adveesrrress could undermine trust in
Notably, however, the proposed methods are simply algorithms with which to sample from existing, already released diffusion models; hence safeguards paired with those models remain in place. For instance, models like those in the Stable Diffusion family are paired with NSFW content checkers that will filter out images with harmful content that might be generated with algorithms, helping mitigate harm from cases where users try to optimize objectives that reward for such content.
J Discussion and Next Steps
While our method yields sfate-of-the-ee class-conditional and text-to-image generation results, is not without its limitations:
J.1 Compute Efficiency
For N noise candidates and K local search iterations per timestep. the e-greedy approach requires NK times the NFEs compared vanilla sampling for generating a single image. While this approach remains computationally linear with respect the number of timesteps T. thus making it significantly less costly than more computationally intensive alternatives like MCTS, it nonetheless poses a computational burden. This additional computational cost could limit practical deployment scenarios, particularly where resources or latency constraints are critical factors.
To mitigate this, recalling Figure 2, as proof-of-concepl with 20 for only {t 0.01 I a, < 1}. K otherwise. This yields rewards within =0.04 of the original results (where K 20 is used at all denoising steps), but cuts the NFEs by more than half. We recognize exploration of similar approaches to increase computational efficiency as an important direction for future research.
run zero-order search and e-greedy
J.2 Dopulation-Basee Metrics
Unlike traditional population- based metrics such as FID or Inception Score [Salimans al., 2016]. our evaluation protocol focused single-image rewards (e.g., classifier confidence, brightness, compressibility). Evaluating on population-based metrics part of our future work. However, is worth noting that the type of single-objective optimization we employ has become somewhat of a standard practice recent works regarding reward optimization diffusion models; c.g. Yeh et al. [2025], Black et al. [2024] use the exact same experimental setup as we do. In addition, method poses no threat to the population- based guarantees that measures like FID try to make, since all our methods-br virtue of being inference-ting search algorithms that are completely gradient-free-avve avoid mode collapse and ensure distributional guarantees [Yeh et al., 2025].
J.3 Reward Overoptimization
For high enough X, the zero-order and e-greedy searches occasionally over-optimize the brightness and compressibility rewards, generating full-white and single-color images, respectively. This is due to the fact that, for large enough X, the local-search iterations enable finding noises potentially far
33

---

outside of) that no longer remain faithful to the context c. While (1) this only occurs with rewards r that grade samples Xy independent of their corresponding contexts c, indicating that such over-optimization could be construed as extreme efficacy of our method; and (2) this behavior can be easily mitigated by reducing/tuning A; alterations to this method to prevent over-optimization irrespective of hyperparameter settings remains important direction of future work. Examples include the regularization mechanisms proposed by Zhang et al. [2024], Tang et al. [2024].
J.4 Future Work
Directions for future work include:
Reducing computational overhead. Future work should aim lower the cost of per-step search, for instance through adaptive scheduling of NFEs or by learning proposal distributions that make candidate selection more efficient.
Exploring alternative bandit solvers. Hybrid methods, such as combining tree-based search with bandit algorithms, may improve sample efficiency and accelerate convergence.
Mitigating reward overoptimization. Regularization strategies could help reduce overfitting to specific reward signals.
Modeling diffusson-ssep interactions. More sophisticated search approaches should account for correlations between noise candidates across timesteps, rather than treating each step independently.
Learning compute-optimal regions. Identifying points during denoising where search is most valuable could allow resources be allocated selectively, improving efficiency without sacrificing quality.
Using reward ensembles. signals, can provide more stable guidance than relying on a single objective.
Averaging across multiple reward metrics, such different VLMM-baedd
Human-in-the-loop extensions. Future systems could be adapted operate in preference- -assed settings where only human ratings or feedback are available instead of closed-form reward functions.
34


Reviews:
[{"review_id": "0a1GZ9Xkgh", "reviewer": "Reviewer_ACK2", "summary": "This paper presents a practical method for optimizing noise trajectories at test time for diffusion models. By framing sampling as an MDP and relaxing it to a sequence of contextual bandits, the authors propose an efficient epsilon greedy search algorithm that adapts global vs. local exploration based on the denoising step. The method shows improved generation quality on EDM and Stable Diffusion.", "questions": "1. Can the authors comment on whether the effectiveness of this algorithm can be theoretically formalized?\n\n2. How robust is performance to changes in guidance scale or step count T?\n\n3. Can the authors explain the baseline selection and the lack of literature discussion?", "limitations": "NA", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strength:\n\n1. The empirical validations are provided for both class-conditional and text-to-image generation.\n\n2. I find following the flow of this paper smooth, which is good.\n\nWeakness:\n\n1. No theoretically grounded analysis or guarantees for the search algorithm's performance.\n2. To me, the compute cost of this method is not light, which is largely due to its MCTS nature. \n3. The related work part lacks many discussions of current literature on both the alignment of diffusion models and inference-time diffusion model scaling. I would encourage the authors to pay more attention to the literature. Besides, I find the experimental setting a bit weird. Many essential baselines are omitted in Table 2. There is a whole line of SMC-based algorithms that can effectively align DMs in inference-time (also requires some compute but can also scale well, for example https://arxiv.org/abs/2501.09685), whereas neither beam-search nor MCTS appears to be a classical baseline here.", "was_revised": true, "final_justification": "I thank the authors for the detailed rebuttal. I especially appreciate the clarifications on compute efficiency, the new regret bound analysis, and the extended results on robustness and baseline comparisons. These additions improve the paper's clarity and empirical contributions. If the next version nicely includes a more comprehensive baseline selection, I consider this a nice work.", "rebuttal": "We thank the reviewer for their feedback, and address their comments below:\n\n1. > To me, the compute cost of this method is not light, which is largely due to its MCTS nature.\n\nPlease refer to Table 3 in Appendix C (contained in `appendix.pdf` in the supplementary material zipfile), which reports number of function evaluations (NFEs) using $T$ denoising steps for each sampling method. In particular, we note the following:\n\n(1) Given $N$ noise candidates per timestep and $S$ simulations, MCTS requires $(N+S)T^2$ NFEs. In contrast $\\epsilon$-greedy only requires $NKT$ NFEs where $K$ is the number of local search iterations per timestep. In other words, for MCTS, NFEs scales quadratically in $T$, but **only linearly for $\\epsilon$-greedy**.\n\n(2) In particular, MCTS's error in approximating true exhaustive search approaches $0$ as $N, S \\to \\infty$; strong performance from MCTS requires an extremely high constant $N+S$ in front of the $T^2$ term. Indeed, MCTS in practice requires on the order of $1000$ simulations, hence $S >> T^2$ and dominates the cost.\n\n(3) For $\\epsilon$-greedy, the constant in front of $T$ is $\\propto K$ - but note that this assumes $K$ is fixed across timesteps. Letting $K_t$ be the number of local search iterations at timestep $t$, we can significantly bring down $\\mathbb{E}_t[K_t]$ by noting that $K_t$ only need be high at intermediate denoising steps where de-mixing occurs, and can be low at extreme (beginning and end) timesteps. We report $\\epsilon$-greedy performance as a function of $\\mathbb{E}_t[K_t]$ below:\n\n$\\mathbb{E}_t[K_t] = 1$ | $\\mathbb{E}_t[K_t] = 4$ | $\\mathbb{E}_t[K_t] = 8$ | $\\mathbb{E}_t[K_t] = 12$ | $\\mathbb{E}_t[K_t] = 15$ | $\\mathbb{E}_t[K_t] = 18$ | $\\mathbb{E}_t[K_t] = 20$ | \n----------------------- | ----------------------- | ----------------------- | ------------------------ | ------------------------ | ------------------------ | ------------------------ | \n 0.8312 / 0.7013 / 0.3261                  | 0.9998 / 0.7070 / 0.9999                   |  0.9997 / 0.6907 / 0.9998                    | 0.9995 / 0.6918 / 0.9999                  |  0.9983 / 0.6961 / 0.9793                     | 0.9974 / 0.7093 / 0.9997                    | 0.9813 / 0.7208 / 0.9885                         |\n\n***Table R1: EDM results by sampling method, varying the value of $\\mathbb{E}_t[K_t]$**. We use the same class labels and (other than $K$) algorithm hyperparameters as in Table 1 of the main paper. Each cell contains three values, corresponding to scores under the different reward functions: Brightness / Compressibility / Classifier.*\n\n**Indeed, $\\epsilon$-greedy attains strong performance across all reward functions, even with extremely low average $K$** (can retain performance within $0.002$ of the highest value with $\\mathbb{E}_t[K_t]$ as low as $4$), **highlighting our method's computational efficiency compared to MCTS.**\n\n2. > The related work part lacks many discussions of current literature on both the alignment of diffusion models and inference-time diffusion model scaling... I find the experimental setting a bit weird. Many essential baselines are omitted in Table 2. There is a whole line of SMC-based algorithms that can effectively align DMs in inference-time..., whereas neither beam-search nor MCTS appears to be a classical baseline here. Can the authors explain the baseline selection and the lack of literature discussion?\n\nWe appreciate the reviewer pointing out this reference, and will surely cite it in the final paper. That being said, please see **response #2 in the below rebuttal to reviewer NHkD**. This contains a comprehensive justification of our current choice of baselines—which actually captures the *entirety of the literature* on test-time noise trajectory optimization.\n\n3. > Can the authors comment on whether the effectiveness of this algorithm can be theoretically formalized?\n\nTo strengthen our claims on $\\epsilon$-greedy's superiority, we derive its regret bound, showing that—even with only black‑box access to a Lipschitz surrogate—simple regret decays sublinearly in the total number of global samples.\n\n**Theorem 1 (Regret of $\\epsilon$‑greedy search).**  \nLet the noise at diffusion step $t$ be $\\mathbf{z}_t \\sim \\mathcal{N}(0,\\sigma_t^{2}\\mathbf{I}_d)$. \n\nFix a confidence level $\\eta\\in(0,1)$ and truncate noise-space to the high‑probability ball  $\\Gamma_t =$ {$z\\in\\mathbb{R}^d : \\|z\\|\\le \\sigma_t\\sqrt{d\\ln(1/\\eta)}$}. \nAssume the surrogate reward $\\tilde R_t:\\Gamma_t\\to[0,1]$ is $L$‑Lipschitz and satisfies  \n$$\n\\bigl|\\tilde R_t(\\mathbf{z})-r(x_0(\\mathbf{z}), \\mathbf{c})\\bigr|\\le\\delta\n\\quad\\forall \\mathbf{z}\\in\\Gamma_t.\n$$ \n\nRun Algorithm 1 ($\\epsilon$‑greedy local search) for $K$ iterations with $N$ candidate points per iteration and global‑exploration probability $\\epsilon$, and let the final pivot be $p_t$. Then the expected simple regret at step $t$ obeys\n$$\n\\mathbb{E}\\bigl[\\max_{\\mathbf{z}\\in\\Gamma_t}\\tilde R_t(\\mathbf{z})-\\tilde R_t(p_t)\\bigr] \\leq \\delta +C_d(\\epsilon NK)^{-1/d},\n$$\n\nwhere the constant $C_d>0$ depends only on the dimension $d$.  Summing over $T$ diffusion steps yields\n$$\nR_{\\mathrm{tot}}=\\sum_{t=1}^T \\mathbb{E}\\bigl[\\max_{\\mathbf{z}\\in\\Gamma_t}\\tilde R_t(\\mathbf{z})-\\tilde R_t(p_t)\\bigr] \\le T\\Bigl[\\delta + C_d(\\epsilon NK)^{-1/d}\\Bigr].\n$$\n\n*Proof.*  By Gaussian concentration, $\\mathbf{z}_t$ lies in $\\Gamma_t$ with probability at least $1-\\eta$, so we may restrict attention to $\\Gamma_t$ at negligible cost.  Each global‐exploration draw is uniform on $\\Gamma_t$, and in total the algorithm collects $M=\\epsilon N K$ such draws.  It is a well-known result that the best of $M$ uniform samples over a $d$‑dimensional, Lipschitz‐reward domain achieves expected simple regret on the order of $M^{-1/d}$.  Since the $\\epsilon$‑greedy pivot never decreases in surrogate reward, its final surrogate value is at least that of the best global sample, inheriting the same $\\Theta(M^{-1/d})$ rate.  Finally, replacing the surrogate by the true reward adds at most $\\delta$, yielding the per‑step bound; summing over $T$ steps gives the stated total regret. $\\blacksquare$\n\n\n**Note the following key points:**\n- **Polynomial decay in global samples.**  With $M=\\varepsilon N K$ global draws per step, regret scales as $O(M^{-1/d})$, so increasing the total sample budget drives simple regret to zero at the classic random‑search rate.\n- **Bias floor $\\delta$.** In practice we can take $\\delta$ to be quite low as it is simply the approximation error induced by application of Tweedie's formula, $|r(\\mathbf{x}_0, \\mathbf{c}) - r(\\hat{\\mathbf{x}}_0^{(t)}, \\mathbf{c})|$.\n\n4. > How robust is performance to changes in guidance scale or step count T?\n\nSee Tables R2 & R3 below. These results confirm the **superiority of $\\epsilon$-greedy irrespective of guidance scale and $T$**.\n\n| Method                   | cfg = 1.5 | cfg = 3.0 | cfg = 7.5  | cfg = 9.0 | cfg = 12.0 |\n| ------------------------ | --------- | --------- | ---------- | --------- | ---------- |\n| **Naive Sampling**       | 0.2446    | 0.2635    | 0.2457     | 0.2650    | 0.2671     |\n| **Best-of-$4$ Sampling** | 0.2612    | 0.2743    | 0.2776     | 0.2783    | 0.2759     |\n| **Beam Search**          | 0.2673    | 0.2314    | 0.2682     | 0.2801    | 0.2780     |\n| **MCTS**                 | 0.2828    | 0.2712    | 0.2885     | 0.2971    | 0.2827     |\n| **Zero-Order Search**    | 0.2539      | 0.2612      | 0.2689     | 0.2443      | 0.2801       |\n| **$\\epsilon$-greedy**    | **0.2958**      | **0.3092**      | **0.2973** | **0.3055**     | **0.2994**       |\n\n***Table R2: SD results by sampling method, varying the classifier-free guidance scale**. We use the same prompts and algorithm hyperparameters as in Table 2 of the main paper. We evaluate only CLIP reward as the brightness and compressibility rewards should be generally insensitive to classfier-free guidance (but indeed find that CLIP too is largely insensitive to cfg).*\n\n| Method                   | $T=18$          | $T=32$          | $T=64$          | $T=128$         | $T=256$         |\n| ------------------------ | --------------- | --------------- | --------------- | --------------- | --------------- |\n| **Naive Sampling**       | 0.4965 / 0.3563 | 0.3912 / 0.3921       | 0.4621 / 0.4549        | 0.48493 / 0.4052      | 0.4541 / 0.4250       |\n| **Best-of-$4$ Sampling** | 0.5767 / 0.4220 | 0.5666 / 0.4992       | 0.6475 /  0.5084      | 0.5801 / 0.5101        | 0.5923 / 0.4917       |\n| **Beam Search**          | 0.6334 / 0.4679 | 0.6013 / 0.4321       | 0.5742 / 0.5023       | 0.6162 / 0.4517       | 0.6618 / 0.4879       |\n| **MCTS**                 | 0.7575 / 0.5395 | 0.6910 / 0.5153       | 0.6831 / 0.5624            | 0.6957 / 0.5018            | 0.7494 / 0.5789            |\n| **Zero-Order Search**    | 0.6083 / 0.3751 | 0.9999 / 0.7417 | 0.9999 / **0.7607** | 0.9999 / 0.7671 | 0.9999 / **0.7677** |\n| **$\\epsilon$-greedy**    | **0.9813** / **0.7208** | **0.9999** / **0.7424** | **0.9999** / 0.7588 |**0.9999** / **0.7680** | **0.9999** / 0.7672 |\n\n***Table R3: EDM results by sampling method, varying the value of $T$**. We use the same class labels and (other than $T$) algorithm hyperparameters as in Table 1 of the main paper. Each cell contains two values, corresponding to scores under the Brightness / Compressibility reward functions (we omit the classifier reward due to limited compute).*\n\n**We hope that with the provided responses & additional strong results, you consider raising your score to support clear acceptance. We are happy to address any additional questions.**\n\n[1] Uehara et al., Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review, 2025.\n\n[2] Ma et al., Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps, 2025.\n\n[3] Li et al., Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding, 2024."}, {"review_id": "XvQHICYCbZ", "reviewer": "Reviewer_NHkD", "summary": "This work propose an efficient test-time ω-greedy search algorithm, inspired by the contextual bandit problem, for optimizing noisy trajectories in diffusion models under non-differentiable reward constraints  and for efficiently optimizing noisy trajectories in diffusion models by formulating denoising as a contextual bandit problem.  It empirically outperforms Monte Carlo Tree Search  and demonstrates global adaptive exploration and local exploitation  behaviors associated with specific diffusion steps, especially being more localized in intermediate steps involving demixing. This method achieves state-of-the-art performance in image generation and is also the first practical method to optimize non-differentiable rewards during test-time sampling.", "questions": "Refer to weaknesses part.", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This study proposes an efficient test-time ω-greedy search algorithm inspired by contextual bandits for optimizing noisy trajectories in diffusion models under non-differentiable reward constraints. The algorithm empirically outperforms Monte Carlo Tree Search and demonstrates adaptive exploration  and local exploitation behaviors associated with specific diffusion steps, especially in the intermediate steps involving de-mixing. Through extensive comparisons with other search strategies, this study quantitatively demonstrates the critical importance of step-size adaptive search locality and shows the effectiveness of the method on EDM and stable diffusion, achieving sample quality improvement on various non-differentiable rewards such as classifier probability and VLM alignment. ​ \n\nWeaknesses: While I am not very familiar with this specific research area, the authors' presentation and experimental results seem methodologically sound and logically consistent. However, I was unable to find broader implications, limitations, and future work in the appendix. In addition, the authors only provide a small number of baselines to verify the effectiveness of the proposed testing-time strategy. I would like to know the impact of the authors' algorithm on a wider range of baselines.", "was_revised": true, "final_justification": "Thanks for your reponse! I maintain my rating.", "rebuttal": "We thank the reviewer for their feedback, and address their comments below:\n\n1. > ...I was unable to find broader implications, limitations, and future work in the appendix\n\nPlease see Appendices E & F for a comprehensive discussion of broader impacts, limitations, & future work. *Note that the appendix is in `appendix.pdf` in the supplementary material zipfile.*\n\n2. > In addition, the authors only provide a small number of baselines to verify the effectiveness of the proposed testing-time strategy. I would like to know the impact of the authors' algorithm on a wider range of baselines.\n\nWe claim that **our current choice of baselines actually exhausts all comparable methods currently existing in the literature**:\n\nThe main classes of algorithms for diffusion model test-time alignment are: (1) best-of-N sampling, (2) SMC-based methods [1], (3) classifier-guidance, & (4) search-based algorithms [2] (optionally paired with value estimation, e.g. SVDD-MC [3]). Our work specifically targets the regime of test-time methods that **(i) require no additional training** and **(ii) can handle arbitrary, non-differentiable rewards**. This already excludes methods such as classifier-guidance and search with learned value functions (which, we note are *orthogonal to the search algorithm itself*—we expect strong performance from combining our $\\epsilon$-greedy algorithm with soft value functions, and leave this to future work).\n\nRegarding SMC guidance: while useful for reward-weighted sampling, it is conceptually and algorithmically distinct from the noise trajectory optimization problem we focus on. SMC operates by generating forward trajectories under the model’s sampling distribution and resampling them based on reward-weighted importance weights, without intervening in the model’s latent dynamics. In contrast, we explicitly optimize the injected noise $z_t$ at each diffusion step $t = T, \\dots, 1$—allowing our method to *explore counterfactual noise choices at each step* and perform structured, fine-grained search to directly optimize terminal reward. In contrast, SMC operates by implicitly sampling from a reward-weighted posterior, without ever modifying the underlying noise path. Our method is thus more akin to trajectory-level planning, with a focus on explicit reward maximization *at the per-instance level*. **The distinction lies in both control parameterization (noise vectors vs. resampling weights) and objective (per-query reward maximization vs. general posterior sampling).**\n\nFocusing on methods that treat the injected noise vectors as the control parameters, we note, as highlighted in [1], that **tree search algorithms over noise trajectories are severely underexplored in the literature**. To the best of our knowledge, the only existing works in this direction are best-of-N sampling and SVDD-PM [1] which is exactly equivalent to beam search with width 1—**both of which we baseline, in addition to more powerful baselines such as MCTS and zero-order search**. Indeed, to our knowledge we are the first work to explore MCTS and bandit-based searches over the denoising tree. Hence, we don't believe we have omitted any essential comparisons.\n\nThat being said, **we have included a suite of additional experiments that more comprehensively test our proposed method**:\n- Evaluate SD across classifier-free guidance scale values **(Response to Reviewer ACK2, Table R2)**\n- Evaluate EDM across # of denoising steps **(Response to Reviewer ACK2, Table R3)**\n- Evaluate EDM on 300 images, a 10x increase in dataset size **(Response to Reviewer Sj6K, Table R4)**\n- Evaluate SD on harder prompts containing compositional information **(Response to Reviewer Sj6K, Table R5)** \n- Evaluate SDXL [2] on Image-Reward [3] **(Response to Reviewer D5sF, Table R6)**\n\n**We hope that with the provided responses & additional strong results, you consider raising your score to support strong acceptance. We are happy to address any additional questions.**\n\n[1] Uehara et al., Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review, 2025.\n\n[2] Podell et al., SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis, 2023\n\n[3] Xu et al., ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation, 2023"}, {"review_id": "7Zu0MD4HGF", "reviewer": "Reviewer_Sj6K", "summary": "This paper frames test time scaling of diffusion as Markov Decision Process (MDP) with a terminal reward. To avoid excessive computation costs of tree-search methods, they view denoising as a sequence of independent contextual bandits. Then, they introduce $\\epsilon$-greedy search algorithm.", "questions": "1. Refer to weakness above.\n2. How many images are used for evaluation on Stable Diffusion?\n3. I suggest to show qualitative results in the main paper.", "limitations": "1. Refer to weakness above.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 1, "clarity": 1, "significance": 1, "originality": 1, "strengths_and_weaknesses": "Strengths\n1. The paper tackles important problem, test-time scaling of diffusion models.\n\nWeakness\n1. Lack of novelty\n\t- The only novel component is $\\epsilon$-greedy noise search.\n\t- SVDD [1] has already proposed method to deal with non-differentiable rewards unlike the last sentence of the abstract.\n2. Lack of discussion on the proposed method.\n\t- The paper does not discuss why $\\epsilon$-greedy method improves the sampling.\n3. Unclear explanation of overall sampling process\n\t- In Algorithm 1, there is no explanation about how the denoising network $D_\\theta$ is used. It seems the written algorithm is incomplete.\n4. Limited and unconvincing experimental support\n\t- Insufficient scale: The valuation on ImageNet is based on only 36 images, which is too small sample size to support or claim meaningful conclusion.\n\t- Oversimplified tasks: The text prompts used for evaluation are very simple. To demonstrate the true value of a test-time scaling method, the experiments should involve more complex prompts which includes compositional information.\n\t- Quality degradation: The qualitative results (Example 2 of Figure 5, Example 3 of Figure 6) suggest, the proposed method may degrade image quality like color saturation. The authors should discuss regarding tradeoff between sample fidelity and text alignment. I suggest authors to include quantitative metrics regarding this. (measuring FID or IS along with Brightness/ Compressibility/ Classifier)\n[1] Li et al., Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding,", "was_revised": true, "final_justification": "The major concerns I have at the first review was\n1. lack of novelty and justification.\n2. experiments that are conducted on too small set.\n3. experiments that are conducted on too easy task.\n4. no discussion regarding the side effect.\n\n1 is fully resolved, as the concern originates from misunderstanding.\n2 is partly resolved as they have conducted experiments on the 300 images, which is statistically meaningful. \n3 is resolved.\n4 is partly resolved as they explain qualitative result in text and suggest solution for it.", "rebuttal": "We thank the reviewer for their feedback, and address their comments below:\n\n1. > Lack of novelty.\nThe only novel component is eps-greedy noise search.\nSVDD has already proposed method to deal with non-differentiable rewards unlike the last sentence of the abstract.\n\nWe would like to respectfully push back against the reviewer's statement re. lack of novelty; we outline our novel contributions below:\n\n- We are the first to investigate sophisticated step-level search algorithms (e.g. MCTS) for denoising.\n- In claiming that this is the first **practical** method for test-time noise trajectory optimization of arbitrary (non-differentiable) rewards, we mean that we are the first to **outperform MCTS** (perhaps the most generally performant tree search algorithm across domains) and **do so with compute-efficiency** ($\\epsilon$-greedy requires $O(T)$, not $O(T^2)$, NFEs). $\\epsilon$-greedy also achieves the **best scaling laws empirically across all sampling methods** (see Appendix D). SVDD [1] does not perform nearly as well; in fact, SVDD-PM is formally equivalent to beam search with beam width 1, and we show $\\epsilon$-greedy beats a **strictly more performant** variant of SVDD-PM (namely, beam search with higher width).\n- Ours is the first work to demonstrate the importance of *adapting the search strategy – local vs. global – based on the diffusion timestep.*\n- To expand further on the comparison to SVDD [1]: it's worth distinguishing between choice of *reward* and choice of *search algorithm*. We focus our work on developing a highly-performant and compute-efficient **search algorithm** for finding the optimal noise trajectory given a reward. Hence, in our experimental setups, we **fix** a downstream reward function $r$ (one of brightness/compressibility/classifier); **fix** a per-step reward function $\\tilde{R}_t = r(\\hat{\\mathbf{x}}_0^{(t)})$; **and then vary** the choice of search algorithm (rejection sampling, beam search, MCTS, $\\epsilon$-greedy). Indeed, the SVDD search algorithm is beam search (*SVDD-PM is identical to beam search with a Tweedie estimate-based reward*), and **$\\epsilon$-greedy outperforms this baseline across the board.** We treat learning a value function as in SVDD-MC to be orthogonal to choice of search algorithm; we would indeed be interested in testing $\\epsilon$-greedy with soft value functions, and while this is outside our paper's scope, we consider this an important direction of future work.\n\n2. > The paper does not discuss why eps-greedy method improves the sampling.\n\nPlease see Section 5.1.3 of the paper, specifically the subsections \"Leaving the denoising tree\" & \"Explore globally and exploit locally\" - both of which provide intuitions about why $\\epsilon$-greedy improves sampling. We provide a short summary of these sections below.\n\n- **Leaving the denoising tree.** Rejection sampling, beam search, and MCTS all optimize a fixed denoising tree, i.e. at each timestep $t$ we cannot explore “noise space” beyond the $N$ initially determined candidates. Zero-order search and $\\epsilon$-greedy address this, with the amount of exploration in noise space dictated by the number of local search iterations $K$, the maximum step size scaling factor $\\lambda$, and the global exploration factor $\\epsilon$.\n- **Explore globally and exploit locally.** At each time $t$, conditioned on latent $\\mathbf{x}_t$, there is a specific “region” in the noise space (call it $\\mathcal{Z}$) containing noise realizations that best fit the downstream objective. In particular, $\\mathcal{Z}$ is not necessarily a subset of $\\mathcal{N}(0, \\sigma_t^2\\mathbf{I})$. If our initial pivot is far from $\\mathcal{Z}$, methods such as zero-order search only allow for local exploration, i.e. we can move at most distance $K\\lambda \\sqrt{d/2}$ away from the pivot over $K$ local search iterations; but with $\\epsilon$-greedy, we can explore the entire Normal with $\\epsilon$ probability; a few such “random jumps” can move us into $\\mathcal{Z}$, at which point we see large returns from hill-climbing. Furthermore, we empirically confirm that the reward landscape as a function of $\\mathbf{x}_t$ is much smoother at extreme time steps rather than intermediate ones, hence there is the most performance gain from hill climbing at intermediate steps which $\\epsilon$-greedy alone can exploit.\n\nPlease also see **response #3 to Reviewer ACK2**, which derives the $\\epsilon$-greedy regret bound and shows that—even with only black‑box access to a Lipschitz surrogate—simple regret decays *sublinearly in the total number of global samples*. Other sampling methods such as zero-order search *cannot attain this regret bound*.\n\n3. > In Algorithm 1, there is no explanation about how the denoising network is used. It seems the written algorithm is incomplete.\n\nIn the main paper, we define the function $f$ as one step of the chosen sampler (e.g. lines 2-6 of Alg. 2 for EDM, lines 2-3 of Alg. 3 for DDIM - see Appendix A in `appendix.pdf` in the supplementary material zipfile). A call to $f$ thus involves a forward pass through the denoising network; and $f$ is indeed called in line 3 of Alg. 1. \n\nWe agree however that this could be explained more clearly, and will surely do so in the final paper if accepted.\n\n4. > Insufficient scale: The valuation on ImageNet is based on only 36 images, which is too small sample size to support or claim meaningful conclusion.\n\nBelow, we report EDM results using the ImageNet reward on 300 images, **a 10x increase in dataset size**. This **confirms the superiority** of our $\\epsilon$-greedy search:\n\n| **Reward**           | **Naive** | **Best-of-$4$ Sampling** | **Beam Search** | **MCTS** | **Zero-Order Search** | **$\\epsilon$-greedy** |\n| -------------------- | --------- | ------------------------ | --------------- | -------- | --------------------- | --------------------- |\n| **Classifier** |       0.3738                     | 0.6413           |  0.6928   |  0.7311                 |          0.6818     |**0.9790**\n\n***Table R4: EDM results by sampling method using classifier reward**. We use the same algorithm hyperparameters as in Table 1 of the main paper.*\n\n5. > Quality degradation: ...the proposed method may degrade image quality like color saturation. The authors should discuss regarding tradeoff between sample fidelity and text alignment. I suggest authors to include quantitative metrics regarding this. (measuring FID or IS)\n\nUnfortunately, we did not have enough compute to generate the $50K$ images recommended for accurate FID computation. However, a qualitative study revealed that for high enough $\\lambda$, the zero-order and $\\epsilon$-greedy searches occasionally over-optimize the brightness and compressibility rewards, producing the examples cited by the reviewer. This is due to the fact that, for large enough $\\lambda$, the local-search iterations enable finding noises potentially far outside $\\mathcal{N}(\\mathbf{0}, \\sigma_t^2\\mathbf{I})$ that no longer remain faithful to the context $\\mathbf{c}$. While (1) this only occurs with rewards that grade samples independent of their corresponding contexts, indicating that such over-optimization could be construed as extreme efficacy of our method; and (2) this behavior can be easily mitigated by reducing/tuning $\\lambda$; alterations to this method to prevent over-optimization irrespective of hyperparameter settings remains an important direction of future work. E.g. in followup work we hope to incorporate the regularization mechanisms proposed by [2], etc.\n\n5. > Oversimplified tasks: The text prompts used for evaluation are very simple. To demonstrate the true value of a test-time scaling method, the experiments should involve more complex prompts which includes compositional information.\n\nBelow, we provide results evaluating each of our methods on **prompts with compositional information**. Following [3], each prompt is of form \"a [SUBJECT] [PHRASE],\" e.g. \"a lion washing the dishes.\" **$\\epsilon$-greedy is again superior across the board:**\n\n| Method                   | Brightness | Compressibility | CLIP       |\n| ------------------------ | ---------- | --------------- | ---------- |\n| **Naive Sampling**       | 0.4579     | 0.6866          | 0.2578     |\n| **Best-of-$4$ Sampling** | 0.5257     | 0.7639          | 0.2618     |\n| **Beam Search**          | 0.5179     | 0.7690          | 0.2646     |\n| **MCTS**                 | 0.5921     | 0.7717          | 0.2954     |\n| **Zero-Order Search**    | 0.4573     | 0.7128          | 0.2802     |\n| **$\\epsilon$-greedy**    | **0.6312** | **0.7869**      | **0.3053** |\n\n***Table R5: SD results by sampling method, across reward functions, on compositional prompts**. We use the same hyperparameters as in Table 2 of the main paper.*\n\n6. > How many images are used for evaluation on Stable Diffusion?\n\nFollowing [4], we used natural language versions of the same 36 class labels used in EDM generation for the Stable Diffusion experiments. We plan to report results on a **10x larger** set of 300 prompts very soon within the reviewer-author discussion period.\n\n7. > I suggest showing qualitative results in the main paper.\n\nWe appreciate the suggestion, and will definitely add qualitative results to the main paper if accepted.\n\n**We hope that with the provided responses & additional strong results, you consider raising your score to support clear acceptance. We are happy to address any additional questions.**\n\n[1] Li et al., Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding, 2024\n\n[2] Zhang et al., Confronting reward overoptimization for diffusion models: A perspective of inductive and primacy biases, 2024\n\n[3] Black et al., Training Diffusion Models with Reinforcement Learning, 2023\n\n[4] Ma et al., Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps, 2025"}, {"review_id": "iJCk2rYsKZ", "reviewer": "Reviewer_D5sF", "summary": "This is a quite interesting and meaningful work, presenting a novel test-time scaling algorithm for diffusion models. The core idea is to treat each sampling step as an independent contextual bandit problem and then use an ϵ-greedy search algorithm to achieve a more comprehensive search. Compared to some previous test-time scaling algorithms, it is remarkably efficient and performs exceptionally well.", "questions": "No.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strength**:\n\n1. The problem of test-time scaling in diffusion models is notably challenging. The method proposed by the authors is highly novel, and I personally believe its exceptional performance stems from the $\\epsilon$-greedy search algorithm, which ensures that each sampling action is based on the previous iteration, thereby achieving a \"cumulative gain\" effect.\n\n2. The authors' proposed method is simple and efficient, which does not detract from its novelty or its beneficial impact on other works. For instance, this test-time scaling algorithm might enable step distillation algorithms for diffusion models to achieve superior results.\n\n3. The writing and logic of the article is clear and very engaging for the reader.\n\n**Weakness**:\n\nI believe this paper has two main shortcomings.\n\n1. The experiments in the paper are relatively limited. The authors might consider conducting further experiments on SDXL, SD3.5, and FLUX, evaluating them on benchmarks such as HPS v2, Image-Reward, GenEval, and T2I-CompBench.\n\n2. The approach of predicting $x_0$ solely using Tweedie's formula with a base model remains too crude, potentially yielding relatively low benefits, especially during the initial sampling stages. The authors could consider using a few-step model distilled from the corresponding base model to execute Tweedie's formula, which might lead to better results.", "was_revised": true, "final_justification": "Since the author added new experiments that resolved my concerns, I have raised my score to \"strong accept.\"", "rebuttal": "We thank the reviewer for their feedback, and address their comments below:\n\n1. > The experiments in the paper are relatively limited. The authors might consider conducting further experiments on SDXL, SD3.5, and FLUX, evaluating them on benchmarks such as HPS v2, Image-Reward, GenEval, and T2I-CompBench.\n\n**We have included a suite of additional experiments that more comprehensively test our proposed method**:\n- Evaluate SD across classifier-free guidance scale values **(Response to Reviewer ACK2, Table R2)**\n- Evaluate EDM across # of denoising steps **(Response to Reviewer ACK2, Table R3)**\n- Evaluate EDM on 300 images, a 10x increase in dataset size **(Response to Reviewer Sj6K, Table R4)**\n- Evaluate SD on harder prompts containing compositional information **(Response to Reviewer Sj6K, Table R5)** \n- Evaluate SDXL [1] on Image-Reward [2] **(Table R6 below)**\n\n| **Reward**           | **Naive** |  **MCTS** | **Zero-Order Search** | **$\\epsilon$-greedy** |\n| -------------------- | --------- |  -------- | --------------------- | --------------------- |\n| **Image-Reward** [2] | 0.4668      | 0.5152     | 0.5036                  | **0.5830**                  |\n\n***Table R6: SDXL results by sampling method, using Image-Reward [2] as the reward function**. We use the same prompts and algorithm hyperparameters as in Table 2 of the main paper. We normalize Image-Reward score to $[0,1]$.*\n\n**$\\epsilon$-greedy remains the highest-performing sampling method across all these experimental settings, highlighting its superiority across reward functions and diffusion model families/sizes.**\n\n\n2. > The approach of predicting x_0 solely using Tweedie's formula with a base model remains too crude, potentially yielding relatively low benefits, especially during the initial sampling stages. The authors could consider using a few-step model distilled from the corresponding base model to execute Tweedie's formula, which might lead to better results.\n\nPast work has shown that computing reward on the Tweedie-estimate $\\mathbb{E}[\\mathbf{x}_0 \\mid \\mathbf{x}_t, \\mathbf{z}_t, \\mathbf{c}]$ **actually outperforms** using learned soft value functions or few-step models. [3] compares the performance of SVDD-MC (beam search using a learned value function) and SVDD-PM (beam search using reward computed on $\\hat{x}_0$ estimated via Tweedie's formula). They find that SVDD-PM is generally more robust since it does not require additional learning (i.e., it directly utilizes the ground truth reward feedback); in fact, SVDD-PM outperforms SVDD-MC across the majority of domains and experimental settings.\n\nThat being said, while learning value functions (e.g. with distillation or even Q-learning algorithms such as DDPG) is orthogonal to our current work (which focuses on choice of *search algorithm* rather than per-step reward function), we are excited by this direction and consider it important future work.\n\n\n**We hope that with the provided responses & additional strong results, you consider raising your score to support strong acceptance. We are happy to address any additional questions.**\n\n[1] Podell et al., SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis, 2023\n\n[2] Xu et al., ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation, 2023\n\n[3] Li et al., Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding, 2024"}]

Author rebuttal:

Meta-review: This paper proposes an approach to enable test-time scaling of diffusion models. Instead of increasing the number of sampling steps, which yields quickly diminishing returnsl, the idea is to optimize the sequence of injected noise vectors. This nontrivial task is done by framing sampling as an MDP and relaxing it to a sequence of contextual bandits. The effectiveness of the proposed method is demonstrated empirically by extensive experiments. One reviewer initially had concerns, but most of them got resolved during the rebuttal/discussion period. Post rebuttal, both reviewers and I found the paper worth publishing, and I'm happy to recommend acceptance.
Decision: Accept (poster)