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

arXiv: https://arxiv.org/abs/2506.20701 · HTML: https://mlsanity.com/paper/2506.20701/

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 77 of its 101 references that are in the map.
- First posted on arXiv 25 Jun 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 77)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [G-Eval: NLG Evaluation using GPT-4 with Better Human Alignment](https://mlsanity.com/paper/2303.16634/) (2023)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Neural Network Acceptability Judgments](https://mlsanity.com/paper/1805.12471/) (2018)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion](https://mlsanity.com/paper/2407.01392/) (2024)
- [Equivariant Diffusion for Molecule Generation in 3D](https://mlsanity.com/paper/2203.17003/) (2022)
- [Reinforcement Learning with Deep Energy-Based Policies](https://mlsanity.com/paper/1702.08165/) (2017)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://mlsanity.com/paper/2203.02923/) (2022)

## Built on by (most cited first, 30 of 34)

- [Inference-time Physics Alignment of Video Generative Models with Latent World Models](https://mlsanity.com/paper/2601.10553/) (2026)
- [Meta Flow Maps enable scalable reward alignment](https://mlsanity.com/paper/2601.14430/) (2026)
- [TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion](https://mlsanity.com/paper/2509.25171/) (2025)
- [HardFlow: Hard-Constrained Sampling for Flow-Matching Models via Trajectory Optimization](https://mlsanity.com/paper/2511.08425/) (2025)
- [Training-Free Adaptation of Diffusion Models via Doob's h-Transform](https://mlsanity.com/paper/2602.16198/) (2026)
- [From Scale to Speed: Adaptive Test-Time Scaling for Image Editing](https://mlsanity.com/paper/2603.00141/) (2026)
- [Inference-Time Scaling of Diffusion Language Models via Trajectory Refinement](https://mlsanity.com/paper/2507.08390/) (2025)
- [Controllable Graph Generation with Diffusion Models via Inference-Time Tree Search Guidance](https://mlsanity.com/paper/2510.10402/) (2025)
- [Proprio: Latent Self-Scoring and Inference-Time Refinement for Physically Plausible Video Generation](https://mlsanity.com/paper/2605.28230/) (2026)
- [Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function](https://mlsanity.com/paper/2512.04559/) (2025)
- [Diffusion Alignment as Variational Expectation-Maximization](https://mlsanity.com/paper/2510.00502/) (2025)
- [Entropy Aware Reward Guidance for Diffusion Language Model Alignment](https://mlsanity.com/paper/2602.05000/) (2026)
- [Diffusion Large Language Models for Black-Box Optimization](https://mlsanity.com/paper/2601.14446/) (2026)
- [UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching](https://mlsanity.com/paper/2602.04344/) (2026)
- [VASR: Variance-Aware Systematic Resampling for Reward-Guided Diffusion](https://mlsanity.com/paper/2604.06779/) (2026)
- [Rethinking Test Time Scaling for Flow-Matching Generative Models](https://mlsanity.com/paper/2511.22242/) (2025)
- [Early Failure Detection and Intervention in Video Diffusion Models](https://mlsanity.com/paper/2603.14320/) (2026)
- [Offline Preference Optimization for Rectified Flow with Noise-Tracked Pairs](https://mlsanity.com/paper/2605.09433/) (2026)
- [Few-step Cofolding with All-Atom Flow Maps](https://mlsanity.com/paper/2606.08375/) (2026)
- [BFMT: Enhancing Search Capabilities of Tree Sampler via Bootstrap Flow-Map Tree](https://mlsanity.com/paper/2607.02915/) (2026)
- [Accelerating Masked Diffusion Large Language Models: A Survey of Efficient Inference Techniques](https://mlsanity.com/paper/2607.12829/) (2026)
- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [DiffusionBrowser: Interactive Diffusion Previews via Multi-Branch Decoders](https://mlsanity.com/paper/2512.13690/) (2025)
- [Spectral Evolution Search: Efficient Inference-Time Scaling for Reward-Aligned Image Generation](https://mlsanity.com/paper/2602.03208/) (2026)
- [Trust-Region Noise Search for Black-Box Alignment of Diffusion and Flow Models](https://mlsanity.com/paper/2603.14504/) (2026)
- [Multimodal Diffusion to Mutually Enhance Polarized Light and Low Resolution EBSD Data](https://mlsanity.com/paper/2604.22212/) (2026)
- [Sequentially-Controlled Interactive Multi-Particle Flow-Maps for Online Feedback-Driven Search](https://mlsanity.com/paper/2607.01144/) (2026)
- [LATS: Levy Adaptive Tree Sampling for Feedback-Driven Diverse Target Discovery](https://mlsanity.com/paper/2609.06761/) (2026)
- [RIDE: Reference-Anchored Inference-Time Diffusion Editing for Scaffold Hopping](https://mlsanity.com/paper/2609.35623/) (2026)
- [PreviewDiff: Multimodal Critic-Guided Search over Diffusion Latents](https://mlsanity.com/paper/2609.36199/) (2026)

## Abstract (arXiv)

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× 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× 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.