Accelerating Diffusion Transformers with Token-wise Feature Caching
Authors: Chang Zou
Xuyang Liu
Ting Liu
Siteng Huang
Linfeng Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: Diffusion Models
Image generation
Video generation
Model Acceleration
Feature Cache
Primary area: generative models
TL;DR: We propose a training-free acceleration method for diffusion transformers by caching the features of unimportant tokens while still computing the important tokens.

Abstract: Diffusion transformers have shown significant effectiveness in both image and video synthesis at the expense of huge computation costs. To address this problem, feature caching methods have been introduced to accelerate diffusion transformers by caching the features in previous timesteps and reusing them in the following timesteps. However, previous caching methods ignore that different tokens exhibit different sensitivities to feature caching, and feature caching on some tokens may lead to 10$\times$ more destruction to the overall generation quality compared with other tokens. In this paper, we introduce token-wise feature caching, allowing us to adaptively select the most suitable tokens for caching, and further enable us to apply different caching ratios to neural layers in different types and depths. Extensive experiments on PixArt-alpha, OpenSora, and DiT demonstrate our effectiveness in both image and video generation with no requirements for training. For instance, 2.36$\times$ and 1.93$\times$ acceleration are achieved on OpenSora and PixArt-$\alpha$ with almost no drop in generation quality. Codes have been released in the supplementary material and Github.

Published as a conference paper at ICLR 2025
ACCELERATING DIFFUSION TRANSFORMERS WITH TOKEN-WISE WISE FEATURE CACHING
Linfeng Zhang't Shanghaai Jiao Tong University University of Electronic Science & Technology of China Sichuan University *National University of Defense Technology
Chang Zou Xuyang Liu
Ting Liu
Siteng Huang²
*Zhejiang University
Code: https: pettps://githen //tgtttgttttttt//e httg:/tthbb..om/hheeea.... com/Shenya-2//oCC
ABSTRACT
Diffusion transformers have shown significant effectiveness in both image and video synthesis at the expense of huge computation costs. To address this problem, feature caching methods have been introduced to accelerate diffusion transformers by caching the features in previous timesteps and reusing them in the following timesteps. However, previous caching methods ignore that different tokens ex- hibit different sensitivities to feature caching, and feature caching on some tokens may lead to 10x more destruction the overall generation quality compared with other tokens. In this paper, we introduce token-wise feature caching, allowing adaptively select the most suitable tokens for caching, and further enable us to apply different caching ratios to neural layers in different types and depths. Extensive experiments on PixArt-o, OpenSora, DiT and FLUX demonstrate our effectiveness in both image and video generation with no requirements for train- ing. For instance, 2.36x and 1. 1.93x acceleration are achieved on OpenSora and Pix-rt-a with almost no drop in generation quality.
INTRODUCTION
Diffusion models (DMs) have demonstrated impressive performance across a wide range of gener- ative tasks such as image generation (Rombach et al.) 2022) and video generation (Blattmann al.) 2023]. Recently, the popularity of diffusion transformers further extends the boundary of visual generation by scaling up the parameters and computations (Peebles & Xie 2023). However, a sig- nificant challenge for diffusion transformers lies in their high computational costs, leading to slow inference speeds, which hinder their practical application in real-time scenarios. To address this, series of acceleration methods have been proposed, focusing on reducing the sampling steps (Song et al. 202T and accelerating the denoising networks (Bolya & Hoffman
2023
Fang al.. (023).
Among these, cache-based methods (Ma et al., 2024b Wimbauer et al. 2024), which accelerate the sampling process by reusing similar features across adjacent timesteps (e.g. reusing the features cached at timestep /. in timestep / - 1), have obtained abundant attention in the industrial com- munity thanks their plug-and-ple
property. As the pioneering works in this line, DeepCache (Ma et al. 2024b) and Block Caching (Wimbauer et al. (024) were proposed to reuse the cached teatures in certain layers of U-Net-like diffusion models by leveraging the skip connections in the U-Net. However, the dependency on the U-Net architectures also makes them unsuitable for dif- fusion transformers, which have gradually become the most powerful models in visual generation. Most recently, FORA (Selvaraju et al. (2024) and D-DIT (Chen et al. 2024b) have been proposed as a direct application of previous cache methods to diffusion transformers, though still not fully analyzed and exploited the property of the transfrmmrrarcchitectur. To tackle this challenge, this paper begins by studying how feature caching influences diffusion transformers at the token level.
Difference in Temporal Redundancy: Figure I shows the distribution of the feature distance be- tween the adjacent timesteps for different tokens, where a higher value indicates that this token exhibits lower similarity in the adjacent timesteps. It is observed that there exist some tokens that show relatively lower distance (in light blue) while some tokens show extremely higher distance (in dark blue), almost 2.5x larger than the mean distance, indicating caching such tokens can lead to an overlarge negative influence. This observation indicates that different tokens have different redundancy across the dimension of timesteps, (i.e. different temporal redundancy).
"Equal contribution. thenyizou@outlokk.ccom @outlook.com Corresponding author: zhanglinfeng@sjuu.cuu

---

Published as a conference paper at ICLR 2025
0.6
0.016
0.014
0.5
00.012
Frequency
0.4
0.010
0.008
0.3
0.006
0.2
0.004
0.002
0.1
0.0
50
100
150
200
250
300
0.0
10
12
Frobenius norm distance
Frobenius norm distance
Figure I: Temporal Redundancy: Distribution of the distance between the feature of tokens in the previous and the current timestep.
Figure 2: Error Propagation: Distribution of the error the final layer output when the same noise is applied to each token in the first layer.
Difference in Error Propagation: Figure introduces the other interesting perspective of error propagation in diffusion transformers. Specifically, self-attention and cross-attention layers are widely utilized in diffusion transformers to formulate the dependency between different tokens. As a result, the error in one of the tokens may propagate to some other tokens by self-attention and finally result the error in all of the tokens. To understand the error propagation in tokens of diffusion transformers, we apply Gaussian noise with the same intensity to each token and compute the resulting error all the tokens in the output of the final layer. Surprisingly, Figure 2 shows that the same noise in different tokens leads to significantly different propagation errors, with the largest propagation error being more than 10x the smallest one. the context of feature caching. this indicates that the same error introduced by feature cache can result in vastly different errors in the final generation result due to error propagation.
Moreover, we have also investigated the difference of the tokens in layers of different depths and types, which demonstrates significant differences, as introduced in the following sections. In sum- mary, different tokens exhibit significant differences in their sensitivities to feature caching, indicat- ing that they deserve different priorities during the caching process. This motivates us to study the token-wise feature caching strategy, which aims select the maximal number of tokens maxi- mize the acceleration ratio while minimizing the resulting error introduced by the feature caching by selecting the tokens that make the least caching error.
To tackle this challenge, this paper introduces Token-wise feature Caching ToCa for training-fiee frree acceleration of diffusion transformers, which provides a fine-grained caching strategy for differ- ent tokens in the same layer and the tokens in different layers. The core challenge of ToCa is to accurately select the tokens that are suitable for feature caching with the computation-cheapel operations. Consistent with the two previous analyses, we mainly study this problem from the per- spective of temporal redundancy and error propagation by defining four scores for token selection. Specifically, for temporal redundancy, we try to select the tokens with the highest similarity (i.e. lowest difference) with their value in the previous timesteps by considering their frequency of being cached, well their distribution in the spatial dimension of the images. For error propagation, we attempt to cache the tokens which makes the least influence on other tokens based on their at- tendance in self-attention and cross-ates atts-ttttntion layers. Besides, all of these scores can be obtained without any additional computation costs.
Extensive experiments text-to-image, text-to-video and class-to-image generation demonstrate the effectiveness of ToCa on PixArt-o, OpenSora, and DIT over previous feature caching methods. For instance, 2.36x acceleration can be achieved on OpenSora without requirements for training, outperforming halving the number of timesteps directly by 1.56 on VBench. On PariPrompt, ToCa even leads 1.13 improvements on CLIP Score, indicating higher consistency the text conditions.
In summary, the contributions of this paper are as follows:
1. We propose Token-wise Caching (IoCa) as a fine-grained feature caching strategy tailored to acceleration for diffusion transformers. To the best of our knowledge, ToCa first introduces the perspective of error propagation in feature caching methods.
2. We introduce four scores to select the most suitable tokens for feature caching in each layer with no additional computation costs. Besides, ToCa enables us to apply different caching ratios in layers of different depths and types, and also bring bag of techniques for feature caching.
3. Abundant experiments on PixArt-a. OpenSora, and DiT have been conducted, which demon- strates that ToCa achieves a high acceleration ratio while maintaining nearly lossless generation quality. Our codes have been released for further exploration in this domain.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
Transformers in Difflusion Models
2015), which iteratively denoise an Diffusion models (DMs) (Ho et al. 2020 Sohl- et al.
initial noise input through series of diffusion steps, have achieved remarkable success across vari- ous generation applications (Rombach et al. 2022 Balaji et al. (022). Early DMs (Ho et al 2020 Rombach et al. 2022) are based on the U-Net architecture (Ronneberger et al., 2013), consistently achieving satisfactory generation results. Recently, Diffusion Transformer (DiT) (Peebles & Xie 2023) has emerged as a major advancement by replacing the U-Net backbone with a Transformer architecture. This transition enhances the scalability and efficiency of DMs across various gener- ative tasks (Chen et al., 2024a; Brooks et al. 2024). For example, PixArt-a (Chen et al.) 2024a) utilizes DiT as scalable foundational model, adapting it for text-to-image generation, while Sora [Brooks et al. 2024) demonstrates DiT's potential in sigh-fidelity video generation, inspiring a se- 2024 Lab & etc. 2024) Despite their success, the
ries of related open-source projects (Zheng et al.
iterative denoising process of these DMs is significantly time-consuming, making them less feasible for practical applications.
Acceleration of Diffusion Models
To improve the generation efficiency of DMs, numerous diffusion acceleration methods have been proposed, falling broadly into two categories: (1) reducing the number of sampling timesteps, and (2) accelerating the denoising networks. The first category aims achieve high-quality generation results with fewer sampling steps. DDIM (Song et al.
202I) introduces a deterministic sampling process that reduces the number of denoising steps while preserving generation quality. DPM- Solver (Lu et al.) 2022a) and DPM-Solver++ (Lu et al. 2022b) propose adaptive high-order solvers for a faster generation without compromising on generation results. Rectified flow (Liu et al., 2023)
optimizes distribution transport in ODE models to facilitate efficient and high-quality generation,
enabling sampling with fewer timesteps. Step-distillation (Salimans & Ho 2022, Meng et al. 2023) minimizes the number of timesteps with knowledge distillation from multiple timesteps to fewer ones. Consistency models (Song et al. (023) accelerate generative modeling by mapping noise di- rectly to data and enforcing self-consistency across steps. In the second category, various efforts have been paid to token reduction (Bolya & Hoffman! 2023] Zhang et al. 2025; 2024a), knowledge distillation (Li et al., 2024),, and weight quantization (Li al.
allaal122220 [2223b]] [Shang et al. 2023) and prun- ing (Fang et al. (023) on the denoising networks. Additionally, recent cache- based methods reduce redundant computations accelerate inference for DMs. These rache-based methods have obtained abundant attention since they have no requirements for additional training. DeepCache (Ma et al.) 2024b) eliminates redundant computations in Stable Diffusion (Rombach et al. intermediate features of low-resolution layers in the U-Net. Faster Diffusion (Li et al. 20023a) accel- erates the sampling process of DMs by caching U-Net encoder features across timesteps, skipping encoder computations at certain steps. Unfortunately, DeepCache and Faster Diffussin are designed specifically for L-NNetbased denoisers and can not be applied to DiT (Chen et al. 2024b). Recently, FORA (Selvaraju et al.) (024) and DDDiT (Chen et al. 2024b) have been proposed to cache the features and the residual of features for DIT. Learning-tor Coarning-to-Cccheee Cache (Ma et al., 2024a) learns an optimal cache strategy, which achieves a slightly higher acceleration ratio but introduces the requirements of training. However, these methods apply the identical cache solution to all the tokens and even all the layers, which leads to significant performance degradation in generation quality.
2022) by reusing
3 METHODOLOGY
In the Methodology section, we briefly introduce Diffusion Models and the feature eaching accel- eration method for Diffusion Transformers, followed by the ToCa workflow and important token selection. The importance of token x; is determined by 8 interaction strength with other tokens, 82 association with global textual information, 8y accumulated cache error, which, if excessive, can cause image collapse and should be controlled, and 84 spatial uniformity of selected tokens.
3.1 PRELIMINARY
Diffusion Models Diffusion models are formulated to contain two processes, including forward process which adds Gaussian noise to a clean image, and a reverse process which gradually denoises standard Gaussian noise real image. By denoting L the timestep and B, as the noise variance
3

---

Published as a conference paper at ICLR 2025
Xt-1
Xt-2
Other Layers
Cache after timestop t:
Other Layers
Update token 2,4,5
Block 7
:
Init.
Block 3
Neurai
layer
Fetch token 2,3 Irom Cacha
Neural 1ayer
Cache ofter timostep D-]
Peature Computed at
Cosputed
Compute Cache 5coxe 0.0.100.....00
X
Compute
Cache
Xt-1
Righ Cache Ratio
timestep
Limestep é-1
(b) ToCa in layers different depth
(a) Overview ToCa in timestep t and t-1
Figure 3: The overview of ToCa on the example of the first layer with caching ratio R =40%. (a) In the first timestep of the cache period, ToCa computes all the tokens and stores them in the cache for initialization. Then, in the next timestep, ToCa first computes the caching score of each token and selects the tokens for cache based on them. Then, ToCa fetches the features of cached tokens from the cache while performing real computations in the other tokens. Then, the features of tokens that have been computed are utilized to update their value in the cache. (b) ToCa applies higher cache ratio in the deep layer and relatively lower cache ratio in the shallow layers.
schedule, then the conditional probability in the reverse (denoise) process can be modeled as 1- v1 04
Po( x²) N It-1; Vac It
E0(22,)) 3..I
where 0₃2 11 3r, ae = 1 II a, and T denotes the number of timesteps. Importantly, cg denotes denoising network with its parameters 0 that takes I, and & as the input and then predicts the corresponding noise for denoising. For image generation with T timesteps, eo is required to infer for T times, which takes most of the computation costs in the diffusion models. Recently, fruitful works demonstrate that formulating as a transformer usually leads to better generation quality.
Diffusion Transformer Diffusion transformer models are usually composed of stacking groups of self-attention layers JSA, multilayer perceptron fMLP, and cross attes-nttention layers fca (for conditional generation). It can be roughly formulated as 91 o 92 o ... 9L where g' ff f(m.)) The upper script denotes the index of layer groups and L denotes its maximal number. We omit the other components such as layer norm and residual connections here for simplicity. For diffusion transformers, the input data X, is sequence of tokens corresponding to different patches in the HxW generated images, which can be formulated as { where H and W denote the height and width of the images or the latent code of images, respectively.
3.2 NAIVE FEATURE CACHING FOR DIFFUSION TRANSFORMERS
We follow the naive scheme for feature caching adopted by most previous caching methods Ma et al. (024b) for diffusion denoisers. Given a set of N adjacent timesteps (t, +....,,+A.. 1), native feature caching performs the complete computation at the first timestep and stores the in- termediate features in all the layers, which can be formulated as C(x;) == f(x)), for V/ E L). where ":=" indicates the operation of assigning the value. Then, in the next 1 timesteps, fea- ture caching avoids the computation of self-attention, cross-attention, and MLP layers by reusing the feature cached at timestep t. By denoting the cache as C and the expected feature of the input x, in the 4n layer F(x)), then for VI E [1, L]. the naive feature caching can be formulated as F(x|T1) F(x[=b)
F(x
C(x)).
(2)
In these NN timesteps, naive feature caching avoids almost all the computation in A - timesteps, leading to around - times acceleration. After the / timesteps. the feature cache then starts new period from initializing the cache as aforementioned,, again. The effecciivnness of feature caching can be explained by the extremely low difference between the tokens in the adjacent timesteps. However, as N increases, the difference between the feature value in the cache and their correct value can be exponentially increased, leading degeneration in generation quality, which motivates study more fine- grained methods for feature caching.
3.3 TOKEN- WISE FEATURE CACHING
The naive feature caching scheme caches all the tokens of the diffusion transformers with the same strategy. However, demonstrated in Figure [ and 5 feature cache introduces significantly
4

---

Published as a conference paper at ICLR 2025
Image
Quary:liaage Tokens
Kaye Test Takens
find the putch
2.0. L.5.5,00.
O p.p. S 00:1
Self-Attention (i))
Cross-Attention.
Spatial
Bonus
(xck 22.2
#2 (ai)) -1
Encourage to Nut Cached
Key: Image
Query: Image Tokenss
frequency of being cathed
value
Sq +3g
Reshape
#:(4)= -
sy(r)
(ngtog)
Cache
Compute 01)
NG (G) (5(0)) + s(6) ((t)
Influence Other Tokens (II) ControlAbiliy
(III) Cache Frequeney
(TV) Uniform Spatial Distribution
Figure 4: The computation of caching scores in ToCa, where token with lower cache score is encouraged to be cached: (I) Self-attention weights are utilized to measure the influence of each token on the other tokens, where token with higher influence is considered not suitable for caching. (II) Cross-attention weights are utilized to measure the influence of each image token on the text (condition) tokens, where image token with higher entropy is considered not suitable for caching. (III) Tokens that have been cached multiple times are encouraged to not be cached in the following layers. (IV) We increase the cache score for the token with the largest cache score in neighboring pixels to make the cached tokens distributed uniformly in the spatial dimension.
different influence to different tokens, motivating to design more fine- fine-grained caching method in the token-level. In this section, we begin with the overall framework of ToCa, then introduce our strategy for token selection and caching ratios.
3.4 OVERALL FRAMEWORK
Cache Initialization Similar to previous caching methods, given a set of adjacent timesteps {t,t 11...,,++. - 1}. our method begins with computing all the tokens at the first timestep t, and storing the computation result (intermediate features) of each self-attention, cross-attention, and MLP layer in cache, denoted by C, as shown in the left part in Figure This can be considered as the initialization of C, which has no difference compared with previous caching methods.
Computing with the Cache In the following timesteps, portant tokens by re-using their value in the cache C. We firstly pre-define the cache ratio R of tokens in each layer, which indicates that the computation of R% of the tokens in this layer should be skipped by using their value in the cache, and the other (1 12%) tokens should still be com- puted. To achieve this, a caching score function S is introduced decide whether token should be cached, which will be detailed in the next section. Then, with S, we can select set of cached tokens as I Cache and the other set of tokens for real computation as I [2 computation of the layer f' for ita token 2y can be formulated as F(xx) = ff(;) - (1 7s (xss), where 7
skip the computation of some unim-
I Cache- Then, the
for E I Cache and 7/2- = for i I Compute- C(x:) denotes fetching the cached value of x; from C, which has computation costs and hence leads to overall acceleration in f.
Cache Updating As a significant difference between traditional cache methods and ToCa, tradi- tional cache methods only update the feature in the cache at the first timestep for each caching period while ToCa can update the feature the cache at all the timesteps, which helps to reduce the error introduced by feature reusing. For the tokens F2 E I Compute which are computed, we update their corresponding value in the cache C, which can be formulated as C(x,) == F(x,) for E I Compuse.
3.5 TOKEN SELECTION
Given a sequence of tokens x, = [xx) token selection aims to select the tokens that are suitable for caching. To this end, we define caching score function S(x,) to decide whether the it token I; should be cached, where token with higher score has lower priority for caching and higher priority to be actually computed. The S(x;) is composed of four sub-scores {s1, 82, S3, $1), corresponding to the following four principals.
(I) Influence to Other Tokens: If a token has a significant contribution to the value of other to- kens, then the error caused by token caching this token casily propagate to the other tokens, ultimately leading to discrepancies between all tokens and their correct values. Consequently, we consider the contribution of each token to other tokens one of the criteria for defining whether it should be cached, estimated with an attention score in self-attention. Recall that the self-attention QK"
can be formulated as O AV, where A Softmax
E R denotes the normalized
Vd
5

---

Published as a conference paper at ICLR 2025
0.32
Layer
Layer
230
Self
0.20
Layer 15
Layer
175
Cross Attn
Trequenc
Layer 28
Layer 25
125
MLP
0.66
190
0.04
0.82
Probenius noom
Frobenius
Frobenius norn
(a)
(b)
(c)
Figure 5: (a) The distance between features at the last timestep and the current timestep for features in different layer depths. (b) The distribution of errors in the output of the final layer when the same
(c) The distribution of errors in the Gaussian noise is applied to tokens in different layer depths. output of the final layer when the same Gaussian noise is applied to tokens in different layer types.
attention map. Q, K, V and O Rxxx are query, key, value and output tokens respectively; d is the hidden size of each token and N is the total number of tokens. More specifically, the ith output token is obtained through o, = j=1 Okij f/y, where ai is the (i,j) element of the attention map A, denoting the contribution of value token 1y to the output token 0j. With these notations, we define 81 to measure the contribution of z; to other tokens $1 (2:) X1 Σ Gij as shown in Figure D.
(II) Influence to Control Ability: The control ability of diffusion models in the text-to-image generation is usually achieved with cross- layer which injects the control signal (e.g. text) into the image tokens. Hence, the cross-attention map reflects how each image token is influenced by the control signal. In this paper, we define the image tokens that are influenced by more tokens in the control signal as the tokens that are not suitable for caching, since the caching error on these tokens leads to more harm in the control ability. Specifically, by denoting Ciy as the (i, element in QK
the cross attention score C Softmax /d where Kiest denotes the keys of text tokens (control tokens). Then, as shown in Figure II), we employ the entropy H (x;) of the cross-attention weight for each image token x; as its influence on the control-ability of diffusion models, which can be formulated as (2(x)) (1;, = شソ= Σド G3 log(c;)).
We observe that when token cached across multiple adjacent layers, the error introduced by feature caching in this token can be quickly accumulated, and the difference between it and its correct value can exponentially amplified, which significantly degrades the overall quality of images. Hence, we define recently cached tokens as unsuitable for cache in the next layers and time steps. Conversely, the tokens that have not been cached for multiple layers and timesteps encouraged to be cached. As shown in Figure III), this selection rule is achieved by recording the times of being cached for each token after their last real computation, which can be formulated as sa(2;) n. where n; represents the number of times that x; has been cached after its last computation. A is the number of timesteps in each feature caching cycle.
(III) Cache Frequency:
The pixels in the neighboring patch of the images usually contain similar information. As discussed in previous works, overwhelmingly pruning the infor- mation in local spatial region may result in significant performance degradation in the whole images (Bolya & Hoffman 2023). Hence, to guarantee that the errors introduced by caching are not densely distributed in the same spatial region, we define the following scoring function: sx(z;) = I(x;)) (A1 $1 (xz) ) + A2 (2(;)) + As ss(zz)), where I(z;) is an indicator function which equals if x; has the highest score of X $1 s1(;) + X2 (2(2;) + >x sz(2;) in its neigh- to balance each score.
(IV) Uniform Spatial Distribution:
boring h x k pixels and 0 in the other settings, and A are
In summary, the overall caching score of x; can be formulated S(x;) =ド== sy(tt), where Xy hyper-parameters to balance each score. Then, with the cache ratio R. the index set for the cached tokens is obtained in the following form:
arg min I Cache (11.22 50% (1.22.....))
(, ...
",)),, ,S(2..))
3.6 DIFFERENT CACHE RATIOS IN DIFFERENT LAYERS
Figure shows the difference in feature caching of different layers, where (a) shows that the output features of different layers have different distances compared with their value in the last step (i.e. different temporal redundancy). (b) and (c) show that when errors in the same density are applied to

---

Published as a conference paper at ICLR 2025
A clese-up of hawk soaring the sky with wings spread wids
A close-up of leopard crouchted ia the tall grass, eyes locked on its prey
close-ug of welf howling at moon in forest clearing
*A misty morning over with rowboat tied a wooden dock
person paddling a canoe through peaceful lake with bug rising from the wates"
*A tranguil beach soft while and clear turquoise waters
Original x 1.0
ToCa
x 1,9
steps 2,0
FORA 2.0
Figure 6: Visualization examples for different acceleration methods on PixArt-a.
layers in different depths and types, the resulting error in the final layer exhibits extremely different magnitudes. Specifically, in the three studies, the disparity between the maximum and minimum error can be several orders of magnitude. Fortunately, ToCa enables us to apply different caching ratios for layers in various depths and types. By denoting the overall caching ratio for all the layers and timesteps as Ro, then two factors and are introduced to adjust the caching ratios. Then the final caching ratio of the layer in depth and type can be written as Rigee = R x Ti X Thype-
n: As introduced in Figure 5ia) and (b) although the features of the shallow layers tend exhibit lower differences than the deeper layers, the error introduced by the cached tokens in the shallow layers can be propagated lowing layers, resulting in much larger caching error. Based on this observation, we set larger and smaller cache ratios for deeper and shallower layers, respectively, by setting ri 0.5 + (1/c- where 0.5 is utilized for 1-center and L denotes the maximal depth and A controls the slope.
the other tokens and amplified during the computation in all the fol-
Typp:: As shown in Figure 5jc). layers of different types have different sensitivities to feature caching. Especially, the error on the token in self-attention layers can quickly propagate to other tokens, due to the property that each token in self-attention layers can attend to the tokens. naive solution is to set lower cache ratios for self-attention layers. However, we observe that even if only smaller ratio of tokens cached, the error introduced by these tokens still quickly propagates to all other tokens, and has almost the same negative influence as caching all the tokens. Based on this fact, we propose to cache all tokens in self- elffaatennionn layers. For MLP and cross-attention layers, Ttype is set to the ratio of their computation costs over the overall computation costs. This strategy encourages layers with more computation costs to have a higher cache ratio.
4 EXPERIMENT
4.1 EXPERIMENT SETTINGS
Model Configurations We conduct experiments on three commonly-used DiT-based models across different generation tasks, including PixArt-a (Chen et al.) (024a) for text-to-image generation,
OpenSora (Zheng et al. 2024) for text-to-video generation, and DiT-XL/2 (Peebles & Xie (023) for class-conditional image generation with NVIDIA A800 80GB GPUs. Each model utilizes its default sampling method: DPM-Solver++ (Lu et al. 2022b) with 20 steps for PixArt-a, rflow (Liu et al. 2023) with 30 steps for OpenSora and DDPM (Ho et al. 2020) with 250 steps for DiT-XL/2.
For each model, we configure different average forced activation cycles A and average caching ratios R. for ToCa follows: PixArt-o: N 3 and R 70%, OpenSora: N = for temporal
7

---

Published as a conference paper at ICLR 2025
Table Qualitative comparison of text-to-image generation on MS-COC02017 and Par- tiPrompts with PixArt-a and 20 DPM++ sampling steps by default.
MS-COC02017
PartiPrompts CLIP
Method
Latency(s) | FLOPs | Speed
FID-30k J CLIP
PixArt-a CChnn al.
2024a)
0.682
11.18
1.00x
28.09
16.32
16.70
50% steps
0.391
5.59
2.00x
37.46
13.85
16.37
FORA(N =2) Seelaraauu et al.
2024)
0,416
5.66
1.98x
29.67
16.40
17.19
FORA(N =3) (Selvaraju al.
2024)
0.342
4.01
2.79x
29.88
16.42
17.15
ToCa (N -3, 60%)
0.410
6.33
1.77x
28.02
16.45
17.15
ToCa (N = 3, 70%)
0.390
5.78
1.93x
28.33
16.44
17.75
ToCa (N = 3, R 80%)
0.370
5.05
2.21
28.82
16.44
17.83
ToCa (N 3, R 90%)
0.347
4.26
2.62x
29.73
16.45
17.82
Table 2: Quantitative comparison in lext-to-video generation VBench. *Results are from PAB
Zhao et al. 2024). PAB1-3 indicate PAB with different hyper-parameters.
Method
Latency(s)
FL.OPs(T) 4
Speed
VBench %)
OpenSora
Zheng al.
2024)
81.18
3283.20
1.00x
79.13
A-DIT
Chen et al.
2024b
79.14 14
3166.47
1.04x
78.21
T-GATE
(Zhang et al.
2024b
67.98
2818.40
1.16x
77.61
PAB 1x
Zhao et al.
2024
60.78
2657.70
1.24x
78.51
PAB 2x
Zhao et al.
2024
59.16
2615.15
1.26x
77.64
PAB
Zhao et
2024)
56.64
2558.25
1.28x
76.95
50% steps
42.72
1641.60
2.00x
76.78
FORA al. 2024)
49.26
1751.32
1.87x
76.91
ToCa I = 80%)
43.52
1439.70
2.28x
78.59
ToCa R 85%)
43.08
1394.03
2.36x
78.34
attention, spatial attention, MLP, and N = 6 for cross-attention, with R 85% exclusively for MLP, and DIT: A = and R 93%. Please refer to the appendix for more implementation details.
Evaluation and Metrics For text-to-image generation, we utilize 30,000 captions randomly se- lected from COCO-2017 (Lin et al., 2014) to generate equivalent number of images. FID-30k computed to assess image quality, while the CLIP Score (Hessel et al. (02T) is used to evaluate the alignment between image content and captions. In the case text-to-video generation, we leverage the VBench framework (Huang et al. 2024), generating 5 videos for each of the 950 benchmark prompts under different random seeds, resulting in a total of 4,750 videos. The generated videos comprehensively evaluated across 16 aspects proposed in VBench. For class-conditional image generation, we uniformly sample from 1,000 classes in ImageNet (Deng et al. 2009) to produce 50,000 images resolution of 256 x 256, evaluating performance using F1D)-50k (Heusel al., 2017). Additionally, we employ sFID, Precision, and Recall as supplementary metrics.
4.2 RESULTS ON TEXT-TO-IMAGE GENERATION
In Table compare ToCa configured with parameters achieve acceleration ratio close to 2.0, against two other training-free acceleration approaches: FORA (Selvaraju et al. 2024), recent cache-based high-acceleration method, and the 10-step DPM-Solver++ sampling (Lu al. 2022b). In terms of generation quality, the quantitative results demonstrate that ToCa achieves the lowest FID among the compared acceleration methods while maintaining a high acceleration ratio. Figure also illustrates that our generated results most closely resemble those of the original PixArt-a. Regarding generation consistency, Table demonstrates that ToCa achieves the highest CLIP score on both MIS-COCO201700 7 et al.) 2014) and the PartiPrompts Yu et al. (022). Figure 6 shows that ToCa generates images that align more closely with the text descriptions compared to other methods. This is particularly evident in the fourth case, where only ToCa successfully generates an image matching "a small rowboat tied to a wooden dock", while other methods fail to generate the content of "a wooden dock". This may be caused by cross-attention
82 in ToCa that ensures the frequent refreshing of tokens that are semantically relevant to the text descriptions, resulting in generated images with enhanced semantic consistency the text prompts.
8

---

Published as a conference paper at ICLR 2025
staps 1.0
ToCa
FORA 1.9
30% steps 2.0
steps 1.0
* 24
x 1.9
50% steps
x 2.0
Figure 7: Visualization examples for different acceleration methods on OpenSora. Please kindly refer the supplementary material or our web page for viewing these videos.
Notably, Table shows that under similar acceleration ratios, ToCa exhibits very marginal de- crease in generation quality. In contrast, directly halving the number timesteps leads to 9.37 increments in FID, indicating significant performance drop. This observation indicates that when the number of sampling steps is already relatively low (e.g., 20 steps). further reduction in the num- ber of sampling steps may severely compromise the generation quality. In contrast, ToCa remains effective, demonstrating the distinct advantage of ToCa in the situation low sampling steps.
4.3 RESULTS ON TEXT-TO- VIDEO GENERATION
Austhetic Quality
We compare ToCa with adjusted rflow sampling steps from 20 to 10, alongside other acceleration methods including FORA. PAB (Zhao al., 2024).
Imaging Quality
Speed Up
Overall
Object
A-DiT (Chen et al. 2024b), and T-GATE ZZhang al. 2024b) using OpenSora (Zheng et al. 2024) for text-to-video generation. As presented in Ta- ble [ the experimental results show that ToCa achieves an impressive VBench score offering the lowest computational cost and highest inference speed among all methods tested. The VBench score of the 2.36x accelerated ToCa scheme drops by only 0.79 compared the non-accelerated scheme, while FORA's score decreases by 2.22, resulting in a 64.4% reduction in quality loss. Additionally, more VBench metrics results are presented in Figure ( which illustrate that ToCa significantly speeds up the original OpenSora with only slight performance degradation on few metrics. Notably, ToCa stands out as the sole acceleration method achieving nearly overall consistency performance with the original OpenSora, clearly outperforming another cache- acceleration method FORA. This again highlights the effectiveness of our proposed cross-attention-bal suring that the generated videos highly aligned with the text descriptions. We further present some video generation results in Figure 7 where we observe that the visual fidelity, and overall consistency of ToCa are closest the original OpenSora. Please kindly refer to our video demos in the supplementary material or the web page for viewing the video demo
Consistency
x1.9
Class
x1.00
Temponal Style
Multiple
Objects
Appearence Style
Ohiginal Model
FORA Poca (Ouss)
Human
50% Timalipp
Action
Soana
Spatial Relationship
Figure VBench metrics and acceleration ratio of proposed ToCa and other methods
token selection strategy, en-
4.4 RESULTS ON CLASS- CLASS-CONDIIIONAL IMAGE GENERATION
Quantitative comparison between ToCa with other training- freee DiT acceleration methods is shown in Table which demonstrates that ToCa outperforms other methods in terms of both FID and sFID by a clear margin under the similar acceleration ratio. For instance, ToCa leads to 0.39 and 0.22 lower values in sFID and FID compared with FORA with similar acceleration ratio, respectively.

---

Published as a conference paper at ICLR 2025
Table 3: Quantitative comparison on class-to-image generation on ImageNet with DIT-XL/2.
Method
Latency(s)
FL.OPs(T) 4 Speed
sFID 4
FID +
Precision
Recall
DIT-XL/2-G (cfg 1.50)
2.012
118.68
1.00x
4.98
2.31
0.82
0.58
33% steps
0.681
39,40
3.01x
6.31
2.76
0.81
0.57
37% steps
0.748
44.15
2.69x
6.04
2.64
0.81
0.58
FORA(N = 3)
0.807
39.95
2.97
6.21
2.80
0.80
0,59
FORA(N = 2.8)
0.815
43.36
2.74x
6.13
2.80
0.80
0.59
ToCa (AV R 93%)
0.820
43.22
2.75x
5.74
2.58
0.81
0.59
Table Ablation studies with DiT-XL/2-G (cfg- 1.50). $1 is used in all experiments. is not used since DiT does not have cross-attention layers.
Ablation studies of token selection based on different attention scores (31 and 82) with PixArt-a. "Random" indicates re- placing attention scores with random values.
Table 5:
R Schedule SenfouuleUniffrrm
Uniform Spatial Cache
ImageNet
$3, 84, f2, Tiype are used in all three settings.
Tupe
Distribution 84 Frequency FID-5k
Token Selection MS-COC02017 Pato0oooo0PPartiPomopopo PartiPrompts Methods
9.32
x
/
9.60
FID-30k
CLIP 1
/
x
/
9.67
Cross-Attention $2
28.33
17.75
Self-Attention $1
28.21
17.13
/
X
9.35
/
9.65
Random
28.46
17.08
Ablation Study Table 4 presents the effect of the two factors on adjusting the caching ratio dif- ferent layers, where applying different caching ratios to layers in different types (Tiype) and different depths (ri) leads to 0.28 and 0.35 FID reduction, respectively. Besides, using the score of uniform spatial distribution (sy) and cache frequency (s,) reduces FID by 0.02 and 0.33, respectively. Table compares the influence of selecting tokens with the self-attention weights ($1) and the cross-attention weights (82). The other ToCa modules including 83, 84, Tt, "type are utilized in these experiments. It observed that 81 tends to achieve lower FID while 82 tends to reach a higher CLIP score, which reasonable since self-attention is mainly utilized for the generation of the overall images while cross-attention utilized to inject the conditional signals. summary, these results demonstrate that all the cache scores in ToCa have their benefits in different dimensions.
5" 150
w/o with
e 125
x 100
1000
46t21 2000 4000
5000
(c) Distributions the eache times of each token (with and w/o
(a) wio Cache Frequency
with Cache Frequency Score S,
Figure 9; Visualization of cached tokens selected with and without Sy (cache frequency). The pixel with darker color indicates the corresponding tokens are more frequently cached, (c) The distribution of the number of being cached for each token w/ and w/o
Visualization on the Cached Tokens Figure (a-b) show the number of times that each token cached during generation, where darker colors indicate more frequent caching. It is observed that both the two schemes perform more cache in the unimportant background tokens while performing more real computations in the tokens of the Arctic fox. However, the image without 83 has bad quality in the background since the background tokens have been cached too many times. contrast, applying the score of cache frequency 83, which aims stop caching the tokens that have been cached in the previous layers, can reduce the gap between the important and unimportant tokens, and prevent the background tokens from overlarge caching frequency. This observation has also been verified in Figure c) that reduces the number of tokens cached by more than 4.5k times.
CONCLUSION
Motivated by the observation that different tokens exhibit different temporal redundancy and dif- ferent error propagation, this paper introduces ToCa, token-wise feature caching method, which adaptively skips the computation of some tokens by resuing their features in previous timesteps. By leveraging the difference in different tokens, ToCa achieves better acceleration performance compared with previous caching methods by a clear margin in both image and video generation, providing insights for token-wise optimization in diffusion transformers.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, al. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01322 2022.
Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffusion: Scaling latent video diffusion models to large datasets. arXiv preprint arXiv:231 1./5127, 2023.
Daniel Bolya and Judy Hoffman. Token merging for fast stable diffusion. Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 4599-4600 2023.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Erie Luhman, Clarence Ng. Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. URL attts:://ppeaae com/research/
v.deo--eeerraiion-mmeee.s -atp::/oooeeeee-------aaa--aaaa- simulratsssmmmmorr
Junsong Chen, Jincheng Yu, Chongjian Ge. Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang. James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-image synthesis. In International Conference on Learning Representa- fions, 2024a.
Pengtao Chen, Mingzhu Shen, Peng Ye, Jianjian Cao, Chongjun Tu, Christos-Savvar Bouganis, Yiren Zhao, and Tao Chen. 8-dit: A training-free acceleration method tailored for diffusion transformers. arXiv preprint arXiv:2406.00 125, 2024b.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scall scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248 255. leee, 2009.
Gongfan Fang, Xinyin Ma, and Xinchao Wang. Structural pruning for diffusion models. arXiv preprint 2rXx 2023.
Jack Hessel, Ari Holzzaan, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning. arXiv preprint arXiv:2104.08718, 2021.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-68.. 2020.
Ziqi Huang, Yinan He, Jiashuo Yu, Fan Zhang, Chenyang Si, Yuming Jiang, Yuanhan Zhang, Tianx- ing Wu, Qingyang Jin, Nattapol Chanpaisit, et al. Vbench: Comprehensive benchmark suite for video generative models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 21807-2188 2024.
PKU- Yuan Lab and Tuzhan AI etc. Open-sora-plan, April 2024. URL https:://bi..og/////
5281/ 10948109
Black Forest Labs. Flux. http::///gtthhaa..ccm///lheeeess
2024.
Senmao Li, Taihang Hu, Fahad Shahbaz Khan, Linxuan Li, Shigi Yang, Yaxing Wang, Ming-Ming Cheng, and Jian Yang. Faster diffusion: Rethinking the role of unet encoder in diffusion models. arXiv preprint 2XXv:2312.09608. 2023a.
Xiuyu Li, Yijiang Liu, Long Lian, Huanrui Yang, Zhen Dong, Daniel Kang, Shanghang Zhang, and Kurt Keutzer. Q-diffusion: Quantizing diffusion models. In Proceedings of the IEEF/CVF International Conference Computer Vision, pp. 17535-17545, 2023b.
Yanyu Li, Huan Wang, Qing Jin, Ju Hu, Pavlo Chemerys, Yun Fu, Yanzhi Wang, Sergey Tulyakov, and Jian Ren. Snapfusion: Text-to-iman -mage diffusion model on mobile devices within two seconds. Advances in Neural Information Processing Systems, 36, 2024.
11

---

Published as a conference paper at ICLR 2025
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740-755. Springer, 2014.
Xingehao Liu, Chengyue Gong, et al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022a.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:221..01095, 2022b.
Xinyin Ma, Gongfan Fang, Michael Bi Mi, and Xinchao Wang. Bearning-to-ccchee Accelerating diffusion transformer via layer caching, arXiv preprint arXiv:2406.01733, 2024a.
Xinyin Ma, Gongfan Fang. and Xinchao Wang. Deepcache: Accelerating diffusion models for free. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 15762-15772, 2024b.
Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 14297-14306, 2023.
William Peebles and Saining Xie. Scalable diffuusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195 4205, 2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and computer-assisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 234-241. Springer, 2015.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.005: 00512, 2022.
Pratheba Selvaraju, Tianyu Ding, Tianyi Chen, Ilya Zharkov, and Luming Liang. Fora: Fost-forward caching in diffusion transformer acceleration. arXiv preprint 2070v:2407.01425, 2024.
Yuzhang Shang, Zhihang Yuan, Bin Xie, Bingzhe Wu, and Yan Yan. Post-training quantization on diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1972-1981, 2023.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2021.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In International Conference on Machine Learning. pp. 32211-2222222 -22522. PMLR, 2023.
Felix Wimbauer, Bichen Wu, Edgar Schoenfeld, Xiaoliang Dai, Ji Hou, Zijian He, Artsiom Sanakoyeu, Peizhao Zhang, Sam Tsai, Jonas Kohler, et al. Cache me if you Accelerat- ing diffusion models through block caching. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 621 1-220, 2024.
12

---

Published as a conference paper at ICLR 2025
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang. Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al. Scaling autoregressive models for content- rich text-to-image generation. arXiv preprint ar arXiv:2206. 10789, 2(3):5, 2022.
Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, and Linfeng Zhang. Token pruning for caching better: 9 times acceleaation on stable diffusion
for free. ArXiv, abs/2501 .00375, 2024a. URL https: /atps:///ppi.. www.nticcccooooe..e........c..
CorpusID:27521133
Evelyn Zhang, Jiayi Tang, Xuefei Ning, and Linfeng Zhang. Training free and gardware-friendll acceleration for diffusion models via similarity- -based token pruning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2025.
Wentian Zhang, Haozhe Liu, Jinheng Xie, Francesco Faccio, Mike Zheng Shou, and Jürgen Schmid- huber. Cross-attention makes inference cumbersome in text-to-image diffusion models. arXiv preprint arXiv:2404.0274770 2024b.
Xuanlei Zhao, Xiaolong Jin, Kai Wang, and Yang You. Real-time video generation with pyramid attention broadcast. arXiv preprint arXiv:2408. 12588, 2024.
Zangwei Zheng, Xiangyu Peng, Tianji Yang, Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all, March 2024. URLhttpps //tptt/tt.tt././///// thub. http://////......ocoo/oaaaaaaaee cowheepcaaaeeeccchheee----0
A APPENDIX
A.I ENGINEERING DETAILS
This section introduces some engineering techniques our work.
A.1.1 STEP-WISE CACHING SCHEDULING
In section 3.6 we propose method for dynamically adjusting the caching ratio R based on the time redundancy and noise diffusion speed across different depths and types of layers, which constitutes key part of our contributions. In the following section, we further explore the dynamic adjustment of R along the timestep dimension, as well as strategies for dynamically adjusting the forced activation cycle N.
At the initial stages of image generation, the model primarily focuses generating contours, while in the later stages, it pays more attention to details. In the early contour generation phase, it is not necessary for too many tokens to be fully computed with high precision. By multiplying by term rt, we achieve dynamic adjustment of along the timestep dimension, where r = 0.5 (0.5 A is a positive parameter controlling the slope, l is the number of timesteps already processed, and T' is the total number of timesteps. By adjusting R in this way, we shift some of the computational load from earlier timesteps to later ones, improving the quality of the generated images. Finally, the caching ratio is determined Riy R x ri Tyype X Tt.
Similarly, we set larger forced activation cycle N during the earlier stages, while a smaller N used during the later detail generation phase to enhance the quality of the details. To ensure that the adjustment of / has minimal impact on the theoretical speedup, we define it follows: N₂ = No/ (0.5 (t/T where No corresponds to the expected theoretical speedup induced by AV, and w, is hyperparameter controlling the slope.
A. 1.2 PARTIALLY COMPUTATION ON CELF-ATTENTION
In the previous section, mentioned that partial computation in the Self-Attention module can lead to rapid propagation and accumulation errors. Therefore, we considered avoiding partial computation in the Self-Attentin module, meaning that during the non-forced activation phase, the Self-Attention module has Ttype 0. In the subsequent Sensitivity Study, we explored trade-off scheme between Attentitnn and MLP modules, with the corresponding formulas for allocation
13

---

Published as a conference paper at ICLR 2025
being Ptype 0.4Xtype for the Self-Attention module, and Ptyee 1 0.6Atype for the MLP module. The factors 0.6 and 0.4 are derived from the approximate computational ratio between these two modules in the DiT model.
A.1.3 TOKEN SELECTION FOR CFG AND NON-CFG
In the series of DiT-based models, the tensors of cfg (class-free guidance) and non-cfg are concate- nated along the batch dimension. A pertinent question token selection is whether the same token selection strategy should be applied to both the cfg and non-cfg parts for the same image (i.e., if a token is cached in the cfg part, should also be cached in the corresponding non-cfg part). We have observed significant sensitivity differences among models with different types of conditioning regarding whether the same selection strategy is used. For instance, the text-to-image and text- to-video models, such as PixArt-a: and OpenSora, if independent selection schemes are applied for the cfg and non-cfg parts, the model performance degrades substantially. Thus, it is necessary to enforce consistent token selection scheme between the cfg and non-cfg parts.
However, in the class-lo-image DiT model, this sensitivity issue is considerably reduced. Using independent or identical schemes for the cfg and non-cfg parts results in only minor differences. This can be attributed to the fact that, in rext-conditional models, the cross-attention module injects the conditioning information into the cfg and non-cfg parts unevenly, leading to significant disparity in attention distribution between the two. Conversely, in class-conditional models, the influence on both parts is relatively uniform, causing noticeable changes in token attention distribution.
A.2 MORE IMPLEMENTATION DETAILS ON EXPERIMENTAL SETTINGS
For the DiT-XL/2 model, we uniformly sampled from 1,000 classes in ImageNet (Deng et al. 2009) and generated 50,000 images with resolution of 256 x 256. We explored the optimal solution for DiT-XI//2 using FID-5k (Heusel et al.] (017) and evaluated its performance with FID-50k. Addi- tionally, sFID, Inception Score, and Precision and Recall were used as secondary metrics. For the PixArt-a model, we used 30,000 captions randomly selected from COCO-2017 (Lin et al. 2014) to generate 30,000 images. We computed FID-30k to assess image quality and used the CLIP Score between the images and prompts to evaluate the alignment between image content and the prompts. For the OpenSora model, we used the VBench framework (Huang al. (024), generating videos for each of the 950 VBench benchmark prompts under different random seeds, resulting in total of 4,750 videos. These videos have a resolution of 480p, an aspect ratio of 16:9, a duration of 2 seconds, and consist of 51 frames saved at frame rate of 24 frames per second. The model was comprehen- sively evaluated across 16 aspects: subject consistency, imaging quality, background consistency, motion smoothness, overall consistency, human action, multiple objects, spatial relationships, ob- ject class, color, aesthetic quality, appearance style, temporal flickering, scene, temporal style, and dynamic degree.
We set the average forced activation cycle of ToCa to N = 2, supplemented with a dynamic adjustment parameter Wt 0.1. The parameter X 0.4 adjusts R at different time steps, and the average caching ratio is R = 70%. The parameter rı = 0.3 adjusts R at different depth layers. The module preference weight Ttype 1.0 shifts part of the computation from cross- attention layers to MLP layers.
PixArt-a:
For OpenSora, we fixed the forced activation cycle for temporal attention, spatial atten- tion, and MLP at 3, and set the forced activation eycle for cross-attention to 6. The ToCa strategy ensures that a portion of token computations conducted solely in the MLP, with Rmip fixed at 85%.
OpenSora:
DIT: We set the average forced activation cycle of ToCa to = 3, supplemented with a dynamic adjustment parameter w, = 0.03 to gradually increase the density of forced activations as the sam- pling steps progress. The parameter X, 0.03 adjusts It at different time steps. Additionally, during the sampling steps in the interval E [50, 100). the forced activation cycle is fixed at // = 2 to promote more thorough computation in sensitive regions. The average caching ratio is R = 93%, and the parameter \ 0.06 adjusts R at different depth layers. The module preference weight Trype = 0.8 means that during steps outside the forced activation ones, no extra computations are performed in attention layers, but additional computations are performed in the MI.P layers.
14

---

Published as a conference paper at ICLR 2025
Originel 1.0
ToCa
33.2% steps 30
FORA 3.0
Figure 10: Visualization examples for different acceleration methods on DIT.
All of our experiments were conducted on 6 A800 GPUs, each with 80GB of memory, running CUDA version 1..1. The DiT model was executed in Python 3.12 with PyTorch version 2.4.0, while PixArt-a and OpenSora were run in Python 3.9. The Py version for PixArt-a was 2.4.0, and for OpenSora it was 2.2.2. The CPUs used across all experiments were 84 vCPUs from an Intel(R) Xe) Gold 6348 CPU 2.60GHZ..
A.3 SENSITIVITY STUDY
We explored the optimal parameter configuration for the ToCa acceleration scheme on DiT and analyzed the sensitivity of each parameter. The experiments used FID-5k and sFID-5k as evaluation metrics. From Figure TT (a) to (f), we respectively investigated the effects of the caching ratio weights A₁. Atype. dІ. the weight of Cache Frequency score λ3, Uniform Spatial Distribution λs, and the dynamic adjustment weight for forced activation W2, It is observed that: (a) The optimal parameter X₂ 0.06, where the corresponding cache ratio shows approximately 6% variation 2.5, at which point the
at both the last and first layers. (b) The optimal parameter is Atype Self-Attention layer does not perform any partial computation, with the entire computational load shifted to the MLP layer. is also noted that as the computation load decreases in the Self-Attention layer and increases the MLP layer, the generation quality shows a steady improvement. (c) The optimal parameter in the figure is X 0.03, and at this point, there is little difference between the best and worst methods, suggesting that the model is not particularly sensitive to the adjustment of cache ratio along the timesteps.
(d) The optimal weight for the Cache Frequency score is X 0.25. We observe that as As increases, the model's generation quality initially shows noticeable improvement, but beyond a weight value of 0.25, the fluctuation is minimal. This indicates that the Cache Frequency has reached a relatively uniform state, achieving a dynamic balance in caching among different tokens. (e) We conducted search for the Uniform Spatial Distribution score with grid sizes of 2 and 4, and the experimental results show that the generation quality with a grid size of is generally better than that with a grid size of 4. This suggests that a finer-graine spatial uniformity indeed contributes to improvement in generation quality (f) We explored the impact of dynamically adjusting the forced activation cycle on the model's generation quality and analyzed the effect of fixing the Force activation cycle 3 at 2 for the relatively more sensitive 50-100 timesteps. The experimental results show that enforcing this fixed cycle in the 50-100 timesteps significantly improves generation quality, with the optimal parameter configuration being w; = 0.4.
In summary, these observations indicate that our method is not sensitive to the choice of hyper- parameters. Actually, experiment results demonstrate that stable performance when directly transferring hyper hyper-parameters parameters from one model to another model in the same model family such DiT in different sizes.
be observed
15

---

Published as a conference paper at ICLR 2025
100
370
410
fD
38.3
3.5
3.0
16.0
0.
30
16.00
C.0 Step-depth weight Ay (c)
16.0
Leyer-depth weight A (a)
Layer- type weight Atype (b)
10.0
10.7
den
10.0
paac 10-10
und ao00
1575
16.93
63
02
04
2-
1.0
16.00
0.
02
C4
0. 06
1.0
360
50
52
0.5
Cache Frequency weight Ag (d)
Uniform Spatial Distribution Ag
Force Activation
Cycle weight Wr (f)
(e)
Figure 11: Sensitivity study on different weights. From (a) (f), the caching ratio weights At. Atgpes X the Cache Frequency score weight A the Uniform Spatial Distribution weight λ4, and the dynamic schedule weight for forced activation wy are presented.
A.4 COMPUTATION COMPLEXITY ANALYSIS
A.4.1 MAIN COMPUTATIONS
Complexity of Attention Layer. the Attention layer, tokens are first processed through linear layer to generate queries, keys, and values. Next, the queries and keys are multiplied using dot product, passed through a softmax function, and then multiplied with the values. Finally, the result is passed through another linear layer to produce the output. The computational cost of the Self- Attention layer can expressed as:
FLOPsSA NN x D 3
x
X 2+ 2 X D x H x xxx
=8 N x D
+4x x D+5x x H ~ O(ND*) + O(N2D),
(4)
where N is the number of tokens, D is the hidden state of each token and H is the number of heads, (H < D). The computational cost Cross-Attention layer can be expressed as:
FLOPSSAA N D D x 2 N₂ D x D + N X N₂ x D + N X N₂ X H x 5 + Ni x N₂ x x 2 + Ni x D X D X2
(5)
=4 X (N1 N₂) X D +4x N₁ x N₂ x D + N1 N₂ H
~0(N + D2) O(N, O(N N2) + O(NN,D)
where N1 = N₂ are the number of image and text tokens, D is the hidden state of each token and H the number of heads, (H < D). In the previous computations, the softmax operation approximated as involving 5 floating coati calculations per element.
Complexity of MLP Layer. The computational cost of ML.P layer can be written as:
FLOPS N D1 x D2 x 2 x D2 x N x D1 x D2 x
=4 X N X D1 D2 + 6 D2 x N 16 N 2 24 x N D O(ND²),
(6)
where N is the number of tokens, D D and D2 4D, are the hidden and middle-hidden state in MLP, respectively. The activation function for MLP is approximated to involve 6 floating-poin operations per element.
A.4.2 COMPUTATION COSTS FROM TOKEN SELECTION
score 81. As mentioned in section 3.5 the Self-Attentin score is computed as (1t) Σ Oijy, where 2; is the ith token, and ai is the element in the self-attention map.
16

---

Published as a conference paper at ICLR 2025
Table 6: Quantitative comparison on class-to-image generation on ImageNet with 50 steps DDIM sampler baseline DiT-XL/2.
Method
Latency(s) FL.OP(()) +
Speed t
sFID
FID +
Precision
Recall
DIT-XL/2-G (efg = 1.50)
0.455
23.74
1.00x
4.40
2.43
0.80
0.59
50% steps
0.238
11.86
2.00x
4.74
3.18
0.79
0.58
40% steps
0.197
9.50
2.50x
5.15
3.81
0.78
0.57
34% steps
0.173
8.08
2.94x
5.76
4.58
0.77
0.56
FORA(N 225)
0.219
10.48
2.27x
6,59
3.83
0.79
0.55
FORA(N =3)
0.211
8.58
2.77
6.43
3.88
0.79
0,54
ToCa (A =3. 3, R. 93%)
0.227
10.23
2.32x
4.74
3.04
0.80
0.57
ToCa (AV = =4, R- 93%)
0.209
8.73
2.72x
5.11
3.60
0.79
0.56
Therefore, the computational complexity of the self-attention score only O(N). In a prac- tical case achieving about 2.3x acceleration, the computation cost of the Self-Attention score ac- counts for approximately 0. 2828% of the main components.
Cross-Attention score $2. As mentioned in section 3.5 the Cross-Attention score is computed as 52(x;) = - Ciy log(cig), where the Ct is the element in the cross-attention map. Therefore, the computational complexity the cross-attention score is only 2.V O In a practical case achieving about 2.3x acceleration, the computation cost of the Cross-Attention score accounts for approximately 0.35% the main components.
Cache Frequency score 83 and Uniform Spatial Distribution score 94. The Cache Frequency score is updated at each step, so its update cost per timestep is N. When the Cache Frequency practical applications, the computation cost is 2N. Thus, the total
score is called for summation cost for one layer is 3N = O(N). The Uniform Spatial Distribution score Są is computed by sort- ing within each grid of size G G and weighting the top-scoring tokens. The computation cost given by N the computational complexity of 84 O(N). a practical case achieving about 2.3x accelera- tion, the computation cost of the Cache Frequency score 83 accounts for approximately 0.044% and the Uniform Spatial Distribution score 84 accounts for 0.15% the main computation compo- nents. In addition, the computational cost for sorting N tokens is O(N log N), which accounts for ppproxximately0.1899 of the main computational cost.
G C log(G" (o)) + 2 x GT, where G is the grid size, which is usually small. Therefore
N
In summary, although ToCa introduces additional computations, its computational complexity of O(N) is negligible compared to the main computational modules with complexities of O(N2D) or O(ND²). In practical tests, the time taken for token selection is minimal, typically less than 1% of the main computational cost. At cache steps, taking a caching ratio of R. 90% as an example, the computational cost of terms with complexity of O D2) is reduced 10% of the original, while the computational cost of terms with complexity of O(N2D) D) is reduced to 1%. (However, as mentioned earlier, in practice, it is more efficient shift all computations at cache steps to the MLP. Therefore, all terms with a complexity of 2 at cache steps are ignored.)
A.5 IMPLEMENTED RESULTS ON CLASS-CONDITION IMAGE GENERATION
In addition to the series of experiments conducted using the DDPM(Ho et al., 2020 sampling method on DIT, which have already been included in the main part, we also performed validation with the more practically relevant DDIM et al. 2021) sampling method to further demonstrate the effectiveness of ToCa as shown in Table
For instance, ToCa leads to 1.32 and 0.28 lower values in sFID and FID compared with FORA with a similar acceleration ratio of approximately 2.7x, respectively, and achieves 0.21 lower values in FID compared to the method of directly reducing the sampling steps with acceleration ratio of approximately 2.5 x. As trade-off between acceleration and performance, we selected the scheme N 3, R 93% the final recommended approach for DDIM sampler.
17

---

Published as a conference paper at ICLR 2025
Xt-1
1 In the fresh step, the model performs full computation all tokens in all layers and updates the computed results into the cache.
Cache after timestep t:
Other Layers
Init. Cache
Neural layer
Fetch token 1 From Cache
Neural layer
Cache after Limestep t-1
Feature Computed at
Computed
Compute Cache Score
X
Compute
Cache
Xt=1
tipestep
In the cache step, the model firstly compute the importance score, deciding the takens to be cached. example, token and cached.
Cachi a Lfter timestep t:
OOhhe Layers
Init.
Neural layer
Fetch taken 1,3 From Cache
Neural layer
Cache alter timested
Feature Compu bed at
Computed at
Compute Cache Score
X
Compute
Cache
t-1
timescep
-imestep
Xt-1
Then, the tokens 2,4,and are computed through the neural layer. The output of cached 1,3 fetched from cache. the Caobe
Cache after timestep t:
OOhhe Layers
Init.
Cache
Neural layer
Fetch taken 1,3 from Cache
layer
Cache after timesteg t-1
Feature Computed at e
Feature Computed at
Compute Cache Score 0.3 0.7 0.5 0.8 0.6
Compute
X t
Limestep t
timestep ビ-1
X Xt-2
4
Cache after timeatep
Other Layers
Other Layers
Update token 2,4,5 to the Cache
Init Cache
Neural layer
Fetch Laken 2
Neural layer
Cache after Limestep t-1
From Cathee
Fea bureure Computed at
Then, calculated tokens 2,4,and are used update the cache.
Feature Computed
X
Computa
Cache
timestep
Figure 12: A more detailed workflow for the proposed ToCa.The cache-and-reuse procedure is conducted on the model at all layers and timesteps.
18

---

Published as a conference paper at ICLR 2025
Table Quantitative comparison in text-to-image generation for FLUX on Image Reward.
Method
Latency(s)
FLOPs(T) L
Speed
Image Reward f
FLUX.1-dev Labs 2024
33.85
3719.50
1.00x
1.202
68% steps
23.02
2529.26
1.47 x
1.200)
FORA (Selvaraju et al.
2024
20.82
2483.32
1.51
1.196
ToCa = 2, R 90%)
19.88
2458.06
1.51x
1.202
FLUX.I-schner Labs
2024
2.882
277.88
1.00x
1.133
75% steps
2.162
208.41
1.33x
1.139
FORA
Selvaraju ct al.
2024
2.365
225.60
1.23x
1.129
FORA 2
Selvaraju et
2024
2.365
225.60
1.23x
1.124
FORA 3
Selvaraju al.
2024
2.365
225.60
1,23X
1.123
ToCa
2, R 90%
1.890
181.30
1.53x
1.134
A.6
RESULTS ON
AND MORE ADVANCED TEXT-TO-IMAGE MODELS
As shown in Table 7 we compared the performance of FORA and ToCa in generating high- resolution images (1024 x 1024) using the more advanced text-conditional image generation models FLUX.I-dee and FLUX. LLUXX.-schnell(LL.. ter, as more efficient model, uses only 4 sampling steps. The evaluation generation quality was conducted using Image Reward, metric better suited to measuring human preference. The generated images were based on 1,632 prompts from the PartiPrompts Yu et al. 2022) dataset to comprehensively evaluate the generation quality of the acceleration methods both FLUX.I-dew and FLUX.1 --cchnell.
2024). The former uses 50 sampling steps, while the lat-
In this comparison, FORA (Selvaraju et al. 2024) represents acceleration the FLUX.1-dee model with 50 sampling steps, where caching is performed every other step (i.e., N 2). FORA! FORA?, and FORA* correspond skipping the 2nd, 3rd, and 4th steps, respectively, during the 4-step gen- eration process. ToCa demonstrated nearly lossless performance under a 1.5x acceleration, with Image Reward scores almost identical the non-accelerated scenario in both the 50-step FL.UX.1 dev and the 4-step FLUX. -cchhelll models. In contrast, all configurations of FORA showed quality degradation. For example, in the 50-step FLUX. 1-dev model, the 1.5x acceleration setting and in the 4-step FLUX. 1-schnell model, 1.2x acceleration setting both resulted in noticeable qual- ity degradation. The corresponding visual results for FLUX. 1-shhnell are presented in Figure B3 demonstrating the lossless acceleration capability of ToCa.
A.7 DETAILS FOR DISTRIBUTION FIGURES
In this section, we provide detailed explanations of the various distribution plots mentioned in the main text supplementary information.
Figure [ illustrates the distribution of the Frobenius norm of the differences between the feature maps o and of+1. which are the outputs of the last DiT Block each timestep and the previous timestep f 1, respectively. The figure also presents the corresponding statistical frequency density of these Frobenius norm values for each token, based on 500 randomly generated samples produced by DIT. This analysis reveals the conclusion that different tokens exhibit varying levels of temporal redundancy.
Figure illuustrass the varying rates of сттоr accumulation and propagation across different tokens. Specifically, Gaussian noise with an intensity of 0.5 was independently added to the ith token of the first layer at each step. The Frobenius norm was then computed between the output features of all tokens at the last layer the same step and the corresponding features from the noise-free output. This process was repeated for all steps and all layers. Given that each noise propagation required re-running the inference process, a random subset of 100 samples from the DiT model was selected for this study, and the noise propagation results were recorded for each iteration. This analysis led to the conclusion that different tokens exhibit varying rates of errr accumulation and propagation.
19

---

Published as a conference paper at ICLR 2025
owl standing telephone
"a ceiling fan with four white blades"
"A spaceship made cardboard.
teapor"
"a half sallfppeeed banana
Original x1.00
ToCa x1.53
"a tiny football in front of three yellow tewnis bolls"
"three small yellow boxes on large blue box"
"We painted the wall with cracks."
"a red circle" "
"A fog.
Original x1.00
ToCa
x1.53
"five red balls on table"
"A Vietnam map"
"element"
map Texas"
"an F1
Original x1.00
VETHAA
ToCa x1.53
Figure 13: Visualization examples for original FLUX. Labs 2024) and proposed ToCa with almost lossless acceleration.
20

---

Published as a conference paper at ICLR 2025
Self-Attention
MLP
Compun Prattlo ratto
18cR8 stepe
Layye
40
30
20
50
30
10
TimeStep
TimeStep
(a)
(b)
Figure 14: Computation ratio for different types of computation layers on different timesteps and layers on DiT. The dark blue lines correspond to fully computed fresh steps. (a) The computation ratio distribution of the Self- Belf-Attention layers. As mentioned earlier, performing partial computation on attention layers is less cost-effective compared to MLP layers. Therefore, we do not implement partial computation attention layers; apart from fresh steps, all other steps directly reuse the corresponding cached features. (b) The computation ratio distribution the ML. layers. As shown, the computation ratio inereases with deeper layers and as the number inferred timesteps during the inference phase grows.
Figure 5(a) illustrates the varying temporal redundancy across layers of different depths. For each timestep, the Frobenius norm of the differences between the features of the current timestep and the corresponding features of the previous timestep at a specific layer depth was computed for each token. The resulting Frobenius corresponding statistical frequency densities. To clearly demonstrate the trends, we selected layers 15, and 28 for visualization. The samples used were randomly chosen from 200 DiT samples.
values
then used plot the distribution alongside their
Figure b)) shows the variation in the offset distribution of the output values from the last layer of timestep when Gaussian noise is added single token different layer depths within the same timestep. This is measured by adding Gaussian noise with an intensity of 0.5 to single token in specific layer one timestep and comparing the deviation in the output features the last layer at that timestep with the noise-free scenario. For clarity, this operation was performed on layers 1, 15, and 25 across all timesteps, using 200 randomly selected samples from the DiT model to generate the examples. is worth noting that Figure b]b) may appear at first glance to violate the normalization condition for frequency density distributions. This is due to the large variations Frobenius norm values, which necessitated the of logarithmic scale on the horizontal axis. Figure 5]a) and (b) demonstrate two key conclusions: deeper layers exhibit poorer temporal redundancy, but errors introduced in deeper layers have smaller impact the output at the same timestep.
Figure ((c) illustrates the results on the PixArt-a model by adding Gaussian noise with an inten- sity of 0.5 x ||ж||| to single token IІ in the 10th layer (approximately the middle layer) at each timestep. This process was performed for three different types of layers (self-attention, cross- attention, and MLP). The Frobenius norm of the error induced by the noise measured the output of the last layer and normalized by the average Frobenius norm ||< of tokens of the same type. The resulting distribution was plotted using 200 prompts randomly selected from the MS-CO MS-CCCO2001 dataset. It is important to note that the additional normalization step, based on the norm values of the tokens, was necessary because the norm values of tokens in self-attention, cross-attention, and MIP layers typically vary significantly. Normalization ensures fair compari- son across these layer types. Additionally, the distribution for the MLP layer in Figure appears more dispersed. This is due to the generally larger variations in MLP output values across different prompts and timesteps. In practice, increasing the number of samples can make the distribution vi- sually denser. However, given that each token requires separate inference for the error propagation experiments, using 200 prompts already incurs significant computational cost while remaining sufficient to reveal the trends.
21

---

Published as a conference paper at ICLR 2025
Algorithm ToCa
Input: current timestep 1, current layer id I.
1: if current timestep t is fresh step then
2: Fully compute Fx)).
3:
(x) = F (2:); # Update the cache.
4: else
5: S(x,,) = -(x;)= ビÿ=
8 Compute the cache score for each token.
85;
6: ICompule == TopK(S(x.)) R%); Fetch the index of computed tokens.
7:
for all tokens 2; do
8:
i E Iommpute then
Compute (x,) through the neural layer. 10: (a;,)) F (x;); Update the cache.
9:
11: end if
12: end for
13:
end if
14:
return F()). # return features for both cached and computed tokens for the next layer.
ToCa
5.0
FORA
DDIM
4.5
FID- -50k
4.0
3.5
3.0
2.5
7.5
10.0
12.5
15.0
17.5
20.0
22.5
FLOPs (T)
Figure 15: Pareto curve with FLOPs-FID to better evaluate the performance of ToCa on DIT with 50 ddim sampling steps as baseline.
22


Reviews:
[{"review_id": "mEpgP2r391", "reviewer": "Reviewer_Vwce", "summary": "The paper introduces ToCa (Token-wise feature Caching), a training-free feature caching method tailored to accelerate diffusion transformers. ToCa allows for adaptive token selection, which aims to cache tokens with minimal caching error based on various criterions. Experiments have been conducted on PixArt-α, OpenSora, and DiT models to show the speedup and quality trade-off of ToCa.", "questions": "- In table 2, ToCa does not outperform PAB with 1.24x speedup (78.34 vs 78.51). Although ToCA achieves a more significant speed-up, it is crucial to preserve the generation quality. I am wondering if ToCa can have better generation quality using a more conservative setting that has similar speed-up to PAB?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The visualizations of temporal redundancy and error propagation illustrate the motivation of ToCa in an intuitive way, making the design choices relatively well motivated.\n- The four token selection scoring functions and layer-specific cache ratios are natural design choices that enable the caching strategy to be more fine-grained\n- ToCa achieves more than 2x acceleration ratios on both text-to-image and text-to-video tasks, while having better quality than baselines.\n- ToCa’s training-free nature makes it more practical for real-world application.", "weaknesses": "- While ToCa achieves reasonable benchmark results, some artifacts remain in the generated images compared to the originals. For instance, in Figure 6, the moon is missing from the \"wolf howling at the full moon\" prompt, and the background forests appear blurred in the \"tranquil beach\" prompt. It is necessary to demonstrate how ToCa performs with high-resolution images (1024x1024) generated by more advanced models like FLUX.1-dev [1].\n- Another important direction in accelerating diffusion models involves reducing the required sampling steps. However, this might compromise ToCa's effectiveness since fewer steps may result in greater feature variations between steps, potentially making them less suitable for caching. Please include experimental results using models with reduced sampling steps, such as FLUX.1-schnell [2] with 4 steps.\n- For advanced model evaluation, FID and CLIP scores may not adequately reflect human preferences. Please provide quantitative results using more recent metrics, such as image reward [3].\n- The method involves numerous hyperparameters, including four token selection scoring metrics ($\\lambda_1$, $\\lambda_2$, $\\lambda_3$, $\\lambda_4$) and caching ratios for different layers and timesteps ($\\lambda_l$, $\\lambda_{type}$, $\\lambda_t$). While A.3 provides a sensitivity analysis, it remains unclear how these parameters should be determined for new models or tasks. The authors should elaborate on hyperparameter selection and impacts, particularly addressing whether an automatic selection method (e.g., a calibration procedure) exists. The complexity of hyperparameters could limit the method's generalizability and practical application.\n- The \"Acceleration of Diffusion Models\" section lacks references to key fundamental works. Specifically, step distillation [4, 5] and consistency models [6] should be cited under \"reducing the number of sampling timesteps,\" and Q-Diffusion [7] should be cited under \"weight quantization.\"\n- Notation inconsistencies: s3 is generally defined as cache frequency but appears as spatial distribution in Table 4, which uses s4 for cache frequency.\n\n[1] https://huggingface.co/black-forest-labs/FLUX.1-dev  \n[2] https://huggingface.co/black-forest-labs/FLUX.1-schnell   \n[3] Xu et al. ImageReward: Learning and Evaluating Human Preferences for Text-to-image Generation. NeurIPS 2023.  \n[4] Salimans et al. Progressive Distillation for Fast Sampling of Diffusion Models. ICLR 2022.  \n[5] Meng et al. On Distillation of Guided Diffusion Models. CVPR 2023.  \n[6] Song et al. Consistency Models. ICML 2023.  \n[7] Li et al. Q-Diffusion: Quantizing Diffusion Models. ICCV 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uJTHUad87c", "reviewer": "Reviewer_3XM7", "summary": "Token-wise Caching (ToCa) is introduced as a fine-grained feature caching strategy designed to accelerate diffusion transformers. It uniquely considers error propagation in caching, utilizing four selection scores to determine the best tokens for caching without incurring extra computational costs. ToCa also allows for variable caching ratios across different layers and integrates multiple techniques to enhance caching efficiency. Extensive experiments on models like PixArt-α, OpenSora, and DiT show that ToCa achieves significant speedups while preserving nearly lossless generation quality.", "questions": "1. The introductory paragraph of the methods section should be rephrased to clarify the concept of the \"naive scheme\" for feature caching. This naive scheme is actually no different from existing approaches, such as DeepCache, FORA, and delta-DiT, which should not be \"we propose\". \n\n2. It should be noted that Figure 3(b) does not depict any low-ratio caching scenarios. \n\n3. There is a discrepancy regarding the application of the scores s1, s2, and s3 across the layers. Specifically, s1 is applicable only to self-attention layers, while s2 pertains solely to cross-attention layers. However, the current definition combines these scores, which raises questions about how this aggregation functions across the different types of layers.\n\n4. Clarification is needed regarding the additional time cost associated with computing the selection scores. Given that the scoring relies on rank computation, there remains a need to calculate and distribute the cached tokens among the uncached tokens. This process is somewhat ambiguous, particularly concerning how the acceleration time is achieved. Presenting the actual extra costs incurred for selecting scores alongside the resulting speedup from the caching process would provide valuable insights.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The motivation behind the proposed approach is both technically sound and clearly explained, supported by two informative figures illustrating the varying levels of similarity among different tokens and the accumulation of error across these tokens. These visual aids effectively convey the rationale for the proposed method.\n\n2. The methodology for selecting scores and making decisions for each layer is novel, offering valuable insights into the acceleration of transformer models. This innovative approach not only enhances the efficiency of the models but also contributes to a deeper understanding of feature caching strategies within the context of transformers.", "weaknesses": "1. There appears to be a lack of a complete algorithmic description within the manuscript. Specifically, a detailed explanation is needed for the token selection process at each layer and each timestep, as well as the procedure for redistributing the cached tokens back into the overall framework. This omission could lead to misunderstandings regarding the efficacy and functionality of the proposed method.\n\n2. The manuscript does not adequately illustrate the computation costs associated with the proposed approach, making it challenging to comprehend the additional expenses involved. A clearer breakdown of these costs would enhance the reader's understanding of the method's efficiency.\n\n3. Including a token diagram for specific layers during both early and late timesteps, across various layer types, may provide additional clarity. Such visual representations could effectively illustrate how tokens are managed and utilized within the different stages of the computation, facilitating a better grasp of the overall caching strategy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iBrG1Y0foZ", "reviewer": "Reviewer_sYWB", "summary": "This paper presents a method called ToCa, which utilizes token-wise feature caching to accelerate the denoising process of DiT. Unlike existing DiT acceleration methods, ToCa offers a fine-grained, token-wise acceleration approach, demonstrating improved performance compared to previous methods.", "questions": "Same with my Weakness 1.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to understand. \n- The observation that the similarity and propagated error differ for each token in DiT is very insightful.\n- Various metrics were proposed to measure the importance of tokens on reusing. \n- Experiments were conducted on a variety of datasets, including not only text-to-image but also text-to-video.", "weaknesses": "- It's difficult to understand how token-wise caching leads to actual acceleration. Since the attention layer calculates the similarity among all inputs through matrix multiplication, even if one output token is not computed, there won't be a significant difference in the overall computational cost (similar to unstructured pruning). A more detailed explanation of where this acceleration comes from or breakdown is required. \n- Evaluation is performed with just single Latency/FID point. An evaluation of the Pareto curve with latency/FID would provide more insightful information. \n- There seems to be a color error in Fig 3(b). The yellow arrow indicating low cache ratio is missing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OIs0XwJb33", "reviewer": "Reviewer_R8q5", "summary": "The authors aim to accelerate DiT inference by caching features and reusing them in future steps. The authors observe the limitations of existing methods that ignore the sensitivities to feature caching among different tokens, thus leading to potential degradation of generation quality. Based on this observation, the main idea of this work is to selectively cache suitable tokens and adjust the caching ratios by considering layer types and depths. Evaluations based on three typical DiT models show the performance of model accuracy and overhead saving.", "questions": "Please refer to the questions in the weakness part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.\tThe author provides two insightful observations that improve my understanding of diffusion transformers. Firstly, different tokens exhibit varying levels of temporal redundancy across different timesteps. Secondly, the authors highlight that due to error propagation in transformer architectures, the same caching error can lead to vastly different impacts on the final generation result. These observations are the basis of the proposed token-wise feature caching method.\n2.\tUnlike previous works that primarily focused on U-Net architectures, the authors shift their attention to the transformer architecture. This is meaningful as existing methods have not fully exploited the unique properties of transformers, which are increasingly popular for visual generation tasks.\n3.\tThe experimental results demonstrate the performance improvements with high inference speedups.\n4.\tThe authors select the most suitable tokens for caching by leveraging both temporal redundancy and error propagation. The proposed method adaptively caches the tokens that minimize the resulting error while maximizing the acceleration ratio.", "weaknesses": "1.\tIn Figure 1, the authors use the Frobenius distance to measure the temporal redundancy of tokens across timesteps. However, it is unclear how this distance is exactly calculated. Providing more details on the calculation process would help readers better understand the experiments.\n2.\tI cannot find a clear formulation for the cache score used in token selection. The authors mention the cache score in Figure 4, but no explicit equation is provided. \n3.\tThe authors update the cache at all timesteps to reduce the error introduced by feature reusing. However, it is unclear how this procedure affects the computation complexity, especially about I/O overhead and cache selection efficiency.\n4.\tThe experiments report the inference speedups but lack a comprehensive analysis of the trade-off between speedups and model accuracy. While FID and CLIP scores are provided in Table 1, these metrics do not fully reflect model accuracy. In addition, Table 3 shows that the proposed method has no significant improvements over existing methods in terms of accuracy and speedups.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)