Research map / Image generation / Fast diffusion sampling
Test-Time Scaling of Diffusion Models via Noise Trajectory Search
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.
Proposes epsilon-greedy noise trajectory search at test time, relaxing a diffusion MDP into per-step contextual bandits, to optimize non-differentiable rewards for EDM and Stable Diffusion sampling.
Problem
Test-time scaling of diffusion models by optimizing the injected noise trajectory of SDE-based samplers under arbitrary non-differentiable rewards, where the search space is high-dimensional and evaluation is costly.
Approach
Cast denoising as a terminal-reward MDP, then relax it to independent contextual bandits per timestep and run an epsilon-contaminated zero-order search that mixes global Gaussian draws with local neighborhood hill-climbing, using Tweedie-estimated clean samples for scoring.
Key results
- On EDM with 18 steps, epsilon-greedy (N=4,K=20) reaches 0.9813 brightness, 0.7208 compressibility, 0.9885 classifier vs. 0.3778 naive for classifier reward, matching or exceeding MCTS (0.9666) at 1440 vs. 3888 NFEs.
- On Stable Diffusion v1.5, epsilon-greedy improves over vanilla sampling by 48%/26%/20% for brightness/compressibility/CLIP rewards and exceeds MCTS by up to 18%.
- The paper claims improvements of up to 164% over baselines in sample quality.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use ε-greedy noise trajectory search instead of zero-order search | EDM · ImageNet 64×64 | Brightness: 0.6083 -> 0.9813 | Helped |
| use ε-greedy noise trajectory search instead of zero-order search | EDM · ImageNet 64×64 | Compressibility: 0.3751 -> 0.7208 | Helped |
| use ε-greedy noise trajectory search instead of zero-order search | EDM · ImageNet 64×64 | Classifier probability: 0.6261 -> 0.9885 | Helped |
| use ε-greedy noise trajectory search instead of zero-order search | Stable Diffusion v1.5 · Stable Diffusion prompts | Brightness: 0.5764 -> 0.7110 | Helped |
| use ε-greedy noise trajectory search instead of zero-order search | Stable Diffusion v1.5 · Stable Diffusion prompts | Compressibility: 0.7281 -> 0.8052 | Helped |
| use ε-greedy noise trajectory search instead of zero-order search | Stable Diffusion v1.5 · Stable Diffusion prompts | CLIP: 0.2689 -> 0.2973 | Helped |
| replace the Tweedie-estimate classifier reward with a learned classifier on noisy latents (ablation) | ε-greedy · ImageNet 64×64 | Classifier probability: 0.9885 -> 0.9998 | Helped |
| replace the Tweedie estimate with a one-step distilled model to predict the clean sample (ablation) | ε-greedy · Stable Diffusion prompts | CLIP: 0.2973 -> 0.2820 | Hurt |
| use ε-greedy noise trajectory search instead of MCTS | EDM · ImageNet 64×64 | Brightness: 0.7575 -> 0.9813 | Helped |
| use ε-greedy noise trajectory search instead of MCTS | EDM · ImageNet 64×64 | Compressibility: 0.5395 -> 0.7208 | Helped |
| use ε-greedy noise trajectory search instead of MCTS | Stable Diffusion v1.5 · Stable Diffusion prompts | Brightness: 0.6077 -> 0.7110 | Helped |
| use ε-greedy noise trajectory search instead of MCTS | Stable Diffusion v1.5 · Stable Diffusion prompts | Compressibility: 0.7322 -> 0.8052 | Helped |
Surprising findings
- Vanilla zero-order search lags behind MCTS, while epsilon-greedy with the same budget far exceeds it.
- Reward sensitivity to the noise state peaks at intermediate timesteps, supporting local exploitation there and global exploration at extreme timesteps.
Limitations
- Relaxation to independent contextual bandits ignores inter-step dependencies, so the search is an approximation of the full MDP.
- Non-monotonic scaling with N and K because greedy selection can pick locally good but globally worse noise.
- Rewards are computed on Tweedie estimates rather than true clean samples, introducing approximation bias.
Experimental setup
Datasets: ImageNet 64x64
Benchmarks: ImageNet 64x64 class-conditional generation Stable Diffusion text-to-image prompts compositional prompts counting prompts
Metrics: brightness reward compressibility reward classifier probability CLIP score counting reward ImageReward
Baselines: Naive sampling Best-of-N sampling Beam search MCTS Zero-order search
Builds on: Monte Carlo tree search zero-order search EDM DDIM epsilon-greedy bandit
Closest papers by meaning
- Scaling Inference Time Compute for Diffusion Models (CVPR 2025 · similarity 0.91)
- Test-time Scaling of Diffusions with Flow Maps (NeurIPS 2026 · similarity 0.90)
- O(d/T) Convergence Theory for Diffusion Probabilistic Models under Minimal Assumptions (ICLR 2025 · similarity 0.90)
- Interpreting and Improving Diffusion Models from an Optimization Perspective (ICML 2024 · similarity 0.90)
- Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models (NeurIPS 2025 · similarity 0.90)
- Score-Optimal Diffusion Schedules (NeurIPS 2024 · similarity 0.90)
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.