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

Query Lower Bounds for Diffusion Sampling
Apr 12, 2026Zhiyang Xun, Eric PriceQuery Complexity Diffusion Sampling
UT Austin
Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps
Feb 13, 2026Sangwoo Jo, Sungjoon ChoiDiffusion Model Sampling Diffusion Sampling
Department of Artificial Intelligence, Korea University, Seoul, South Korea.
Noise Schedule Design for Diffusion Models: An Optimal Control Perspective
May 21, 2026Seo Taek Kong, Weina Wang, R. SrikantDiffusion Sampling Diffusion-Based Image Generation
ECE & CSL, University of Illinois Urbana-Champaign · Computer Science Department, Carnegie Mellon University · ECE, CSL & NCSA, University of Illinois Urbana-Champaign

Apr 12, 2026J/K move · Enter open · S save

Query Lower Bounds for Diffusion Sampling

Zhiyang Xun, Eric Price

UT Austin

Query Complexity Diffusion Sampling Information-Theoretic Lower Bounds Diffusion Models

Open paper AskPDF arXiv

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.

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