Reproduce “LearniBridge: Learnable Calibration of Feature Caching for Diffusion Models Acceleration” on OpenResearch

About this paper

Xuyue Huang, Zhe Chen, Wang Shen, Xiao-Ping Zhang

LearniBridge adds low-rank LoRA calibration to feature caching in diffusion Transformers, learning to reconstruct skipped-timestep features from cached intermediate states without retraining the backbone. Across FLUX.1-dev, HunyuanVideo, and WAN 2.1-1.3B, it preserves generation quality at roughly 3.5–6× acceleration and improves robustness to the accumulated errors that affect conventional caching methods.

Key insights

• Singular-value and subspace analyses suggest that cross-timestep feature corrections have concentrated low-rank structure and are partly prompt-invariant, making parameter-efficient calibration feasible.
• The output projection and feed-forward layers contribute more to correction quality than the query, key, and value projections, while adapter ranks above approximately 64 provide limited additional benefit.
• Calibration benefits substantially from increasing the prompt set up to about five prompts, after which additional examples yield smaller gains; prompt length and diversity still affect reconstruction quality.

Summary by alphaXiv.