Inference-time Scaling of Diffusion Models through Classical Search

1 INTRODUCTION

• Diffusion samples may violate physical laws or miss human intent, while the large generative space can require multiple trials for satisfactory outputs.The paper uses verifier functions to measure sample quality, including reward functions, Q-functions, classifiers, and vision-language models.Source: §1 · p. 1
• Global search traverses denoising states with BFS or DFS, evaluates intermediate states using a verifier, and prioritizes high-quality paths through branching and backtracking.Figure 1 depicts global search locating high-score modes within the base distribution.Source: §1 · pp. 1–2, Figure 1 · p. 2
• Local search uses Langevin MCMC with verifier-gradient and diffusion-score guidance to optimize samples beyond the base model’s capabilities.The paper describes this as a theoretically grounded component of the combined framework.Source: §1 · p. 2
• The unified framework jointly scales local and global search, demonstrating superior performance over prior methods that scale them in isolation.The paper presents this as advancing the Pareto frontier of inference-time scaling across diverse domains.Source: §1 · p. 2

Particle-based scaling methods

• Particle filtering scales inference by increasing particle counts and using resampling, while tree-search methods evaluate intermediate nodes and expand promising candidates (Kim et al., 2025; Singhal et al., 2025a; Li et al., 2024; Guo et al., 2025; Ma et al., 2025).The paper treats both families as special cases of its BFS framework, which applies branching and filtering to parallel denoising particles.Source: §2 · pp. 2–3
• Earlier tree-search approaches typically use fixed-schedule BFS; the paper adds adaptive DFS backtracking and reports superior performance over prior methods (Guo et al., 2025).It characterizes DFS as more efficient and adaptive than BFS.Source: §2 · pp. 2–3

Gradient-based guidance methods

• Classifier guidance requires a trained noise-dependent classifier, while training-free guidance on clean samples approximates noisy conditional probabilities and is described as inheriting first-order approximation bias (Dhariwal & Nichol, 2021; Ye et al., 2024; Chung & Ye, 2022; Yu et al., 2023; He et al., 2023).The paper contrasts these approaches with annealed Langevin MCMC, which it describes as asymptotically exact sampling without additional training.Source: §2 · p. 3

3.1 DIFFUSION PROBABILISTIC MODELS

• A diffusion model defines a forward noising process and learns a reverse denoising process, with its denoising objective equivalent to learning the score of the noisy-data marginal.Generation transforms noise into data using reverse transition kernels, with deterministic DDIM or stochastic DDPM samplers described as options.Source: §3.1 · p. 3

3.2 COMPOSITIONAL AND CONTROLLABLE GENERATION OF DPMS

• Controllable generation seeks samples from a composition of the base distribution and an objective, but exact sampling would require training a time-dependent objective function on base-model samples, which may be impractical.The paper therefore adopts optimization-based methods to approximate the target distribution.Source: §3.2 · p. 3
• Annealed Langevin MCMC constructs a sequence of composed distributions and uses their combined scores to sample from the target without extra training.The method follows the annealing path toward the composed data distribution, with asymptotic convergence as the step count grows and step size shrinks.Source: §3.2 · pp. 3–4

4.1 GLOBAL SEARCH FOR MODE IDENTIFICATION

• Global search represents denoising as a fixed-depth tree and evaluates intermediate particles to allocate computation toward promising paths.This uses intermediate denoised estimates rather than waiting to compare only final samples, as in best-of-N sampling.Source: §4.1 · p. 4

4.1.1 UNIFIED BFS-STYLE LINEAR SEARCH

• The improved BFS design, Increase-Max-SSP, consistently outperforms prior methods in efficiency; the ablations identify SSP resampling as key for performance.The framework represents prior methods as special cases of a unified BFS design space spanning schedules, scoring functions, and resampling strategies.Source: §4.1.1 · p. 5

4.1.2 DFS-STYLE NON-LINEAR SEARCH

• DFS denoises one particle until its verifier score falls below a threshold, then backtracks to explore another region of the sample manifold.The threshold lets users balance output quality and compute; verifier-driven backtracking allocates more exploration to difficult or low-quality trajectories.Source: §4.1.2 · pp. 5–6

4.2 SCALING LOCAL SEARCH VIA LANGEVIN MCMC WITH VERIFIER GRADIENT

• Local search targets a compositional distribution by hill-climbing with Langevin MCMC; in the continuous limit, training-free guidance with recurrence is equivalent to Langevin MCMC on annealed distributions.The recurrence corresponds to Langevin MCMC on the diffusion model distribution, while the guidance term defines an annealing path toward high-reward regions beyond base-model modes.Source: §4.2 · p. 6

5.1 ELUCIDATING THE DESIGN SPACE OF BFS

• The improved BFS (Increase, Max, SSP) consistently outperforms previous methods across compute budgets and models.The authors use it as the baseline in subsequent experiments.Source: §5.1 · p. 7, Table 2 · p. 7

5.2 ADAPTIVE AND EFFICIENT INFERENCE-SCALING WITH DFS

• DFS consistently outperforms Best-of-N and improved BFS across compute budgets, with up to 2× less computational cost.Lower thresholds allocate less compute; higher thresholds scale compute for better quality, and harder prompts naturally consume more compute.Source: §5.2 · pp. 7–8, Figure 3 · p. 7

5.3 JOINT SCALING LOCAL AND GLOBAL SEARCH

• Jointly scaling local and global search establishes a new Pareto frontier, while scaling either separately yields suboptimal performance.Compute is measured as total NFEs across local and global search.Source: §5.3 · p. 8

5.3.1 LONG HORIZON PLANNING

• In PointMaze, scaling local-search steps produces successful plans free of violations, unlike naive samples that often violate the maze layout.Figure 5 contrasts a failed trajectory with a successful sample.Source: §5.3.1 · p. 8, Figure 5 · p. 8
• BoN with additional local-search steps improves the Pareto frontier, while TFG fails to scale with increased compute and can become trapped in local optima.Combining local search with BFS or DFS further improves scaling efficiency.Source: §5.3.1 · pp. 8–9, Figure 4 · p. 8

5.3.2 OFFLINE REINFORCEMENT LEARNING

• TTS performs comparably to training-based baselines on D4RL locomotion, while offline distillation of its generated actions outperforms state-of-the-art diffusion-based online RL methods.The results are reported in Tables 3 and 4.Source: §5.3.2 · p. 9, Table 3 · p. 9, Table 4 · p. 9

5.4 MITIGATING REWARD HACKING WITH DOUBLE VERIFIER

• On conditional ImageNet generation, double-verifier significantly improves performance over a single verifier with 2× less compute and reduces OOD samples, as indicated by higher MSP.The evaluation reports FID, class accuracy, and MSP.Source: §5.4 · pp. 9–10, Table 5 · p. 10

6 LIMITATIONS AND CONCLUSION

• Inference-time scaling jointly uses local and global search, but still requires additional hyperparameters.The authors leave evolutionary search for hyperparameter tuning to future work.Source: §6 · p. 10

B.1 DISCUSSIONS ON CONCURRENT WORKS

• ABCD combines BFS and DFS but uses fixed termination, limiting adaptive compute allocation and wasting compute on easier instances.It evaluates particles after full denoising, which the authors say yields smaller score estimation errors than DFS.Source: §B.1 · p. 15
• The framework combines classical BFS and DFS with annealed Langevin MCMC local search, which the authors say is critical for challenging decision-making domains.They describe the local-search approach as the first to scale diffusion local search using annealed Langevin MCMC.Source: §B.1 · p. 15

B.3 INFERENCE SCALING IN DIFFUSION MODELS WITHOUT DIFFERENTIABLE OPTIMIZATION

• Prior diffusion inference-scaling methods include SMC approaches interpreted as BFS, while the authors claim their search design improves compute efficiency.The cited methods differ in scoring and resampling; Ma et al. (2025) explored local and global search but did not systematically evaluate or improve compute efficiency.Source: §B.3 · p. 15

B.4 GRADIENT GUIDANCE OF DIFFUSION MODELS

• Training-free guidance approximates noisy conditional probabilities using denoised outputs, but its theoretical foundations and scaling behavior with recurrence remain largely unexplored.The authors connect recurrence to Langevin MCMC and interpret guidance as defining an annealing path toward high-reward regions beyond base-model modes.Source: §B.4 · p. 16

C.1 LANGEVIN MCMC AS GRADIENT FLOW IN MEASURE SPACE

• For smooth log-concave targets, prior work shows ULA can converge efficiently; for complex multimodal distributions, the stated guarantee is only convergence to a concave vicinity.The paper frames Langevin SDE as gradient flow of KL divergence in measure space, with the target distribution as the minimum.Source: §C.1 · pp. 16–17

C.2 ANNEALED LANGEVIN MCMC SAMPLING

• Annealed Langevin MCMC takes multiple steps at each noise level, unlike reverse diffusion, which denoises across noise levels.Sampling through smoothed intermediate distributions can improve mixing on multimodal distributions, sharing a principle with simulated annealing.Source: §C.2 · p. 17

C.3 ANNEALED LANGEVIN MCMC WITH RECURRENT TRAINING-FREE GUIDANCE

• The section proves that naive recurrence corresponds to Langevin MCMC and that guidance defines an annealing path, then provides convergence analysis.The analysis is divided into recurrence equivalence, the effect of guidance, and convergence.Source: §C.3 · p. 18

C.3.1 EQUIVALENCE BETWEEN LANGEVIN MCMC AND NAIVE RECURRENCE

• As the denoising step size approaches zero, a recurrence step approximates the unadjusted Langevin algorithm and discretizes the Langevin SDE.The recurrence first uses the learned score function, then adds noise; its formulation depends on the sampler.Source: §C.3.1 · pp. 18–19

C.3.2 ANNEALED LANGEVIN MCMC WITH GUIDANCE

• With training-free guidance, the guidance term acts as the score of annealed verifier functions, defining an annealing path that biases sampling toward high-reward regions.The passage describes the combined guidance terms as the gradient of an annealed function.Source: §C.3.2 · pp. 19–20

C.3.3 CONVERGENCE ANALYSIS

• Theorem 1 bounds the Wasserstein distance between the sample after recurrence steps and the annealed target, under bounded-support, strong-log-concavity, smoothness, and Lipschitz-curvature assumptions.The supplied theorem passage states the conditions and bound, but its bound expression is not readable here.Source: §C.3.3 · p. 20

C.4 RELATIONSHIP BETWEEN LANGEVIN MCMC AND GRADIENT ASCENT

• Gradient ascent optimizes verifier reward, whereas Langevin MCMC also optimizes divergence from the base distribution; recurrence can pull samples back toward the original manifold.The authors say recurrence may be omitted for small guidance strength, but helps avoid OOD samples when multiple gradient updates traverse modes.Source: §C.4 · p. 21

C.5 IMPLEMENTING LOCAL SEARCH WITH TFG HYPER-PARAMETER SPACE

• The implementation uses training-free guidance for local search, with recurrence depth equal to the number of local-search steps.The passage says this equivalence also enables efficient hyperparameter search within the training-free guidance framework.Source: §C.5 · p. 21

D GLOBAL SEARCH OF DENOISING DIFFUSION MODELS

• The global-search section presents BFS and permits evaluation and resampling at a fixed subset of denoising steps.This allows experiment-specific hyperparameters for evaluation timing.Source: D GLOBAL SEARCH OF DENOISING DIFFUSION… · p. 22, §D.1 · p. 22

D.1.1 OVERVIEW OF DESIGN SPACE

• BFS design choices include increasing tempering, which gives estimated scores more influence at lower-noise steps where estimation bias is smaller.The authors also adopt SSP resampling as a variance-reduced design choice; other variance-reduced methods perform similarly and all outperform naive multinomial resampling.Source: §D.1.1 · pp. 22, 24

D.1.2 OVERVIEW OF PRIOR METHODS

• SVDD selects the best sample at each step, TreeG ranks or resamples particles before branching, DAS uses increasing or adaptive tempering with SSP resampling, and FK-steering defaults to max scoring.The FK-steering passage notes that its official implementation uses multinomial resampling, which may lead to suboptimal performance.Source: §D.1.2 · p. 24

D.2 DFS-BASED SEARCH

• DFS stores prior results in a buffer and returns its best sample when the budget is exhausted and no further backtracking is allowed.Its procedure also retrieves the best buffered sample during backtracking.Source: §D.2 · p. 25
• Earlier evaluation steps increase uncertainty but enable backtracking, while larger backtracking depth broadens exploration at greater computational cost.A smaller depth reduces computation and favors local search.Source: §D.2 · p. 25
• When backtracking to the terminal noise step, the method initializes with fresh Gaussian noise to avoid cumulative errors from repeated restarts under nonzero terminal SNR.This behavior is reported for many diffusion schedules.Source: §D.2 · p. 25

• Experiments used Nvidia A100 GPU clusters and over 1000 GPU hours.Source: E EXPERIMENT DETAILS · p. 25

E.1 ABLATION OF BFS DESIGN SPACE

• For ImageReward prompts, the experiments followed FK-steering’s temperature and resampling interval settings and used TreeG with a fixed branch-out size of 2.Source: §E.1 · pp. 25–26

E.2 TEXT-TO-IMAGE COMPOSITIONAL GENERATION WITH DFS

• Text-to-image experiments used SSD-1B with 50 DDIM steps, evaluated DFS and BFS at selected timesteps, and swept BFS temperature, reporting the best performance.DFS and BFS were evaluated at steps {25, 35, 45}, with backtrack depth ΔT = 25; BFS temperatures ranged from 0.5 to 8.Source: §E.2 · p. 26
• Example prompts and images before and after inference scaling are shown for CompBench.Source: §E.2 · p. 26

E.3 LONG HORIZON MAZE PLANNING

• Long-horizon PointMaze tasks were difficult for the diffusion model without extra inference compute, and the evaluation navigated from bottom left to top right in a 20x20-cell maze.The model was trained for 1.2M steps and used 16 denoising steps, where performance was reported to saturate.Source: §E.3 · p. 26
• The maze verifier uses the ground-truth layout to score trajectory-point violations, but a trajectory with no violating points may still fail because segments between points can cross walls.The authors state that the verifier alone cannot generate a successful plan.Source: §E.3 · pp. 26–27
• The passage reports that using time steps {4, 3, 2, 1} helps scale up performance with more compute, in the context of hyperparameter-search results shown immediately afterward.Source: §E.3 · p. 28

E.4 OFFLINE RL

• Offline RL experiments reused Lu et al.’s pretrained diffusion model and Q-function without its time-dependent energy function, and used 15 DDIM steps.Source: §E.4 · p. 28
• Offline RL global search used only 4 particles on Medium-Expert and Medium and 2 on Medium-Replay, so the experiments used Best-of-N rather than BFS or DFS.The paper notes these counts are below the 50 and 32 particles used in cited prior work.Source: §E.4 · p. 28
• The offline RL comparison included IQL and diffusion-based baselines, while reproduced TFG and DAS baselines received approximately twice the hyperparameter space and computational cost.The best baseline performance was reported, and the method was evaluated on seeds different from those used for hyperparameter search.Source: §E.4 · pp. 28–29

E.4.1 OFFLINE POLICY DISTILLATION

• Policy distillation replaced Medium-dataset actions with test-time-search samples, finetuned the model with early stopping, then evaluated the finetuned diffusion model without a verifier.Source: §E.4.1 · p. 29
• One epoch of policy finetuning converged rapidly in Hopper and Walker2d and took about 4 minutes on one Nvidia RTX 4090, versus several hours for DPPO.The reported epoch time includes inference-scaling data generation, and this offline approach requires no online data collection.Source: §E.4.1 · p. 29

• The conditional ImageNet setup used separate classifiers for global and local verification, reporting FID on 256 generated samples and class accuracy from another classifier.The global verifier was applied sparsely, which the authors say added negligible computational cost.Source: §E.5 · p. 30
• Higher maximum softmax probability indicates greater classifier confidence that an image belongs to a dataset class and is less out-of-distribution.Source: §E.5 · p. 30
• Adversarial patches exploited a single verifier, and examples with high predicted target-class probability were attributed to reward hacking.Source: §E.5 · p. 30, Figure 9 · p. 30