LearniBridge: Learnable Calibration of Feature Caching for Diffusion Models Acceleration

Diffusion Transformers (DiTs) have driven substantial progress in image and video generation but suffer from prohibitive computational costs. Feature caching accelerates inference by reusing intermediate representations. Existing methods rely on historical features for implementation simplicity, yet suffer from severe error accumulation at high acceleration ratios. To address this limitation, we investigate the nature of the requisite feature correction. We demonstrate that the optimal calibration update is characterized by a shared low-rank subspace across diverse prompts. Guided by this structural insight, we propose LearniBridge, a learnable calibration mechanism for feature caching that bridges multiple timesteps through lightweight LoRA updates. This mechanism enables effective calibration requiring only $3-5$ training samples. Extensive experiments on image and video generation show that LearniBridge achieves up to $5.87\times$, $5.75\times$, and $4.10\times$ acceleration on FLUX, HunyuanVideo, and WAN 2.1, respectively. On WAN 2.1, it improves VBench by 1.28\% over the previous SOTA at $4.10\times$ acceleration.

Diffusion-based AI models have become very powerful tools for generating high-quality images and videos, but they usually require a large amount of computation and are therefore slow and expensive to run. A common way to speed them up is to reuse information computed in earlier steps, rather than recalculating everything from scratch. However, when this reuse is too aggressive, small errors can build up and noticeably reduce the quality of the generated results. In this paper, we study how to correct these accumulated errors efficiently. We find that the needed corrections share a simple underlying structure across many different text prompts. Based on this observation, we propose LearniBridge, a lightweight method that learns how to adjust reused features so that the generation process remains accurate while becoming much faster. Importantly, LearniBridge only needs a few training examples to work effectively. Experiments on both image and video generation models show that LearniBridge can greatly reduce the computation needed for generation while maintaining strong visual quality. It achieves significant speedups on several widely used models and improves over previous methods on video generation benchmarks. This makes high-quality image and video generation more efficient and practical.