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

2026-07-31

AI Summary

Diffusion model inference suffers from high latency due to iterative denoising, and existing static caching strategies struggle to accommodate the dynamic demands across varying prompts and timesteps. This work proposes OnlineCache, the first framework that jointly optimizes dynamic caching and error correction: it employs a lightweight policy network to adaptively determine when to cache intermediate features and introduces a learnable corrector to compensate for approximation errors. A bilevel optimization strategy is adopted to balance global generation quality with local fidelity. Evaluated on FLUX.1-dev, OnlineCache achieves nearly 3× acceleration without compromising output fidelity, and consistently outperforms state-of-the-art methods on DiT and CogVideoX architectures.

Abstract

Diffusion models have revolutionized generative tasks but incur high latency due to iterative denoising. While cache-based strategies accelerate inference by reusing intermediate features, they largely rely on static, sample-agnostic schedules. We argue that this rigidity overlooks two facts empirically validated in this paper: (i) generation difficulty varies across prompts, requiring adaptive resource allocation--complex inputs demand more computation while simpler ones require less; (ii) error sensitivity fluctuates across timesteps, where static policies may cache high-error steps or waste computation on low-error ones. We therefore propose OnlineCache, a dynamic caching framework that jointly learns when to cache and how to correct approximation errors. We leverage policy gradient to train a lightweight network for adaptive speed-quality trade-offs, and incorporate a learnable corrector to mitigate caching-induced errors. Both modules are jointly optimized under a bilevel optimization framework, with the policy targeting global generation quality and the corrector minimizing local errors. Our method automatically allocates computational resources across both samples and timesteps, improving overall generation quality. Extensive experiments demonstrate clear superiority. On FLUX.1-dev model, OnlineCache achieves nearly 3 speedup while preserving generation fidelity. On DiT and CogVideoX, it similarly delivers competitive acceleration without compromising quality; across all scenarios, it consistently outperforms existing cache-based acceleration baselines.

Problem

• diffusion models
• caching policies
• adaptive inference
• error sensitivity
• computational efficiency

Innovation

• dynamic caching
• error correction
• policy gradient
• bilevel optimization
• diffusion inference acceleration