Test-Time Scaling of Diffusion Models via Noise Trajectory Search

• URL: http://arxiv.org/abs/2506.03164v1
• Date: Sat, 24 May 2025 19:13:29 GMT
• Title: Test-Time Scaling of Diffusion Models via Noise Trajectory Search
• Authors: Vignav Ramesh, Morteza Mardani,
• Abstract summary: 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.
• Score: 7.243632426715941
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: The iterative and stochastic nature of diffusion models enables test-time scaling, whereby spending additional compute during denoising generates higher-fidelity samples. Increasing the number of denoising steps is the primary scaling axis, but this yields quickly diminishing returns. Instead optimizing the noise trajectory--the sequence of injected noise vectors--is promising, as the specific noise realizations critically affect sample quality; but this is challenging due to a high-dimensional search space, complex noise-outcome interactions, and costly trajectory evaluations. We address this by first casting diffusion as a Markov Decision Process (MDP) with a terminal reward, showing tree-search methods such as Monte Carlo tree search (MCTS) to be meaningful but impractical. To balance performance and efficiency, we then resort to a relaxation of MDP, where we view denoising as a sequence of independent contextual bandits. This allows us to introduce an $\epsilon$-greedy search algorithm that globally explores at extreme timesteps and locally exploits during the intermediate steps where de-mixing occurs. Experiments on EDM and Stable Diffusion reveal state-of-the-art scores for class-conditioned/text-to-image generation, exceeding baselines by up to $164\%$ and matching/exceeding MCTS performance. To our knowledge, this is the first practical method for test-time noise trajectory optimization of arbitrary (non-differentiable) rewards.

Related papers

• Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models [13.312007032203857]
  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.
  arXiv  Detail & Related papers  (2025-06-25T17:59:10Z)
• Noise Conditional Variational Score Distillation [60.38982038894823]
  Noise Conditional Variational Score Distillation (NCVSD) is a novel method for distilling pretrained diffusion models into generative denoisers.<n>By integrating this insight into the Variational Score Distillation framework, we enable scalable learning of generative denoisers.
  arXiv  Detail & Related papers  (2025-06-11T06:01:39Z)
• Sampling Binary Data by Denoising through Score Functions [2.9465623430708905]
  Tweedie-Miyasawa formula (TMF) is key ingredient in score-based generative models.<n>TMF ties these together via the score function of noisy data.<n>We adopt Bernoulli noise, instead of Gaussian noise, as a smoothing device.
  arXiv  Detail & Related papers  (2025-02-01T20:59:02Z)
• Robust Representation Consistency Model via Contrastive Denoising [83.47584074390842]
  randomized smoothing provides theoretical guarantees for certifying robustness against adversarial perturbations.<n> diffusion models have been successfully employed for randomized smoothing to purify noise-perturbed samples.<n>We reformulate the generative modeling task along the diffusion trajectories in pixel space as a discriminative task in the latent space.
  arXiv  Detail & Related papers  (2025-01-22T18:52:06Z)
• Score-based Generative Models with Adaptive Momentum [40.84399531998246]
  We propose an adaptive momentum sampling method to accelerate the transforming process. We show that our method can produce more faithful images/graphs in small sampling steps with 2 to 5 times speed up.
  arXiv  Detail & Related papers  (2024-05-22T15:20:27Z)
• Noisy Pair Corrector for Dense Retrieval [59.312376423104055]
  We propose a novel approach called Noisy Pair Corrector (NPC) NPC consists of a detection module and a correction module. We conduct experiments on text-retrieval benchmarks Natural Question and TriviaQA, code-search benchmarks StaQC and SO-DS.
  arXiv  Detail & Related papers  (2023-11-07T08:27:14Z)
• DiffSED: Sound Event Detection with Denoising Diffusion [70.18051526555512]
  We reformulate the SED problem by taking a generative learning perspective. Specifically, we aim to generate sound temporal boundaries from noisy proposals in a denoising diffusion process. During training, our model learns to reverse the noising process by converting noisy latent queries to the groundtruth versions.
  arXiv  Detail & Related papers  (2023-08-14T17:29:41Z)
• Diffusion Model Based Posterior Sampling for Noisy Linear Inverse Problems [14.809545109705256]
  This paper presents a fast and effective solution by proposing a simple closed-form approximation to the likelihood score. For both diffusion and flow-based models, extensive experiments are conducted on various noisy linear inverse problems. Our method demonstrates highly competitive or even better reconstruction performances while being significantly faster than all the baseline methods.
  arXiv  Detail & Related papers  (2022-11-20T01:09:49Z)
• Denoising Distantly Supervised Named Entity Recognition via a Hypergeometric Probabilistic Model [26.76830553508229]
  Hypergeometric Learning (HGL) is a denoising algorithm for distantly supervised named entity recognition. HGL takes both noise distribution and instance-level confidence into consideration. Experiments show that HGL can effectively denoise the weakly-labeled data retrieved from distant supervision.
  arXiv  Detail & Related papers  (2021-06-17T04:01:25Z)
• Learning based signal detection for MIMO systems with unknown noise statistics [84.02122699723536]
  This paper aims to devise a generalized maximum likelihood (ML) estimator to robustly detect signals with unknown noise statistics. In practice, there is little or even no statistical knowledge on the system noise, which in many cases is non-Gaussian, impulsive and not analyzable. Our framework is driven by an unsupervised learning approach, where only the noise samples are required.
  arXiv  Detail & Related papers  (2021-01-21T04:48:15Z)
• Simultaneous Denoising and Dereverberation Using Deep Embedding Features [64.58693911070228]
  We propose a joint training method for simultaneous speech denoising and dereverberation using deep embedding features. At the denoising stage, the DC network is leveraged to extract noise-free deep embedding features. At the dereverberation stage, instead of using the unsupervised K-means clustering algorithm, another neural network is utilized to estimate the anechoic speech.
  arXiv  Detail & Related papers  (2020-04-06T06:34:01Z)

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.