🤖 machine learning

OnlineCache: Learning Dynamic Caching Policies with Error Correction for Efficient Diffusion Inference

OnlineCache is a dynamic caching framework that employs a learnable policy and error corrector to adaptively allocate computational resources across varying prompts and timesteps, achieving significant inference speedups for diffusion models while preserving generation quality.

Original authors: Zhikang Xie, Xichen Ye, Yifan Wu, Haoshen Yu, Li chenan, Peizhu Gong, Weizhong Zhang, Cheng Jin

Published 2026-08-03

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhikang Xie, Xichen Ye, Yifan Wu, Haoshen Yu, Li chenan, Peizhu Gong, Weizhong Zhang, Cheng Jin

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 forced to repaint the entire canvas from scratch for every single brushstroke. That is essentially how modern AI image generators, known as diffusion models, work today. They start with a chaotic cloud of digital noise and slowly, step-by-step, refine it into a clear picture. While this process creates stunningly realistic images, it is incredibly slow and computationally expensive, like trying to build a skyscraper by laying one brick at a time and then checking your work by rebuilding the whole tower. To make these AI tools faster and more useful for real-time applications, scientists have been looking for ways to skip unnecessary steps. The most popular idea so far has been "caching," which is like keeping a copy of a previous brushstroke and reusing it if the picture hasn't changed much yet. However, the old way of doing this was rigid: it used a fixed schedule, deciding to skip steps based on a simple timer, regardless of whether the current image was easy or hard to draw.

This paper introduces a smarter, more flexible approach called OnlineCache. Instead of following a rigid rulebook, the authors trained a tiny, lightweight "coach" (a policy network) to watch the painting process in real-time. This coach decides, moment by moment, whether it's safe to reuse an old brushstroke or if it needs to do the full, heavy calculation. The paper argues that the difficulty of generating an image varies wildly depending on the prompt (some ideas are easy, some are complex) and that some moments in the painting process are more sensitive to errors than others. By teaching the AI to adapt its strategy on the fly and even include a "corrector" to fix small mistakes made by skipping steps, OnlineCache achieves massive speedups. On the FLUX.1-dev model, it nearly triples the speed (a 3× speedup) while keeping the image quality just as high as the slow, careful method. The authors demonstrate that this dynamic, learning-based approach consistently outperforms previous static methods across different image sizes, datasets, and even video generation tasks.

The Problem: The "One-Size-Fits-All" Trap

To understand why OnlineCache is a big deal, we first need to look at the problem it solves. Diffusion models are like artists who start with a blurry, noisy sketch and gradually sharpen it. To get a final image, they might take 30 or more steps. In the past, researchers tried to speed this up by saying, "Hey, steps 10, 11, and 12 look similar, so let's just reuse the result from step 10." This is called caching.

However, the existing methods were like a strict teacher who says, "You must skip steps 10, 11, and 12 for every student, no matter what." This is a static policy. The paper points out two major flaws with this approach:

Different Prompts, Different Needs: Some image requests are simple (like "a red ball"), while others are complex (like "a cyberpunk city in the rain"). A static teacher treats them the same, wasting time on the easy ones and rushing the hard ones.
Different Moments, Different Risks: In the painting process, some steps are critical for getting the shape right, while others are just adding texture. If you skip a critical step, the whole image might look wrong. Static rules don't know the difference; they might skip a crucial step or waste time on a boring one.

The authors argue that this rigidity is the bottleneck. They observed that the "difficulty" of the task changes from prompt to prompt and from step to step, yet old methods ignored this.

The Solution: A Smart Coach That Learns

OnlineCache changes the game by turning the caching decision into a learning problem. Instead of a fixed rule, the system uses a tiny neural network (the "policy") that acts like a smart coach watching the painting process.

How the Coach Decides:
The coach looks at a few key clues before deciding whether to skip a step:

The Current State: What does the image look like right now? (Is it still a mess, or is it almost done?)
The Cached State: What did the last saved step look like?
The Time: How far along are we in the process?

Based on these clues, the coach asks: "Is it safe to reuse the old result, or do we need to do the hard work?" If the coach thinks it's safe, it skips the heavy computation (saving time). If it thinks the image is tricky or the step is critical, it forces the AI to do the full calculation.

The "Bilevel" Twist (The Coach and the Fixer):
The paper introduces two versions of this system. The first is just the coach. The second, more advanced version (called Bilevel Optimization or BLO) adds a second character: a Corrector.

The Coach decides when to skip.
The Corrector is a tiny tool that fixes the small errors that happen when the coach skips a step.

Think of it like a fast-forwarding video player. The coach decides when to fast-forward. If the fast-forward skips too much, the picture might look blurry. The corrector is like a "sharpen" button that instantly fixes the blur, ensuring the fast-forwarded video still looks crisp. The system trains both the coach and the corrector together: the coach learns to skip only when the corrector can handle the fix, and the corrector learns to fix whatever the coach skips.

What the Experiments Showed

The authors tested this system on several powerful AI models, including FLUX.1-dev, DiT, and CogVideoX (for video). Here is what they found:

Speed vs. Quality: On the FLUX.1-dev model, OnlineCache achieved nearly a 3× speedup. This means it generated images three times faster than the standard method. Crucially, the quality didn't drop; in fact, the images were often sharper and more accurate than those made by other fast methods.
Adaptability: The system proved it could handle different scenarios. When asked to generate simple images, it skipped more steps. When the prompt was complex, it worked harder. It also figured out that some specific moments in the painting process were too important to skip, while others were safe to ignore.
Beating the Competition: The paper compared OnlineCache against other top-tier methods like ERTACache and TeaCache. In almost every test, OnlineCache produced better images with less error. For example, on one test, it reduced the visual error (L1 error) by 37.48% compared to a static method that used the same amount of skipping.
Generalization: The system was trained on one set of images (MSCOCO) but worked perfectly on completely different images (Parti-Prompts) and even different resolutions (from 512x512 to 1024x1024) without needing to be retrained. It even worked on video generation, showing that the "smart coach" idea isn't just for pictures.

Why This Matters

The paper suggests that the future of fast AI isn't about building bigger, faster computers, but about being smarter about how we use the computers we have. By moving away from rigid, "one-size-fits-all" rules and embracing a dynamic, learning-based approach, we can make AI generators significantly faster without sacrificing the quality that makes them amazing. The authors show that by letting the AI decide when to take a shortcut and how to fix the mistakes, we can get the best of both worlds: speed and perfection.

In short, OnlineCache teaches the AI to be a flexible artist rather than a rigid robot, resulting in faster, higher-quality creations that adapt to the task at hand.

Technical Summary: OnlineCache

Problem Statement

Diffusion models have established themselves as the dominant paradigm in generative AI across images, video, and audio. However, their reliance on iterative denoising processes incurs substantial computational overhead and high latency, hindering real-time deployment. While cache-based acceleration strategies have emerged to reuse intermediate features and exploit temporal redundancy, existing approaches suffer from significant rigidity. Most current methods rely on static, sample-agnostic schedules (e.g., fixed caching ratios) or heuristic-driven rules (e.g., greedy local thresholds).

The authors argue that these rigid policies fail to address two empirically validated sources of heterogeneity in the denoising process:

Sample-level heterogeneity: Generation difficulty varies significantly across different prompts. Static policies allocate computation uniformly, failing to assign more resources to complex inputs and fewer to simple ones.
Timestep-level heterogeneity: Error sensitivity fluctuates throughout the denoising trajectory. Static policies may inadvertently cache high-error steps or waste computation on redundant, low-error steps, leading to suboptimal quality.

Methodology

To address these limitations, the paper proposes OnlineCache, a dynamic caching framework that formulates the caching schedule as a sequential decision-making problem. The framework learns when to cache and how to correct approximation errors, adapting to both specific samples and timesteps.

Core Components

Policy Network ( $\pi_\theta$ ):
- Architecture: A lightweight Multi-Layer Perceptron (MLP) that outputs a binary decision ( $a_t \in \{0, 1\}$ ) for each timestep: whether to reuse a cached residual or perform full computation.
- State Representation ( $s_t$ ): To ensure efficiency and generalization, the policy does not ingest raw latent tokens. Instead, it utilizes statistical sufficiency and temporal awareness:
  - Channel-wise mean ( $\mu$ ) and standard deviation ( $\sigma$ ) of the current latent $x_t$ and the cached residual $r_{cached}$ .
  - The maximum absolute value of the cached residual ( $\max(|r_{cached}|)$ ) to capture salient localized updates.
  - Time embedding ( $emb_{time}(t)$ ) to provide phase information.
- Optimization: The policy is trained using Policy Gradient (REINFORCE). The objective maximizes a trajectory-level reward that balances visual fidelity against a target acceleration rate.
Reward Design:
- Quality Reward ( $R_q$ ): Based on perceptual similarity (LPIPS) and structural similarity (SSIM) between the generated image and the ground truth.
- Acceleration Constraint ( $R_a$ ): A penalty term enforcing a target cache ratio (e.g., 50% caching).
- Stability Penalty ( $R_s$ ): A penalty applied if consecutive skips exceed a tolerance threshold to prevent error explosion.
Bilevel Optimization (BLO) Variant:
- To mitigate errors introduced by skipping computation, the authors introduce a learnable Error Corrector ( $C_\phi$ ), implemented as a lightweight MLP.
- Outer Loop: Optimizes the policy $\theta$ to maximize global generation quality.
- Inner Loop: Optimizes the corrector $\phi$ to minimize local errors (Mean Squared Error) on the cached steps.
- Training Strategy: Due to the intractability of computing exact implicit gradients (caused by discrete action sampling and high dimensionality), the authors employ a First-Order Approximation. They alternate between updating the policy (assuming the corrector is fixed) and updating the corrector (via supervised fine-tuning on trajectories generated by the new policy). This coordinate descent approach converges to a joint equilibrium.

Key Contributions

Dynamic, Instance-Aware Framework: OnlineCache moves beyond fixed rules, dynamically allocating computational resources based on prompt difficulty and timestep sensitivity.
Joint Optimization of Policy and Correction: The BLO variant uniquely couples the caching decision with an error corrector, ensuring that the policy only skips computation when the corrector can effectively repair the induced trajectory deviations.
Efficient and Generalizable Design:
- The policy network is lightweight (negligible inference latency).
- The use of statistical summaries (mean/std) rather than raw tensors makes the approach resolution-agnostic and capable of zero-shot transfer across resolutions and closely aligned model architectures.
Comprehensive Evaluation: The authors provide extensive experiments across image (FLUX.1-DEV, DiT) and video (CogVideoX) generation, demonstrating superiority over both static and learning-based baselines.

Experimental Results

The paper reports significant performance gains across multiple models and datasets:

FLUX.1-DEV (Image Generation):
- On MSCOCO (512×512), OnlineCache (BLO variant) achieves a 2.96× speedup with a cache ratio of 72.22%, outperforming ERTACache (2.68×) and TeaCache (2.60×) in both speed and visual quality (lower LPIPS, higher SSIM/PSNR).
- On Parti-Prompts, it achieves up to 5.25× speedup with graceful degradation in quality, maintaining competitive metrics even at extreme acceleration levels.
- Generalization: The method trained on 512×512 images successfully generalizes to 1024×1024 without retraining. It also transfers effectively to the distilled FLUX.1-Schnell model and the Parti-Prompts dataset.
DiT-XL/2 (ImageNet-256):
- Compared against learning-based methods (FastCache, L2C) and training-free methods (TeaCache, AdaCache), OnlineCache achieves a 1.88× speedup (under DPM++ 50 steps) while improving FID from 3.47 to 3.28. The authors attribute this quality gain to the suppression of sampling noise via selective caching and error correction.
CogVideoX-2B (Video Generation):
- On video generation tasks, OnlineCache outperforms the TeaCache baseline, achieving higher VBench scores and lower LPIPS at comparable speedups (e.g., 1.79× speedup with 78.09% VBench score vs. TeaCache's 77.40%).

Significance and Claims

The paper claims that OnlineCache establishes a principled, globally optimized paradigm for cache-based acceleration. By treating caching as a sequential decision-making problem rather than a static rule, the framework effectively navigates the trade-off between efficiency and fidelity.

The authors emphasize that their approach:

Automatically allocates resources across both samples and timesteps, addressing the intrinsic heterogeneity of diffusion processes.
Requires no architectural modifications to the backbone diffusion models, making it a "plug-and-play" solution.
Delivers superior Pareto frontiers between speed and quality compared to existing state-of-the-art methods, validated across diverse modalities and resolutions.

The work suggests that future acceleration strategies should move beyond heuristic or static scheduling toward adaptive, instance-aware policies that jointly optimize for both speed and error correction.