PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models

Vineet Jain, Kusha Sareen, Mohammad Pedramfar, Siamak Ravanbakhsh

arXiv:2506.20701v1cs.LGcs.AIstat.ML

TL;DR

Inference-time alignment must guide diffusion using rewards available only at the end, while existing approaches leave prior-run information unused. DTS uses a tree of denoising trajectories to refine value estimates, and its search variant targets high-reward samples. The paper reports improved sampling results and describes a limitation when sample counts greatly exceed tree size.

  • Problem

    Inference-time alignment must guide denoising from rewards available only at the end, and prior sampling runs are not systematically reused to improve future samples.

    Source: §1 · p. 2, §3.2 · p. 5

  • Method

    DTS builds a tree over diffusion states and uses terminal-reward backups to refine intermediate soft-value estimates across rollouts.

    Source: §1 · p. 2, §4.2 · p. 5, §6 · p. 11

  • Results

    DTS achieves the lowest FID and CMMD in all three class-conditional settings after 10^6 NFEs, and its sample quality improves as compute increases.

    Source: §5.1 · p. 9, §6 · p. 11

  • Takeaways & Limitations

    Soft-value backups provide global credit assignment, while reusing computation makes sample quality improve as additional compute is spent.

    Source: §3.2 · p. 5, §6 · p. 11

    When the number of samples significantly exceeds tree size, repeated samples can correspond to high-probability leaves.

    Source: §6 · p. 11

Abstract

from arXiv · show

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, resulting in 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$^\star$), 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 $10\times$ less compute. In text-to-image generation and language completion tasks, DTS$^\star$ effectively searches for high reward samples that match best-of-N with up to $5\times$ less compute. By reusing information from previous generations, we get an anytime algorithm that turns additional compute into steadily better samples, providing a scalable approach for inference-time alignment of diffusion models.

1 Introduction

Inference-time alignment must guide denoising using rewards revealed only at the end, while existing reward estimates can bias guidance. The paper frames this challenge as tree search, proposing DTS for reward-aligned sampling and DTS⋆ for high-reward search.

1 Introduction

  • Existing gradient-based, SMC, and search-based methods approximate unseen terminal rewards, and those approximations can bias decisions and degrade sample quality.The paper motivates tree search by viewing the pretrained diffusion model as a deterministic policy and reverse-process Gaussian steps as stochastic environment transitions.Source: §1 · pp. 1–2
  • Inference-time alignment seeks to guide denoising when rewards are available only at the trajectory end, a credit-assignment challenge compounded by the need to reuse information across samples.The paper asks how prior inference-time samples can improve future samples through a sequential, scalable process.Source: §1 · pp. 1–2
  • DTS uses rollouts in a finite-horizon denoising tree to refine intermediate-state value estimates, while DTS⋆ searches denoising trajectories for modes in high-volume target-density regions.DTS is for sampling from the target density; DTS⋆ is for optimization toward high-reward samples.Source: §1 · p. 2
  • Up to 10× less compute matches baseline performance on class-conditional image generation, while up to 5× less compute matches best-of-N on text-to-image alignment and language completion.The authors also report asymptotically exact target-distribution samples for DTS in the limit of infinite rollouts.Source: §1 · p. 2

2 Preliminaries

A diffusion model defines a generative chain whose learned reverse process denoises from a simple prior. Alignment reweights the pretrained model by reward, and its optimal policy can be characterized through a recursively defined soft value function.

2 Preliminaries

  • The learned reverse process iteratively denoises samples from a simple prior, such as a standard Gaussian, according to a noise schedule.The reverse transitions induce the model’s generative distribution.Source: §2 · p. 3
  • Alignment targets the posterior proportional to the pretrained density times the probability of a desirable property, commonly modeled as exponentiated reward with inverse temperature λ.The paper assumes λ = 1 unless otherwise stated and considers sampling from the target distribution or finding its mode.Source: §2 · p. 3
  • The pretrained generative process can be viewed as a policy, and the target distribution as its optimal policy under an objective regularized by divergence from the pretrained model.This formulation is related to maximum-entropy reinforcement learning, with KL divergence to the pretrained model replacing entropy regularization.Source: §2 · p. 3
  • The soft value function obeys a recursive relation analogous to the soft Bellman equation and exactly characterizes the optimal policy.Its terminal condition is the reward, and the formulation connects optimal sampling with soft-value estimation.Source: §2 · p. 3

3 Inference-time adaptation of diffusion models

Inference-time alignment methods guide denoising toward terminal rewards, but their value estimates rely on approximations that become unreliable at high noise. Existing sampling procedures also discard intermediate evaluations rather than reusing them across runs.

3 Inference-time adaptation of diffusion models

  • Existing inference-time alignment methods use gradient guidance, sequential Monte Carlo, or search, but rely on approximations to estimate unseen terminal rewards.Fine-tuning with reinforcement learning is incompatible with reward functions that are unknown until inference time.Source: §3.1 · p. 4, §3 · p. 4

3.1 The value estimation problem

  • At high noise, Tweedie-based one-step value estimates become increasingly inaccurate and can be essentially random, because they substitute a learned score for the true intermediate score.The expected reward is approximated by evaluating the reward at the one-step posterior mean.Source: §3.1 · pp. 4–5

3.2 Scaling axes

  • Existing SMC and search methods discard intermediate evaluations between runs, so extra compute cannot cumulatively improve value estimates or sample quality.This is especially relevant because value estimation is difficult at high noise levels.Source: §3.2 · p. 5

4 Diffusion Tree Sampling and Search

DTS organizes denoising as a tree and reuses rollout rewards to refine values, supporting both target-distribution sampling and high-reward search. In illustrative experiments, it approximates target densities more accurately than the compared methods, while sample quality improves with compute and value estimates have lower bias and variance.

4.1 Denoising tree

  • The denoising process forms a finite-horizon tree whose nodes are noisy states and whose edges are diffusion steps, allowing information and soft-value estimates to be tracked across trajectories.DTS samples from the target density, while DTS⋆ is a search variant for high-reward samples.Source: §4.1 · p. 5

4.2 Tree-based sampling

  • DTS repeatedly selects by estimated values, expands a child, rolls it out to a terminal state, then backs up its reward through the tree.The rollout path is added to the tree, and terminal rewards update parent values using the soft Bellman equation.Source: §4.2 · pp. 5–6
  • With bounded rewards and λ > 0, DTS's empirical terminal-state distribution converges to the optimal policy as tree iterations approach infinity.The tree can also be used for approximate target sampling while it is being built.Source: §4.2 · pp. 5–6

4.3 Design choices for diffusion alignment

  • DTS⋆ selects the child with the largest soft-value estimate, while DTS samples by values and uses no explicit exploration bonus in practice.Progressive widening governs branch expansion, and UCT guides exploration for DTS⋆.Source: §4.3 · pp. 6–7

4.4 Illustrative experiments

  • In the illustrative experiments, the prior is either an equally weighted eight-Gaussian mixture or has checkerboard support; rewards unevenly weight modes or favor proximity to the origin.Methods are benchmarked by diffusion-model function evaluations.Source: §4.4 · pp. 7–8
  • In both 2D settings, DTS approximates the target density more accurately than the compared methods, and its sample quality improves with additional compute.Gradient-based methods DPS-RG and TDS-RG suffer from instability and require gradient clipping to stabilize denoising.Source: Figure 5 · p. 7, §4.4 · p. 8
  • DTS reduces value-estimate bias using reward information and variance by aggregating multiple rollouts; one-step denoising and single-rollout estimates have high bias and variance.The compared methods' errors generally worsen at higher diffusion timesteps.Source: §4.4 · p. 8, Figure 6 · p. 8

5 Experiments

Experiments evaluate DTS for class-conditional posterior sampling and DTS⋆ for text-to-image and text generation. Across these settings, the methods show strong posterior fidelity or high-reward search, with performance improving as compute increases.

5 Experiments

  • The experiments evaluate faithful sampling from high-dimensional posteriors and search for high-reward samples in high-dimensional image space.The class-conditional tasks use unconditional MNIST and CIFAR-10 diffusion priors; text-to-image tests use Stable Diffusion v1.5 on DrawBench and animal prompts.Source: §5 · p. 8, §5.1 · p. 8, §5.2 · p. 10

5.1 Class-conditional posterior sampling

  • Across all three class-conditional settings, DTS achieves the lowest FID and CMMD, and its FID remains low with better scaling as compute increases.The evaluation covers MNIST single digits, MNIST odd/even digits, and CIFAR-10 single classes; the CIFAR-10 classifier has about 85% accuracy, and TDS/SMC show mode-collapse characteristics.Source: Table 1 · p. 9, §5.1 · p. 9, Figure 7 · p. 9

5.2 Text-to-image generation

  • On DrawBench prompts, DTS⋆ outperforms the baselines; on aesthetic scoring, FK/SMC earns the highest rewards but severely over-optimizes, while DTS⋆ balances reward with faithfulness to the prior.The authors hypothesize that soft-value backups aggregate posterior mass, so reward spikes in vanishing-probability prior regions contribute little to value estimates.Source: §5.2 · p. 10

5.3 Text generation

  • As function evaluations increase, DTS⋆ consistently achieves the highest text rewards while maintaining greater output diversity than FK/SMC.The text reward is grammatical-acceptability probability from a BERT-based classifier, and diversity is measured by distinct trigrams.Source: §5.3 · p. 11, Figure 10 · p. 11

6 Discussion

The discussion identifies high-dimensional search, repeated leaves, and sequential control flow as practical constraints, while noting ways pretrained priors and batching help. The related-work discussion positions DTS as connecting entropy-regularized MCTS to consistent posterior sampling with pretrained diffusion kernels.

6 Discussion

  • DTS uses soft-value backups to propagate terminal rewards through a finite-horizon tree, yielding global credit assignment and improved sample quality as compute grows.In high dimensions, an uninformed tree grows exponentially, so a good pretrained model is important for pruning the effective search space.Source: §6 · p. 11
  • When the number of samples substantially exceeds the tree size, high-probability leaves can produce repeated samples; a learned value network for a fixed reward is proposed as future work.The current work focuses on zero-shot inference-time alignment for unseen rewards.Source: §6 · p. 11
  • Tree-based control flow is less parallelizable than particle-based SMC, although batching enables per-step parallelism and repeated draws from a constructed tree require no further model calls.The authors suggest asynchronous rollout actors as a possible scaling direction.Source: §6 · p. 11

A Extended related work

  • DTS applies soft-value backups and Boltzmann selection with a pretrained diffusion kernel, and proves consistency for sampling from the KL-regularized posterior rather than selecting one high-reward action.The authors describe this as bridging entropy-regularized MCTS for control and unbiased posterior sampling for inference-time alignment.Source: A Extended related work · p. 17

B Sequential Monte Carlo for diffusion sampling

Sequential Monte Carlo (SMC) approximates reward-aligned diffusion sampling with particles and estimated values, but value errors and high-variance weights can limit sampling quality. DTS shares a soft Bellman recursion with GFlowNets while using search rather than learning sampler parameters.

B Sequential Monte Carlo for diffusion sampling

C Connection with Generative Flow Networks

  • DTS uses Monte-Carlo log-sum-exp backups to estimate soft values and samples children proportionally to their exponentiated estimated values.Repeated rollouts make its terminal distribution converge to the reward-tilted posterior, as with a perfectly trained GFlowNet.Source: C Connection with Generative Flow… · p. 18
  • DTS uses a fixed pretrained diffusion kernel without parameter updates, suiting new unseen rewards; GFlowNets learn a sampler for a fixed reward.DTS is positioned for large inference budgets, whereas GFlowNets amortize training across repeated queries.Source: C Connection with Generative Flow… · p. 18

D Proofs and derivations

The derivations establish the recursive soft value relation and the transitions of the optimal policy. The value recursion uses a terminal reward condition, while the policy follows from the target and marginal densities.

D Proofs and derivations

  • The derivations provide the recursive soft-value relation and the expression for the optimal policy used in the paper.The paper presents these derivations for completeness.Source: §D.1 · p. 18
  • The soft-value derivation factors the reward-weighted trajectory through the reverse transition, then applies terminal condition V0(x0) = r(x0) to obtain Equation (5).The recursion is analogous to the soft Bellman equation in maximum entropy reinforcement learning.Source: §D.1 · p. 19
  • The optimal-policy transitions are obtained by dividing the target joint density by the marginal density of xt.The derivation identifies the marginal and then uses the ratio to obtain the transitions.Source: §D.1 · p. 19

D.2 Proof of Proposition 1

DTS is asymptotically consistent: as rollouts accumulate, its terminal-state distribution converges to the target. The algorithm alternates tree traversal, expansion, terminal reward evaluation, and value backups, with experiments using scheduled branching.

D.2 Proof of Proposition 1

  • With bounded reward and λ > 0, DTS terminal states converge empirically to the optimal policy π* as tree iterations M approach infinity.Infinite rollouts make value estimates approach the true soft values, so DTS's sampling distribution matches the target.Source: §D.2 · pp. 19–20

E DTS and DTS⋆algorithm

  • Each iteration selects a path, expands a node with a base-model sample, rolls out to a terminal state, evaluates its reward, and backs up soft values and visit counts.DTS selects children probabilistically; DTS⋆ selects by UCT.Source: E DTS and DTS⋆algorithm · p. 21

F.1 Tree structure

  • Experiments branch only every few timesteps, with little observed performance difference from branching at every step at equal function evaluations.The reported progressive-widening settings use α = 0.8 and C = 2 for two-dimensional and image experiments, and α = 0.7 and C = 2 for text.Source: §F.1 · pp. 21–22, Table 2 · p. 21

F.2 Experiment details

The experiments use different pretrained or trained diffusion models and reward functions across two-dimensional data, image generation, text-to-image, and language generation.

F.2 Experiment details

  • The two-dimensional experiments use an MLP denoiser with four 128-dimensional hidden layers, trained for 500 epochs on 10,000 samples.The noise schedule is linear, with βmin = 0.001 and βmax = 0.07.Source: §F.2 · p. 22
  • MNIST generation uses a trained Unet and classifier-based class log-likelihood rewards, while CIFAR-10 uses a pretrained pixel-space Unet and a ResNet-18 classifier reward.Both image denoisers use stochastic DDIM sampling with η = 1.0.Source: §F.2 · p. 22
  • Text-to-image uses Stable Diffusion v1.5 with ImageReward and LAION aesthetic rewards, while language generation uses MDLM with a CoLA classifier reward.The CoLA reward is the log probability that a string is acceptable.Source: §F.2 · p. 23

F.3 Compute

The appendix reports experiment runtimes and describes how the baselines were implemented and configured.

F.3 Compute

  • Class-conditional image experiments used approximately 3 GPU hours per class, totaling approximately 66 GPU hours across 22 tasks.Text-to-image experiments required approximately 123 GPU hours, while the text-generation experiments required approximately 22.5 GPU hours.Source: §F.3 · p. 23

G Details of baselines

  • All baselines were re-implemented in a unified codebase and validated by reproducing quantitative results from their original papers.The authors state that the complete source code, including all baselines, will be released upon publication.Source: G Details of baselines · p. 23
  • DPS adapts inverse-problem guidance by replacing the objective gradient with the reward gradient; its reward is a classifier log likelihood for class-conditional images.The experiments use Tweedie’s formula for the clean-sample estimate, and the diffusion model’s predicted mean is also used.Source: G Details of baselines · pp. 23–24
  • SMC/FK-Steering uses the pretrained diffusion model as its proposal, with fixed-interval resampling and adaptive resampling to increase sample diversity.Sampling experiments use the ‘diff’ potential with λ = 1.0, while search experiments use the ‘max’ potential with λ = 10.0.Source: G Details of baselines · p. 24
  • DAS tempers reward terms with a monotone schedule, while TDS uses a twisted proposal with SMC to sample from the target posterior.A separate value-estimate equation uses differences between reward estimates at successive diffusion states and a terminal reward.Source: G Details of baselines · p. 24

H.1 Class-conditional image experiments

Across three class-conditional settings, DTS generally matches the target distribution more accurately than other methods at most evaluated compute levels. Additional examples illustrate class-conditional samples, text-to-image outputs, and text completions.

H.1 Class-conditional image experiments

  • DPS can produce unstable, unnatural images, whereas SMC-based methods show mode-collapse signs; DTS balances diversity and high rewards while matching the posterior.The authors describe this trend in random samples across the MNIST and CIFAR-10 settings.Source: §H.1 · p. 24, Figure 12 · p. 26, Figure 13 · p. 27
  • Across three class-conditional settings, DTS has the lowest FID and CMMD for most function-evaluation counts.Figure 11 reports FID, CMMD, average log rewards, and average diversity across MNIST digit, MNIST odd/even, and CIFAR-10 class-generation settings.Source: §H.1 · p. 24, Figure 11 · p. 25

H.2 Text-to-image examples

  • Text-to-image examples show samples as inference-time compute increases and cover artistic style, spatial relationships, and object counts.Figure 14 measures compute in diffusion-model function evaluations; Figures 15–17 present randomly selected text-image pairs for the listed prompt concepts.Source: §H.2 · p. 28, Figure 14 · p. 28, Figure 15 · p. 29

H.3 Text completion examples

  • The appendix provides additional randomly selected text completions from MDLM, FK-Steering, and DTS⋆ using a CoLA classifier reward.The examples are shown in Figure 18.Source: §H.3 · p. 30, Figure 18 · p. 30
Loading 2506.20701v1…