Forecast Then Calibrate: Feature Caching as ODE for Efficient Diffusion Transformers

Authors

  • Shikang Zheng Shanghai Jiao Tong University, China South China University of Technology, China
  • Liang Feng Shanghai Jiao Tong University, China
  • Xinyu Wang Shanghai Jiao Tong University, China Tsinghua University, China
  • Qinming Zhou Shanghai Jiao Tong University, China Tsinghua University, China
  • Peiliang Cai Shanghai Jiao Tong University, China
  • Chang Zou Shanghai Jiao Tong University, China
  • Jiacheng Liu Shanghai Jiao Tong University, China
  • Yuqi Lin Shanghai Jiao Tong University, China
  • Junjie Chen Shanghai Jiao Tong University, China
  • Yue Ma The Hong Kong University of Science and Technology, China
  • Linfeng Zhang Shanghai Jiao Tong University, China

DOI:

https://doi.org/10.1609/aaai.v40i16.38349

Abstract

Diffusion Transformers (DiTs) have demonstrated exceptional performance in high-fidelity image and video generation. To reduce their substantial computational costs, feature caching techniques have been proposed to accelerate inference by reusing hidden representations from previous timesteps. However, current methods often struggle to maintain generation quality at high acceleration ratios, where prediction errors increase sharply due to the inherent instability of long-step forecasting. In this work, we adopt an ordinary differential equation (ODE) perspective on the hidden-feature sequence, modeling layer representations along the trajectory as a feature-ODE. We attribute the degradation of existing caching strategies to their inability to robustly integrate historical features under large skipping intervals. To address this, we propose FoCa (Forecast-then-Calibrate), which treats feature caching as a feature-ODE solving problem. Extensive experiments on image, video generation, and super-resolution tasks demonstrate the effectiveness of FoCa, especially under aggressive acceleration. Without additional training, FoCa achieves near-lossless speedups of 5.50× on FLUX, 6.45× on HunyuanVideo, 3.17× on Inf-DiT, and maintains high quality with a 4.53× speedup on DiT.

Downloads

Published

2026-03-14

How to Cite

Zheng, S., Feng, L., Wang, X., Zhou, Q., Cai, P., Zou, C., … Zhang, L. (2026). Forecast Then Calibrate: Feature Caching as ODE for Efficient Diffusion Transformers. Proceedings of the AAAI Conference on Artificial Intelligence, 40(16), 13449–13457. https://doi.org/10.1609/aaai.v40i16.38349

Issue

Section

AAAI Technical Track on Computer Vision XIII