Research map / Image generation / Fast diffusion sampling
Diffusion Tree Sampling: Scalable inference‑time alignment of diffusion models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Casts inference-time alignment of diffusion models as tree search with soft-value Monte Carlo backups, yielding DTS (asymptotically exact sampling) and DTS* (greedy search variant) that reuse past rollouts.
Problem
Inference-time alignment of pretrained diffusion models to new reward functions is hard because terminal rewards are only available at the end of denoising, and value estimates (Tweedie one-step or single rollout) are biased and high-variance at high noise; past runs are also not reused.
Approach
Build a denoising tree where nodes are noisy states, selection samples children by exp(λ·soft value), rollouts are added to the tree, and terminal rewards are soft-backed-up (log-sum-exp) to refine value estimates across all runs; the greedy variant selects by max soft value (with UCT) for search.
Key results
- On MNIST and CIFAR-10 class-conditional sampling, DTS matches the FID of the best baseline with up to 10x less compute and achieves the lowest FID and CMMD in all three settings at 10^6 NFEs (e.g., CIFAR-10 FID 0.195 vs 0.241 for DAS).
- In text-to-image and language completion, DTS* matches best-of-N performance with up to 5x less compute.
- Bias-variance analysis on 2D shows the tree-based soft value estimates reduce both bias and variance relative to Tweedie one-step and single-rollout estimates, especially at high noise.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DTS instead of DAS | MNIST class-conditional sampling · MNIST | FID: 0.039 ± 0.017 -> 0.014 ± 0.005 | Helped |
| Use DTS instead of SMC/FK | MNIST class-conditional sampling · MNIST | CMMD: 0.177 ± 0.142 -> 0.068 ± 0.030 | Helped |
| Use DTS instead of SMC/FK | MNIST even/odd sampling · MNIST even/odd | FID: 0.027 ± 0.009 -> 0.007 ± 0.003 | Helped |
| Use DTS instead of DAS | MNIST even/odd sampling · MNIST even/odd | CMMD: 0.079 ± 0.011 -> 0.036 ± 0.029 | Helped |
| Use DTS instead of DAS | CIFAR-10 class-conditional sampling · CIFAR-10 | FID: 0.241 ± 0.037 -> 0.195 ± 0.041 | Helped |
| Use DTS instead of DAS | CIFAR-10 class-conditional sampling · CIFAR-10 | CMMD: 0.822 ± 0.203 -> 0.745 ± 0.201 | Helped |
| Branch only at selected denoising timesteps rather than every timestep (ablation) | DTS · Diffusion alignment experiments | performance | No clear effect |
| Use soft-value backup rather than max-backup (ablation) | DTS for text generation · Text generation | reward | Hurt |
Surprising findings
- SMC/FK achieves the highest aesthetic-score rewards but shows severe over-optimization, while DTS* balances reward and faithfulness.
- Tweedie-based one-step value estimates are essentially random at high noise levels even in simple 2D settings.
Limitations
- Tree-based control flow is sequential and less parallelizable than particle methods like SMC.
- In high dimensions an uninformed tree grows exponentially, so pure tree search relies on a good pretrained prior; repeated samples appear when sample count exceeds tree size.
- Value function is not learned, so no amortization across fixed rewards.
Experimental setup
Datasets: MNIST CIFAR-10 DrawBench ImageNet-1000 animal prompts CoLA 15 text completion prompts
Benchmarks: MNIST CIFAR-10 DrawBench ImageNet animal prompts (aesthetic) CoLA text completion
Metrics: FID CMMD expected log reward diversity (CLIP pairwise cosine) max ImageReward max aesthetic score distinct trigrams MMD bias and variance of value estimates
Baselines: DPS SMC/FK TDS DAS best-of-N SMC-Rollout
Builds on: Monte Carlo Tree Search sequential Monte Carlo DPS TDS DAS UCT
Closest papers by meaning
- Scaling Inference Time Compute for Diffusion Models (CVPR 2025 · similarity 0.91)
- A General Framework for Inference-time Scaling and Steering of Diffusion Models (ICML 2025 · similarity 0.90)
- Test-time Alignment of Diffusion Models without Reward Over-optimization (ICLR 2025 · similarity 0.90)
- Breaking AR’s Sampling Bottleneck: Provable Acceleration via Diffusion Language Models (NeurIPS 2025 · similarity 0.90)
- Accelerating Diffusion Models with Parallel Sampling: Inference at Sub-Linear Time Complexity (NeurIPS 2024 · similarity 0.90)
- Inference-Time Alignment of Diffusion Models with Direct Noise Optimization (ICML 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.