TLDR: OmniCache is a novel training-free method that accelerates Diffusion Transformer (DiT) models by intelligently reusing cached computations. Unlike previous approaches that focus on later sampling stages, OmniCache analyzes the entire sampling trajectory to strategically distribute cache reuse, particularly in early stages where the model can self-correct noise. It also dynamically estimates and filters cache-induced noise, achieving significant speedups (2-2.5x) with minimal quality loss, even on challenging distilled models where other methods fail.
Diffusion models, especially those built with Transformer architectures, have become incredibly powerful for generating images and videos. However, this power comes at a significant computational cost, making real-time deployment a challenge. A new research paper introduces OmniCache, a clever, training-free method designed to dramatically speed up these Diffusion Transformer (DiT) models by rethinking how we reuse cached computations.
Existing methods for accelerating DiT models often focus on reusing computations in the later stages of the denoising process, primarily because model outputs tend to be more similar then. However, the authors of OmniCache argue that this approach can actually degrade the quality of the generated content. They found that while later stages show higher output similarity, they also have weaker denoising strength, meaning any introduced noise from caching is harder to correct.
OmniCache takes a fundamentally different approach. Instead of focusing on local similarities, it adopts a “trajectory-oriented global perspective.” The core idea is to analyze the entire sampling trajectory—the path a diffusion model takes as it transforms noise into a clear image or video. The researchers discovered that if cache reuse is applied in the early stages of sampling, the diffusion model’s inherent self-correction capabilities can effectively filter out any noise introduced by the cached computations. In contrast, applying cache reuse too late leaves insufficient time for correction, leading to irreversible deviations and poorer quality.
How OmniCache Works
The method operates in two main stages: a calibration stage and an inference stage.
During the **calibration stage**, OmniCache first simplifies and estimates the sampling trajectories of various diffusion models. It then uses a curvature-based approach to identify the optimal points in the trajectory for cache reuse. Regions with minimal curvature indicate stable directional progression, making them ideal candidates for caching without compromising quality. The calibration stage also estimates the correlation of cache-induced noise between steps.
In the **inference stage**, when a predetermined cache reuse step is reached, OmniCache replaces the original computations with cached attention and MLP outputs, saving significant computational effort. Crucially, it dynamically estimates and corrects the noise introduced by this reuse. Recognizing that diffusion models handle different types of signals at various stages (low-frequency for overall structure early on, high-frequency for fine details later), OmniCache applies either high-pass or low-pass filtering to the estimated noise. This ensures that only relevant noise is addressed, preventing unnecessary alterations to the hidden states.
Also Read:
- TADSR: Achieving Controllable and Realistic Image Super-Resolution in a Single Step
- Entropy-Driven Efficiency: Quantizing Vision Transformers by Exploiting Attention Redundancy
Impressive Results
Experiments demonstrate OmniCache’s effectiveness. It achieves a 2 to 2.5 times speedup on popular models like OpenSora and Latte, with virtually no loss in generative quality. Even on more challenging, less redundant distilled models, such as CogVideoX-5b-I2V-distill (which uses only 16 sampling steps), OmniCache delivers a 1.45 times speedup without performance degradation. This is particularly noteworthy because existing cache reuse methods often cause these distilled models to fail or “collapse.”
An ablation study further highlighted the importance of OmniCache’s noise correction and filtering modules, showing how they significantly improve video quality and lead to nearly lossless performance compared to the original models.
While OmniCache offers substantial improvements, it does have a minor limitation: cache reuse cannot be applied for three consecutive steps. This constraint is in place to ensure the reliability of the cache-induced noise estimation, which might slightly limit its maximum acceleration potential.
This innovative approach, detailed in the paper OmniCache: A Trajectory-Oriented Global Perspective on Training-Free Cache Reuse for Diffusion Transformer Models, marks a significant step forward in making powerful generative AI models more efficient and accessible for real-time applications. The research was carried out by Huanpeng Chu, Wei Wu, Guanyu Fen, and Yutao Zhang.
