OTCache: Optimal Transport for Geometry-Aware Caching in Diffusion Models

Abstract: We propose OTCache, a training-free framework for accelerating diffusion sampling via caching schedule prediction. Existing graph-based caching methods reduce redundant computation by optimizing shortest-path objectives, but rely on an additive independence assumption, which often breaks down in the low NFE regime. To address this issue, OTCache models caching schedules across inference budgets as a smooth evolution in policy space, inspired by Optimal Transport (OT). The framework consists of three stages: (1) obtaining a high-fidelity \textbf{reference schedule} using a graph-based caching method under a conservative budget; (2) performing a lightweight anchor search under an extreme low-budget setting via Optuna optimization with an end-to-end perceptual objective; and (3) predicting schedules for target budgets via quantile interpolation between the reference and anchor policies using continuous warping representations. Experiments on FLUX.1 [dev], Qwen-Image, and HunyuanVideo show that OTCache achieves 4.5x, 4.7x, and 3.66x acceleration, respectively, while consistently improving generation fidelity over state-of-the-art caching baselines. This work provides a new perspective on accelerating diffusion models through Optimal-Transport-inspired schedule modeling. Code:https://github.com/UnicomAI/OTCache

• The paper presents OTCache, which leverages optimal transport to model caching schedules as smooth geodesics in policy space, overcoming limitations of additive graph-based methods.
• The three-stage framework—comprising reference schedule extraction, anchor policy discovery, and transport-based interpolation—achieves up to 4.70× acceleration with minimal quality loss.
• Experimental results on text-to-image and text-to-video tasks validate significant improvements in LPIPS, PSNR, and SSIM, highlighting OTCache's potential for real-time generative applications.

OTCache: Optimal Transport-based Geometry-Aware Caching for Efficient Diffusion Model Inference

Introduction and Motivation

Diffusion models, particularly flow matching (FM) and its transformer-based variants, have demonstrated state-of-the-art quality in generative modeling across image, video, and multimodal domains. However, their pronounced computational footprint and inference latency remain substantial barriers to deployment, especially in interactive or resource-constrained settings. Traditional acceleration methods such as distillation, pruning, and quantization deliver efficiency but require costly retraining or architectural modifications, imposing limitations on adoption flexibility.

Training-free, schedule-based caching approaches have thus emerged to reutilize intermediate latent representations, primarily via scheduling policies defined over the diffusion or ODE trajectory. Existing graph-based caching methods, notably MeanCache, formulate the schedule as a shortest-path problem over a multigraph with additive local error surrogates. However, the fundamental assumption of additive independence becomes invalid in aggressive (low-step/NFE) acceleration regimes, resulting in suboptimal fidelity and significant perceptual quality deterioration.

Figure 1: OTCache surpasses traditional additive surrogates by recovering fidelity under ultra-low NFE, and exploits smooth structural regularity across optimal schedules as the NFE budget varies.

OTCache addresses these limitations by modeling caching policy evolution geometrically in schedule space, inspired by optimal transport (OT). The method leverages the smooth, structured transition of optimal policies across different NFE budgets, architecting a solution that interpolates between high-fidelity and aggressively pruned endpoints. This enables robust, geometry-aware caching schedules for any target computational budget.

Methodology

Limitations of Additive Graph-based Scheduling

MeanCache and similar graph-based frameworks rely on the assumption that scheduling errors propagate additively and independently across the denoising trajectory. Empirical analysis reveals this surrogate objective rapidly diverges from true generation quality in the low-NFE regime due to strong non-linear cumulative effects and long-range interactions among caching decisions. Furthermore, optimal paths for different budgets reveal consistency and smoothness, indicating an underlying trajectory on policy space rather than discrete, unrelated optima.

OTCache hypothesizes that optimal schedules for varying budgets are different projections of a shared geodesic in policy space. This motivates leveraging OT’s Wasserstein geometry to morph between budget-specific scheduling distributions, retaining structural continuity and globally-informed cache decisions.

Three-Stage OTCache Framework

OTCache operates in three stages:

1. Reference Schedule Extraction: For a conservative, high-NFE budget, apply a standard graph-based caching schedule (e.g., MeanCache) to obtain a high-fidelity reference policy, providing a stable structural prior.
2. Anchor Policy Discovery (Low-NFE End): Directly search for an optimal low-NFE caching schedule under an ultra-constrained budget using end-to-end perceptual loss (LPIPS). Black-box optimization (Optuna + CMA-ES) is employed in the discrete, first-order (gap) parameter space, with warm-start initialization from MeanCache to improve efficiency and robustness.
3. Optimal Transport-based Interpolation: Both endpoints are represented as continuous, monotonic warping curves via PCHIP splines over normalized progress coordinates. Target-budget schedules are then predicted via quantile interpolation (Wasserstein geodesic), weighted by locality and reference confidence (via log-conditioning), and discretized with a power-law warping that emphasizes conservative early steps.

Figure 2: Schematic of OTCache—leveraging graph-based reference, anchor search, and optimal-transport interpolation to predict robust schedules for arbitrary NFE budgets.

This approach eliminates the independent-budget assumption, reducing search complexity by offloading expensive optimization to a single anchor, then interpolating for all intermediate budgets.

Experimental Results

Text-to-Image Generation

Extensive evaluation on FLUX.1 $1024\times1024$" title="" rel="nofollow" data-turbo="false" class="assistant-link">dev and Qwen-Image ($1664\times928$) demonstrates that OTCache consistently achieves new Pareto-optimal trade-offs. For FLUX.1, OTCache attains $4.50\times$ speedup ($\mathcal{B}=10$), preserving ImageReward at $0.996$ and delivering significant improvements in LPIPS and PSNR over all baselines. At $3.04\times$ ($\mathcal{B}=15$), it further reduces LPIPS from $0.142$ (MeanCache) to $0.126$ and boosts PSNR to $26.03$.

Figure 3: Quality comparison at high acceleration ratios on FLUX.1 [dev], demonstrating OTCache's visual and structural superiority under extreme compression.

For Qwen-Image, OTCache achieves up to $1664\times928$0 acceleration at $1664\times928$1 without pronounced loss in structural similarity (SSIM $1664\times928$2) and attaining LPIPS of $1664\times928$3, outperforming all graph- and threshold-based methods. Content-alignment is robust to acceleration, with object spatial arrangement and rare-word semantics preserved more reliably than in existing caching approaches.

Figure 4: High-acceleration visual comparison on Qwen-Image. OTCache delivers better spatial and semantic consistency at minimal NFE.

Text-to-Video Generation

On HunyuanVideo, OTCache delivers $1664\times928$4 speedup ($1664\times928$5) with VBench scores of $1664\times928$6 and LPIPS $1664\times928$7, outperforming strong baselines (MeanCache, TeaCache, DiCache) in perceptual fidelity and temporal coherence. The qualitative analysis indicates sustained robustness to flickering and artifact introduction, with critical structural aspects (e.g., motion, object stability) intact even under aggressive acceleration.

Figure 5: Frame-level text-to-video comparison at high acceleration, showing that OTCache preserves temporal and spatial consistency.

Ablation and Search Analysis

Ablation studies reveal that OTCache's power-law warping parameter $1664\times928$8 ($1664\times928$9) optimally redistributes NFE density to early, high-volatility steps, improving both LPIPS and PSNR. Evaluations on the anchor search stage confirm rapid convergence (median $4.50\times$0 queries for Top-1 schedules) and significant average gain ($4.50\times$1 LPIPS improvement over MeanCache anchors). OTCache demonstrates solid initialization robustness, with policy search always seeded from MeanCache for best practical outcomes.

Figure 6: Anchor search convergence—the median Top-1 optimum is identified within 50 trials.

Figure 7: OTCache consistently yields superior quality-latency profiles relative to state-of-the-art baselines.

Implications and Future Directions

OTCache’s geometry-aware, training-free schedule modeling dramatically lowers the adoption cost for diffusion model acceleration. It offers a favorable speed-quality trade-off without retraining or model modification, establishing it as a compelling candidate for deployment in latency-sensitive, real-time, and large-scale generative scenarios.

The optimal transport-based interpolation methodology is broadly applicable beyond diffusion models—potentially informing schedule design for other iterative generative frameworks, caching strategies in temporal/structured models, or even adaptive numerical solvers.

Further research could probe:

• Extension of caching geometry modeling to support adaptive, on-the-fly schedule adaptation, e.g., prompt- or modality-conditioned interpolation.
• Integration with online reinforcement frameworks to continually refine both anchor and reference endpoints.
• Exploration of higher-dimensional OT interpolation for multi-branch models or multi-modal tasks.

Conclusion

OTCache delivers a principled, training-free framework for diffusion model acceleration by modeling schedule evolution as an optimal-transport geodesic in policy space. The framework sidesteps additive surrogate limitations and the combinatorial intractability of independent per-budget optimization. Empirical results confirm that OTCache sets a new performance frontier in text-to-image and text-to-video generative modeling, offering significant acceleration with minimal impact on perceptual or structural fidelity (2606.31026).