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

• URL: http://arxiv.org/abs/2506.20701v1
• Date: Wed, 25 Jun 2025 17:59:10 GMT
• Title: Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models
• Authors: Vineet Jain, Kusha Sareen, Mohammad Pedramfar, Siamak Ravanbakhsh,
• Abstract summary: Adapting a pretrained diffusion model to new objectives at inference time remains an open problem in generative modeling.<n>We introduce a tree-based approach that samples from the reward-aligned target density by propagating terminal rewards back through the diffusion chain.<n>By reusing information from previous generations, we get an anytime algorithm that turns additional compute into steadily better samples.
• Score: 13.312007032203857
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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, 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.

Related papers

• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [62.640128548633946]
  We introduce a novel inference-time scaling approach based on particle Gibbs sampling for discrete diffusion models.<n>Our method consistently outperforms prior inference-time strategies on reward-guided text generation tasks.
  arXiv  Detail & Related papers  (2025-07-11T08:00:47Z)
• Faster Diffusion Models via Higher-Order Approximation [28.824924809206255]
  We propose a principled, training-free sampling algorithm that requires only the order of d1+2/K varepsilon-1/K $$ score function evaluations.<n>Our theory is robust vis-a-vis inexact score estimation, degrading gracefully as the score estimation error increases.
  arXiv  Detail & Related papers  (2025-06-30T16:49:03Z)
• Test-Time Scaling of Diffusion Models via Noise Trajectory Search [7.243632426715941]
  We introduce an $epsilon$-greedy search algorithm that globally explores at extreme timesteps and locally exploits during the intermediate steps where de-mixing occurs.<n>Experiments on EDM and Stable Diffusion reveal state-of-the-art scores for class-conditioned/text-to-image generation.
  arXiv  Detail & Related papers  (2025-05-24T19:13:29Z)
• Distributional Diffusion Models with Scoring Rules [83.38210785728994]
  Diffusion models generate high-quality synthetic data.<n> generating high-quality outputs requires many discretization steps.<n>We propose to accomplish sample generation by learning the posterior em distribution of clean data samples.
  arXiv  Detail & Related papers  (2025-02-04T16:59:03Z)
• O(d/T) Convergence Theory for Diffusion Probabilistic Models under Minimal Assumptions [6.76974373198208]
  We establish a fast convergence theory for the denoising diffusion probabilistic model (DDPM) under minimal assumptions.<n>We show that the convergence rate improves to $O(k/T)$, where $k$ is the intrinsic dimension of the target data distribution.<n>This highlights the ability of DDPM to automatically adapt to unknown low-dimensional structures.
  arXiv  Detail & Related papers  (2024-09-27T17:59:10Z)
• Active Diffusion Subsampling [15.028061496012924]
  In maximum entropy sampling, one selects measurement locations that are expected to have the highest entropy, so as to minimize uncertainty about $x$.<n>Recently, diffusion models have been shown to produce high-quality posterior samples of high-dimensional signals using guided diffusion.<n>We propose Active Diffusion Subsampling (ADS), a method for designing intelligent subsampling masks using guided diffusion.
  arXiv  Detail & Related papers  (2024-06-20T15:05:06Z)
• Rejection via Learning Density Ratios [50.91522897152437]
  Classification with rejection emerges as a learning paradigm which allows models to abstain from making predictions.<n>We propose a different distributional perspective, where we seek to find an idealized data distribution which maximizes a pretrained model's performance.<n>Our framework is tested empirically over clean and noisy datasets.
  arXiv  Detail & Related papers  (2024-05-29T01:32:17Z)
• Projection Regret: Reducing Background Bias for Novelty Detection via Diffusion Models [72.07462371883501]
  We propose emphProjection Regret (PR), an efficient novelty detection method that mitigates the bias of non-semantic information. PR computes the perceptual distance between the test image and its diffusion-based projection to detect abnormality. Extensive experiments demonstrate that PR outperforms the prior art of generative-model-based novelty detection methods by a significant margin.
  arXiv  Detail & Related papers  (2023-12-05T09:44:47Z)
• Towards Faster Non-Asymptotic Convergence for Diffusion-Based Generative Models [49.81937966106691]
  We develop a suite of non-asymptotic theory towards understanding the data generation process of diffusion models. In contrast to prior works, our theory is developed based on an elementary yet versatile non-asymptotic approach.
  arXiv  Detail & Related papers  (2023-06-15T16:30:08Z)
• Convergence for score-based generative modeling with polynomial complexity [9.953088581242845]
  We prove the first convergence guarantees for the core mechanic behind Score-based generative modeling. Compared to previous works, we do not incur error that grows exponentially in time or that suffers from a curse of dimensionality. We show that a predictor-corrector gives better convergence than using either portion alone.
  arXiv  Detail & Related papers  (2022-06-13T14:57:35Z)
• Local policy search with Bayesian optimization [73.0364959221845]
  Reinforcement learning aims to find an optimal policy by interaction with an environment. Policy gradients for local search are often obtained from random perturbations. We develop an algorithm utilizing a probabilistic model of the objective function and its gradient.
  arXiv  Detail & Related papers  (2021-06-22T16:07:02Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.