Leveraging Inference-Time Compute for Diffusion Models via Global Scheduling of Denoising Trajectories

Yuan Cao; Yifu Tang; Hangqi Li; Zeyu Zheng

cs.LGSep 19, 2026

Leveraging Inference-Time Compute for Diffusion Models via Global Scheduling of Denoising Trajectories

Authors: Yuan Cao, Yifu Tang, Hangqi Li, Zeyu Zheng

Organizations: School of Mathematical Sciences, Peking University · Department of Industrial Engineering and Operations Research & BAIR Lab, University of California, Berkeley

Abstract

Diffusion models generate a sample by traversing a denoising trajectory, a sequence of stochastic noise-reduction steps that transforms pure noise into a draw from a target distribution. At deployment time, additional computation can improve sample quality without retraining: at each step, the sampler draws several candidate noise samples, scores the resulting predictions with a quality criterion called the verifier, and retains the best candidate at the cost of one network evaluation per candidate. This raises a resource allocation question: given a fixed budget of function evaluations, how should search effort be distributed across the steps of the denoising trajectory? We formulate this as a computational budget allocation problem. First, we show that, to leading order in the step size, the expected gain from evaluating $K$ candidates at a step factorizes into an endogenous, step-specific sensitivity parameter times a universal sample-size factor equal to the expected best of $K$ standard-normal draws. Second, for a fixed sensitivity profile, the optimal allocation solves a separable concave integer program with water-filling structure; at fixed total sensitivity, its advantage over uniform allocation increases with sensitivity dispersion in the majorization order. Third, we prove that when sensitivities vary across instances, no adaptive policy can avoid worst-case regret that grows linearly in the trajectory length, which motivates a design that anchors the allocation offline and adapts online only to recover instance-specific slack. We extend the analysis from independent random search to a broader family of local search operators, and instantiate it as an implementable algorithm. Experiments on three families of diffusion samplers show that the proposed allocation attains the quality of the uniform benchmark with 20 to 50 percent fewer function evaluations.

Figures & tables

Figure 1: The two-level decomposition of noise trajectory search. The global scheduler

\mathcal{G}

(top) receives the budget

B

and distributes per-step budgets

\{K_{t}\}

across the trajectory. At each step, the local search operator

\mathcal{L}_{t}

(bottom) evaluates

K_{t}

candidate noise realizations and retains the best. The bars depict a non-uniform allocation, in which steps where replacing the noise candidate yields larger quality gains receive larger budgets.

Operator $\mathcal{L}$	Gain sequence $\phi_{K}^{\mathcal{L}}$	Growth rate	Status
Random search	$a_{K}\sim\sqrt{2\log K}$	$O(\sqrt{\log K})$	Exact
$\epsilon$ -greedy	$\geq\mathbb{E}[a_{N+1}]$ , $N\sim\mathrm{Bin}(K{-}1,\epsilon)$	$\Omega(\sqrt{\log K})$	Lower bound
Local perturbation	$c_{d}(\lambda)\,(K-1)$	$O(K)$	Exact in the idealization

Table 1: Operator-specific gain sequences in the projected-score idealization. All three operators satisfy Assumption 2 ; derivations and verifications appear in Appendix B .

Figure 2: The online stopping decision in the

(\sigma_{t},\Delta a_{K})

plane. The horizontal threshold detects small recent gains; the vertical threshold detects low sensitivity through small empirical variance. Stopping is triggered only when both criteria fire (shaded region). The contour

\sigma_{t}\cdot\Delta a_{K}=\lambda^{*}

depicts the common marginal threshold from Theorem 2 ; the rectangular rule operationalizes the corresponding two-threshold decision in GAINS .

		Stable Diffusion (budget $B$ )
Verifier	Policy	400	500	800
Brightness	Uniform	$0.7025\pm 0.047$	$0.7094\pm 0.063$	$0.7511\pm 0.062$
Brightness	GAINS	$\mathbf{0.7248\pm 0.056}$	$\mathbf{0.7346\pm 0.060}$	$\mathbf{0.7832\pm 0.046}$
Compressibility	Uniform	$0.8833\pm 0.017$	$0.8825\pm 0.017$	$0.8892\pm 0.020$
Compressibility	GAINS	$\mathbf{0.8946\pm 0.016}$	$\mathbf{0.9004\pm 0.015}$	$\mathbf{0.9008\pm 0.012}$
		EDM (budget $B$ )

Table 2: Verifier scores under matched local-search budgets. Each entry is the mean over 20 replications, plus or minus one standard deviation. Boldface marks the better policy at each budget.

Verifier	Uniform	Offline only	GAINS (offline + online)
Brightness	$0.7025\pm 0.047$	$0.7158\pm 0.065$	$\mathbf{0.7248\pm 0.056}$
Compressibility	$0.8833\pm 0.017$	$0.8873\pm 0.018$	$\mathbf{0.8946\pm 0.016}$

Table 3: Decomposition of the two stages on Stable Diffusion at fixed budget

B=400

. Offline only applies the profiled allocation without online control.

Policy	Brightness	Compressibility
Uniform	$0.6949\pm 0.076$	$0.7938\pm 0.089$
GAINS	$\mathbf{0.7213\pm 0.075}$	$\mathbf{0.8076\pm 0.096}$

Table 4: Generalization to a broader prompt distribution: 20 prompts, 10 replications each, on Stable Diffusion.

	Uniform		GAINS
Verifier	Zero-order	Random	Zero-order	Random
Brightness	$0.4920\pm 0.072$	$0.7382\pm 0.039$	$\mathbf{0.5080\pm 0.075}$	$\mathbf{0.7457\pm 0.045}$
Compressibility	$0.5727\pm 0.066$	$0.8417\pm 0.027$	$\mathbf{0.5832\pm 0.061}$	$\mathbf{0.8454\pm 0.025}$

Table 5: Global scheduling combined with two qualitatively different local operators on Stable Diffusion. GAINS improves upon the uniform benchmark under both operators.

		Budget $B$
Verifier	Policy	80	100	150
Brightness	Uniform	$0.5492\pm 0.003$	$0.5499\pm 0.003$	$0.5515\pm 0.004$
Brightness	GAINS	$\mathbf{0.5507\pm 0.003}$	$\mathbf{0.5519\pm 0.003}$	$\mathbf{0.5526\pm 0.003}$
Compressibility	Uniform	$0.8135\pm 0.005$	$0.8156\pm 0.004$	$0.8157\pm 0.006$
Compressibility	GAINS	$\mathbf{0.8162\pm 0.003}$	$\mathbf{0.8170\pm 0.006}$	$\mathbf{0.8193\pm 0.005}$

Table 6: Flow-based sampler with stochasticity injected at selected steps. Margins are smaller than in Table 2 , consistent with the narrower sensitivity dispersion among the eligible stochastic steps.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Procedure	Actions used	Backtracks?	Schedule
Plain sampling	Fwd	No	$K_{t}=1$ for all $t$
Best-of- $N$ ( Ma et al. 2025 )	Fwd	No	$N$ full trajectories
$\epsilon$ -greedy ( Ramesh and Mardani 2025 )	EpsG, Fwd	No	Uniform $K_{t}=K$
Direct noise optimization ( Tang et al. 2024 )	Zero, Fwd	No	Uniform $K_{t}=K$
Rollover budget forcing ( Kim et al. 2025 )	Rand, Fwd	No	Online rollover
T-SCEND ( Zhang et al. 2025b )	Rand, Fwd	No	Fixed two-phase

Appendix

Table 7: Existing inference-time search procedures as policies in the noise-trajectory-search MDP. The second column lists the action types each procedure invokes (Fwd equals forward move; Rand, EpsG, Zero, and Lang denote random,

\epsilon

-greedy, zero-order, and Langevin search; Any denotes any operator satisfying Assumption 2 ). The third column indicates whether the policy ever revisits an earlier step. The fourth column describes how the per-step search budget is set.

Inference-Time Compute Allocation

Inference-time compute allocation decides where, when, and how much computation models spend during prediction to improve quality under resource limits.

See all 82 papers

Apr 12, 2026

Query Lower Bounds for Diffusion Sampling

Diffusion models generate samples by iteratively querying learned score estimates. A rapidly growing literature focuses on accelerating sampling by minimizing the number of score evaluations, yet the information-theoretic limits of such acceleration remain unclear. In this work, we establish the first score query lower bounds for diffusion sampling. We prove that for

d

-dimensional distributions, given access to score estimates with polynomial accuracy

\varepsilon=d^{-O(1)}

(in any

L^p

sense), any sampling algorithm requires

\widetildeΩ(\sqrt{d})

adaptive score queries. In particular, our proof shows that, within any polynomial total-query budget, successful sampling requires searching over

\widetildeΩ(\sqrt{d})

distinct noise levels, providing a formal explanation for why multiscale noise schedules are necessary in practice.

Zhiyang Xun, Eric Price

UT Austin

Query Complexity Diffusion Sampling Information-Theoretic Lower Bounds Diffusion Models

Ask PDF

Feb 13, 2026

Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps

Diffusion-based generative models have achieved remarkable performance across various domains, yet their practical deployment is often limited by high sampling costs. While prior work focuses on training objectives or individual solvers, the broader sampling design problem, specifically solver selection and scheduling, remains largely governed by static heuristics. We propose SDM, a principled, training-free sampling framework that adapts both the numerical solver and the timestep schedule to the intrinsic properties of the diffusion trajectory. By analyzing the PF-ODE dynamics, we show that velocity variation is small in high-noise stages and increases near the data manifold, identifying intervals where solver order is most consequential. In parallel, we introduce an offline-calibrated adaptive scheduling method that explicitly controls the local Wasserstein discretization error and projects the calibrated trajectory to a prescribed NFE budget. We further extend the formulation to a mixed-transition Wasserstein error bound, providing a unified error-propagation view of adaptive scheduling and solver selection within the overall SDM framework. Across standard benchmarks, with extensions to modern ODE samplers, high-resolution synthesis, and text-to-image generation, SDM achieves improved sample quality compared to baseline methods, attaining an FID of 1.93 on CIFAR-10, 2.41 on FFHQ, and 1.98 on AFHQv2, with a reduced number of function evaluations compared to existing samplers. Our code is available at https://github.com/aiimaginglab/sdm.

Sangwoo Jo, Sungjoon Choi

Department of Artificial Intelligence, Korea University, Seoul, South Korea.

Diffusion Model Sampling Diffusion Sampling Adaptive Sampling Diffusion Timestep Scheduling Diffusion Model Acceleration

Ask PDF

May 21, 2026

Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

We develop a principled framework for analyzing and designing noise schedules in diffusion models. We show that one can recast this design problem as an optimal control problem, whose state is the Fisher information of the diffusion process, which evolves according to an ODE, and whose control input is the noise schedule. The objective of the optimal control problem is a functional involving the Fisher information, which is shown to be an upper bound on the Kullback-Leibler sampling error. By solving this optimal control problem, we obtain sufficient conditions on noise schedules under which

\tilde{\mathcal{O}}(d/n)

sampling error is achievable, where

d

is the data dimension and

n

is the number of discretization steps. While existing theoretical works also prove that

\tilde{\mathcal{O}}(d/n)

sampling error bounds are achievable, these results hold for specific noise schedules that differ from those used in practice. Under additional assumptions on the data distribution, we derive a family of noise schedules that recovers the functional forms of schedules widely used in practice. We show that these schedules achieve the

\tilde{\mathcal{O}}(d/n)

error bound when their parameters are chosen according to our analysis. Our theory offers a principled approach to noise schedule design, yielding explicit functional forms and guidance for selecting their parameters.

Seo Taek Kong, Weina Wang, R. Srikant

ECE & CSL, University of Illinois Urbana-Champaign · Computer Science Department, Carnegie Mellon University · ECE, CSL & NCSA, University of Illinois Urbana-Champaign

Diffusion Sampling Diffusion-Based Image Generation Diffusion Timestep Scheduling Optimal Control Diffusion Models

Ask PDF

Abstract

Inference-Time Compute Allocation

Explore similar work

Query Lower Bounds for Diffusion Sampling

Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps

Noise Schedule Design for Diffusion Models: An Optimal Control Perspective