LearniBridge: Learnable Calibration of Feature Caching for Diffusion Models Acceleration

Abstract: 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 WAN2.1, respectively. On WAN2.1, it improves VBench by 1.28% over the previous SOTA at $4.10\times$ acceleration. Our code is available at https://github.com/Iiiiiiirene/LearniBridge.

• The paper presents a LoRA-based calibration method that leverages a low-rank, prompt-invariant subspace to correct cached features, achieving up to 6x acceleration in diffusion model inference.
• It employs a two-phase process where minimal prompt training optimizes adapters in the final DiT block, reducing redundant computations while preserving output quality.
• Experimental results show significant speedups in both image and video generation with minimal quality degradation, outperforming existing cache-based approaches in fidelity and temporal consistency.

LearniBridge: Learnable Calibration of Feature Caching for Diffusion Models Acceleration

Motivation and Structural Insights

Diffusion Transformers (DiTs) offer high-quality image and video generation but are constrained by excessive computational burdens, particularly due to repeated backbone evaluations across all diffusion timesteps. Feature caching strategies have emerged to address these inefficiencies, enabling inference acceleration by reusing intermediate tensor representations. However, existing cache-based approaches often suffer severe error accumulation at elevated acceleration ratios, compromising generative fidelity and temporal consistency.

The paper "LearniBridge: Learnable Calibration of Feature Caching for Diffusion Models Acceleration" (2606.26778) systematically investigates the structure of caching errors and demonstrates that the optimal feature correction admits a low-rank, prompt-invariant subspace. SVD analysis of aggregated input matrices ($X_t^l$) across diverse prompts reveals rapid decay in singular values, indicating concentration of spectral energy in a few principal components.

Figure 1: SVD analysis demonstrates a pronounced low-rank structure in cached intermediate activations, and small inter-group angles confirm the prompt-invariance of the optimal correction subspace.

This structural insight implies that the weight correction required for feature calibration ($\Delta W^l$) is inherently low-rank and stable across prompt variations. The empirical finding provides a principled basis for parameter-efficient learnable calibration and suggests broad generalization after training on a minimal prompt subset.

Methodology: LoRA-Based Calibration Architecture

LearniBridge operationalizes these observations by employing Low-Rank Adaptation (LoRA) to implement lightweight, learnable corrections exclusively within the final DiT block. The calibration workflow consists of a two-phase process:

1. Training Phase: A pre-calibration pass records final block inputs at fully computed timesteps, and LoRA adapters are optimized to map cached features to their target representations at skipped timesteps. Only a few (3–5) prompts are required for effective calibration, given the prompt-invariant low-rank structure.
2. Inference Phase: Instead of recomputing the entire backbone for each timestep, cached features from full-compute steps are fed directly into the LoRA-augmented final block, yielding a calibrated output without redundant computation.

   Figure 2: Overview of LearniBridge, detailing the pre-calibration and LoRA finetuning workflow for plug-and-play feature correction.

By restricting adaptation to the final block—covering all linear layers including query, key, value, output projections, and feed-forward weights—LearniBridge achieves negligible memory and latency overhead. The method bypasses auxiliary blocks during inference and operates as a modular, plug-and-play calibration mechanism.

Experimental Results

Image Generation

Quantitative evaluations on FLUX.1-dev demonstrate that LearniBridge achieves up to $6\times$ speedup with minimal quality degradation. Under aggressive acceleration, LearniBridge maintains superior ImageReward, CLIP, PSNR, and SSIM scores compared to all baselines, evidencing robust fidelity and semantic alignment.

Figure 3: Visual comparison on FLUX.1-dev; LearniBridge preserves content fidelity and visual sharpness at nearly $6\times$ acceleration, outperforming baselines that exhibit blurring and color artifacts.

Video Generation

On HunyuanVideo and WAN 2.1-1.3B, LearniBridge attains speedup ratios of $5.75\times$ and $4.10\times$, respectively, while consistently outperforming the previous SOTA in VBench score and perceptual metrics. Baseline methods suffer from content drift, abnormal motion, and degraded color reproduction at high acceleration, while LearniBridge preserves both spatial and temporal coherence.

Figure 4: LearniBridge maintains high motion fidelity and content consistency under high-speed acceleration for HunyuanVideo; baseline outputs exhibit pronounced artifacts and visual instability.

Figure 5: On WAN 2.1-1.3B, LearniBridge outputs closely match original video quality, avoiding blurring and inconsistent color seen in baseline approaches.

Ablation and Architectural Analyses

The impact of varying the LoRA rank is evaluated extensively. Optimal calibration is achieved with a moderate rank ($r=64$) beyond which performance saturates, confirming that the feature correction task is efficiently solvable with constrained parameter space.

Removal of LoRA adapters from individual linear layers leads to notable performance drops, especially from output and feed-forward projections, indicating these as critical for effective calibration.

Figure 6: (a) PSNR improves up to $r=64$, then plateaus; (b) removing adapters from feed-forward and output projections yields the most perceptual degradation.

Further ablations affirm that calibration generalizes over both prompt length and prompt number, with only a handful of moderately complex prompts sufficing for robust adaptation.

Practical and Theoretical Implications

LearniBridge leverages the theoretical result that temporal feature shift calibration in DiTs is governed by a shared, low-dimensional subspace. By exploiting this property via LoRA-based updates, the method achieves robust acceleration across diverse generative tasks without sacrificing model quality.

Practically, LearniBridge enables substantial reduction in inference cost for state-of-the-art diffusion models, making deployment feasible for high-throughput or real-time generative scenarios in both image and video domains. Its parameter-efficiency and prompt-invariance mark it as readily adaptable across diverse diffusion model architectures.

Theoretically, the work suggests that structural analysis of intermediate representations can be leveraged for principled model acceleration. The confirmed universality of low-rank corrections may inform further research into meta-learning and transfer architectures for generative models.

Conclusion

LearniBridge introduces a LoRA-based calibration technique for cached features in diffusion models, formally grounded in empirical low-rank, prompt-invariant subspace analysis. The approach systematically outperforms prior cache-based and extrapolation methods at higher acceleration ratios, preserving generative fidelity even at $5.87\times$ speedup. Its lightweight training, broad generalization, and plug-and-play calibration architecture position LearniBridge as a practical and theoretically informed solution for diffusion model acceleration, with implications for future research on efficient generative modeling and adaptive feature reuse in Transformer architectures.