Rethinking Token-wise Feature Caching: Accelerating Diffusion Transformers with Dual Feature Caching

TL;DR

This work questions the necessity and effectiveness of always computing 'important' tokens in token-wise caching for diffusion transformers. It introduces DuCa, a dual feature caching scheme that alternates aggressive and conservative caching cycles and employs random token selection to promote token diversity. Across DiT-based models and tasks (image, video, and class-conditional generation), DuCa achieves substantial speedups (often >2×) while preserving near-original generation quality and alignment, aided by compatibility with FlashAttention. The study highlights the importance of token diversity over fixed importance scores and provides practical guidelines for cache cycle parameters, offering a scalable path to accelerate diffusion models in real-world deployments.

Abstract

Diffusion Transformers (DiT) have become the dominant methods in image and video generation yet still suffer substantial computational costs. As an effective approach for DiT acceleration, feature caching methods are designed to cache the features of DiT in previous timesteps and reuse them in the next timesteps, allowing us to skip the computation in the next timesteps. Among them, token-wise feature caching has been introduced to perform different caching ratios for different tokens in DiTs, aiming to skip the computation for unimportant tokens while still computing the important ones. In this paper, we propose to carefully check the effectiveness in token-wise feature caching with the following two questions: (1) Is it really necessary to compute the so-called "important" tokens in each step? (2) Are so-called important tokens really important? Surprisingly, this paper gives some counter-intuition answers, demonstrating that consistently computing the selected ``important tokens'' in all steps is not necessary. The selection of the so-called ``important tokens'' is often ineffective, and even sometimes shows inferior performance than random selection. Based on these observations, this paper introduces dual feature caching referred to as DuCa, which performs aggressive caching strategy and conservative caching strategy iteratively and selects the tokens for computing randomly. Extensive experimental results demonstrate the effectiveness of our method in DiT, PixArt, FLUX, and OpenSora, demonstrating significant improvements than the previous token-wise feature caching.

Figures (7)

• Figure 1: Comparison between conservative caching (a) and aggressive caching (b): Conservative caching skips some tokens instead of all tokens in each layer to maintain the generation quality, but also leads to lower acceleration ratios. $t$ and $t-1$ denote the refreshing and caching timesteps, respectively.
• Figure 2: The caching error and acceleration ratio of three feature caching strategies. The caching error is defined as the $L2$-norm distance between $x_t$ computed with and without feature caching.
• Figure 3: The overview of DuCa. (a) In the first timestep of a cache cycle, DuCa computes all tokens and stores them in the cache as initialization. Then, DuCa performs a conservative caching step followed by an aggressive caching step, and then another conservative caching step, alternating in this manner to form a complete cache cycle. (b) The proposed method selects tokens for DuCa by setting random scores for all tokens to diversity. The tokens with the largest random scores are selected as cached tokens. The computation results for these tokens are directly retrieved from the cache, while the remaining tokens are computed through the neural layer, with their updated results stored back into the cache.
• Figure 4: Visualization results for different acceleration methods on FLUX.1-dev model.
• Figure 5: Visualization results for different methods on OpenSora (videos available in the supplementary material).