Diffusion Tree Sampling: Scalable inference‑time alignment of diffusion models
Authors: Vineet Jain
Kusha Sareen
Mohammad Pedramfar
Siamak Ravanbakhsh
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion models
inference-time alignment
Monte Carlo Tree Search
test-time sampling and search
generative models
Primary area: deep_learning
TL;DR: We build a Monte Carlo Tree over the diffusion denoising process that can be used for scalable, compute-efficient, inference‑time alignment of pretrained diffusion models to new reward functions

Abstract: Adapting a pretrained diffusion model to new objectives at inference time remains an open problem in generative modeling. Existing steering methods suffer from inaccurate value estimation, especially at high noise levels, which biases guidance. Moreover, information from past runs is not reused to improve sample quality, leading to inefficient use of compute. Inspired by the success of Monte Carlo Tree Search, we address these limitations by casting inference-time alignment as a search problem that reuses past computations. We introduce a tree-based approach that _samples_ from the reward-aligned target density by propagating terminal rewards back through the diffusion chain and iteratively refining value estimates with each additional generation. Our proposed method, Diffusion Tree Sampling (DTS), produces asymptotically exact samples from the target distribution in the limit of infinite rollouts, and its greedy variant Diffusion Tree Search (DTS*) performs a robust search for high reward samples. On MNIST and CIFAR-10 class-conditional generation, DTS matches the FID of the best-performing baseline with up to $5\times$ less compute. In text-to-image generation and language completion tasks, DTS* effectively searches for high reward samples that match best-of-N with $2\times$ less compute. By reusing information from previous generations, we get an _anytime algorithm_ that turns additional compute budget into steadily better samples, providing a scalable approach for inference-time alignment of diffusion models.

Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models
Vineet Jain
Kusha Sareen
Mohammad Pedramfar Siamak Ravanbakhsh
Mila Québec Al Institute, McGill University
Jain.vineed kusha. sareen, mohammad pedramf ar, si amak. ravanbakhah))mm aimia.quubbcc
Abstract
Adapting a pretrained diffusion model to new objectives at inference time remains open problem in generative modeling. Existing steering methods suffer from inaccurate value estimation, especially at high noise levels, which biases guidance. Moreover, information from past runs is not reused to improve sample quality, resulting in inefficient use of compute. Inspired by the success Monte Carlo Tree Search, we address these limitations by casting inference-ting problem that reuses past computations. We introduce a tree-based approach that samples from the reward-aligned target density by propagating terminal rewards back through the diffusion chain and iteratively refining value estimates with each additional generation. Our proposed method, Diffusion Tree Sampling (DTS), produces asymptotically exact samples from the target distribution in the limit of infinite rollouts, and its greedy variant, Diffusion Tree Search (DTS"), performs a global search for high reward samples. On MNIST and CIFAR-10 class-conditional generation, DTS matches the FID of the best-performing baseline with up to 10x less compute. In sotttoo-magee generation and language completion tasks, DTS* effectively searches for high reward samples that match best-of-N with up to 5x less compute. By reusing information from previous generations, we get anytime algorithm that turns additional compute into steadily better samples, providing a scalable approach for inference-time alignment of diffusion models. Project page: htttps:/diiffie tree- wwwre--smppll.. samplismppssasesaa gi github.is
alignment as search
Introduction
Diffusion models have emerged one of the most powerful frameworks for generative modeling, achieving state-of-the-an results across a wide range of modalities, including image synthesis [30 (75) 64), molecule conformer generation [32 89], and text generation [68; 48). Despite their success, adapting a pretrained diffusion model satisfy new, user-defined objectives at inference time without expensive retraining or fine-tuning remains major challenge [79]
Most objectives can be cast as reward function, turning alignment into posterior sampling problem where the target is to sample from the pretrained model density weighted by exponentiated reward. The key challenge is that rewards are only available at the end of the denoising trajectory. So, inference-tine time alignment seeks guide the denoising process based on unseen terminal rewards.
range of different methods have been proposed - gradient-basee guidance uses reward gradients to bias the denoising process; sequential Monte Carlo (SMC) 84 77
[7:10
3 where one
7
I9 39)
which maintains a population of particles and resamples them during denoising based on an estimate of terminal rewards; or more recently, search-basee methods [45] 46 50] that perform a local greedy search based on approximate rewards. The common issue undermining all of these methods that they rely on certain approximations to estimate the unseen terminal rewards. As we demonstrate in Section these approximations bias decisions and degrade sample quality.
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

emagl baby
making
licking littting
Bowen the that of triengli
latte
Bohwwroooogorr
start. blue
"Diifuuiio"
Stable Diffusion 1.5
Diffaisoon
Stable Diffusion 1.5 +
FK-SStering
Drufimaon
Stable Diffusion 1.5
DTS (Ours)
DiFusion
Figure 1: Sample text-image pairs using Stable Diffusion v1.5 [64] and ImageReward [8] as the guiding function, with generated samples picked at random for each method and prompt.
We therefore address the following challenges questions in this work: (1) how to guide the diffusion process at inference-ting when rewards are available only at the end? This also known as the credit assignment problem in reinforcement learning (RL) literature [5]]; (2) inference-time samples can potentially inform and improve future samples how to systematically use this information in a sequential yet scalable sampling process?
Fortunately, RL also provides solution that has been historically quite successful in addressing both challenges- Monte Carlo Tree Search (MCTS) 6 We therefore ask: can we leverage MCTS for steering diffusion models? We observe that during denoising, the pretrained diffusion model can be viewed as a deterministic policy, while the reverse process Gaussian step can be viewed as stochastic environment transition. This is exactly the classical setting for MCTS [41, suggesting could use tree search to solve the problem of inference-time alignment.
Our proposed algorithm, Diffusion Tree Sampling (DTS) a novel inference-time alignment method that casts the denoising process as finite-horizon tree, where similar to MCTS, rollouts are used continuously improve value estimates for intermediate noisy states. For applications that require optimization, rather than sampling from the target density, we propose search variant Diffusion Tree Search (DTS*)- - that performs principled search in the space of denoising trajectories to identify the modes within high-volume regions of the target density.
Our contributions can be summarized as follows:
We formulate inference-time alignment of diffusion models as tree search problem for sampling from the reward-aligned distribution or optimizing for high reward samples.
We develop a general tree-based algorithm that yields asymptotically exact samples from the target distribution in the limit of infinite rollouts.
We demonstrate that DTS significantly reduces bias and variance in value estimation compared common approximations used by many existing methods.
We show that both DTS and DTS* scale more favorably compared to leading baselines and match their performance with up 10x less compute class-conditional image generation, and up to 5x less compute on eext-to-image alignment and language completion tasks.
2 Preliminaries
Diffusion models. Diffusion models 30 751 define a generative process via Markov chain that
/a, X VI a, E, <~ (0, I), where E {1, (1..... T}} indexes discrete time steps. and {az 1T, defines a noise schedule. The noise schedule chosen such that at t = T. the marginal distribution of the samples resembles simple fixed distribution, such standard Gaussian, p(xT) (0 I).
progressively adds noise to data Paota (x), referred to as the forward process,

---

A learned reverse process iteratively denoises samples: Po(x x₂) =N (o t), t) I). where o; the posterior variance calculated from the forward noise schedule and μs is parameterized typically by neural networks and optimized by minimizing the variational bound on the data likelihood or via denoising score matching. The generative process induces a distribution:
T
po (X ................................ xT-1, p(xx) II
po xt). p(x)) (O.
(I)
1=1
Alignment of diffusion models. Consider pretrained diffusion model po and optimality
variable O E (0, 1) which denotes whether a sample x ~ Po(x) satisfies some desirable property.
This is equivalent sampling from the posterior distribution p(x O 1) x px(x) p(O 1 x)- typical assumption is that p(O = 1 x) x cxp(Ar(x)) where r is some reward function and is the inverse temperature. For the rest of this paper, assume that À = 1 unless otherwise stated and we define the alignment problem as sampling from the target distribution or finding its mode, where Z is the normalization constant:
1
π (x) = px(x) exp(Ar(x)).
(2)
Reinforcement learning approach. Since the generative process in diffusion models defines Markov chain, we may consider the model pe as a policy. The target distribution '² can be seen as the optimal policy for the following objective:
1 " (x) == argmax (x [r(x)] - DKL. (- |po).
This is closely related to the maximum entropy RL objective 993 [2],
except that the entropy regularization is replaced by the KI. divergence with the pretrained model pa- We define the soft value function at timestep t as the expected exponentiated reward starting from x, and following pot
(x,,)
log (xo lxx)) (exp (Ar(x.))]
This soft value function satisfies the following recursive relation, analogous to the soft Bellman equation and exactly characterizes the optimal policy n²:
Vi(x;) = = log Epo (xx.) (exp A Vx(xx) = r(xo).
(5)
po ( x1) exp (AVi-1
o. (x xx) J ( x, ) exp ( dx;-1
(6)
This formulation explicitly connects optimal sampling with soft value estimation, motivating various practical approximations and sampling methods discussed subsequent sections. For completeness, we derive Equations (5) and (6) in Section In the rest of the paper, use V22 to denote the true soft value function and 0e to denote estimates.
Inference-time adaptation of diffusion models
One option to obtain the optimal policy in Equation 66 is to fine-tune the diffusion model using RL 21 5 8T [8]. Sampling from unseen reward functions would require guiding the denoising process during inference to align with the optimal policy without modifying the prior pretrained model. We discuss some of the most relevant works below, and works in related areas in Section A
Gradient-based guidance. One way sample from the optimal policy I² to use the first-order Taylor expansion of V2-1 around the pretrained mean (x t). This yields the gradient- based denoising step x ~N (Nr (x,, t) + A o Vx._, Vi-1 (x, 11,, af 7). This can be considered a form of classifier guidance [17] and is used in many proposed inference-ting steering methods [1] 3 27
The gradient approximation can be improved by using Monte Carlo samples for estimation 74
Sequential Monte Carlo. Particle-based methods are another very popular approach, where a population of samples is maintained approximately sample from the desired distribution. Sequential Monte Carlo (SMC) [15] uses potential functions, which usually approximate the soft value function, to assign weights particles and resample them at every step. Different variations of SMC have been proposed for diffusion model alignment [84] 77 7 19 39 (72). Classical SMC guarantees exact sampling in the limit of infinite particles and exact value estimation. In practice, however, the repeated sampling procedure can reduce diversity due to weight variance and inaccurate value estimates. We provide detailed background SMC for diffusion sampling in Section
3

---

2-
Sw50
75 T
One-step prediction x,(X,) using Tweedie's 's formula for different time steps, along with average mean squared error with the ground truth data samples. Close to / = 0, the predictions are fairly accurate, towards the maximum timestep = 99, they devolve into random predictions.
Figure 2:
Search-based methods. Recently, there has been a growing interest in using search-based methods to align diffusion models [50]. Most of these methods propose doing a local search 45 46) by
obtaining multiple denoising candidates at each step and selecting the best one based on their value. More recently, tree search has been combined with best-of-N [92], and an MCTS- based approach [9]] has been applied in the specific context of diffusion foreing [9] over sequences for planning. However, these methods either do not use an explicit backup mechanism, resulting in a limited local search 45 46 ?Σ [50, or they rely on inaccurate value estimates
[9]]. DTS, on the other hand, performs global credit assignment using all trajectories for asymptotically exact sampling.
3.1 The value estimation problem
Most existing methods use soft values from Equation
5 either by taking their gradient [11 74 as
91). Since estimating these soft values intractable, they employ the following set of approximations. The first step is to apply Jensen's inequality:
SMC potentials (84 (39,, or for search [45 46
Vx(x,) log Epo(xon- (xx) (exp (Ar(x)))] Epo(xo.- 1|x,) [x(x)]]
Estimation of this expectation is computationally expensive since it requires multiple rollouts from the current sample at timestep & to the clean sample. The expected reward is further approximated as
E Ene(xos [xx) [r(x))] x r(xo (x:)), where &o is the posterior mean obtained using Tweedie's 's formula 20 in single step:
(x;)) Ep (Xoss- [Xo]
(x+ (1 C. V log (x()))
(8)
Var
The posterior mean is an approximation because the true score function for intermediate marginal densities Vx, log Pr (x;) replaced by the learned score function. We investigate the effect this approximation for a diffusion model trained on a mixture of Gaussians in Figure 2 where we see that the prediction based on Tweedie's 'ss formula gets increasingly inaccurate for higher noise levels. Therefore, despite the wide adoption of this approximation, the value estimates used for guidance are essentially random higher noise levels even simple 2D settings.
3.2 Scaling with compute
Efficient utilization of available compute critical for any inference-time alignment algorithm. Existing SMC or search-based methods treat each sampling procedure as independent event, and all intermediate evaluations are discarded. Consider a streaming or repeated sampling setting. There is no mechanism to assimilate information from prior runs to improve sample quality. This could be particularly useful for correcting errors value estimation, which, as we saw above, is difficult at high noise levels. In other words, these methods scale parallelly by increasing particle count, but do not scale sequentially by turning extra compute into cumulative improvements estimate quality.
4 Diffusion Tree Sampling and Search
The pitfalls above suggest two complementary desiderata for an effective inference- time sampler:
(D1) Use information from low-noise timesteps, where the reward signal is reliable, to refine decisions made at high-noise timesteps, rather than treating every step in isolation.
(D2) Reuse information from previously explored trajectories that additional compute improves sample quality instead of merely increasing parallel particle count (this property is characteris- tic of anytime algorithm).

---

Best-of-N
SMC
DTS
Dendlding
1=0
selected
expandied/
terminal
denoosnng
backup
node
rollout
node
Figure 3: Illustration of various inference-time steering methods, where size of the node represents the associated values. Left: Best-of-N denoises multiple samples using the base diffusion model and selects the one with the highest reward. Center: SMC maintains a population of particles and resamples based an estimate of the value function. Right: DTS and DTS maintain tree that accumulates information across multiple rollouts and backs up the terminal reward to refine value estimates. The diagram illustrates the four phases: selection, expansion, rollout, and backup.
To address these issues, in this section we develop solution by first interpreting the denoising process as a tree in Section 4.1 We then introduce a general tree-based algorithm to sample from the target density in Section 4.2 and describe its application to diffusion model alignment in Section 4.3 Finally, in Section 4.4 we empirically evaluate our method on 2D datasets to validate, a very clear and controlled setting, whether DTS satisfies the desiderata mentioned above.
4.1 Denoising tree
The Markov property of the reverse diffusion chain naturally induces a finite horizon tree in Rdd where d is the dimensionality of the space over which we are diffusing. Here, the nodes at depth t represent noisy states x, and the edges represent a denoising step. Each node x, can be stochastically denoised into multiple children X;-1 ~ po(- xt).
This framing allows us to keep track of information across multiple denoising trajectories, including estimates of the soft value function, which helps with global credit assignment. Using the tree structure gives us the flexibility to sample from the target density or search for the highest reward sample with minimal changes the underlying algorithm. We call the sampling variant Diffusion Tree Sampling (DTS) and the search variant Diffusion Tree Search (DTS*).
4.2 Tree-based sampling
Similar MCTS, we construct a tree T, where nodes represent states x, and edges represent transitions pe p=(x;-1 xx) following the base diffusion model. Each node maintains the current state and timestep (x,, f), an estimate of the soft value function o(x;), and the visit count N N(x)). Since we do not have a fixed starting state, we introduce a dummy state as the root Xy that transitions to the prior in our diffusion model i.e., p(Xx xy+1) No0, I). Additionally, use C(x,) to denote the set of children of node X4.
The goal is to expand this tree while improving value estimates as we expand it, so that it can be used for approximate sampling from the target distribution at any time during the construction. The resulting tree sampling process provably samples from the target distribution TT in the limit infinite rollouts. The tree-building procedure of DTS repeats the following steps iteratively:
Selection. Starting from the root XT sample a child x E C(x.) from Boltzmann distribu- tion exp (Xx(xx)))) recursively until either unexpanded node is reached or / - O.
2. Expansion. If we reach node x, such that the number of children is less than the maximum allowed value and > O, we create a new child node x,-1 ~ po xx) and initialize 6(x-11) =
0, -
=1
3. Rollout. From the newly created node, perform rollout till terminal states x by recursively
sampling from (- x.) for t' = 1,..., O. An important distinction from traditional MCTS is that we add the rollout path to T.
5

---

4. Backup. Evaluate the terminal node using the reward function 0(x)) = r(xo) and use soft Bellman equation (Equation (5) to update parent node values using the children node values recursively for t = ... The visit counts for all nodes in the path are also updated.
Each traversal of the tree, from the root to the backup of the value function, constitutes one tree- building iteration. For sampling from T, we simply start from the root and perform selection steps until we reach a terminal node. A formal algorithm is provided in Section B
Proposition 1 (Asymptotic consistency). Let r be bounded and X > 0, then DTS produces sequence of terminal states whose empirical distribution converges to the optimal policy n' as the number of tree iterations M + CC.
Proof sketch. By construction, the tree policy selects X with unnormalized probability po(x++1 x₂) exp(A6(xx+1)), which is the optimal policy defined in Equation (6. By telescoping the product over t, we obtain the final samples at t = T are sampled from pe (x) exp(r(xx)). A more detailed proof is given in Section D
4.3 Design choices for diffusion alignment
The algorithm discussed above can be applied to any Markov chain. However, in this work, we apply it to the problem of inference-time alignment of diffusion models. We discuss various considerations and design choices below, with more implementation details in Section
Sampling or Search. DTS designed to sample from the target distribution xr" but for settings where a single high-quality sample is required, we introduce search variant, DTS*. It keeps the same soft value backup but modifies the selection step by always selecting the child with the largest soft value estimate instead of Boltzmann sampling. Since DTS* uses soft values, this is different from standard MCTS- implements a marginal-MAP (max-sum) inference scheme (63) over the tree, At every noise step, the algorithm selects the branch whose subtree carries the greatest mass under オ* and, once 0 reached, returns the highest-density leaf inside that dominant region. As we will see in the Section S this volume-based selection helps avoid reward over-optimization.
Branching. Extensions of MCTS to continuous spaces commonly use progressive widening [14] to decide the maximum number of branches (xx;) allowed per node based on the number of visits: B(x))= C (Nx))" C' > 0, often should be expanded more, since they represent more promising directions for denoising. We adopt the same strategy and during tree traversal, if we encounter a node such that [C(x,) | < B(x,)
(0, 1). The high-level intuition is that nodes that are visited more
and >> will always expand.
Exploration. There is rich literature on search methods for classical MCTS, and the most popular approach, UCT [11, is an application of upper-confidence bounds (2) to trees. We employ this
exploration strategy for DTS", i.e. we choose the child x,-1 E C(x,) the UCT estimate:
with the maximum value of
log N(x) N (x
UCT(x:-1) = (xx-11) + Cust
Cuct > 0.
(9)
For DTS, we do not employ explicit exploration, because, in practice observe that sampling
obviates the need for an exploration bonus or handcrafted mechanism. This also supported theoretically by recent work in the bandit setting [57].
Efficient implementation. The main computational cost is incurred when using the diffusion model proposal to sample new children or perform rollouts. We implement an efficient batched version of the algorithm by collecting nodes in a batch and performing single batched denoising rollout per iteration. The selection and backup steps involve simple tensor operations and pointer manipulation with negligible cost. Therefore, while the control flow of our method is sequential, the practical algorithm can be parallelized. Note that once the tree has been built, sampling is near instantancous by repeatedly selecting children without any model calls.
4.4 Illustrative experiments
In this section, we perform experiments on simple 2D settings answer the following questions:
Does DTS sample accurately from the target distribution?
.
Does reward backup in DTS result in more accurate value estimates (desideratum D1)?
Does sample quality of DTS improve with more inference-ting time compute (desideratum D2)?

---

Prior
Target
DPS-RG
SMC
TDS-RG
DAS
SMC-Rollout
DTS (Ours)
Figure Samples from the prior p(xo), target p(x.) exp(r(xo)) and different sampling methods at 10 NFEs. Top: The prior is an qqual-weighter mixture of Gaussians, and the reward function dis- tributes mass unevenly. Bottom: The prior has support on alternate square regions in a checkerboard pattern, and the reward function r(z, y) -0.5(z² + y²) is negative distance from the origin. More
details the experimental setup provided in Section
E22
Gaussian Mixture
Checkerbeard
Gaussian Mixture
Checkerbeard
0.0
100
Biant())
0.5
-0.3
0.5
20-1
10"1
3.5
Vertót
1.0
0.5
1
10
0.00
0.0
10"
10
10
10
10
10f
20
40
60
80
20 40
60
Inference Budget (NFEs)
Inference Budget (NFEs)
Diffusion timestop (t)
Diffusion timestep (t)
DTS (Qurs)
DPS-RG
-
SMC
TDS RG
DAS
+
SMC SMC--Rolouut
discrepancy (MMD) between generated samples and target ground truth samples as a function of number of function evaluations of the prior diffusion model. Right: Bias and variance of value estimates for different approaches at 10 NFRs.
Figure 5: Left: Maximum
Setting. We compare DTS with several inference-time steering methods including some which were originally proposed for posterior sampling in inverse problems. and adapt them to the reward- soaudd-guddncee gradient-based guidance only, (2) SMC gradient guidance. (4) DAS 39 SMC gradient guidance + tempering. We also implement a version of SMC, which we call
setting: (1) DPS-RG (our reward awward-guiedd version of DPS [
SMC (71) (3) TDS-RG (reward-guided version of TDS 84
SMC-Rollout, where the values are estimated via one full DDIM (73 rollout. For fair comparison, benchmark with respect the number of function evaluations (NFEs) of the diffusion model.
Results. Figure plots the samples obtained using different methods for two different settings. In accurately, with other methods distributing mass inaccurately to different areas of the support. In particular, gradient gradient-based methods like DPS-RG RG and TDS-RG suffer from instability and require gradient clipping to stabilize denoising steps. Figure [5 shows the Maximum mean discrepancy (MMD) between generated samples and ground truth samples for an inereasing amount of inference budget, measured in terms of the number of function evaluations of the diffusion model. The results empirically validate that the sample quality of DTS improves with more compute, satisfying desideratum D2.
both cases, DTS approximates the ground truth target density
analysis of value estimates. To verify DTS leads to more accurate soft value estimates, we perform biri-nari analysis. We estimate ground truth soft value estimates based on 1000 rollouts from noisy states and then compute the bias and variance of the value estimates. Figure 5 compares DTS, Tweedie's 's formula (SMC + variants), and single full DDIM rollout (MMC-Rollout) Both one-step denoising and single rollout have high bias and variance, which generally get worse for higher timesteps. Our tree- tree-based approach reduces bias by using accurate reward information and reduces variance by aggregating information from multiple rollouts. This empirically validates that DTS satisfies desideratum DI.
5 Experiments
We validate the efficacy of DTS and DTS* for image and text generation. Our experiments show that DTS draws faithful samples from high-dimensional posteriors, and DTS* efficiently searches high-dimensional image space to discover high-rewar samples. We also test our method on image
7

---

Table 1: Comparison of inference-time posterior sampling methods. We report the sdd of each
metric across the relevant classes and highlight +5% values from the best experimental value.
Datasct
MNIST
MNIST evencedd
CIFAR-10
Algoitian
1
(ID
CMMD (lugrts)) (T's Divenity G) (MMD(() Klkarla)) (1) Diversity (7) FID(D) CMMD ()) Ellaris)) (T? Diveskly
DPS
0.359 202 0,64 coa
-0.23 sor 0.434 sarer 0.123 1009 0.293 201s
-0002 suet
0,572 c009 1i1 2,609 2114 0,002 some
0.551 2000
SMC/FK
0.00
00002
0.22
0.027 9.123
-0003
0.312
1,09
0.457
TDS
0087
0.463 grae
.0000 soon
0.04 200 0.053 4an0 0,200 +00
-0001 4am
0.576
0.887 earo 2.675 =aws
.0.06 eeas
0.459 -nro
DAS
0,09
pe31 0,179 coas -0006 godte 5040 0.011 zoom 0.079 2001 -0015
200
0,003 0.24) saor 0,822 2130 +0.584 2330
00500 20re
DTS (ours) 0.014 0,068 coat
40023 04422 0.007
0.036
-0.010
0,597
2000 0,195 0,765
+0,305
(0.42
MNIST single class
MNIST even/odd
CIFAR-10 single class
109
10a
109
10-2
10a
10
10"
10
10a
10a
10
10*
10 Inference Budget (NFEs)
10
Inference Budget (NFEs)
Inference Budget (NFEs)
DTS (Ours)
DPS
SMC/FK
TDS
DAS
Figure 6: FID (lower better) versus number of function evaluations for different methods on MNIST single digit generation averaged over all 10 digits (left), MNIST odd and even digit generation (center), and CIFAR-10 single class generation averaged over all 10 classes (right). All methods were evaluated with 5000 generated samples per class.
inpainting task in Section I provide detailed memory/wall-clock time analysis in Section H and some additional results in Section A detailed description of each experimental setting is provided in Section F22 and details of baseline implementations are in Section G
5.1 Class-conditional posterior sampling
We evaluate DTS the task of sampling from class-conditioned posterior distribution p(x c) x po())p(c x) where (x) is a pretrained unconditional diffusion model and p(c x) is a classifier. This would correspond to setting r(x) log p(c x) in Equation (
Setting. We use MNIST and CIFAR-10 10 datasets, each with 10 classes. In both cases, the priors unconditional diffusion models in pixel-space
train one from scratch on MNIST and an off-the-shelf model for CIFAR-10. 10. For MNIST, we consider two settings: sampling from individual digits, and sampling from even/odd digits. The latter is multimodal posterior with reward function r(x) = (, loc = ? x) for the even digits and similarly for odd digits. For CIFAR-10, 10, we sample from individual classes.
Prior
DTS (Car)
DPS (Car)
SMC (Car)
TDS (Car)
Prior
DTS (five)
DTS (even)
DPS (three)
TDS (even)
7 5 3
4
5 5 5
2 0 4 2 6 2
3 3 3
3
2 8
2
0
1 2
7 8
3
5 5 5
8
4 6
6
3
3
3
0
0
0
a 6
4
9
$ 5
5
5
8
6
3
3
3
3
0
2
2
9
5 41
5
5
5
5
8
6
4
6
3
3
3
3
8
0
6
2
Figure 7: Samples generated from the CIFAR-10 (top) and MNIST (bottom) base diffusion models,
and posterior samples using different methods at 10th NFEs. Gradient-based guidance such as DPS can be unstable, leading to samples that lie outside the support of the prior. SMC-hased methods struggle to accurately sample from multi-modal distributions for MNIST even digits, TDS oversamples from the digit two and undersamples from the digit four, and for CIFAR-10 car, both SMC and TDS suffer from mode collapse.
8

---

Baselines. We compare the performance of DTS with DPS I.I SMC/FK 77I TDS [4 and DAS [9]. DPS, TDS, and DAS use reward gradients, while SMC/FK is derivative-free We report two distribution- based metrics - Fréchet Inception Distance (FID) and CLIP maximum mean discrepancy (CMMD) 37 that compare generated samples with ground truth samples from the dataset, in addition to average rewards and CLIP diversity (pairwise cosine distance).
Results. Table I reports the mean=sid of various metries for different methods after 10 NFEs. In all three settings, DTS achieves the lowest FID and CMMD by considerable margin, indicating it closely matches the true posterior. We observe that the margin of improvement on CIFAR-IO narrows slightly. We attribute this to reward noise: the CIFAR-10 10 classifier achieves an accuracy of ~ 855, so its logits provide noisier signal than the near- pear--t classifier used on MNIST. Even so, DTS still outperforms all baselines. TDS and SMC in particular show characteristics of mode collapse with very high average rewards and low diversity, whereas DPS often generates samples that lie outside the support of the base model. Figure shows that DTS achieves very low FID across different NFEs and has better scaling properties with more compute compared to existing methods. Figure Sesureeppr SigureZreeeenss example outputs from each method, highlighting their specific characteristics. We present samples for all classes and plots of additional metrics as a function of NFEs in Section J.2
presents
5.2 Text-to-image generation
Setting. We use Stable Diffusion v1.5 [64], a latent diffusion model, as the prior over 512 x 512 images x (x y) where y denotes the text prompt. We evaluate on two different settings:
(a) Brach [67], which is comprehensive benchmark of 200 prompts, with ImageReward [88] y(x,y) that encodes prompt accuracy as well as human preferences; and (b) following Black al. [5).
we use 45 common animals from the ImageNet-1000 1000 dataset as prompts, with the LAION aesthetics predictor [69] r(x) that encodes aesthetic quality of an image but does not check for prompt accuracy.
Baselines. A strong baseline for high-reward generation is best-of-N, which draws N samples from the base model and keeps the one with the highest reward. SMC has also been applied this problem [71], but (a) as discussed in Section 3.1 it relies on one-step value estimates that become inaccurate at high noise, and (b) Section 5. shows that it often collapses onto narrow modes. We compare DTS* with best-of-N and FK- Kteerrering (SMC) [7T]. We compare all methods in a compute-matched matched setting, measured in terms of the number of diffusion model calls (NFEs).
Figure& the maximum ImageReward and aesthetic score versus inference compute. DTS* outperforms the baselines for DrawBench prompts, see Figure | for examples. One strong
Results.
feature DTS* is its favorable scaling properties with more compute. In the aesthetic score setting, FK/SMC achieves the highest rewards across NFEs, but we see severe over-optimization prompts. DTS* manages to strike the right balance between achieving high rewards while still
all
3.8
SD 1.5 + ImageReward
Base
FK-stterrng (SMC)
DTS (Ours)
Reward
1.4
1.2
6.15
0.32
8.15
8.01
7:49
SestheticCcore
wold
1.5
7.89
7.83
7.82
7.37
7.0
dolphin
6.5 10
30
Inforence Budget INFEs)
0755 (Dur))
of -d-
FROSC
Figure 8: Left: Maximum ImageReward 88 vs. compute (NFEs) per prompt, averaged over 200
prompts from Draw 677, and maximum aesthetic score 69 vs. compute (NFEs) per prompt, averaged over 45 common animal prompts. Right: Samples generated using SD-v1.5 [64 for simple animal prompts and aesthetic score as the reward at 100k NFEs. For each prompt, DTS* faithfully matches the prompt while achieving high reward, whereas SMC samples score higher but visibly over-optimize. Numbers in the corner show aesthetic scores.

---

0.0
MDLM + CoLA
Base
Reward -1.454
The city centre workers are blamed cribicsed also for border management in recent years. has involved taking millions of dollars (nominally Medicare Deperances) for months. The package reduced to $1.50 million, the...
Reeward 1
-0.2
Treasury
-0.3
-0.4
FK/SMC
Reward -0.059
212
The city defended by Royal Jamaican Squadron. city is defended by the English Dafance Squadron, The clefanded by the Continental Guards. The navy defanded by the Catalonian Squadron. The city is defended by the...
Distrct Trirams
0.95
0.90
DTS (Ours)
Reward -0.032
0.85
The city's Public Safety Committee has also promised to consider ways discourage residents from buying private homes their But tax-banefit analysts oppese such stop because they fear it would dramatically reduce the local tax base. Taxpayers...
0.80 212
212
2
2 2
Inferance Buriget (MFES)
DTS (Ouns)
Figure 9: Left: Maximum reward and distinct trigrams vs. compute (NFEs) per prompt, across 15 simple prompts with MDLM (68) using a classifier trained CoLA (54) the reward. DTS obtains the highest reward while maintaining diverse outputs. Right: Typical samples generated by the base model, FK/SMC, and DTS* for the prompt 'The city" at 2 NFEs.
maintaining faithfulness to the prior model. We hypothesize that since DTS* backs up soft values, each node aggregates posterior mass rather than peak density. Consequently, reward spike that lies taiisiigg-robbbllity probability region of the prior contributes negligibly to the value estimates, resulting in an implicit KL-regularization effect (cf. Section 2
in a vanishing
5.3 Text generation
Setting. We evaluate DTS* on text generation using MDLM [68], a discrete diffusion language model. We generate three text completions of length 64 for each of 15 prompts introduced by Han et al. [23] The reward is defined as the log probability that the text is classified grammatically 'acceptable" by a BERT-based classifier [54] trained on the Corpus of Linguistic Acceptability (CoL.) [83] We
also report diversity by computing the number of distinct trigrams in each generated sequence. For decoding, we find that using DTS* with max-back max-backup (A - oc) yields the best performance. We compare our method against two baselines: FK/SMC 71 and best-of-N.
Results. As shown in Figure 9 DTS* consistently achieves the highest rewards as the number of function evaluations (NFEs) increases. Notably, reward functions in text domains particularly susceptible to over-optimization [3 25 8), leading to the less diverse outputs observed when using FK/SMC. By contrast, DTS* produces outputs that have both high rewards and high diversity. We perform evaluation using an LLM in Section J.4, since it aligns better with human judgement 47
Discussion
We have introduced novel framework that casts inference-tine alignment of diffusion models as finite-horizon tree search. By propagating terminal rewards via a soft value backup, our approach achieves global credit assignment and improved sample quality as compute increases. Below, we highlight practical considerations, point out limitations, and suggest directions for further work.
High-dimensions and the role of pretrained model. In high dimensions, an uninformed search tree grows exponentially with dimension, rendering pure tree search infeasible. A good quality pretrained model acts as a powerful prior, significantly pruning the effective search space. Even then, without gradient information, such methods struggle if the posterior is very sharp, such in certain inverse problems (Section D.
In several applications of MCTS in game play, such as AlphaZero 770, deep neural networks approximate both policy and value. While our current work focuses on zero-shot inference- -imencee-imeeeeee alignment for any unseen reward, an exciting future direction would be to integrate a learned value network for fixed reward.
Learning the value function.
Compute cost. The control flow of brse- methods is sequential, which makes them less parallelizable than particle-basee based methods such SMC. However, discussed in Section 4.3] we implement an efficient version by batching calls to the diffusion model. Moreover, once the tree is constructed, sampling incurs no further model calls, making repeated draws effectively free.
10

---

Acknowledgments and Disclosure of Funding
This research is in part supported by CIFAR Al Chairs and the NSERC Discovery program. Mila and the Digital Research Alliance of Canada provided computational resources.
References
[1] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023.
[2] Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47:235-256, 2002.
[3] Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 843-852 2023.
[4] Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in 2021.
Neural Information Processing Systems,
[5] Kevin Black, Michael Janner, Yilun Du Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024.
[6] Cameron Browne, Edward Powley, Daniel Whitehouse, Simon M Lucas, Peter Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. survey of monte carlo tree search methods. IEEE Transactions on Computational Intelligence and AI in games, 4(1):1-43, 2012.
[7] Gabriel Cardoso, Sylvain Le Corff, Eric Moulines, et al. Monte carlo guided denoising diffusion models for bayesian linear inverse problems. In The Twelfth International Conference on Learning Representations, 2024.
[8] Asli Celikyilmaz, Elizabeth Clark, and Jianfeng Gao. Evaluation of text generation: A survey, arXiv preprint arXiv:2006. 14799, 2020.
[9] Boyuan Chen, Diego Martí Monsó, Yilun Du Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusson forcing: Next-token prediction meets full-sequence diffusion. Advances in 2024.
Neural Information Processing Systems, 37:24081-2444
[10] Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, page 027836492411236668 2023.
[11] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul
Ye. Diffusion posterior sampling for general noisy inverse problems. The Eleventh Interna- tional Conference on Learning Representations, 2023.
[12] Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models differentiable rewards, In The Twelfth International Conference Learning Representations, 2024.
[13] Florentin Coeurdoux, Nicolas Dobigeon, and Pierre Chainais. Plug-and-phl split gibbs sam- pler: embedding deep generative priors in bayesian inference. IEEE Transactions on Image Processing. 33:3496 3507, 2024.
[14] Adrien Couëtoux, Jean- Ban-tBttstte Hoock, Nataliya Sokolovska, Olivier Teytaud, and Nicolas Bonnard. Continuous upper confidence trees. In Learning and Intelligent Optimization: 5th International Conference, LION 5, Rome, Italy, January 17-21, 2011. Selected Papers 5, pages 433 445. Springer, 2011.
[15] Pierre Del Moral, Amaud Doucet, and Ajay Jasra. Sequential monte carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411-433 2006.

---

[16] Tristan Deleu, Padideh Nouri, Nikolay Malkin, Doina Precup, and Yoshua Bengio. Discrete probabillistic inference as control in multi-path environments. In The 40th Conference on Uncertainty in Artificial Intelligence, 2024.
[17] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8794, 2021.
[18] Carles Domingo-Eniin Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint arXiv:2409.08856 2024.
[19] Zehao Dou and Yang Song. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations, 2024.
[20] Bradley Efron. Tweedic's formula and selection bias. Journal of the American Statistical Association, 106(496):1602-1614, 2011.
[21] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning eox-tto-image diffusion models. Advances in Neural Information Processing Systems, (16799588-78887 2023.
[22] Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International Conference on Machine Learning, pages 1352-1361. PMLR, 2017.
[23] Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. Ssd-Im: Semi-auttoregressvee simplex- based diffusion language model for text generation and modular control, 2023. URL https
//arxiv. orgoogg////o Bbrrg/abs///n 1243
[24] Philippe Maasen-Estrucl Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idgl: Implicit q g-learning as an actor-critic method with diffusion policies. arXiv preprint a334v:2304. 10573, 2023.
[25] Tatsunori B Hashimoto, Hugh Zhang, and Percy Liang. Unifying human and statistical eval- uation for natural language generation. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Tech- nologies, Volume I (Long and Short Papers), pages 1689-1701, 2019.
[26] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[27] Yutong He, Naoki Murata, Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Dongjun Kim, Wei-Hsiang Liao, Yuki Mitsufuji,, Zico Kolter, Ruslan Salakhutdinov, et al. Manifold preserv- ing guided diffusion. In The Twelfth International Conference on Learning Represeniations, 2024.
[28] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv.1606.08415. 2016.
[29] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.
[30] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:68-006851, 2020.
[31] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020.
[32] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning. pages 8867-8887. PMLR, 2022.
12

---

[33] Yujia Huang, Adishree Ghatare, Yuanzhe Liu, Ziniu Hu, Qinsheng Zhang, Chandramouli S Sastry, Siddharth Gururani, Sageev Oore, and Yisong Yue. Symbolic music generation with non non-differentiable rule guided diffusion. In Proceedings of the 41st International Conference on Machine Learning, pages 19772-19797, 2024.
[34] Vineet Jain and Siamak Ravanbakhsh. Learning to reach goals via diffusion. In International Conference on Machine Learning, pages 2170-121195 PMLR, 2024.
[35] Vineet Jain, Tara Akhound-Sadegh, and Siamak Ravanbakhsh. Sampling from energy-based policies using diffusion. arXiv preprint arXiv:2410.013/2, 2024.
[36] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. International Conference on Machine Leurning, pages 9902-9915. PMLR, 2022.
[37] Sadeep Jayasumana, Srikumar Ramalingam, Andreas Veit, Daniel Glasner, Ayan Chakrabarti, and Sanjiv Kumar. Rethinking fid: Towards a better evaluation metric for image generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9307-9315, 2024.
[38] Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning, Advances in Neural Information Processing Systems, 36: 67195-67212, 2023.
[39] Sunwoo Kim, Minkyu Kim, and Dongmin Park. Test-time alignment of diffusion models without reward over-optimization. In The Thirteenth International Conference on Learning Representations, 2025.
[40] Diederik P Kingma. Adam: A method for stochastic optimization.
arXiv preprint
arXiv: 1412.6980 2014.
[41] Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In European conference on machine learning, pages 282-293. Springer, 2006.
[42] Piotr Kozakowski, Mikolaj Pacek, and Piotr Miloš. Planning and learning using adaptive entropy tree search. In 2022 International Joint Conference on Neural Networks (ICCNN), pages 1 1-8. IEEE, 2022.
[43] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436-444, 2015.
[44] Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302. 12192, 2023.
[45] Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, Gokcen Eraslan, Surag Nair, Tom- maso Biancalani, Shuiwang Ji, Aviv Regev, Sergey Levine, et al. Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding, arXiv preprint arXiv:2408.0m 08252, 2024.
[46] Xiner Li, Masatoshi Uehara, Xingyu Su, Gabriele Scalia, Tommaso Biancalani, Aviv Regev, Sergey Levine, and Shuiwang Ji. Dynamic search for inference-ting time alignment in diffusion models. arXiv preprint arXiv:2503.0200 02039, 2025.
[47] Yang Liu, Dan Iter, Yichong Xu, Shuohang Wang, Ruochen Xu, and Chenguang Zhu. G-eval: Nlg evaluation using gpt-4 with better human alignment. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 251 1-2522, 2023.
[48] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In International Conference on Machine Learning, pages 32819 18819-28449 PMLR, 2024.
[49] Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li. and Jun Zhu. Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In International Conference on Machine Learning, pages
PMLR, 2023.
13

---

[50] Nanye Ma, Shangyuan Tong, Haolin Jia, Hexiang Hu, Yu-Chuan Su, Mingda Zhang. Xuan Yang. Yandong Li, Tommi Jaakkola, Xuhui Jia, al. Inference- time scaling for diffusion models beyond scaling denoising steps. arXiv preprint MxXv::2501.09332, 2025.
[51] Marvin Minsky. Steps toward artificial intelligence. Proceedings of the IRE, 491):8-30, 1961.
[52] Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, and Pierre-Luu Bacon. Maximum entropy gflownets with soft q-learning. In International Conference on Artificial Intelligence and Statistics, pages 2593-2601. PMLR, 2024.
[53] Nikita Morozov, Daniil Tiapkin, Sergey Samsonov, Alexey Naumov, and Dmitry Vetrov. Im- proving gflownets with monte carlo tree search. arXiv preprint 1xXi::2406.136555 2024.
[54] John X. Morris, Eli Lifland, Jin Yong Yoo, Jake Grigsby, Di Jin, and Yanjun Qi. Textattack: framework for adversarial attacks, data augmentation, and adversarial training in nlp, 2020. URL bttpss/////ax org/ /bs//////// 2005. 05909
[55] Michael Painter, Mohamed Baioumy. Nick Hawes, and Bruno Lacerda. Monte carlo tree search with boltzmann exploration. Advances in Neural Information Processing Systems, 36: 78181-78192, 2023.
[56] Kushagra Pandey, Farrin Marouf Sofian, Felix Draxler, Theofanis Karaletsos, and Stephan Mandt. Variational control for guidance in diffusion models. In Forty-second International Conference on Machine Learning, 2025.
[57] Mohammad Pedramfar and Siamak Ravanbakhsh. Multi-armed sampling problem and the end of exploration, 2025.
wwtp:/wwwwwrgaab.....a/a wtt.:rggggbb...../////////////.. 10797
[58] Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text- to-image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.03739 2023.
[59] Michael Psenka, Alejandro Escontrela, Pieter Abbeel, and Yi Ma. Learning diffusion model policy from rewards via q-score matching. International Conference on Machine Learning, pages 41163-41182 PMLR, 2024.
[60] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pages $748-8763. PMLR, 2021.
[61] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36:53728-53111 2023.
[62] Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. arXiv preprint arXiv:2304.02532. 2023.
[63] Christian Robert, Arnaud Doucet, and Simon Godsill. Marginal map estimation using markov chain monte carlo. In 1999 IEEE International Conference Acoustics, Speech, and Signal Processing. Proceedings. ICASSP99 (Cat. No. 99CH36258), volume 3, pages 1753-1756. IEEE, 1999.
[64] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[65] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention- 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings. part III 18, pages 234 241. Springer, 2015.
[66] Litu Rout, Yujia Chen, Nataniel Ruiz, Abhishek Kumar, Constantine Caramanis, Sanjay Shakkot- tai, and Wen-Sheng Chu. Rb-modulation: Training free personalization of diffusion models using stochastic optimal control. arXiv preprint 2705::2405.17441, 2024.
14

---

[67] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, al. Photorealistic ext-tto-mmage diffusion models with deep language understanding, Advances in Neural Information Processing Systems, 35:36479-36494 2022.
[68] Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin. Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models, 2024. URL httthtttp///rxttti sarxiv.org/aann
wtt.:orr/abb/400005555555 2eabbs/2400007555 07524
[69]
Chrisoph Schuhmann. laion aesthetics/
Laion aesthetics, 2022.
URL https://laion: ietiaaiinaaibblg... ai/blog/
[70] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, (02(6419):114--111 2018.
[71] Raghav Singhal, Zachary Horvitz, Ryan Techan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. A general framework for inference-time scaling and steering of diffusion models. arXiv preprint arXiv:2501.06888, 2025.
[72] Marta Skreta, Tara Akhound-Sadegl Viktor Ohanesian, Roberto Bondesan, Alan Aspuru- Guzik, Arnaud Doucet, Rob Brekelmans, Alexander Tong, and Kirill Neklyudov. Feynman- kac correctors in diffusion: Annealing, guidance, and product of experts. In Forty-second International Conference on Machine Learning, 2025.
[73] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. The Eighth International Conference on Learning Representations, 2020.
[74] Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming- Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-guided diffusion models for plug-and- -unyaa controllable generation. In International Conference Machine Learning, pages 2023.
PMLR,
[75] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[76] Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry P Vetrov. Generative flow networks entropy--regularized rl. In International Conference Artificial Intelligence and Statistics, pages 4213-4221. PMLR, 2024.
[77] Brian L Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi S Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif- scaffolding problem. In The Eleventh International Conference on Learning Representations. 2023.
[78] Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine Fine- tuning of continuous-time diffusion models as entropy-rege regularized control. arXiv preprint wwXi:22402.15144 2024.
[79] Masatoshi Uehara, Yulai Zhao, Chenyu Wang, Xiner Li, Aviv Regev, Sergey Levine, and Tom- maso Biancalani. Inference-time alignment in diffusion models with reward-guided generation: Tutorial and review. arXiv preprint arXiv:2501. 009655, 2025.
[80] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[81] Siddarth Venkatraman. Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, et al. Amortizing intractable inference in diffusion models for vision, language, and control. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
15

---

[82] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
[83] Alex Warstadt, Amanpreet Singh, and Samuel R. Bowman. Neural network acceptability judgments. Transactions of the Association for Computatioonl Linguistics, 7:625 641, 2019. doi: 10.1 10.1acl/ a 00290.
Stg/Q19-1044
[84] Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practi- cal and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36:3 0es:33372-31403 2023.
[85] Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score: Better aligning text-to-image models with human preference. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2096-2105, 2023.
[86] Zihui Wu, Yu Sun, Yifan Chen, Bingliang Zhang, Yisong Yue, and Katherine Bouman. Princi- pled probabilistic imaging using diffusion models as saug-and-play priors. Advances Neural Information Processing Systems, 177118889-118427, 2024.
[87] Chenjun Xiao, Ruitong Huang, Jincheng Mei, Dale Schuurmans, and Martin Müller. Maximum entropy monte-carlo planning. Advances in Neural Information Processing Systems, 32, 2019.
[88] Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong, Imagereward: Learning and evaluating human preferences for texl-to-image generation. Advances in Neural Information Processing Systems, 36: 5903-15995, 2023.
[89] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
[90] Long Yang, Zhixiong Huang. Fenghao Lei, Yucun Zhong, Yiming Yang, Cong Fang, Shiting Wen, Binbin Zhou, and Zhouchen Lin. Policy representation via diffusion probability model for reinforcement learning. arXiv preprint arXiv:2305.13122, 2023.
[91] Jaesik Yoon, Hyeonseo Cho, Doojin Baek, Yoshua Bengio, and Sungjin Ahn. Monte carlo tree diffusion for system 2 planning. arXiv preprint arXiv:2502.07202, 2025.
[92] Tao Zhang, Jia-Shu Pan, Ruiqi Feng, and Tailin Wu. T-scend: Test-time scalable mets-enhanced diffusion model. arXiv preprint arXiv:2502.00190 2025.
[93] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pages 1433-1438. Chicago, IL, USA, 2008.
[94] Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford. Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv: 190900899990893 2019.
16

---

A Extended related work
We discussed the main approaches that have been proposed for inference- time alignment of diffusion models in Section Below, we briefly review additional steering works and three tangentially related areas: fine-tuning of diffusion models, their use in reinforcement learning, and antropy-regulanee variations of Monte Carlo Tree Search.
Additional steering methods. Some gradient-steering works treat guidance as stochastic control: Pandey et al. [56] learn KL-regularised drifts for differentiable rewards but break when high-rew high-reward modes lack coverage; Rout et al. [66] propose a training- maining-free method to modify the drift for style transfer in vision; Huang et al. [33] extend to non-differentiable music rules via high-variance REINFORCE. Our search-based alignment does not require differentiability or dense coverage assumptions and is domain-agnostti A different line of work focuses on plug-and-play samplers [3] which also use Tweedic's formula to get biased values and use compute- heavy MCMC/Gibbs chains for inverse problems. Skreta et al. Our principled search method produces high-quality samples by using unbiased values and works with arbitrary rewards.
Fine-tuning of diffusion models. To sample from the target distribution x. for a fixed, known reward function, one option is to amortize the posterior sampling problem and update the model parameters via fine-tuning. The paradigm mirrors the trajectory of large-language-motee alignment (94, 61]. Supervised preference finetuning trains directly on synthetic pairs scored by reward model [4 85]. Some early methods exploit differentiable objectives to back-propagate a single scalar all the way to the noise prediction network [12 [58, whereas more traditional reinforcement learning approaches cast each reverse step as an action and optimize expected reward [5]. To avoid
over-optimization of the reward, recent works use KL regularization [21 [8] 81.
Diffusion models in reinforcement learning. Since the introduction of diffusion models as power- ful frameworks for generative modeling, they have become popular for sampling actions or future states in RL. The earliest successes were in offline imitation learning, where some approaches model trajectories 36. or expert policies 10 from offline datasets. Other works maximize a Q-function in addition to behavior cloning [82 38]. employ an explicit actor-critic scheme [24], or treat the critic as an energy function guide the denoiser 49). Some goal-conditioned extensions have also
been proposed 662 34. Recent works have explored similar ideas in the online setting 10, [5] (35]
Those methods aim maximize return for control tasks, while we aim to draw unbiased samples from the reward-tilted distribution for any chosen reward.
Entropy-regularize MCTS. Monte- Carle Tree-Search (MCTS) has recently been extended to soft-value objectives that incorporate an entropy bonus (87, which uses a log-sum-exp value update and samples actions from Boltzmann distribution, guarantecing improved exploration at the cost of converging the soft rather than the standard optimum. Follow-up work proposed adapt the entropy term to a predefined value 42 and decay the entropy term (55] Very recently, Morozov et al.
53 used soft-backup MCTS to improve planning in Generative Flow Networks 4 Our Diffusion
Tree Sampling (DTS) follows the same Boltzmann selection and soft value backup pattern, it is the first to embed pre-trained diffusion kernel inside the tree and to prove consistency for sampling from the Ke-reaulrd posterior, not just selecting a single high-reward action. In this sense, DTS bridges the gap between antropy-regulariee MCTS used for control and unbiased posterior sampling required for inference-time alignment of generative models.
B Sequential Monte Carlo for diffusion sampling
Many existing methods for inference-ting diffusion alignment 77 39] apply sequential
Monte Carlo (SMC) [15] to the reverse diffusion chain. SMC maintains a population of K particles to approximately sample from a sequence of intermediate targets ( (x.r)) 1/rr culminating in the desired x²" (xx) x (xo) exp(Ar(Xo)). In diffusion alignment, one usually sets
(xx:)) p(x))
Po (x,-1 xx)) exp (A (x,), s=t+1
II
(10)
where út is a potential approximating the soft value V,, Each SMC iteration for = T,T - 1,.... has three steps:
(k) 1. Propagation. Sample particles x ~ ge(. xx()),, for k = 1,.... K where qe is the proposal distribution, often set to be the diffusion transition po xx).
17

---

2. Weighting,
Assign importance weights
pa( (8:)
(k) x exp (x8--(2(21)
X
(11)
qe [x x (1t))
importance rabio
The first factor corrects for using proposal and the second tilts weights toward high estimated value.
)K to obtain an equally weighted 3. Resampling. Resample (x )i proportional to {wi-1)
particle set (x(*) 1K for the next iteration.
(x--1Jk=1
Classical SMC guarantees that, as K xo and if the potentials are exact, the empirical measure (R) Ex X converges to the target distribution x", where 8(x) the Dirac delta at æ. In practice, however, this repeated sampling procedure can reduce the diversity of samples, especially when the weights have high variance. This results an effective sample size which is much lower than K.
Another major issue when applying SMC to diffusion models is that estimating the soft value function V. is not straightforward and errors in the approximation can lead to inaccurate sampling. The next subsection discusses the velue-estimation problems in more detail.
C Connection with Generative Flow Networks
Diffusion Tree Sampling can viewed as an on-the-fly, non- non-parametric realization of the ideas behind Generative Flow Network (GFlowNet) [4. Both frameworks ultimately seek to sample from an unnormalised density:
1 x ='xx)= = [(x), 2 2- =
but they do so with different machinery and at different points in the learning inference pipeline. GFlowNets assume a unique initial state So and define probability over complete paths (so ST = x) through
T
T
PO(T) = II Po(se 8t-1) = PO(ST) II (
and train the parameters 9 so that the forward flow, leaving every non-terminal state equals the backward flow, entering it plus injected terminal reward r(x) = log fx). For the special case of tree-structured graph, this constraint in log form is soft Bellman equation [76] [6 52]:
F(s) = log Σ
Pa(s (') exp (AF(s))))
s EChhldd(s)
with F(s) the learned log-flow function.
DTS satisfies the same soft Bellman recursion (cf. Equation (5 but does without learning parameters. During tree construction, DTS estimates the soft value V by Monte-Carlo log-sum-exp backups; selection then samples children proportionally to (xp(\);1), where 6- 1 is the estimated soft value. Repeated roll-outs make the empirical terminal distribution converge to the reward-tilted posterior n² (xu) x Pe (xn) exp (Ar Axx))), just perfectly trained GelowNet would.
The key differences between DTS and GFlowNets summarized below.
Proposal. DTS uses a fixed, pretrained diffusion kernel Pa proposal, whereas GFlowNets learn their forward policy Pg.
Learning search. DTS performs pure inference without updating any parameters, whereas GFlowNets learn the parameters of the sampler to amortize future sampling.
Computational regime. DTS excels when has strong prior and large inference budget for new rewards; GFlowNets shine when the reward fixed and repeated queries amortize the training cost.
18

---

Because DTS is search procedure, it is ideal for adapting pretrained diffusion model to different unseen reward functions without retraining. GFlowNets, in contrast, learn fast parametric sampler for single reward.
D Proofs and derivations
D.1 Derivation of Equations 5 and
We derive the recursive relation satisfied by the soft value function well as the expression for the optimal policy in Section for completeness.
Soft value function. This recursive relation is analogous to the soft Bellman equation in maximum entropy RL [93 22). Starting from the definition of (xxx)):
Vx(
log Epo(xor- (exp (Ar(xo))]
log / X1 1 xx,) exp (Ar(xo)) dxndx; ... dxp-1
= log / p(xo,xI.... 2 ((x exp (Ar(xo)) dxodx1 ... dxx-1
X log / (p(x--1xx)) (f p(x0,x1........ X1, X1 exp (Ar(xo)) dxndxı ... dxl-2 dxi-1
1
-exp(AVi-
1 log log / p(xx-1(xx) exp (x dx, log R. 1xx (exp (AV (x
1
-
The above relation combined with the terminal condition Vo(xo) r r(xo) gives Equation (5) Optimal policy.
The joint target density over the full chain (Xo, (No...., 1, xx) given by: exp (Ar(xo)),
x .xo...... XI-1, xx) = Σ Pa
where 2 represent the normalization constant of this joint density.
The marginal joint density of (x,, Xx-1) under x² is:
/ po (xo,.... xt xx) exp (Ar(xo)) dxo ......... .xxx-2
(X,, X Xx-1)
z Po (xx))oo(Xx11 xx) (f po (Xo, .....,.,,....................... Xx-2 xx-1) exp (Ar(xo)) dxo dxc-2
1
Pe pe (xx)))(xxx11 x:) exp (x
Similarly, the marginal density of x, under nt" is:
(x:) Po (x:) exp (AVi(x;)
By dividing these two marginals, we get the transitions under the optimal policy:
(x,, x1 (x xx) exp
(x:-1 x1)
n (x,))
exp (>V)(x,,)
px 1 x.) exp (AV,-
p x1) exp (AV,-1 (x,-1) dx
The above relation gives the optimal policy from Equation (
D.2 Proof of Proposition
Proposition (Asymptotic consistency). Let be bounded and X 0, then DTS produces sequence of terminal states whose empirical distribution converges to the optimal policy T² as the number of tree iterations M +
19

---

Proof. We use p(- X²) to denote general proposal distribution. For application to diffusion aligmment, this would correspond transitions under the pretrained model po(: xt). Additionally, we use q(- x²) to denote the transition density of DTS.
Step 1: Transition probability under DTS. Recall that under DTS, given node X,, create each child by sampling from the base model p(- xt). During tree traversal, we select the next state X;-1 proportional to the exponentiated soft value function. Thus, the transition probability of DTS from X₁ to X is given by:
x xx)exp exp (A6(x,-1))
p(x:-1 xx) exp (Ar((x-1))
x(x xx)
(12)
J p(xx11 xx) exp (A6((4-1)) dx;-1
exp (Ao(x:))
where the second equality follows from the definition of the soft Bellman equation:
1 i(x:) X log (x--1x((1xx.) (exp (Ao(x;-1))
Step Joint density trajectory. Recall that the root node of DTS contains a dummy state XT that transitions to the diffusion process prior q(xr XT+1) = (0, Then, the joint density of full trajectory (XT, XI- Xo} under DTS given by:
T+1 II (xx x()
T+1
pxx, x1) exp (A6(x(-1))
((XT, XT- xo)
exp (Ao(xe))
exp (X0(x)))
T+1
II
p(x--1 Xx)
exp (Xo(XT+)
t=1
exp (X6(x)))
exp (Ao(XT+) XT Xo).
Step 3: Marginalizing, Marginalizing over intermediate states x,....y Xy, we get the distribution of terminal state No:
.xxdxxxxTTT ((xo) / (IXT;XT-I:..... dxxdxx-1.....
exp (A0(xo))
P(xT, xo) dx .d.d
exp (Ao (Ao(XT+1))
exp (Aύ(Xo))
exp (Ao(XT+1)) p(xo).
By definition, the soft value function at the terminal node is ((x)) = r(xo). Plugging this and using the definition of value function from Equation (4 we have;
exp (Ar(xo)) p(xo)
((xo) = f'p(
11. XT exp (Ar(x))) p(Xo)
X0 XT+1) exp (Ar(xo)) dx ... dx dxco
p'((x) exp (Ar(xo)) dxo
This has the form of the target distribution in Equation 6 except that it uses the value estimates that are calculated based on rollouts starting from each state x,. the limit of infinite rollouts, these value estimates approach the true soft values, confirming that the sampling distribution from DTS exactly matches the target distribution =².
Therefore, DTS is consistent, as correctly generates samples from the desired target distribution asymptotically.
20

---

E DTS and DTS" algorithm
Algorithm Diffusion Tree Sampling (DTS) and Diffusion Tree Search (DTS*)
1: Input: base policy po, reward function number of iterations M, inverse temperature A,
parameters a, Cuct
2: Initialize root node X7+1 with dummy value, ((x++1) 0, N(x++1)= =1
3: Initialize tree T with root node Xy+1
4: for m = -,..... M do
5: P (x++1}
6: -
7: // Selection
8:
while [((x)) C Cx,)x and / >> do
9:
[DTS] select child probabilistically: x Execco exp(Ao(xx))
exp(Ae(xx 1))
10:
[DTS*] select child maximizing UCT: x;-1 = magx f(x() + cwa N(x))
((x))
P (x
12:
f t-t-1
end while
14:
// Expansion: expand X by sampling a new child
if > and [((x)) C (x,)) then
16:
// Rollout: from new node Xt-1 sample rollout path to terminal X
while 0 do
18:
X,-1 ~ (- xx), 8(x,1)==0, N(xx-))= =1
19:
P PU (x-1)
20:
t t
21:
end while
22:
end if
23:
Evaluate terminal reward: 0(xo) r(xo)
24:
// Backup: update value along path P
25:
for t = O,..., T do
26:
Soft backup: ((x++1) < 1 log [((x+11] Exec(oan) exp()((x:)))
27:
Update visits: (xx-1)) +- N(xx+1) +1
28: end for
29:
end for
30: return T
Algorithm Diffusion Tree Sampling (DTS) and Diffusion Tree Search (DTS*) inference
Input: Constructed tree T, number of samples N, inverse temperature
2: Initialize population of samples S <-
3: for 1,..., N do
4: Set f < +
5: while do
6:
[DTS] select child probabilistically: x Extecoo) exp(A6(xx))
exp(Ae(x; 1))
7:
[DTS*] select
select child maximizing value: Xt-1 arg maxx o(x')
8:
1 1-t-1
9:
end while
S-SU (xo)
11: end for
return Samples
Implementation details for DTS and DTS*
F1 Tree structure
The algorithm presented in Section 4.2 and Section allows every state x, along the denoising trajectory to be considered for branching. However, in practice only branch every few timesteps.
21

---

We noticed very little difference in performance between the two cases for the same number of function evaluations, however, we expect branching at every step to outperform for a very high compute budget. We match the tree branching schedule with the resampling schedule for all baselines with SMC, similar to the setting from Singhal al. [7]]. The exact setting for each experiment is presented in Table 2 where we always branch at the root node corresponding to T.
Table 2: Branching schedule for DTS and DTS* which is also the resampling schedule used for SMC-based methods SMC/FK (7].. TDS [84
DAS (9].
Domain
Total denoising steps
Branching schedule
Two-dimensional
100
100(root), 80, 60, 40, 20 50(root), 40, 30, 20,
Image pixels (MNIST, CIFAR-10)
50
Image latents (SD-v1.5)
50
50(root), 40, 30, 20, 10 64(root), 54, 44, 34, 24, 14
Text tokens (MDLM)
64
Apart from this, we have hyperparameters associated with progressive widening that control the maximum number of branches at any node. We used a = 0.8 and C = 2 for all two-dimensional and image experiments and 0.7 and C' for text generation. There is scope of improving the performance of DTS and DTS* further by tuning these parameters for specific tasks.
F.2 Experiment details
Illustrative 2D
Base diffusion model: The denoising network an MLP that takes as input the 2-dimensional data and the timestep and outputs a 2-dimensional noise prediction. The timestep is transformed using sinusoidal embeddings [80]. The network has four hidden layers of 128 dimension each with the sigmoid linear unit (SiLU, (28)) activation. We used the linear noise schedule with 3min 0.001 and 8max 0.07 and the score matching objective. The optimizer used for training was Adam [40] with a learning rate of 3 x 10-3 We train the model for 500 epochs training set of 10000 samples.
Reward function: Gaussian mixture: The reward function is:
r(x) = log Σ
wi exp(-lIx - 14 1=112/200)
i=1
where w, = 2x)().5.), p1; cos ا sin ) i 1,...,8,.. with 0.3.
Checkerboard: The reward is negative distance from the center r(x) -/
Class-conditional MNIST
Base diffusion model: The denoising network is a Unet architecture 665 that operates on images
of size 32 x 32 x 1 (upscaled from 28 28 x 1) with block channels {32, 64, 128, 256). We use the DDIMScheduleef] from diffusers library with default parameters, except we set = 1.0 the inference process stochastic like DDPMs 30 We the AdamW optimizer with learning rate 10-4 for 100 epochs on the MNIST training set.
We train classifier p(c x) on the MNIST training set. The classifier is a convolutional neural network (43) using two 5 5 kernels with (16, channels followed by 2 2 max pooling operation with ReLU activations. The features are then flattened and followed by a linear layer with 10 outputs corresponding
Reward function:
trained using Adam optimizer with learning rate 10 The reward function for single class generation is the log likelihood of the class rx) log p(c i x) for í {0, 1,..., 9}. For the even odd generation, is defined as r(x) = maxes log p(c p(cc= = x), where S = (0. 2, 4, 6, 8} for even digit generation and S = (1,3.5,7.3 for odd digit generation.
the classes. The network
22

---

Class-conditional CIFAR-10
We used the pre-trained diffusion model adpm-cifar10 from Hugging Face, which uses Unet architecture and diffuses over 32 x 32 x 3 images in pixel- space. We use the DDIMScheduler with 1.0 for stochastic denoising.
Base diffusion model:
Reward function: We train a classifier p(c x) on the CIFAR-I0 10 training set. The classifier uses a ResNet-18 [26] backbone that outputs an embedding which is average pooled, flattened, and passed to single linear layer with 10 outputs. The network is trained using Adam optimizer with learning rate 10-3. Similar to MNIST single class generation, the reward function is the log likelihood of the class (x) log p(c
x) for {0, (0.1....,;}
Text-to-image
Base diffusion model: We use Stable Diffusion v1.5 from Hugging Face, which is latent diffusion model [64] The diffusion process defined over 64 64 4 latent variables, which are obtained by encoding 512 x 512 x 3 images using a variational autoencoder. The model uses CLIP (60] to encode text prompts into embeddings which then used condition the generative process via classifier-free guidance [29]. We use the DDMMScheduler with 1.0.
Reward function: We use pre-trained models as reward functions including ImageReward" r(x, y) that encodes prompt accuracy as well as human preferences the LAION aesthetic score predicto rxx) (xx) that encodes aesthetic quality of an image.
http:///gtthbb.ooo////////ggggnn wtom//HUDD//mmgeeeaaarat. ImageReward
https://gathub.comm
predicccor
Conditional text
Base diffusion model: We use MDLM for our text generation experiments. This a discrete diffusion model with 110M parameters that directly predicts the tokens. We define the diffusion process over a context length of tokens with 64 sampling steps and use the standard discrete unmasking update for stochastic denoising.
Reward function: We use BERT-based classifie trained on the Corpus of Linguistic Accepat- bility (CoLA) {3
This reward function r(x) encodes the linguistic acceptability of a given string The reward is the log probability of the text being "acceptable". We find this model to be more robust reward-hasd hacking than alternatives.
https://hugginggaace httpss//haggiiggaae....//uuloo wttp://wwww.../p..oooooee/eeeeee htwpa//leseeooorrrruurooomoooooo
wttp://faaccccccocoeeeeeeeeeecce wtt.t/////ttttttttttecooco http:///oobbrraabaa-oooll... CoLe
F.3 Compute
We report execution times single A100 GPU with 80 gigabytes of memory.
Each 2D experiment including all methods in 15 minutes. Adding up the time five seeds and two different datasets, the combined run time approximately 2.5 GPU hours.
The MNIST and CIFAR-10 10 class-condition
approximately GPU hours per class including all methods. Over all 22 tasks (10 MNIST single digit + 2 MNIST even/odd 10 CIFAR-10 classes) equals approximately 66 GPU hours.
experiments
The text-to-image experiments using Stable Diffusion v1.5 require roughly 30 minutes per prompt across all methods. Adding up all 200 prompts from DrawBench and 45 animal prompts, reproducing all experiments requires approximately 123 GPU hours.
The text generation experiments using MDLM requires roughly 30 minutes per prompt. Thus, generating completions per prompt the prompts requires roughly 22.5 GPU hours.
23

---

G Details of baselines
We re-implemented all baseline methods in our unified codebase since most of them use SMC as backbone and share the same underlying infrastructure. Each implementation was validated by reproducing the quantitative results reported in its original paper. Section B provides a concise primer on SMC for reference. The complete source code including all baselines will be released publicly upon publication of this work.
DPS. Diffusion Posterior Sampling [1] was originally proposed for noisy inverse problems such as image super-resolution and de-blurring using the gradient of the final objective. To adapt this method for general reward functions, we make minor modification by replacing the gradient the inverse problem objective with the gradient of the reward function:
x N ())( + oo Vx, T(x((x)) I).
(13)
where &o is obtained using Tweedie's formula (ef. Section 3), 14 is the predicted mean of the base diffusion model, and r(x) is the reward function the two-dimensional experiments and classifier log likelihoods log p(c = i x) for class-conditional image experiments. The official implementation provided here
OMC//K-Steeringg In our paper, SMC refers to the simplest variant, FK-Steering 71.. which defines different weighting schemes and uses the pre-trained diffusion model as the proposal distribu- tion. As per the setting in Singhal et al. (71), we perform the resampling step at fixed intervals during denoising (given in T and use adaptive resampling to increase diversity of generated samples. diw--dimesiioal and class- cands-ooniiiinall image generation) use the 'diff'
Our sampling experiments potential with À 1.0, whereas the search experiments texxt-t--imge and text generation) use the 'max' potential with À = 10.0. The weights for resampling are given by Equation (II) where the proposal equal to the pre-trained diffusion transition and the value estimates are equal to:
(1)) { diff (xx-1) (8)((x11)) (xx(x:)), i
adiff 'T
(XT) (8o(XT))-
1mma (3)-1) max (1(o0(x: 1)), mi
mi (k) maxx
(xo(x)! (
$2!
We adapted the official implementation provided here
TDS. Twisted Diffusion Sampler 84 comprises of a "twisted" proposal which is used along with SMC to sample from the target posterior distribution. For general reward functions, the twisted proposal is the same the used in Equation (13) and the final weights are obtained using Equation (I] after plugging in the twisted proposal and the value estimates:
(x,-1 xx) N (5 (x), + oo (xxr(x((x))) of i ( (So (x - (xx(xx)) (T(XT) r ((oo(x)))
The official implementation is provided here
DAS. Diffusion Alignment as Sampling 39 re-uses the twisted proposal of TDS but multiplies the
reward term by a monotone tempering schedule 0 = Pr I ----1 S % = to reduce the bias from inaccurate value estimates at high noise levels. The weights given by Equation (11) after plugging in the tempered proposal and value estimates:
qe ($ xx) / (x t) A of (xtx(((( 1). 8 (x = 7 (xx Tt (8o(x:)), (T(xx) = 7T r (x0(xT)).
The official implementation is provided here
H Memory and wall-clock time comparison
Below, we provide peak memory usage for each method in the text-to-image generation experiment using Stable Diffusion 1.5 and ImageReward. We perform this experiment with 40 different prompts and average the results. The numbers are reported for a single NVIDIA A100 GPU with 40GB of memory using the PyTorch profiler.
24

---

SD 1.5 ImageReward
SD 1.5 ImageReward
SD 1.5 + ImageReward
1.7
30000
5000
Menrory Usage (MB)
1.6
40000
20000
1.5
3000
2000
1.3
10000
1.2
10a
104
10'
10
10a
10'
Inference Budieet (NFEs)
Inference Budget (NFES)
Time (5)
DTS (Ours)
Best-of-N
FK/SMC
Figure 10: Comparison of methods using SD-v1.5 and ImageReward Left: Peak memory usage (MB) for different NFEs. DTS* is comparable best-of-N, whereas SMC/FK requires much more memory. Center: Wall-clock time (seconds) for different NFEs. Right: Reward vs. wall-clock time, showing that DTS* scales more favorably in terms of runtime when accounting for its performance. At higher levels of compute, DTS* can match the performance of best-of-N with roughly 2x less wall-clock time. The wall-clock time for DTS* can be greatly reduced by using larger batch sizes, at the expense of slightly more memory and possibly slight degradation in performance. We used batch size of here keep in line with the rest of the experiments.
Peak Memory Usage. We observe that the memory usage of DTS is much lower compared to FK/SMC, and is more comparable best-of-N. Initially. DTS* uses slightly more memory than best-of-N, but with more compute, it tilts in favor of DTS*, showing that it actually scales better in terms of memory.
The main memory usage DTS* is the tree as noted by the reviewer, which stores x 64 x latents. Note that currently the entire tree is kept on the GPU, and we can save GPU memory by offloading the tree to the CPU, For bfst-of-N, the memory usage peaks at the end when all N candidates need to be decoded using the VAE into 512 x 512 x 3 images and evaluated using the reward function pick the best sample. This problem does not occur with DTS* since the rewards are calculated when traversal reaches the terminal state, and multiple high-resolution images do not need to be kept in memory. For FK/SMC, this cost is very high since all N particles need be decoded into 512 x 512 x 3 images, then evaluated using the reward function at each resampling step. general, for images and text, the reward function itself is large pretrained model, and computing the reward requires a non-trivial amount of compute.
Wall-clock time. In terms of wall clock time, DTS* approximately 2.5 slower than the other methods for the same number of function evaluations (NFEs), since it is harder to parallelize. Note that the runtime reflects tree building from seratch once the tree is built, repeated sampling is just pointer chasing and is near instantaneous (whereas for other methods repeated sampling would require starting the entire procedure again).
We discuss two important considerations regarding the runtime:
When taking into account their performance, DTS* is still preferable in terms of wall clock time since it outperforms the basclines and needs much fewer NFEs to reach the same level of performance. As noted in the paper and the table below, for the text-to-image task, it needs roughly 5x fewer NFEs to match the reward of the best performing baseline (best-of-N) at 100% NFEs. In other words, to match the performance of best-of-N at 100k NFEs, DTS* is approximately 2x faster in terms of wall clock time.
We can significantly reduce the runtime of DTS/DTS* by using multiple actors in parallel that traverse and update a shared tree (similar to AlphaZero [70] training). This will give a roughly linear speedup in runtime at the cost of increased memory for copying the diffusion model (the tree is not copied since it can be shared among all actors, hence increase in memory will be sublinear in terms of the number of actors). This represents an inherent trade-off between memory usage and runtime. We did not implement distributed pipeline since it requires some engineering effort, but a viable option.
25

---

CIFAR-10 Inpainting
CIFAR-10 Inpainting
0.200
0.020
0.150
MSE
0.100
0.016
0.014
0.050
0.012
10
10
10a
10
Inference Budget (NFES)
Inference Budget (NFEs)
DTS (Ours)
DPS
SMCIFK
D (Ours)
DPS
Figure 11: Mean square error (MSE) between generated images and reference image on 20 randomly sampled images from CIFAR- 10. Methods that do not gradient information (SMC. DTS) during inference struggle due to the very sharp posterior. Incorporating the gradient into DTS allows it to outperform baselines and continue improving with more compute.
I Image inpainting experiments
As added application, test the feasibility of our method solving inverse problems. We run an inpainting experiment on CIFAR-10 images with a Gaussian forward channel as in Chung et al. [Α We run this experiment 20 images (Iwo images sampled randomly from each class) and average the results. For each 32 x 32 image, we mask the 10 x 10 patch from the center and use the negative error between the generated sample and the unmasked pixels from the original image as the reward. We report the final average mean squared error (MSE) between 16 generated images and the reference image.
Figure 11 shows that methods that do not use gradient information, such as SMC 7 and DTS, do not seem to perform well on this task. We hypothesize this because the inpainting task has very sharp posterior where the sampled image must exactly match the unmasked pixels (compared to class conditional sampling where the posterior represents a semantic concept and has wider mode). This is very hard search problem, and without the forward channel gradient to guide the sampling process, there is very little signal in most regions of the pretrained model distribution.
We then incorporate the gradient into DTS, where we use the gradient term in addition to the score model to obtain the next state when constructing the tree. The backup and selection procedure remains the same. With the gradient- guided proposal, DTS improves upon DPS [[]] and even SMC + gradient (which is the same as TDS 844. We also observe that DTS allows the performance to consistently improve with more compute, whereas vanilla DPS does improve with more compute. At 500k NFEs, DTS + gradient improves 32.96% compared to vanilla DPS.
J Additional experimental results
J.1 Comparison between DTS and DTS*
We compare DTS and DTS* on the text-to-image setting using Stable Diffusion 1.5 with ImageRe- ward on 40 prompts from DrawBench. We picked complex prompts where the base model struggles to generate high reward images to demonstrate the difference more clearly. For DTS, we sample 16 images from the tree and pick the one with the highest reward. The results show that there is noticeable gap, since DTS spends more compute expanding suboptimal paths, whereas DTS* aggressively searches only the highest reward regions except for the small exploration bonus.
Table 3: DTS and DTS* with SDv-1.5 and ImageReward on 40 prompts from DrawBench.
Algorithm NFEs +
1000
2500
5000
10000
25000
50000
DTS
0.774
1.025
1.064
1.112
1.214
1.238
DTS*
0.825
1.086
1.254
1.342
1.459
1.552
26

---

J.2 Class-conditional image experiments
We supplement Table and Figure with additional results and samples.
We plot all four metrics FID, CMMD, average log rewards and average diversity across different number of function evaluations (NFEs) for the three settings considered in Section [51] The plots show that across the three settings for most values of NFEs, DTS matches the target distribution more accurately compared to other methods (lowest FID and CMMD).
We also present random samples for each method and setting in Figures 13 to 14 We observe the same trend as noticed in Figure 7 gradient-base guidance like DPS can be unstable leading to unnatural images, while SMC-based methods show signs of mode collapse with low average diversity and high average rewards. DTS balances both diversity and high rewards effectively by closely matching the true posterior distribution.
MNIST
MNIST
CIFAR-10
109
10
10a
F10
10
10
10
10
109
109
10
10
10
10-2
10
10
10a
0.7
0.7
0.7
Averoge diversity
Average bgg rewardd
0.2
-0.6
-0.6
-0.6
10a
10a
10a
10a
10a
30
10"
Inference Budget (NFEs)
Inference Budget (NFEs)
Inference Budget (NFEs)
(Ours)
SMC/FK
TDS
Figure 12: Distribution level metrics vs. number of function evaluations for different methods on MNIST single digit generation averaged over all 10 digits (left), MNIST odd and even digit generation (center), and CIFAR-10 10 single class generation averaged over all 10 digits (right). All methods were evaluated with 5000 generated samples per class. Metries reported: FID (lower is better), CMMD (lower is better), Average log rewards (higher is better), and average diversity (higher is better).
27

---

DTS (Ours)
DPS
SMC/FK
TDS
DAS
0
0
0
/
/
a
2 2
2 2 2 2 2
2
2
2
2 2 2 2 2 3 3 3
2 2 2
2
2 2 2
2
2 2 1 2
2
2
2 2 2 3 3 3 3
2
2
2 2 2 3
2
2 2 2
~
2 2 3 3
2
2
2
2 2
2
2
2
2 2 2
2
2
2
2
3
3
3
3 1333
3
3 3 3 3
3
3
3
3
3 3
0 3
3
3
3 3 3 3
3
3 3
3 3
3
3
33
3 3 3
3
3
3 33
3
3 333
3
3
33
3
4 4 4
4 4 4
4 4
4
4 4
4 4
4 4
4
4-
4 4
4 4
4
4 4
4
444
4 4444
4
44 4 4
4 444
4
4 4
444
4 4
4
4444
4444
444
4 4
5 55
S 5
5
5555
5 5555
5 5553
5 5 5
5 5855
5 5
5555
$
5
5
6
5 5555
5 5 5
5565
555 5 5
5 5 5
S 1555
6 6 6
6 6626 6
6 6 9
9
9999
9 3 999 9
6 6 <
6 6 6 6
6 6
7 6
6
6 6
> 6 9
9 6 6 6
666
6 6
6
6
9
6
9
6 6
6 6 66
6 6
6
6
6
6
6 6
777 7
7
7 7779
7
7 7 7
7 7 7 7
7
7
7
4
7 7 7 7 7 8 & 8 8 8 8 9
7 7
7 7 7 7 2 8 8 8 8 8 ? 8 S 8 8 F 8 8 /
7
7
F 7
7 7 T 8 8 8 8 9 9 9 9 4 6 6 0 1 1 / 1
7 7 7 6 9 6 8 9 9 9 9 6 8 6 6
7
7
7 7 8 8 8 8
7 7
7
7
7 7 8
7
7 2 3
7 7 7 7 7 8 8 8 8 8 9 9 9 9 9 4 4 9 9 9 2 2
7 8
7
7
7 7
7 8
8 r % 4
8
8 8 8
8 8 8
8 8 8 8
8
8
8
6
8 8
8
8
S 8 3 9 9 + 6 2 0
8
8
8
8
8 9
g
9 9 9
9 9 9 9
9 4 9 9
9 9 9 9 9 6 6 4 6 5- / / /
9
9 9 9 4 O
9 9 9 9 8
9
9 9 9 6 4
9 9
9 9 2 6 4 40 0 2 6000
9 9 2 6 0 O 6 / 7 3 7 9
9
9 8
9
9
2 8 6 6 8
2
4
2 A 6
2
2
O 0 2 2 6 8 O 6 2
2
G 4 6 5
2 6
a
4 2
0
2
2 2
8
a
2
3
9
3
7 1 /
9 9
3 9 3 9
/
3
3
5
9
3
9
/ 7 /
/ 3
1
/ 7
5
/ 5
/
4
3
5
5 /
7
8 3 3 5
/
5 5 7
3
7
Figure 13: MNIST posterior samples generated using different methods for digits 0-9, even and odd.
28

---

DTS uurss)
DPS
SMC/FK
TDS
DAS
airplane
car
bird
deer
dog
frog
horse
ship
truck
Figure 14: CIFAR-10 posterior samples generated using different methods for all classes.
29

---

J.3 Text-to-image examples
We present more samples for qualitative analysis. Figure 13 shows how samples change with increasing amount of inference-time compute, providing visual evidence for the quantitative results from Figure & Figures 16 to 18 shows text-image pairs testing different concepts such as artistic style, spatial arrangement and object count.
An eoncy al bully parda trnartin a red bas, bios gloves, Steen attirs, blue pautts
Leiang claet piupple flomen pot. A puuppl Hoome pot the shape of Scrunglle
6:mmm macro abook a kitton Hiching baby mrock, atotto Tetning
panda makind
Legh Arnold Bil warsicmeoor
segit aago
K
DIFILSTION
DIFTSMONN
2.K
DIFILIUN
Compute budget (NFEs)
Dificion
10K
fittion
29K
Dusion
50K
100K
Figure 15: Text-image pairs from Figure 1 with increasing amount of inference-time compute,
measured number of function evaluations (NFEs) the diffusion model.
30

---

tighn in lab onaz. Mlami vibe, turning well clled sciencs content machine, digital art
An amolent #gyptiann painting depleting argunent over whose turn is ls to take out the trash
painting by Wood of astronaut
photogrsph 19800 airship shaped like pig. foeting over whout Seld
painting portrait ropal Burjer King posing with Whopper
couple, Amerioan gothio style
Stable Diffusion 1.5
Stable Diffusion 1.5 FK-Steering
Stable Diffusion 1.5 + DTS (Ours)
Figure 16: Sample text-image pairs using Stable Diffusion v1.5 and ImageReward as the guiding function for prompts requiring specific artistic style. Samples are picked random for each method and prompt.
lagtop top teddy bear
A giraffe underneath
An umbrolls top of spoon
banans the left of apple
bub
Stable Diffusion 1.5
Stable Diffusion 1.5
PK-Steering
Stable Diffusion 1.5
DTS* (Ours)
Figure 17: Sample text-image pairs using Stable Diffusion v1.5 and ImageReward as the guiding function for prompts requiring specific spatial relationships between objects. Samples are picked at random for each method and prompt.
31

---

One oat and two dogs atsting the grass
Three ears on the stzeet
Pour dogs the slrves
Pive dogls the street
Five ahe strved
Stable Diffusion 1.5
Stable Diffusion 1.5 FK-Steering
Stable Diffusion 1.5 + DTS (Ours)
Figure 18: Sample text-image pairs using Stable Diffusion v1.5 and ImageR as the guiding function for prompts requiring specific object counts. Samples are picked random for each method and prompt.
J.4 Text completion examples
We present additional text completions for the base MDLM model, FK-Steering and DTS* in
Figure 19 We also evaluate the text samples using ChatGPT-4o, since reward this case does necessarily capture quality, and previous works 47 suggest that LLM evaluation aligns better with human judgement.
Base
Reward -1.827
Base
Reward -0.721
Base
Reward -1.709
Once a time, the Quentum constant was an excel lent dream and people believed that
The pizza was also shown to be some Alaska's Crust Filety, and later grilled
The prasidentt of the country must navigate through thicket job that they're getting with them. Having the president have thicket a job that they have to work with them. They get to
megical. light and some kind of magical systolized just prior
2012 Alaska Fim Festival in fried eggplent. the pizza restaurant was making cameo on Pizza Planet 3. This...
magical being...
FK/SMC
Reward -0.075
FK/SMC
Reward -0.072
FK/SMC
Reward -0.059
The president the country commented. AAP/ABC News Topies:
Once upon a time women would commit themselives into fleets, spawning "Stupid", STupid", "Stupid", "ST", "Stupid", "Stupid", "Stupid", "St
pizza
8,298 50.57.1.61 16 99, 1...2222.666.11111 8.8.7....8........... 8.069.5.11 13.155111.16.555 19.1116.69911111661 051.0944131...3200000000 18.36.43
dadaaalgoveemmet,, government-and poliics, elections, united-stater First Trump
Trump on Wednesday was in some way...
(Ours)
(Ours)
(Ours)
The pizza was topped with two teriyali These huge, tomatoey balls looked like they came from cougg was kind soft and they balanced nicely. vegetable
The president of the country supposedly spent some debating should stay armed." " he argued: "Your military officers bring their heavy wespons with you and you
Once upon toxted you and said, "It' going really well say, suppose you would surprised, wouldn't you?" fow days were texting me, and
box. The
them."
Figure 19: Sample text completions using MDLM and a CoLA classifier as reward. Samples are picked at random for each method and prompt.
32

---

Table Text samples generated for 2 NFES with the CoLA reward (Section 5.3) are scored by ChatGPT-4o between 0 and 1 following rubric. Results are reported as mean=sid by averaging over 3 completions for each of the prompts. DTS improves 13.7% over the best baseline.
Algorithm
Best-of-N
FK/SMC
DTS* (Ours)
ChatGPT-40 score
0.283 1.000
0.216
1001
0.322 0.002
Prompt for LLM evaluation
You are an expert language evaluator. Your task is to assess the overall quality of a text sample based on the following criteria:
1. Fluency: Is the text generally grammatically correct and understandable?
Diversity: Does the text show some variety in vocabulary or sentence structure, even if modest?
3. Plausibility: Is the content mostly coherent, logical, and believable in its context?
Please be constructive and fair in your evaluation - highlight the strengths of the sample as well areas that could be improved. Some minor errors simplicity acceptable if the overall message is clear.
For each sample, provide brief explanation touching on each criterion, followed by quality score between 0.0 (very poor) and 1.0 (excellent). Format the final score using \boxed(score),
Example 1 (High Quality)
Sample: "h" shimmering lake reflected the golden hues of the setting sun birds glided across the surface, their wings slicing through the calm evening air." ..n
Evaluation: Fluency: The sentence flows smoothly with proper grammar and elegant phrasing. Diversity: The vocabulary is rich ("shimmering", "glided", "golden hues"), "), and sentence structure is varied. Plausibility: The description is vivid and realistic. \boxed(0.95)
Example 2 (Low Quality but Understandable)
Sample: """Lake is shiny. Birds is fly. The sky is color like orange. is lake.
Evaluation: Fluency: Some grammatical issues (" ("birds is fly". "it is lake"), but basic meaning is understandable. Diversity: Limited vocabulary and repetition. Plausibility: The message simple but not nonsensical. (boxed(0.25)
Now evaluate the following sample:
Sample: """ (sample)"" Evaluation:
33

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We make three main claims in the paper: (1) our tree-based method yields asymptotically exact samples from the target distribution (supported by Proposition [.: (2) DTS reduces bias and variance in value estimation compared to existing methods (supported by experiments in Section 4.4 Figure 5): (3) DTS and DTS* achieve improved sample quality with more compute and scale more favorably compared to existing methods (supported by experiments in Section
Figure 7
Figure Figure
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in 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. It fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the limitations of our method in Section iincluding the dependence on good pretrained model, limitations when dealing with sharp posteriors inverse problems, and the fact that it might be harder to parallelize compared to existing methods. Our experiments cover different modalities and applications empirically validate the performance of our method.
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 create a separate "Limitations" section in 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 on 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 a few datasets or with a few runs. In general, empirical results often depend 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 or images are taken in low lighting. Or a speech-to-tes, system might not be used reliably to provide closed captions for online lectures because it fails 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
34

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We include one theoretical result (Proposition Ty. We make mild assumptions that are clearly stated in the proposition. We provide brief proof sketch in the main text and a full proof in Section
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proof's in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's 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 to 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: We include full details all experimental setups in the main text (Section 5) We also provide formal algorithm in Section [ and implementation details, including hyperparameters used for DTS. DTS* and the baselines in Section and Section G
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to 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.
If the contribution
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 novel architecture, describing the architecture fully might suffice, if the contribution is 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 is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to hosted model (c.g.. in the case of a 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 all submis-
sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be a way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-soure models, it may be that access to the model is limited in some
35

---

way (e.g., s. to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We have released the code publicly and provided the link in the abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
Seliddes/CodeSubmisse
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
ciaips.cc/pbbll.
aassionPollii onPolicy) 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 reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is 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 (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All experimental settings/details available in Section and in Section F For each set of experiments, we mention the datasets used, the metrics reported, and the baselines for comparison.
Fnn
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should presented in the core of 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, appendix, or 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: The paper reports error bars and information about staticical significance (eg. repetitions etc.) wherever relevant. These can be found in most plots in the paper and the main results table in Section 5
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
36

---

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
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It 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 2-sigma enror 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 show in tables figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide estimate of the number of GPU hours all the experiments in the paper and a description the hardware in Section F3 We also compare memory usage and wall-clock time of different methods in Section
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' ' make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The paper conforms with the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there a special consider- ation 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: [No]
Justification: Our contribution is entirely methodological and implication-aanostte it pro- vides generic inference-time time steering framework without targeting any particular domain or use case, and thus raises no specific societal or ethical concerns that require highlighting in this context.
Guidelines:
The answer NA means that there is societal impact of the work performed.
37

---

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.
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 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 is 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 (c.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 of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper does not currently release new data or models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines 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 not 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 (c.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: We have cited all works, datasets, and assets that are mentioned in the paper. We also include links to the publicly available implementations of baselines that were used in the paper in Section
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 or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from particular source (c.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, patasets.pppprrswitte wwwpppppppphhoooooopooppoo thcode. com datasets
38

---

has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset'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 are released in this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was 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
Question: 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)?
Answer: [NA]
Justification: There no crowdsourcing part of this study.
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 as much detail as possible should be included in the main paper.
According to the NeurIPS 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 rescarch 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: This study did not have any participants.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with
human subjects.
Depending on 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 NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
39

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development does not involve LLMs.
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 our LLM policy (https://neuripp: cc//onffreencc//225//M for what should or should not be described.
40


Reviews:
[{"review_id": "uzFB5WNAVM", "reviewer": "Reviewer_1bSF", "summary": "This paper proposes Diffusion Tree Sampling (DTS) and its greedy variant Diffusion Tree Search (DTS*), for better adapting a pretrained diffusion model to new objectives at inference time. This method adopts a tree-based approach to sample from the reward-aligned target density and refine value estimates with each additional generation. From theory proof to experimental results, DTS and DTS* have shown a promising way for new objectives with less computation and higher performance.", "questions": "- Please provide more details on the backup process. And the comparison about the value estimation should be added.\n- For text generation, the results do not express the advantages of the proposed method. Please add more evaluation metrics for a full analysis to showcase the superiority of the proposed method.\n- Please provide experiments with different numbers of nodes and different numbers of denoising steps.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n- The paper combines RL with the diffusion model and tackles the credit assignment problem during the denoising process. Therefore, DTS can assign not only terminal rewards but also processing rewards.\n- The authors provide theoretical guarantees to ensure that DTS produces a sequence of terminal states whose empirical distribution converges to the optimal policy.\n- The illustrative experiments, class-conditional posterior sampling, and image and text generation experiments demonstrate the effectiveness of the proposed method. \nWeaknesses\n- Lack of comparison in value estimation problem: The core idea of the tree-based sampling approach is to use the value of children nodes to update the value of parent nodes, as mentioned in the backup process. However, there are no experiments about how much the differences is between the estimated value from the proposed method and the ground-truth value. And the comparison between other baselines concerning the value estimation.\n- Insufficient baselines in terms of image and text generation: The proposed method is only compared with FK/SMC and best-of-N for image and text generation. What about the results of other baselines? The results in Figure 8 between the proposed method and best-of-N in the text generation are almost the same. What is the advantage of the proposed method in this setting?\n- The setting of the number of nodes and denoising steps: How to set these two variants, is it an empirical value, or else? Whether the experiment section should include a sensitive analysis.", "was_revised": true, "final_justification": "In the author's response, the clarification about baselines addresses some of my concerns. Results for CoLA reward function are also very helpfu. Therfore, I adjust my rating and lean towards acceptance.", "rebuttal": "We thank the reviewer for their comments and suggestions. We address your concerns below, and also draw attention to our response to other reviewers for added experiments - image inpainting (reviewer ULya), comparison between DTS and DTS* (reviewer T4KC),  and comparison of memory and runtime (reviewer t7aZ). Note that this year, NeurIPS does not allow posting any links or documents in the rebuttal. \n\n> There are no experiments about how much the differences is between the estimated value from the proposed method and the ground-truth value\n\nIndeed, we believe this is an important point and we do have experiments to support it; based on reviewer’s comment we will highlight this more clearly. We have a bias-variance analysis of the value estimates with respect to the ground truth values in Figure 5. Here, we perform 1000 Monte Carlo rollouts from noisy states using the base model to get a reasonably good approximation of ground truth. Figure 5 shows that the value estimates of DTS have lower bias and variance compared to the baselines, including DPS and SMC + variants like TDS, DAS (which all use Tweedie’s formula) and SMC + rollout (using a single DDIM rollout). We limit this error analysis to the 2D experiments, since estimating the ground truth values is prohibitively expensive in higher dimensions, both due to model size and also because we require exponentially more MC rollouts to reliably estimate the ground truth values. Please let us know if you have an additional kind of evaluation in mind.\n\n> The proposed method is only compared with FK/SMC and best-of-N for image and text generation. What about the results of other baselines?\n\nThere are two different settings considered in our paper: (1) posterior sampling (illustrative 2D and class-conditional image generation), and (2) what one may call “high reward search” (text-to-image and text diffusion). The nature of these two distinct tasks and prior literature motivate the choice of algorithms.\n\nIn the first setting, there is a well-defined and tractable ground truth distribution. DTS and the baselines in these experiments are designed to sample approximately from a target distribution; hence, we report distribution-based metrics.\n\nIn the second setting, the goal is to produce the single best sample that achieves both high likelihood under the pretrained model and high reward [1,2]. To generate an image for a given caption, for example, we care about having a single good image. DTS and the other baselines from the previous setting are not designed to produce a single good sample (and the corresponding papers also consider the sampling setting only). Therefore, we only use baselines from prior works that report results for a single best sample. The reviewer may be interested in our comparison of DTS and DTS* on text-to-image generation in our response to reviewer T4KC.\n\n> The results in Figure 8 between the proposed method and best-of-N in the text generation are almost the same. What is the advantage of the proposed method in this setting?\n\n> For text generation, the results do not express the advantages of the proposed method. Please add more evaluation metrics for a full analysis to showcase the superiority of the proposed method.\n\nFor text generation, even small differences in reward values can result in qualitatively distinct text samples. The difficulty of text evaluation has been discussed in previous works [3,4], which show that metrics generally correlate poorly with human judgments and do not necessarily follow a linear scale. Below, we score different methods with ChatGPT 4o, which is again not a perfect metric, but aligns closer with human evaluation [5].\n\nThe fact that FK/SMC produced overly repetitive outputs is also primarily attributed to issues with these reward functions. DTS* uses soft values that accumulate posterior mass and can balance likelihood under the pretrained MDLM with optimizing the reward, producing better quality samples due to the influence of MDLM. Following the reviewer’s observation, the difference between DTS* and FK/SMC in the text experiment also seems small in magnitude, but the samples in Figure 8 differ significantly. \n\nWe ran an additional experiment using a reward function trained on Corpus of Linguistic Acceptability (CoLA) [6], and we observed a similar phenomenon, see results below. The absolute difference between DTS* and baselines is again small, but the difference in text quality is significant. \n\n**Results for CoLA reward function**\n\n*CoLA rewards, higher is better*\n| **Method / NFEs** | $2^9$ | $2^{11}$ | $2^{13}$ | $2^{15}$ | $2^{17}$ |\n| :--- | ---: | ---: | ---: | ---: | ---: |\n| **DTS\\*** | **-0.262202** | **-0.141963** | **-0.074698** | **-0.045357** | **-0.035025** |\n| **Best-of-N** | -0.382968 | -0.180514 | -0.100434 | -0.063361 | -0.049272 |\n| **FK/SMC** | -0.296245 | -0.200835 | -0.120380 | -0.088619 | -0.056350 |\n\n\n**Prompt: “Once upon a time”**\n- **DTS\\* (reward = -0.027):** Once upon a time, I texted you and said,”It’s going really well.” I’d say,”I suppose you would be surprised, wouldn’t you?” Now, a few days ago, you were texting me, and\n- **Best-of-N (reward = -0.037):** Once upon a time, Samuel Adams was a famous launch merchant. Now I am afforded opportunity to observe the financial earnings of our Founders, President, and reminds me, at the @Beauty Post. I find myself very excited about & very interested in game mechanics, character models, jurists, movements\n- **FK/SMC (reward = -0.075):** Once upon a time women would commit themselves into fleets, spawning “Stupid”, STupid”, “Stupid”, “ST”, “Stupid”, “Stupid”, “Stupid”, “St\n\n\n**Prompt: “The pizza”**\n- **DTS\\* (reward = -0.027):** The pizza was topped with two teriyaki balls. These huge, tomatoey balls looked like they came from a pizza box. The dough was kind of soft and they balanced nicely. The vegetable sauce was voluminous. Another topping feature was ricotta. Ricotta uses a soft crust\n- **Best-of-N (reward = -0.038):** The pizza was made primarily of a cheese seasoned with sugar and spices. Image caption Cardiffers were first to sip their new beer on Boxing Day in 2016. Cardiffers were first to taste its brand new beer when it opened Boxing Day in 2016. Its doors opened when the 3200-seat\n- **FK/SMC (reward = -0.072):** The pizza is described as follows: B.298.50.57.16.99.1202.7.4601.8.069.5.1105.15.1113.16.5916914513.352.28017.28018.36.43\n\n\n**LLM evaluation**\n\nWe provide text samples to ChatGPT 4o and asking it to score them between 0 and 1 following a rubric (the prompt is omitted due to the character limit; we can provide it if the reviewer wishes). Results are averaged over 3 completions for each of the 15 prompts for $2^{17}$ NFEs. We observe that DTS\\* improves **31.3%** over Best-of-N when using inifini-gram reward, and **13.7%** when using CoLA reward.\n\n*LLM evaluation of text samples, higher is better*\n| **Method / Reward** | **infini-gram** | **CoLA** |\n|:---|---:|---:| \n| **DTS\\*** | **0.302 ± 0.023** | **0.322 ± 0.022** |\n|  **Best-of-N** | 0.230 ± 0.022 | 0.283 ± 0.020 |\n| **FK/SMC** | 0.111 ± 0.014 | 0.216 ± 0.019 |\n\n\n> The setting of the number of nodes and denoising steps\n\nWe do not explicitly specify the number of nodes for the tree. Instead, we specify a maximum number of diffusion model evaluations (NFEs reported in all of our figures) for a fair comparison across methods. For a concrete example, at $10^6$ NFEs with 50 denoising steps per sample, best-of-N and the SMC-based methods would use 20000 particles, and DTS/DTS* would have $10^6$ total nodes. All of our experiments compare these methods across different NFEs to understand how they scale with more compute (or equivalently, how they compare when varying the number of particles/nodes). \n\nThe number of denoising steps is orthogonal to the method used for inference-time alignment, and generally increasing the denoising steps results in higher quality samples from the base model, which is common across all methods. We use typical values of denoising steps for all methods: 50 is the default for pre-trained image and text-to-image models, and for text diffusion, it depends on the sequence length.\n\n\n> Please provide more details on the backup process\n\nThe backup process follows Equation 5, where the value of a state $x_t$ is the log-mean-exp of the value at states $x_{t-1}$, where the mean is over samples from the pretrained model $p_\\theta$. In DTS/DTS*, this translates to the value of each node being the log-mean-exp of its children values; since the children are sampled using $p_\\theta$, the expectation is over $p_\\theta$ as in equation 5. During tree construction, we traverse the tree until we reach a terminal node. We set the value of the terminal nodes equal to the exponent of the reward, and then iteratively update the values of the parents along this path. We hope this explanation helps clarify the backup, we will be happy to answer any further questions.\n\n\n*[1] Singhal et al. \"A General Framework for Inference-time Scaling and Steering of Diffusion Models.\" In Forty-second International Conference on Machine Learning, 2025.*\n\n*[2] Ma et al. \"Inference-time scaling for diffusion models beyond scaling denoising steps.\" arXiv preprint arXiv:2501.09732 (2025).*\n\n*[3] Hashimoto et al. \"Unifying Human and Statistical Evaluation for Natural Language Generation.\" In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics, Volume 1 (Long and Short Papers), pp. 1689-1701. 2019.*\n*[4] Celikyilmaz et al. \"Evaluation of text generation: A survey.\" arXiv preprint arXiv:2006.14799 (2020).*\n\n*[5] Liu et al. \"G-Eval: NLG Evaluation using Gpt-4 with Better Human Alignment.\" In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, 2023.*\n\n*[6] Warstadt et al. \"Neural network acceptability judgments.\" Transactions of the Association for Computational Linguistics 7 (2019): 625-641.*"}, {"review_id": "xua1eHuSZ2", "reviewer": "Reviewer_ULya", "summary": "This paper presents Diffusion Tree Sampling (DTS) a inference time method for guidance in diffusion models. DTS adopts MCTS for diffusion models by constructing a tree of possible rollouts using the diffusion chain and backing up rewards for more accurate value estimation of intermediate noisy states. The authors also present a search variant of the method, DTS*, for generating single high quality samples.", "questions": "See weaknesses", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**:\n\n1. The method builds on top of MCTS and thus inherits the advantages of MCTS as an anytime algorithm.\n2. The method is relatively straightforward to understand.\n3. Empirical results on Stable Diffusion are good to have but I have some concerns on other experimental designs (see weaknesses).\n\n**Weaknesses**:\n\nI don't have much criticism around the methodology since the method is pretty clear from the main text and the authors have clarified a lot of expensive design choices in Section 6 (However I think the authors should emphasize the runtime per prompt from App F.3 in more details in this section). I have some followup questions around the experimental design instead which I highlight as follows.\n \n1. How significant are the results on class conditional synthesis like MNIST/CIFAR-10. I feel like these benchmarks are heavily deprecated and do not reveal the complete picture.  Can the authors choose some other tasks. For instance, I would be interested in comparing DTS for inverse problems like super-resolution or in-painting? On these tasks DPS could still be considered a competitive baseline. \n\n2. Second, for the experiments on Stable Diffusion, can the authors explain why Best-of-N is a strong baseline since it feels like the most naive algorithm for any kind of inference time steering for diffusion models? Moreover, based on my understanding, unlike methods which tilt the diffusion guidance towards a pre-specified reward, Best-of-N does not do it. Therefore this observation seems a bit odd.\n\n3. How is exploration controlled in this algorithm? Do you always create a new node when the |C_t| < maximum value or with some probability traverse a chain with a high value estimate?\n\n4. Minor issues: Line 171: Do you mean x_T+1 instead of x_0\n5. Lastly there is some existing work [1,2,3] on gradient based steering of diffusion models which relies on similar ideas of tilting the reward distribution while simultaneously learning the guided posterior directly. I guess the paper can benefit from additional discussion around the limitations of these works.\n\n[1] Variational Control for Guidance in Diffusion Models - Pandey et al.\n\n[2] RB-Modulation: Training-Free Personalization of Diffusion Models using Stochastic Optimal Control - Rout et al.\n\n[3] Symbolic Music Generation with Non-Differentiable Rule Guided Diffusion - Huang et al.", "was_revised": true, "final_justification": "Based on the response from the authors, I will maintain my rating in favor of acceptance.", "rebuttal": "We thank the reviewer for their thorough review and positive assessment of our work. We address your concerns below, and also draw attention to our response to other reviewers for new experiments - text diffusion with LLM evaluation (reviewer 1bSF), comparison between DTS and DTS* (reviewer T4KC), and comparison of memory and runtime (reviewer t7aZ). Note that this year, NeurIPS does not allow posting any links or documents in the rebuttal. \n\n> I think the authors should emphasize the runtime per prompt from App F.3 in more details in this section\n\nWe refer to our response to reviewer t7aZ for a comparison of memory usage and runtime per prompt for text-to-image generation using Stable Diffusion 1.5 with ImageReward. We will include these results in the final version.\n\n> How significant are the results on class conditional synthesis like MNIST/CIFAR-10. I feel like these benchmarks are heavily deprecated and do not reveal the complete picture. \n\nWhile these benchmarks are saturated for tasks like image classification, they are still used in recent works for diffusion inference-time alignment or fine-tuning [1,2] since it is not trivial to sample from class-conditioned distributions using a base model that samples uniformly from all classes. We believe the results are significant, since the performance gap between DTS and the other methods is substantial, especially as we increase inference-time compute. The qualitative samples in Figure 6 and Appendix H.2 also showcase the specific characteristics of samples obtained using different methods. In particular, existing methods seem to miss certain modes or show characteristics of mode collapse, which shows that this task is not trivial for these methods. \n\n> Can the authors choose some other tasks. For instance, I would be interested in comparing DTS for inverse problems like super-resolution or in-painting? On these tasks DPS could still be considered a competitive baseline.\n\nFollowing the reviewer’s request, we ran an inpainting experiment on CIFAR10 images with a Gaussian forward channel as in the DPS paper. Due to time limitations, we run this experiment on 20 images (two images sampled randomly from each class) and average the results. For each 32x32  image, we mask the 10x10 patch from the center and use the negative error between the generated sample and the unmasked pixels from the original image as the reward. We report the final average MSE between 16 generated images and the reference image. \n\nThe results below show that DPS, which uses the gradient of the reward, is the only method that obtains good results on this task. We hypothesize this is because the inpainting task has a very sharp posterior where the sampled image must exactly match the unmasked pixels (compared to class conditional sampling where the posterior represents a semantic concept and has a wider mode). This is a very hard search problem and without the gradient, there is very little signal in most regions of the pretrained model distribution.\n\n*Avg. MSE between generated images and reference image, lower is better*\n| **Method / NFEs**   | **10000** | **25000** | **50000** | **100000** | **250000** | **500000** |\n|:---------|----------:|----------:|----------:|----------:|----------:|----------:|\n| **DPS**      | **0.01806** | **0.01563** | **0.01502** | **0.01501** | **0.01607** | **0.01773** |\n| **SMC/FK**   | 0.14587  | 0.19196  | 0.14906   | 0.19285  | 0.17591  | 0.19706  |\n| **DTS (Ours)**  | 0.16829 | 0.15937 | 0.15717 | 0.14121 | 0.13013 | 0.13018 |\n\n\nWe then incorporate the gradient into DTS, where we use the gradient term in addition to the score model to obtain the next state when constructing the tree. The backup and selection procedure remain the same. With the gradient-guided proposal, DTS improves upon DPS and even SMC + gradient (which is the same as TDS). We also observe it allows the performance to consistently improve with more compute, whereas vanilla DPS does not improve with more compute. At 500k NFEs, DTS + gradient improves **32.96%** compared to vanilla DPS.\n\n*Avg. MSE between generated images and reference image, lower is better*\n| **Method / NFEs**   | **10000** | **25000** | **50000** | **100000** | **250000** | **500000** |\n|:---------|----------:|----------:|----------:|----------:|----------:|----------:|\n| **DPS**      | 0.01806 | 0.01563 | 0.01502 | 0.01501 | 0.01606 | 0.01773 |\n| **SMC+gradient/TDS**     | 0.01415 | 0.01331 | 0.01916 | 0.02086 | 0.01981 | 0.01591 |\n| **DTS+gradient (Ours)**    | **0.01302** | **0.01263** | **0.01227** | **0.01194** | **0.01153** | **0.01189** |\n\n\n> Why is best-of-N a strong baseline?\n\nThere are two different types of experiments considered in our paper: (1) posterior sampling experiments (the illustrative 2D and class-conditional image generation experiments), and (2) what one may call “high reward search” experiments (text-to-image and text diffusion experiments).\n\nIn the first setting, there is a well-defined and tractable ground truth distribution (the reward-tilted distribution in the 2D case and the class-conditioned image distribution in the image case). Here, we can compare with the ground truth samples to evaluate how well methods like DTS and baselines match the target distribution.\n\nIn the second setting, the goal is to produce the single best sample that achieves both high likelihood under the pretrained model and high reward [3,4]. To generate an image for a given caption, for example, we care about having a single good image. In this setting, prior work [3,4] shows that this best-of-N approach is surprisingly strong because the pretrained model already draws from a rich, high-entropy distribution and best-of-N simply cherry-picks the luckiest draw. \n\nRecent work [5] provides theoretical guarantees for best-of-N for test-time scaling of LLMs: under some coverage conditions (if the base model assigns non-negligible probability to every high-reward region), there exists a finite N for which best-of-N achieves optimal performance. This makes best-of-N both a powerful baseline and a useful stress-test for any new inference-time alignment method. Yet the guarantees are brittle - when coverage is imperfect (the realistic case), the bounds loosen and best-of-N may miss entire reward modes; and as N grows large, the best sample increasingly exploits imperfections in the reward model, leading to reward hacking. We believe that our principled search method using unbiased soft values should help with both of these issues, since soft value functions can help guide denoising towards undersampled regions as well as avoid reward hacking.\n\n> How is exploration controlled in this algorithm? Do you always create a new node when the |C_t| < maximum value or with some probability traverse a chain with a high value estimate?\n\nWe always expand from a node if the number of children is less than the maximum allowed value. The trade-off between exploration and exploitation is controlled with this maximum value, which depends on the number of visits (this branching technique is called progressive widening in the literature [6]). Since our selection step (whether sampling as in DTS or picking the max as in DTS*) prioritizes higher value nodes, these nodes end up being visited more often and hence are allowed more children during expansion.\n> Minor issues: Line 171: Do you mean x_T+1 instead of x_0\n\nWe believe the reviewer is referring to line 177, and it is a typing error - it should be x_{T+1} as suggested. Thanks for pointing that out!\n\n> Lastly there is some existing work [A,B,C] on gradient based steering of diffusion models which relies on similar ideas of tilting the reward distribution while simultaneously learning the guided posterior directly. I guess the paper can benefit from additional discussion around the limitations of these works.\n\nThank you for suggesting these works. We will add the following discussion in the related works section:\n \n“Gradient-steering works treat guidance as stochastic control: Pandey et al. [A] learn KL-regularised drifts for differentiable rewards but break when high-reward modes lack coverage; Rout et al. [B] propose a training-free method to modify the drift for style transfer in vision; Huang et al. [C] extend to non-differentiable music rules via high-variance REINFORCE. Our search-based alignment does not require differentiability or dense coverage assumptions and is domain-agnostic.”\n\n*[1] Wu et al. \"Practical and asymptotically exact conditional sampling in diffusion models.\" Advances in Neural Information Processing Systems 36 (2023): 31372-31403.*\n\n*[2] Venkatraman et al. \"Amortizing intractable inference in diffusion models for vision, language, and control.\" Advances in neural information processing systems 37 (2024): 76080-76114.*\n\n*[3] Singhal et al. \"A General Framework for Inference-time Scaling and Steering of Diffusion Models.\" In Forty-second International Conference on Machine Learning, 2025.*\n\n*[4] Ma  et al. \"Inference-time scaling for diffusion models beyond scaling denoising steps.\" arXiv preprint arXiv:2501.09732 (2025).*\n\n*[5] Huang et al. \"Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment.\" In Forty-second International Conference on Machine Learning, 2025.*\n\n*[6] Couëtoux et al. \"Continuous upper confidence trees.\" In International conference on learning and intelligent optimization, pp. 433-445. Berlin, Heidelberg: Springer Berlin Heidelberg, 2011.*\n\n*[A] Pandey et al. \"Variational Control for Guidance in Diffusion Models.\" In Forty-second International Conference on Machine Learning, 2025.*\n\n*[B] Rout et al. \"Rb-modulation: Training-free personalization of diffusion models using stochastic optimal control.\" arXiv preprint arXiv:2405.17401 (2024).*\n\n*[C] Huang et al. \"Symbolic Music Generation with Non-Differentiable Rule Guided Diffusion.\" In International Conference on Machine Learning, pp. 19772-19797. PMLR, 2024.*"}, {"review_id": "qJmrYAGabd", "reviewer": "Reviewer_T4KC", "summary": "The paper presents a Monte Carlo Tree Search algorithm for inference-time alignment of diffusion models. Authors cast inference-time alignment as a search problem that reuses past computations, and intoroduce an algorithm that samples from reward-aligned target density by propagating rewards through the search tree. Authors present a theoretical motivation for their method and experimentally show that the method achieves better quality of the samples than considered baselines, as well as better scales with compute.", "questions": "0) See Weaknesses.\n\n1) Line 270: does $10^6$ denote the number of function evaluations used to generate each image, or the whole set of samples? If it is the NFE for a single image, this seems unreasonably high for such simple datasets. Did the authors measure the GPU time needed to generate one image? The same question for the NFEs denoted in plots in Figures 6, 7, 8. \n\n2) Algorithm 1 describes how the search tree is built. Are final samples used for evaluation the same as the leaves created during the construction of the tree, or are they obtained by resampling paths inside the tree after construction? I did not find the description of how the final samples for evaluation are obtained in the paper, I think it would be better to clarify explicitly.\n\n3) At first I struggled to understand the motivation and message of Section 3.1 in the text, as it feels a bit disconnected from the rest of the work. I think that explicitly explaining in the text how the value estimates by themselves can help with the problem at hand (e.g. how they can be used for guidance) could improve the readability of this section.", "limitations": "No concern.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**: \n\nThe paper is generally well-structured and easy to read (except for some theoretical parts and evaluation details, see weaknesses and questions below). The idea of applying MCTS for inference-time alignment of diffusion models is novel and, in my opinion, very interesting and promising. The experimental evaluation and results are solid. I also did enjoy the illustrative synthetic experiments in Section 4.4, although I would suggest moving Figure 9 from Appendix H.1 to the main text as it would better help illustrate the performance of the proposed method (in Figure 4 the difference with some baselines is visually less obvious).\n\n**Weaknesses**: \n\nMy main concern for this paper is the presentation of theoretical background and analysis for the proposed method. \n\nFirst of all, up to my knowledge, the definition of soft value in Section 2 is not the most common one used in KL-regularized RL problems. Usually, the soft optimal value is defined through the expected reward minus KL term, where the expectation is taken with respect to the optimal policy, not the reference policy (which is $p_\\theta$ in this paper). In the context of RL-based diffusion fine-tuning see, e.g., Section 2.2 of [1]. Although, from my understanding, the aforementioned definitions should produce equivalent functions, authors should either reference some previous literature that studies the exact same definition, or elaborate and provide more context for the reader. Authors reference two papers on entropy-regularized RL, which is a related setting, but not exactly the same, and study a different problem. \n\nSecondly, theoretical derivations in Appendix D seem valid to me, however, whether the Algorithm 1 in Appendix E directly follows them is quite unobvious. E.g., optimal policy in Equation 6 is proportional to $p_\\theta$ times the exponent of value, while in line 9 of Algorithm 1 the next state is sampled proportionally to just the exponent of value. Similarly, the mentioned soft Bellman equation in Equation 5 expresses the value through the expectation of exponentiated values in the next states with respect to the reference policy, while soft backup in Algorithm 1 (line 26) just takes the sum of exponentiated values in the next states, not the expectation with respect to $p_\\theta$. Do I understand correctly that this is due to the fact that the expectation with respect to $p_\\theta$ is thought to be estimated via Monte Carlo, since paths in the tree are sampled using $p_\\theta$? If this is the case, the authors should definitely elaborate on this more in their theoretical derivations. If not, then I do not really understand the theory behind the proposed algorithm. Also, if Monte Carlo estimate is used in line 26, shouldn't there be the logarithm of the mean, not the logarithm of the sum?\n\nAddressing these points will help to substantially improve the theoretical part of the work and its clarity in my opinion.\n\nThe presented connection to GFlowNets in Appendix C is interesting and insightful, however, there is a caveat in how it is written in the text. If it is the usual GFlowNet task of sampling from the target distribution, and $F$ denotes the log-flow, then the soft Bellman equation in Appendix C should use the backward policy $P_B(s \\mid s')$, not the trainable forward policy $P_\\theta(s' \\mid s)$. See [2] for the analysis of connections to entropy-regularized RL and how soft Bellman equations are expressed through GFlowNet flows and policies. However, [3] puts the task of fine-tuning of diffusion models in the GFlowNet framework, and in such case the reference policy will replace the backward policy in equations. In this case the approach proposed in this paper will indeed be very related to the one presented in [4], which is another insightful observation and connection made by the authors.\n\nI also have some minor points related to the Experiments section. I think that putting some direct comparison of DTS and DTS* would help the reader to better understand the benefits of each. It looks a bit odd that some experiments in the paper present the results only for DTS, and others only for DTS*. In addition, it would be good to put some references to previous literature near the statement \"reward functions in text domains are particularly susceptible to over-optimization, often resulting in repetitive outputs\", as making such observation from just one experiment in the paper is not well enough validated.\n\nReferences:\n\n[1] Uehara et al. Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review. 2024\\\n[2] Tiapkin et al. Generative Flow Networks as Entropy-Regularized RL. AISTATS 2024\\\n[3] Venkatraman et al. Amortizing intractable inference in diffusion models for vision, language, and control. NeurIPS 2024\\\n[4] Morozov et al. Improving GFlowNets with Monte Carlo Tree Search. ICML 2024 SPIGM Workshop", "was_revised": true, "final_justification": "Despite the weaknesses in presentation I outlined in my review, I believe that this is a solid work. I'm glad that my comments and suggestions about the clarity of theoretical results and comparisons between DTS and DTS* were helpful to the authors for the purpose of improving the paper. I decided to maintain my score and remain positive towards accepting this work.", "rebuttal": "We thank the reviewer for their detailed review and positive assessment of our work. We address your concerns below, and also draw attention to our response to other reviewers for additional experiments - text diffusion with LLM evaluation (reviewer 1bSF), image inpainting (reviewer ULya) and comparison of memory and runtime (reviewer t7aZ). Note that this year, NeurIPS does not allow posting any links or documents in the rebuttal. \n\n> Definition of soft value function.\n\nThe understanding of the reviewer is correct, these two definitions are equivalent. Indeed, we use the same definition of the value function as in Uehara et al. [1] (Equation 6 in Section 2.1). We will add a reference to this work in Section 2 as suggested by the reviewer.\n\n> Sampling and expected value in backup\n\nThe reviewer understands correctly. Since the nodes are expanded using $p_\\theta$, when we perform selection based on the exponentiated values, we are effectively sampling from the product $p_\\theta \\exp(V_t(x_t))$. The same reasoning is applied in the backup, since the Monte Carlo estimate uses samples drawn from $p_\\theta$. We mention that children are drawn from $p_\\theta$ in step 1 of the proof, but we agree with the reviewer that it should be made more explicit, especially when explaining the backup operation. We will update this in the paper.\n\nThank you for pointing out the typing error in Algorithm 1, line 26 - instead of sum, it should be mean over the children (dividing by the number of children $|C_t|$). \n\n> Connection to GFlowNets\n\nThank you for pointing this out. The soft Bellman equation in Appendix C should indeed use the backward policy. In the case of diffusion, $P_B(s | s’)$ is the noise process that moves from data to noise space. We think the reviewer may find Figure 2 from He et al. [2] helpful. We will correct this in the paper.\n\n> Some experiments in the paper present the results only for DTS, and others only for DTS*\n\nThere are two different types of experiments considered in our paper: (1) posterior sampling experiments (the illustrative 2D and class-conditional image generation experiments), and (2) what one may call “high reward search” experiments (text-to-image and text diffusion experiments). The nature of these two distinct tasks and prior literature motivate the choice of algorithms. DTS is designed to tackle problem (1) and DTS* for problem (2).\n\nIn the first setting, there is a well-defined and tractable ground truth distribution (the reward-tilted distribution in the 2D case and the class-conditioned image distribution in the image case). DTS and the baselines considered in these experiments are designed to sample approximately from this target distribution; hence, we report distribution-based metrics.\n\nIn the second setting, the goal is to produce the single best sample that achieves both high likelihood under the pretrained model and high reward [3,4]. To generate an image for a given caption, for example, we care about having a single good image. DTS and the other baselines from the previous setting are not designed to produce a single good sample, unlike DTS*. Therefore, we only consider each of these methods in its respective setting. We will explain the difference between the two settings more clearly in the paper. \n\nFollowing the reviewer’s request, we compare DTS and DTS* on the text-to-image setting using Stable Diffusion 1.5 with ImageReward. Due to time limitations, we limit this experiment to 40 prompts from DrawBench. We picked complex prompts where the base model struggles to generate high reward images to demonstrate the difference more clearly. For DTS, we sample 16 images from the tree and pick the one with the highest reward. The results show that there is a noticeable gap, since DTS spends more compute expanding suboptimal paths whereas DTS* aggressively searches only the highest reward regions (except for the small exploration bonus).\n\n*ImageReward, higher is better*\n\n| **Method / NFEs**   | **1000** | **2500** | **5000** | **10000** | **25000** | **50000** |\n|:---------|----------:|----------:|----------:|----------:|----------:|----------:|\n| **DTS**      | 0.774 | 1.025 | 1.064 | 1.112 | 1.214 | 1.238 |\n| **DTS\\***   | **0.825**  | **1.086**  | **1.254**   | **1.342**  | **1.459**  | **1.552**  |\n\n\n> References for supporting the statement about text reward functions\n\nFollowing the reviewer’s suggestion, we will add the following references to support the statement:\n\nHoltzman et al. [5] highlight that text generation strategies that optimize for perplexity of some learned model often lead to dull, repetitive, and unnatural text. Hashimoto et al. [6] observe that metrics often fail to capture the trade-offs that human evaluators make, such as quality vs. diversity. The excellent survey by Celikyilmaz et al. [7] discusses that statistical evaluation of text generally correlates poorly with human judgments and does not necessarily follow a linear scale. We add evaluations of the text samples using ChatGPT 4o in our response to reviewer 1bSF, following previous work which shows these evaluations tend to have higher correlation with human judgement [8].\n\n> Meaning of NFEs\n\nThe NFE values in all our experiments refer to the total number of function evaluations used in the algorithm. We test all methods for different values of NFEs to study how these methods scale with more compute. For a concrete example, if we are comparing all methods at 10^6 NFEs with 50 denoising steps per sample, best-of-N and the SMC-based methods would use 20000 (=10^6 / 50) particles, and DTS/DTS* would have 10^6 total nodes in the tree. Each individual sample is still obtained using the same number of denoising steps (50 in this example). We refer to our response to reviewer t7aZ for comparison of runtime and memory usage of different methods.\n\n> How are the final samples obtained?\n\nOnce the tree is constructed, we perform a final selection process (depending on DTS/DTS*) by traversing the tree from the root to obtain the final samples. These would be a subset of the leaves constructed in the tree. Following the reviewer’s suggestion, we will add a separate Algorithm 2 that, given a constructed tree, performs selection to obtain the final samples.\n\n> Clarifying motivation and message of Section 3.1\n\nThank you for the feedback. We will add some motivation for studying value functions to connect it with the rest of the paper and improve clarity. We will add the text below to the beginning of Section 3.1:\n\n“The central problem for inference-time alignment of diffusion models is accurate credit assignment during denoising to evaluate which noisy states have the potential to lead to higher reward samples. This notion is captured by value functions as defined in Equation 5. The discussion above explains how most existing methods use the value function in some form, however, since they are generally intractable to compute, certain approximations are commonly used in literature.”\n\nWe will add the following text to the end of Section 3.1 to connect the results of Figure 2 with the rest of the paper:\n\n“These inaccurate value estimates lead to biased guidance during denoising, resulting in samples that may not approximate the target distribution accurately.” \n\n> Moving Figure 9 from Appendix H.1 to the main text\n\nWe had to move Figure 9 to the appendix due to space constraints. The final version of the paper allows one additional page, and if accepted, we plan to move Figure 9 to the main paper.\n\n\n*[1] ​​Uehara, Masatoshi, Yulai Zhao, Chenyu Wang, Xiner Li, Aviv Regev, Sergey Levine, and Tommaso Biancalani. \"Inference-time alignment in diffusion models with reward-guided generation: Tutorial and review.\" arXiv preprint arXiv:2501.09685 (2025).*\n\n*[2] He, Haoran, Emmanuel Bengio, Qingpeng Cai, and Ling Pan. \"Random Policy Evaluation Uncovers Policies of Generative Flow Networks.\" In Forty-second International Conference on Machine Learning, 2025.*\n\n*[3] Singhal, Raghav, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. \"A General Framework for Inference-time Scaling and Steering of Diffusion Models.\" In Forty-second International Conference on Machine Learning, 2025.*\n\n*[4] Ma, Nanye, Shangyuan Tong, Haolin Jia, Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang et al. \"Inference-time scaling for diffusion models beyond scaling denoising steps.\" arXiv preprint arXiv:2501.09732 (2025).*\n\n*[5] Holtzman, Ari, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. \"The Curious Case of Neural Text Degeneration.\" In International Conference on Learning Representations, 2020.*\n\n*[6] Hashimoto, Tatsunori B., Hugh Zhang, and Percy Liang. \"Unifying Human and Statistical Evaluation for Natural Language Generation.\" In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 1689-1701. 2019.*\n\n*[7] Celikyilmaz, Asli, Elizabeth Clark, and Jianfeng Gao. \"Evaluation of text generation: A survey.\" arXiv preprint arXiv:2006.14799 (2020).*\n\n*[8] Liu, Yang, Dan Iter, Yichong Xu, Shuohang Wang, Ruochen Xu, and Chenguang Zhu. \"G-Eval: NLG Evaluation using Gpt-4 with Better Human Alignment.\" In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2023.*"}, {"review_id": "W75CQ1Xc9I", "reviewer": "Reviewer_t7aZ", "summary": "This paper introduces Diffusion Tree Sampling (DTS), a novel method for inference-time alignment of diffusion models that is inspired by Monte Carlo Tree Search (MCTS). The authors frame the iterative denoising process as a search problem on a tree, where nodes represent noisy states. By performing rollouts to the terminal state and backing up the observed rewards through the tree, DTS iteratively refines the value estimates of intermediate states. This approach addresses a key weakness of prior methods: the reliance on inaccurate, one-step value estimations, especially at high noise levels. DTS is designed as an \"anytime\" algorithm that reuses past computations to progressively improve sample quality. A greedy search variant, DTS*, is also proposed for optimization tasks.", "questions": "- The key concern on this method is clearly its scalability with memory. Could the authors provide an analysis comparing the memory footprint of DTS to baselines like SMC for a fixed task? \n- In the ordinary DTS without UCT exploration, is the stochasticity from the Boltzmann sampling sufficient to ensure the tree's exploration, if the modes of the reward is quite diverse?\n- I would suggest including a brief discussion on other related methods on this topic, e.g. SGS [1,2], and recent advances in SMC [3,4], etc..\n\n\n### Minor Issues\n- It seems in the proofs in Appendix D.1, several $V$ lacks their corresponding subscripts, making it hard to follow.\n- On line 178, the root should be changed from $x_0$ to $x_t$.\n\n[1] Wu, Zihui, et al. \"Principled probabilistic imaging using diffusion models as plug-and-play priors.\" Advances in Neural Information Processing Systems 37 (2024): 118389-118427.\\\n[2] Coeurdoux, Florentin, Nicolas Dobigeon, and Pierre Chainais. \"Plug-and-play split Gibbs sampler: embedding deep generative priors in Bayesian inference.\" IEEE Transactions on Image Processing (2024).\\\n[3] Skreta, Marta, et al. \"Feynman-kac correctors in diffusion: Annealing, guidance, and product of experts.\" arXiv preprint arXiv:2503.02819 (2025).\\\n[4] Chen, Haoxuan, et al. \"Solving inverse problems via diffusion-based priors: An approximation-free ensemble sampling approach.\" arXiv preprint arXiv:2506.03979 (2025).\\\n[5] He, Jiajun, et al. \"RNE: a plug-and-play framework for diffusion density estimation and inference-time control.\" arXiv preprint arXiv:2506.05668 (2025).", "limitations": "Limitations are adequately addressed.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n- The concept of applying MCTS-like method to diffusion model alignment is novel and provides a new paradigm for guidance. One key feature of the algorithm is \"anytime\", that is reusing information from all past simulations to improve value estimates. \n- The paper provides comprehensive empirical evidence of DTS's effectiveness. It is shown to match the performance of leading baselines on a variety of tasks, including image generation, text-to-image, etc.\n\n### Weaknesses\n- The method requires a large amount of memory, including the entire search tree and intermediate states. This could lead to prohibitive computational cost for large-scale tasks.\n- It would be beneficial to include more details on DTS*, the proposed variant of DTS, discussing on how the modified algorithm fosters exploration.", "was_revised": true, "final_justification": "I am updating the score to 5 (Accept) upon their promised revision, in light of the overall quality and novel idea of the paper.", "rebuttal": "We thank the reviewer for their careful review and positive assessment of our work. We address your concerns below, and also draw attention to our response to other reviewers for additional experiments - text diffusion with LLM evaluation (reviewer 1bSF), image inpainting (reviewer ULya), and comparison between DTS and DTS* (reviewer T4KC). Note that this year, NeurIPS does not allow posting any links or documents in the rebuttal. \n\n> Comparing memory footprint and wall clock runtime\n\nBelow, we provide peak memory usage for each method in the text-to-image generation experiment using Stable Diffusion 1.5 and ImageReward. We perform this experiment with 40 different prompts and average the results. The numbers are reported for a single NVIDIA A100 GPU with 40GB of memory using the pytorch profiler.\n\n**Peak Memory Usage (MB):**\n\nWe observe that the memory usage of DTS* is much lower compared to FK/SMC, and is more comparable to best-of-N. Initially, DTS* uses slightly more memory than best-of-N, but with more compute, it tilts in favor of DTS*, showing that it actually scales better in terms of memory. \n\nThe main memory usage of DTS* is the tree as noted by the reviewer, which stores 64x64x4 latents. Note that currently the entire tree is kept on the GPU, and we can save GPU memory by offloading the tree to the CPU. For best-of-N, the memory usage peaks at the end when all N candidates need to be decoded using the VAE into 512x512x3 images and evaluated using the reward function to pick the best sample. This problem does not occur with DTS*, since the rewards are calculated when traversal reaches the terminal state, and multiple high-resolution images do not need to be kept in memory. For FK/SMC, this cost is very high since all N particles need to be decoded into 512x512x3 images, then evaluated using the reward function *at each resampling step*. In general, for images and text, the reward function itself is a large pretrained model, and computing the reward requires a non-trivial amount of compute.\n\n*Peak memory usage (MB), lower is better*\n\n| **Method / NFEs** | 1000 | 2500 | 5000 | 10000 | 25000 | 50000 | 100000 | \n|:----------|-----:|-----:|-----:|------:|------:|------:| ------:| \n| **DTS\\* (Ours)** | 8059 | 9140 | 9460 | 10103 | **12025** | **15226** |  **17318** |\n| **Best-of-N** | **8056** | **8331** | **8789** | **9705** | 12453 | 17033 | 19782 | \n| **FK/SMC** | 27167 | 27348 | 27559 | 27974 | 29220 | 31296 | 33948 |\n\n**Wall clock time (seconds):**\n\nIn terms of wall clock time, DTS* is approximately 2.5x slower than the other methods for the same number of function evaluations (NFEs), since it is harder to parallelize. Note that the runtime reflects tree building from scratch - once the tree is built, repeated sampling is just pointer chasing and is near instantaneous (whereas for other methods repeated sampling would require starting the entire procedure again). \n\nWe discuss two important considerations regarding the runtime:\n\n1) When taking into account their performance, DTS* is still preferable in terms of wall clock time since it outperforms the baselines and needs much fewer NFEs to reach the same level of performance. As noted in the paper and the table below, for the text-to-image task, it needs roughly 5x fewer NFEs to match the reward of the best performing baseline (best-of-N) at 100k NFEs. In other words, to match the performance of best-of-N at 100k NFEs, DTS* is approximately 2x faster in terms of wall clock time.\n\n2) We can significantly reduce the runtime of DTS/DTS* by using multiple actors in parallel that traverse and update a shared tree (similar to AlphaZero training). This will give a roughly linear speedup in runtime at the cost of increased memory for copying the diffusion model (the tree is not copied since it can be shared among all actors, hence increase in memory will be sublinear in terms of the number of actors). This represents an inherent trade-off between memory usage and runtime. We did not implement a distributed pipeline since it requires some engineering effort, but it is a viable option.\n\n*Wall clock time (seconds), lower is better*\n\n| **Method / NFEs** | 1000 | 2500 | 5000 | 10000 | 25000 | 50000 | 100000 | \n|:----------|-----:|-----:|-----:|------:|------:|------:| ------:| \n| **DTS\\* (Ours)** | 52 | 132 | 267 | 543 | 1372 | 2763 | 5513 | \n| **Best-of-N** | 21 | 50 | 96 | 188 | 461 | 926 | 1910 |\n| **FK/SMC** | 23 | 55 | 107 | 213 | 535 | 1073 | 2204 |\n\n*ImageReward score, higher is better*\n\n| Method | 1000 | 2500 | 5000 | 10000 | 25000 | 50000 | 100000 |\n|:----------|---------:|---------:|---------:|---------:|---------:|---------:| ---------:| \n| **DTS\\* (Ours)** | 1.148 | 1.309 | 1.407 | 1.506 | 1.605 | 1.665 | 1.704 |\n| **Best-of-N** | 1.140 | 1.269 | 1.350 | 1.417 | 1.478 | 1.530 | 1.564 |\n| **FK/SMC** | 1.182 | 1.221 | 1.324 | 1.338 | 1.423 | 1.506 | 1.528 |\n\n> It would be beneficial to include more details on DTS*, the proposed variant of DTS, discussing on how the modified algorithm fosters exploration.\n\nDTS* involves changing the selection rule, where instead of sampling children proportional to exponentiated values, we pick the child with the highest value at each node. We discuss some considerations in Section 4.3.  Given the reviewer’s suggestion (and also those of reviewer ULya), we will elaborate on some of the details a bit more with the extra page allowed in the final version of the paper. DTS* inherently does not foster exploration, which is why we use UCT to add some exploration in the tree traversal. \n\n> In the ordinary DTS without UCT exploration, is the stochasticity from the Boltzmann sampling sufficient to ensure the tree's exploration, if the modes of the reward is quite diverse?\n\nThe reviewer raises an excellent point. In the sampling case, since DTS selects children using Boltzmann sampling, it inherently has exploration ability since there is a non-trivial probability of sampling suboptimal children. Empirically, we found no performance difference when using UCT with DTS over a range of exploration coefficients. This is also supported theoretically by recent work in the bandit setting, where if the goal is to sample from the Boltzmann distribution of the reward, simply picking arms based on the Boltzmann distribution of the estimated reward without explicit exploration leads to optimal regret bounds [1].\n\n> I would suggest including a brief discussion on other related methods on this topic, e.g. SGS [A,B], and recent advances in SMC [C,D], etc..\n\nWe thank the reviewer for pointing out these works. We will include the following discussion in the related works section. \n\n“Recent plug-and-play samplers [A, B] also use Tweedie’s formula to get biased values and use compute-heavy MCMC/Gibbs chains for inverse problems. Skreta et al. [C] use SMC correctors to re-weight particles each step which inherits the same problems of diversity collapse and low effective sample size when the weights have high variance. Chen et al. [D] propagate a diffusion-PDE ensemble, which is accurate but can be compute-intensive; while He et al. [E] steer paths via high-variance density-ratio estimates. Our principled search method produces high quality samples by using unbiased values and works with arbitrary rewards.”\n\n> It seems in the proofs in Appendix D.1, several V lacks their corresponding subscripts, making it hard to follow.\n\n> On line 178, the root should be changed from x_0 to x_T\n\nThank you for pointing out these typing errors. The root should indeed be $x_{T+1}$ in line 178, and we will add subscripts to the value function in the appendix. \n\n*[1] Pedramfar, Mohammad, and Siamak Ravanbakhsh. \"Multi-Armed Sampling Problem and the End of Exploration.\" arXiv preprint arXiv:2507.10797 (2025).*\n\n*[A] Wu, Zihui, Yu Sun, Yifan Chen, Bingliang Zhang, Yisong Yue, and Katherine Bouman. \"Principled probabilistic imaging using diffusion models as plug-and-play priors.\" Advances in Neural Information Processing Systems 37 (2024): 118389-118427.*\n\n*[B] Coeurdoux, Florentin, Nicolas Dobigeon, and Pierre Chainais. \"Plug-and-play split gibbs sampler: embedding deep generative priors in bayesian inference.\" IEEE Transactions on Image Processing 33 (2024): 3496-3507.*\n\n*[C] Skreta, Marta, Tara Akhound-Sadegh, Viktor Ohanesian, Roberto Bondesan, Alan Aspuru-Guzik, Arnaud Doucet, Rob Brekelmans, Alexander Tong, and Kirill Neklyudov. \"Feynman-Kac Correctors in Diffusion: Annealing, Guidance, and Product of Experts.\" In Forty-second International Conference on Machine Learning, 2025.*\n\n*[D] Chen, Haoxuan, Yinuo Ren, Martin Renqiang Min, Lexing Ying, and Zachary Izzo. \"Solving inverse problems via diffusion-based priors: An approximation-free ensemble sampling approach.\" arXiv preprint arXiv:2506.03979 (2025).*\n\n*[E] He, Jiajun, José Miguel Hernández-Lobato, Yuanqi Du, and Francisco Vargas. \"RNE: a plug-and-play framework for diffusion density estimation and inference-time control.\" arXiv preprint arXiv:2506.05668 (2025).*"}]

Author rebuttal:

Meta-review: This paper proposes an inference-time steering method for diffusion models based on Monte Carlo Tree Search (MCTS). The proposed method is straightforward, principled, and demonstrates strong empirical performance across multiple modalities, including 2D data, image, and text. As acknowledged by the reviewers, the paper is clearly written and well-structured. The discussion primarily addressed clarifications regarding the method, theoretical aspects, and experimental setup. All the reviewers found the response from the authors satisfactory. Overall, this paper is a solid contribution, and thus I recommend acceptance.

One important note is the existence of very recent related work [1], which also applies MCTS to diffusion models. Although there are clear differences in implementation and application scenarios, it would be helpful to mention and discuss these distinctions in the revised manuscript.

[1] Jaesik Yoon, Hyeonseo Cho, Doojin Baek, Yoshua Bengio, Sungjin Ahn. Monte Carlo Tree Diffusion for System 2 Planning. ICML 2025.
Decision: Accept (poster)