🤖 machine learning

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

OTCache is a training-free framework that accelerates diffusion model sampling by modeling caching schedules as a smooth evolution in policy space using Optimal Transport, achieving significant speedups and improved fidelity across various models by interpolating between a high-fidelity reference schedule and an optimized low-budget anchor.

Original authors: Huanlin Gao, Fang Zhao, Qiang Hui, Fuyuan Shi, Shaoan Zhao, Yantao Li, Chao Tan, Ting Lu, Yuren You, Kai Wang, Shiguo Lian

Published 2026-07-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Huanlin Gao, Fang Zhao, Qiang Hui, Fuyuan Shi, Shaoan Zhao, Yantao Li, Chao Tan, Ting Lu, Yuren You, Kai Wang, Shiguo Lian

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to paint a masterpiece, but you are only allowed to make a very limited number of brushstrokes. If you try to paint the whole picture in just a few strokes, you might miss important details or make the colors look wrong. This is the challenge facing modern AI image and video generators: they usually need hundreds of "brushstrokes" (called steps) to create a high-quality image, which takes a long time and uses a lot of computer power.

To speed things up, researchers have tried a trick called caching. Think of this like a painter who, instead of mixing new paint for every single stroke, reuses a bucket of paint they mixed a moment ago. If the scene hasn't changed much, reusing the old paint saves time.

However, the old way of deciding when to reuse the paint (the "caching schedule") had a major flaw. It assumed that every decision to reuse paint was independent of the others. It was like a GPS that calculates the shortest path by adding up the distance of each tiny road segment, ignoring the fact that taking a shortcut early on might force you into a terrible traffic jam later. In the "low-stroke" regime (where you are trying to be extremely fast), this old method breaks down, leading to blurry or distorted images.

Enter OTCache: The "Smart Navigator"

The authors of this paper propose a new system called OTCache. Instead of treating every decision as an isolated math problem, they view the entire painting process as a smooth, flowing journey. They use a mathematical concept called Optimal Transport, which you can think of as a "smoothie blender" for strategies.

Here is how OTCache works in three simple steps:

The Safe Map (Reference): First, the system looks at a "safe" scenario where it has plenty of time (a high number of steps). It uses a standard, reliable method to figure out the best way to paint in this comfortable setting. This gives them a high-quality "map" of how the painting should evolve.
The Extreme Test (Anchor): Next, they look at the opposite extreme: a scenario where they have almost no time (very few steps). Here, the old "add-up-the-distances" method fails. So, OTCache uses a smart, lightweight search tool to find the absolute best way to paint in this rush, even if it has to break the old rules. This becomes their "anchor" point for the extreme speed.
The Smooth Blend (Interpolation): Now, imagine you want to paint at a speed between the safe map and the extreme test. Instead of guessing, OTCache takes the "Safe Map" and the "Extreme Test" and blends them together smoothly. It treats the painting schedule like a flexible rubber band, stretching and warping it to fit the exact speed you need. Because the "Safe Map" and the "Extreme Test" are connected by a smooth mathematical curve, the middle speeds work perfectly too.

The Results

The paper tested this on three powerful AI models (FLUX.1, Qwen-Image, and HunyuanVideo). The results were impressive:

Speed: OTCache made the AI 3.6 to 4.7 times faster than before.
Quality: Unlike other fast methods that made the images look weird or distorted (like a goat with extra legs or a giraffe facing the wrong way), OTCache kept the images looking sharp and true to the original description.
Reliability: It solved the problem where other methods would fail completely when asked to go too fast.

In a Nutshell

Think of OTCache as a smart travel planner. Old planners would tell you, "Take the fastest road for the first mile, then the fastest road for the second mile," without realizing that the first road might lead to a dead end. OTCache looks at the whole trip from start to finish, understands how the journey changes as you speed up, and creates a smooth, optimized path that gets you to your destination (the perfect image) quickly without crashing. It is a "training-free" solution, meaning it doesn't need to relearn how to paint; it just learns how to paint faster by being smarter about when to reuse its tools.

Technical Summary: OTCache

Problem Statement

Diffusion models, particularly those based on Flow Matching (FM) and Rectified Flow, have achieved state-of-the-art results in image and video synthesis. However, their practical deployment in resource-constrained or interactive scenarios is hindered by the massive computational overhead and memory footprints associated with iterative denoising processes.

While training-free caching strategies have emerged to accelerate inference by reusing intermediate features, existing graph-based methods (e.g., MeanCache, LeMiCa) face significant limitations in low-NFE (Number of Function Evaluations) regimes. These methods typically formulate cache scheduling as a constrained shortest-path problem on a directed graph, relying on an additive independence assumption. This assumption posits that the end-to-end degradation of a schedule can be approximated by summing independent local edge errors. The authors argue that this surrogate breaks down under extreme acceleration (low NFE), where large integration steps cause earlier caching decisions to interact non-linearly with subsequent steps, leading to substantial fidelity losses that graph-based optimizers fail to recover. Furthermore, existing approaches treat each budget independently, ignoring the structural regularities and smooth evolution of optimal schedules across different inference budgets.

Methodology: OTCache

To address these limitations, the authors propose OTCache, a training-free framework that models cache scheduling as a smooth evolution in policy space, inspired by Optimal Transport (OT). The framework operates in three distinct stages:

Stage 1: Reference Policy Acquisition (High-Budget)
Under a conservative, high-budget setting (e.g., $B_{ref} = 20$ ), a standard graph-based caching method (instantiated as MeanCache) is used to generate a high-fidelity reference schedule ( $\pi_{ref}$ ). In this regime, the additive shortest-path surrogate remains reliable, providing a robust structural prior for the underlying ODE trajectory.
Stage 2: Anchor Policy Search (Low-Budget)
In the extreme low-budget regime (e.g., $B_{anc} = 8$ ), where graph-based surrogates fail, OTCache performs a lightweight black-box search to find a near-optimal anchor schedule ( $\pi_{anc}$ ).
- Objective: Instead of minimizing local edge errors, the search directly minimizes an end-to-end perceptual objective (LPIPS) between the cached output and the full-budget ground truth.
- Optimization: The search utilizes Optuna with CMA-ES, initialized with a warm start from the MeanCache schedule. The optimization occurs in a "gap" (first-order difference) space to preserve monotonicity and improve search landscape conditioning.
Stage 3: Optimal Path Prediction (Target Budget)
To predict schedules for arbitrary target budgets ( $B \in [B_{anc}, B_{ref}]$ ), OTCache models the transition between the anchor and reference policies as a smooth trajectory in the space of probability measures.
- Continuous Representation: Discrete schedules are lifted to continuous, strictly monotonic time-warping curves using Piecewise Cubic Hermite Interpolating Polynomials (PCHIP).
- Quantile Interpolation: The target schedule is generated via Wasserstein geodesic interpolation (linear interpolation of quantile functions) between the anchor and reference curves.
- Weighted Geodesic: The interpolation weights are determined by two priors: Locality (distance to the target budget) and Confidence (trust in the high-budget reference).
- Conservative Realization: The continuous geodesic is discretized back into a schedule using power-law warping ( $\rho \ge 1.0$ ) to bias NFE density toward the early stages of the reverse ODE, where vector field volatility is highest.

Key Contributions

Rethinking Cache Scheduling: The authors identify two fundamental flaws in existing graph-based schedulers: (i) the misalignment of additive shortest-path surrogates under nonlinear error propagation in low-NFE regimes, and (ii) the failure to leverage structural relationships across different budgets.
Budget-Conditioned Schedule Evolution: OTCache introduces a three-stage framework that treats cache schedules as probability measures. By modeling their evolution as a smooth trajectory in policy space, the method leverages Optimal Transport geometry to interpolate between reliable endpoints, avoiding the pitfalls of independent budget optimization.
Training-Free Acceleration: The framework requires no retraining of the diffusion model, making it immediately applicable to existing architectures.

Experimental Results

The authors evaluated OTCache on three state-of-the-art generative models: FLUX.1 [dev], Qwen-Image, and HunyuanVideo.

FLUX.1 [dev]: At a 4.50× acceleration (10 steps), OTCache achieved an LPIPS of 0.254 and an ImageReward of 0.996. This significantly outperformed MeanCache (LPIPS 0.272) and other baselines, which suffered from content inconsistencies and object distortions at similar speeds.
Qwen-Image: OTCache achieved a 4.70× acceleration (10 steps) with an LPIPS of 0.171, demonstrating superior structural fidelity (SSIM 0.864, PSNR 21.48) compared to MeanCache and LeMiCa, which exhibited noticeable content shifts and object repositioning.
HunyuanVideo: For text-to-video generation, OTCache achieved a 3.66× acceleration with a VBench score of 80.37% and LPIPS of 0.252, outperforming MeanCache (3.59×) and avoiding the temporal artifacts seen in methods like DiCache and TeaCache.

Ablation Studies:

Content Consistency: OTCache maintained rare-word semantics and structural details (e.g., specific objects in a scene) at high acceleration ratios where MeanCache exhibited content drift.
Parameter $\rho$ : A non-linear allocation of NFE density ( $\rho = 1.3$ ) was found to be optimal, balancing early-stage stability with late-stage refinement.
Search Efficiency: The anchor search in Stage 2 was shown to be highly efficient, with the median number of trials required to find the Top-1 optimum being approximately 50.

Significance and Claims

The paper claims that OTCache provides a new perspective on accelerating diffusion models by shifting from discrete, independent shortest-path optimization to geometry-aware schedule modeling. By interpreting the evolution of caching policies across budgets as a smooth trajectory in policy space and utilizing Optimal Transport for interpolation, OTCache recovers fidelity losses that traditional graph-based methods cannot address in low-NFE regimes. The work demonstrates that it is possible to achieve substantial acceleration (up to 4.7×) while maintaining or even improving generation fidelity compared to state-of-the-art baselines, all without the need for model retraining.