Budget-Constrained Step-Level Diffusion Caching

Abstract: Step-level caching accelerates diffusion models by exploiting temporal redundancy across denoising steps. Existing methods make per-step cache decisions using threshold-based heuristics, without directly optimizing for final output quality. As a result, their inference latency varies across inputs and is difficult to control at deployment. In this work, we propose BudCache, which inverts this formulation: rather than letting per-step error thresholds dictate the runtime cost, we fix the compute budget in advance and search for the cache policy that best preserves the final output. To tackle the combinatorial complexity of step selection, we combine Simulated Annealing with deterministic Hill Climbing. This offline search identifies high-quality cache policies within minutes and introduces no online search or thresholding overhead during inference. When the compute budget is very tight, we further introduce cache-aware schedule alignment, which adapts the time discretization to the selected cache policy to reduce cache-induced trajectory mismatch. Experiments on FLUX.1-dev and Wan2.1 show that BudCache achieves better generation quality than heuristic caching baselines under the same inference budgets. Code is available at https://github.com/Westlake-AGI-Lab/BudCache

• The paper presents BudCache, a novel budget-constrained caching strategy that deterministically optimizes step-level reuse in diffusion models to meet strict NFE budgets.
• It employs a hybrid two-phase optimization using simulated annealing and hill climbing to search a combinatorial space of caching policies, ensuring enhanced output quality.
• Experiments across image and video synthesis tasks show that BudCache achieves significant speedups and improved fidelity compared to traditional threshold-based caching methods.

Budget-Constrained Step-Level Diffusion Caching: A Technical Analysis

Introduction and Motivation

The computational burden of iterative generative sampling in modern diffusion models and flow-matching frameworks is a persistent challenge for high-fidelity visual synthesis, particularly at scale. Step-level caching has emerged as a promising approach for reducing the Number of Function Evaluations (NFE) by exploiting inter-step temporal redundancy in the denoising process. However, existing caching methods predominantly employ heuristic, threshold-based policies, lacking direct global optimization for output fidelity and introducing input-dependent latency variation.

"Budget-Constrained Step-Level Diffusion Caching" (2606.13496) introduces BudCache, a paradigm shift in step-level caching from threshold-based runtime decisions to a budget-constrained, deterministic, and globally optimized caching strategy. Rather than allowing heuristics to dictate both step reuse and computational expense, BudCache fixes the NFE budget a priori and searches for a static cache policy to maximize output quality. This global discrete allocation avoids unpredictable inference latency and can generalize across images, videos, and varying diffusion backbones.

Figure 2: Overview of the BudCache framework, illustrating the two-step cache-policy search and cache-aware schedule alignment process.

Core Methods

Caching Policy Optimization

BudCache formulates inference acceleration as a constrained combinatorial optimization. For a discretized time schedule given by standard solvers—where each step corresponds to a network evaluation—a binary vector $\mathbf{m} \in \{0,1\}^K$ is introduced. $m_k=1$ signals a fresh computation, $m_k=0$ denotes cache reuse. The constraint $\sum m_k = B$ enforces a strict NFE budget $B$.

The cache policy search objective is to minimize output-level distortion (e.g., MSE) between the cached sampler and full-compute baseline, considering the global effect of each caching decision on trajectory drift and final output quality. The search space is combinatorially large: for $K$ time steps and $B$ active entries, there are $\binom{K}{B}$ possible policies.

To efficiently explore this space, BudCache employs a hybrid two-phase optimization:

• Simulated Annealing (SA): Probabilistic global exploration using Swap/Shift operators enables the search to escape local minima and explore diverse step allocations.
• Deterministic Hill Climbing (HC): Local greedy refinement from the SA output achieves convergence to an efficient local optimum.

This hybrid strategy robustly converges to near-identical minima from diverse initializations, ensuring policy reliability under practical constraints.

Figure 1: Convergence consistency of the proposed hybrid optimization approach across random initializations, confirming robustness and consistent solution quality.

Cache-Aware Schedule Alignment

Standard discretization schedules are suboptimal when aggressive caching modifies the sampler's error accumulation profile. BudCache introduces a second-stage cache-aware schedule alignment: after fixing the cache mask, the continuous time steps $\boldsymbol{\sigma}$ are optimized via backpropagation to minimize output discrepancy (e.g., image MSE) relative to the full solver trajectory. This stage corrects mismatch between the temporal allocation implied by the cache policy and the actual local per-step error, particularly critical in tight-budget (low-NFE) regimes.

This schedule optimization is performed offline, is computationally efficient, and adds zero inference-time overhead.

Experimental Results

BudCache is evaluated on several axes: image synthesis (FLUX.1-dev), video generation (Wan2.1-T2V-1.3B/14B), and compatibility with few-step distilled models (Hyper-SD).

Quantitative Comparisons

BudCache consistently surpasses state-of-the-art caching baselines including TeaCache, MagCache, LeMiCa, and DiCache under matched computational budgets across key metrics: LPIPS, SSIM, PSNR, CLIP-Score, and human preference proxies (ImageReward, HPSv2.1). Notably:

• At 2.51x speedup, BudCache achieves LPIPS = 0.1759 vs. TeaCache's 0.3218.
• In high-speed (3.00x) regimes, BudCache (ultra variant) remains superior despite more aggressive acceleration.
• On Wan2.1-T2V-1.3B, BudCache achieves PSNR = 25.93 vs. TeaCache's 21.52 at substantially lower latency.

Qualitative Comparisons

BudCache demonstrates marked improvements in semantic integrity and high-frequency detail preservation, especially in challenging text rendering and intricate geometry synthesis.

Figure 4: Qualitative comparison with baselines on FLUX.1-dev, highlighting BudCache's fidelity in semantic and fine-grained detail preservation.

On video synthesis, BudCache's cache policies and schedule alignment strategies yield both faster sampling and superior temporal-consistency preservation compared to threshold-driven approaches.

Figure 5: Visualization of videos generated by Wan2.1-T2V-1.3B, emphasizing both acceleration and improved temporal coherence with BudCache.

BudCache remains effective even when transferred across solvers, resolutions, and classifier-free guidance scales, indicating policy robustness and transferability.

Tight Budget and Schedule Alignment Effects

When the compute budget is very tight (e.g., 5–7 NFE), BudCache outperforms both uniform step reduction and existing cache baselines, especially when combined with schedule alignment. For instance, at 5 NFE, BudCache achieves an ImageReward increase of nearly 0.2 over uniform-step solvers.

Analysis and Generalization

BudCache's global cache search yields robust and generalizable policies, as evidenced by:

• High transferability across solvers, resolutions, and guidance scales (no re-search needed).
• Applicability to both standard and distilled few-step models (Hyper-SD), confirming the continued presence of exploitable redundancy.
• Feasibility of searching cache policies on small models and transferring them to larger instances (e.g., Wan2.1-T2V-14B), facilitating cheap adaptation for large-scale models.

Practical and Theoretical Implications

BudCache demonstrates that global, budget-constrained cache policy search is an efficient and deterministic alternative to heuristic, input-dependent caching strategies. Key implications include:

• Deterministic Latency: Deployment is guaranteed to meet pre-specified NFE/latency targets for all inputs.
• Global Quality Optimization: Offline search aligns the caching strategy explicitly with output fidelity, outperforming error-thresholded approaches.
• Zero Overhead at Inference: All optimization is performed offline, with no added online complexity.
• Robust Generalization: Same cache allocation can generalize across settings, models, and tasks, vastly simplifying deployment and model upgrades.

Theoretically, this work decouples step granularity from budget and provides a framework to analyze temporal redundancy and resource allocation in diffusion trajectories.

Figures

Figure 3: Visualizations of 1024x1024 images generated by FLUX.1-dev under BudCache policies, highlighting preservation of fine details even at high acceleration ratios.

Figure 6: Qualitative comparison on Hyper-SD, validating BudCache's applicability to distilled few-step models.

Conclusion

BudCache introduces a globally optimized, budget-constrained paradigm for step-level diffusion caching that deterministically improves sample quality under fixed computational budgets. Its hybrid optimization efficiently navigates the combinatorial policy space, and optional schedule alignment further mitigates errors arising from aggressive caching. The framework is robust, generalizes across domains and models, and offers a practical pathway to deterministic, latency-aware deployment of high-fidelity diffusion models.

Future avenues include developing more transferable cache policies, joint cache-schedule optimization, and deeper analyses of the limits and trade-offs between global and instance-adaptive caching strategies in diverse generative contexts.