DiTFastAttn: Attention Compression for Diffusion Transformer Models

We truly appreciate the reviewer for the insightful and constructive comments.

W1. The authors sample 5K images to evaluate the generation quality while the standard setting is sampling 50K images (following original DiT).

Thanks. Following your suggestion, we have increased the evaluation size of ImageNet to 50K and COCO to 30K. Tables can be found in the global rebuttal (Table 1 & 2). More discussions can be found in the global rebuttal Q1.

W2. What is the adopted compression strategy in each layer and each timestep? It would be better to illustrate the searched compression plan in a figure.

Thanks for your suggestion. We generated figures for adopted compression strategy in different threshold settings, check Q2 in global rebuttal and Figure 1 in the rebuttal PDF.

W3. Can the proposed method be combined with other acceleration methods \eg Deepcache?

We believed that Deepcache benefits from the hierarchical nature of U-Nets so it can not be directly applied on DiT model. However, we do think our method can be combined with other acceleration methods. For example, our method is orthogonal with quantization method, so we can apply both quantization and our method to the model to achieve further acceleration.

W4. Why do the authors only demonstrate the Inception Score in Figure 9? FID is the most common metric in image generation.

Showing only one metric in a figure tends to be more clear. In a future version of our paper, we will include an ablation study plot to further demonstrate the FID results.

W5. What are the FLOPs and Latency in Table 1 exactly? The authors should demonstrate the value rather than only the fraction since there is still enough space.

Thanks. We make a summary table in the global rebuttal Q1 (Table 1 & 2) to contain exactly FLOPS and Latency. Summary tables like that will be included in the future version of our paper.

Thanks for your comment.

Search time varies for different models and parameter settings. We use a greedy method in compression plan search so increase threshold (more compression) will result in a shorter search time. The maximum search time for a model can be determined by setting the threshold to 0. One thing to note is that once the search is completed, the compression strategy will be locally cached, eliminating the need for a repeated search. The cached strategy will be utilized the next time our method is applied for image or video generation. For reference, here we post the search times for DiT and PixArt-Sigma 1K using the default settings.

Table1. Search time for DiT-XL-2-512

Table2. Search time for PixArt-Sigma-XL-2-1024-MS

Thanks for your patience. Here we post the results for DiT after changing our experiment setting to align with DiT paper. From this table, we observe that when the threshold is increased from 0 to 0.10, the IS drops from 210 to 180, and the FID decreases from 3.16 to 3.09, then increases to 4.52. The behavior of the FID in different settings follows the pattern reported in the previous discussion[1]. We observe a higher reduction in FLOPs and Attention FLOPs in this setting. We believe the current results demonstrate the effectiveness of our method across different settings.

So far, we increased evaluation sample size, changed experiments settings and provided results that aligned with the original paper according to W1; provided compression plan and search time that asked in W2 and FLOPs number that asked in W4; explained the combination with other acceleration model for W3; revised our paper according to W5. We are happy to answer additional questions and provide more clarification if needed. We kindly ask you to reconsider the score since we nearly reached the end of the discussion period.

[1] Jayasumana S, Ramalingam S, Veit A, et al. Rethinking fid: Towards a better evaluation metric for image generation[C]// CVPR. 2024: 9307-9315.

We appreciate the reviewer for the insightful and constructive comments.

W1. the evaluation of the effect of parts of the method & design choices seems lacking: there are missing ablations regarding different variants (e.g., naive approaches) of leveraging specific redundancies, and especially the evaluation of the compression plan algorithm is lacking and its hyperparameters (threshold $\delta$, number of samples).

We show the effect of some naive approaches in Figure 9 in our paper, and we have added more images generated (like Figure 5a in the attached pdf) and compression plan (like Figure 5c in the attached pdf) by naive approaches. We conducted an ablation study on number of samples, and we found that adding number of samples had little effect on the compression plan.

DiT n_samples ablation

In global rebuttal Q5, we added more ablation studies.

W2. & Q2. Open questions remain, especially regarding unorthodox design choices and whether the CFG redundancy is actually an attention redundancy as described or a general redundancy already thoroughly described in previous works.

Thanks for your comments. We believe that the issue of conditional guidance (CFG) redundancy is not limited to the attention mechanism alone. There has been some prior work exploring methods to reduce CFG redundancy. However, our approach differs from previous efforts, as we have examined the use of CFG sharing at a finer granularity. Specifically, we have assessed the influence of conditional and unconditional computations across different layers and different timesteps.

Our work applies CFG sharing selectively, only on certain layers and time steps, in order to avoid significantly degrading the model's performance. In this way, our work acknowledges and builds upon existing research, while offering a novel perspective on addressing redundancy within the unique framework of diffusion models. Going forward, we plan to extend our method to other model components in the future.

Q1. While their approach to only computing attention locally does seem to be effective, I am surprised by the fact that the authors chose to apply a 1D window attention on the flattened image token sequence instead of one of the many 2D-aware local attention mechanisms such as (Shifting) Window Attention. These should have been the obvious choice in this situation, and the authors fail to motivate their unorthodox choice.

Please see global rebuttal Q4.

Q3. What is the distribution of which acceleration method is used how often, and where in the generation process?

Please check Q2 in global rebuttal and Figure 1 in the rebuttal pdf attached.

Q4. Can the authors provide an explanation regarding the intuition of FID improving when reducing the attention FLOPS in Fig. 5?

We have thoroughly reviewed the evaluation code and confirmed that there are no issues on our end. We believe the discrepancy may be due to the inaccuracy of the FID metric, as indicated in previous research [1]. The FID does not fully capture the visual quality of the generated images. However, changes in the FID score can still be meaningful, as they reflect shifts in the distribution of the generated images.

[1] Jayasumana S, Ramalingam S, Veit A, et al. Rethinking fid: Towards a better evaluation metric for image generation[C]// CVPR. 2024: 9307-9315.

Q5. Subjectively, more compression results in a loss of contrast, both for generated images (albeit slightly), as well as videos. Do the authors have an explanation as to what might be causing this?

We evaluated the images generated using different methods and found that the high threshold of ASC was likely the cause. We included an example generated by different methods and at varying ASC thresholds in the rebuttal PDF (Figure 5). We suspect that excessive CFG sharing can lead to a loss of detail and less pronounced edges in the generated images. If the trade-off between contract quality and ASC is a significant concern, the user should consider disabling the ASC method.

Q6. For their similarity analyses in Fig. 4, the authors only analyze the cosine similarity of attention outputs. Are the magnitudes similar as well? Wouldn't this be a requirement to enable successful caching?

Thanks. We plot the magnitude of attention output, the plot can be found in the rebuttal pdf (Figure 6). It is similar to cosine similarity.

Q7. Fig. 5 is hard to parse due to the multitude and different scales of the presented metrics. I'd suggest that the authors separate the metrics into separate stacked graphs, which should also render these graphs more accessible for people with color perception impairments. I would also suggest adding arrows to indicate whether lower or higher is better for each individual metric. Similar improvements to other figures could also help make the paper more accessible.

Thanks for your advice. Please see Q3 in global rebuttal.

Q8. For a camera-ready version, I think the paper would benefit from additional uncurated qualitative examples

Thanks. We plan to include more uncurated qualitative examples, ablation studies, plots, and tables in the future version of our paper, including the results presented in the rebuttal PDF.

We appreciate the reviewer for the insightful and constructive comments.

W1 & Q1. How do the authors ensure the generalizability of their method across different datasets and tasks? A more rigorous theoretical analysis could help understand the limitations and potential extensions of the proposed techniques.

About generalizability, our experiments demonstrate the effectiveness of DiTFastAttn across multiple datasets and tasks:

• Image generation: We evaluated on class conditional generation (for DiT) and text conditional generation (for PixArt-Sigma).
• Video generation: We successfully applied DiTFastAttn to OpenSora for video generation tasks.

Our method shows improved compression ratios and quality preservation as resolution increases from 512x512 to 2048x2048, indicating good scalability. This diversity in datasets, tasks, and model architectures provides strong evidence for the generalizability of our approach. However, we acknowledge that further evaluation on additional domains could further strengthen this claim.

For theoretical analysis, we thank the reviewer for this valuable suggestion. While our current work focuses on empirical results, we agree that a more rigorous theoretical analysis would be beneficial.

• We have provided rigorous computational complexity analysis for each proposed technique (WA-RS, AST, ASC) and their combinations. Some directions we could explore in future work include:
• Investigate how the error propagates through the denoising steps when using approximated attention outputs.
• Quantify the information loss when using window attention compared to full attention.

We will consider adding a discussion of these theoretical aspects in an extended version of the paper.

W2 & Q2. What is the model complexity, running speed, and model size under different compression settings?

In Appendix A.4 of our paper, we discussed the complexity of our compression method. Additionally, in the global rebuttal, we have added summary tables (Table 1 and 2) to present the exact FLOPs and latency (running speed). The result shows our method is effective to reduce both the complexity and improve the running speed. The model size will not change.

For model size and VRAM usage, when AST or WA-RS is applied, the hidden states from the previous timestep are stored. The shape of stored hidden states of a layer is [batchsize*2, num_of_head, seqlength, hidden_size]. For example, with the DiT XL/2 512 model, at a batchsize of 1, the VRAM usage of the original model is 11652MB. To store the FP16 hidden states, an additional 2016MB of VRAM is needed, which is 17.3% of the original model's VRAM usage. We will include a VRAM analysis in the limitations section.

W3. Insufficient Comparisons: The paper does not compare the related methods, such as flashattention, and KV cache. This mainly weakens the contribution. A more detailed comparison with existing attention acceleration and model compression techniques would be beneficial.

Thanks. Our baseline (we denoted as raw setting in our method) has used flashattention in their attention computation, and our method can achieve up to 1.6x speedup compared to it. We'll add a note to the paper to avoid confusion.

KV cache is a method used in large language models for auto-regressive generation. In diffusion model, DiT don't need to attend the tokens from previous timesteps. KV cache isn't applicable to diffusion models like DiT. So we cannot compare our method with KV cache. Note that our method is inspired by caching concepts but applied differently:

• Caching activation outputs across denoising timesteps and conditional/unconditional branches.
• WA-RS technique caches residuals to maintain long-range dependencies with efficient local attention.

W4 & Q3. Can the authors offer clearer definitions and detailed illustrations for each redundancy type (spatial, temporal, and conditional redundancy)?

Although we have provided some definitions in our paper, we agree that providing more clarity on these concepts will strengthen our paper. Here we denote $X$ as input, $Y$ as output

Window Attention with Residual Sharing (WA-RS):

• Definition: The original paper Eq.1 and Eq.2.
• Illustration: The original paper Figure.3. We will add a heat map visualization of the attention matrix, clearly showing the concentration of attention values along the diagonal.

Attention Sharing across Timesteps (AST):

• Definition: For step k, $Y_k=W_{o}O_k$ and Y_{k+1} = \\{Y_k if $AST_k = 1$; $W_{o}O_{k+1}$ if AST_k = 0\\}

• Illustration: The original paper Figure.2 is clear  Attention Sharing across CFG (ASC):

• Definition: For step k, Q_{k} = \\{W_{Q}X_{k,:c/2} if $ASC_k = 1$, $W_{Q}X_{k}$ if ASC_k = 0\\} and so as $K_{k}$ and $V_{k}$. The output is Y_{k} = \\{[W_{o}O_k, W_{o}O_k] if $ASC_k = 1$; $W_{o}O_{k}$ if ASC_k = 0\\}

• Illustration: The original paper Figure.2 is clear

W5. Application: The paper mainly focuses on attention reduction, but It is not very effective at low resolutions because the attention mechanism accounts for a small proportion of the computational cost.

While it's true that our method provides greater benefits at higher resolutions, we believe this is actually a strength rather than a limitation:

• Even at lower resolutions, our method still provides measurable improvements. For example, on the 512x512 DiT model, we achieve up to 31% FLOP reduction and 10% latency reduction for attention computation.
• As transformer-based diffusion models continue to scale up in size and target higher resolutions, the relative importance of attention computation increases. Our method is thus well-positioned to provide even greater benefits for future large-scale models.

We truly appreciate the reviewer for the insightful and constructive comments.

Q1. You should cite Sora in the introduction and not OpenSora.

Thanks. We have added the citation of Sora in the introduction.

Q2. How is your windowed attention method different from neighborhood attention (Natten)? This should be cited with local attention.

We adopted 1D window attention, which can be recognized as 1D neighborhood attention, and we have added references to Natten in the article. The evaluation shows our 1D window attention is more efficient (See global rebuttal Q4).

Q3. Including the 1-sigma boundaries (e.g. aggregated over 1k samples) in the MSE plots on Fig 3a would help understand how much variation there is for the attention layers.

We have added the MSE change plot on 1k samples and included the 1-sigma boundaries. We have also added the model type in the figure title and caption. The results indicate that the variation is not significant. Please refer to Figure 4.b in the attached rebuttal PDF.

Q4. You explored the impact with CFG but have you looked at the impact with negative conditioning? Does sharing the conditional attention perform similarly or degrade?

We investigated the impact of negative conditioning. The attentions with positive conditioning and negative conditioning have a high degree of similarity, and our method is also effective in reducing computation with negative conditioning. We provided an example in the attached PDF (Figure 4c, "low quality" as the negative prompt). We also added a plot to show that similarity across CFG with negative prompt behave similar as the one without negative prompt (Figure 4d). We will add more examples and analysis of negative conditioning in our paper.

Q5. There appears to be some interesting behavior in Figure 4, which may be interesting to explore from an explainability perspective.

In Figure 4a, we observe that the middle timesteps exhibit less similarity, which may suggest that the structural transformation occurs primarily in the middle timesteps, while the other timesteps focus on noise removal and refinement. In Figure 4b, we note that the middle layers show more differences between the conditional and unconditional cases, indicating that the middle layers are responsible for the condition-specific processing.

Q6. For your compression plan search, did you consider using other metrics such as LPIPS or SSIM? Would they impact the chosen plan?

Yes, we've considered LPIPS and SSIM for our comparison. LPIPS only works for image comparisons and requires the channel number to be set to 3, whereas we are comparing latent space (with 4 channels). LPIPS also needs to use a network as a feature extractor, which results in a significant inference overhead.

We also tried the SSIM metric, and the results are included in Table 4 of the global rebuttal. The SSIM metric appears to achieve comparable results to our own metric. However, SSIM is more computationally complex to calculate. Therefore, we have decided to use our own metric for this task. We have included the compression plan in Figure 2 of the rebuttal PDF. Interestingly, the SSIM metric selects more lower-level layers for compression, while the existing metric chooses to compress the higher layers. We will research this in the future.

Q7. Evaluation sample number

Please see global rebuttal Q1. Our conclusions remain unchanged at the larger evaluation size.

Q8. When you evaluated FID, IS, CLIP, did you do so on the same 5k samples used for determining the plan? I would have used the ImageNet/COCO validation set for the plan search and then evaluated on the test set to ensure the two distributions were non-overlapping.

For the calibration process, we only use class labels (for DiT) and text prompts (for PixArt-Sigma) to generate a small set of images (8 images for DiT, 6 images for PixArt 1K, and 4 for PixArt 2K). No real images are used in the calibration process, so there is no risk of overlapping.

Q9. What dataset did you use for your OpenSora evaluation?

Since the video generation process is too resource-intensive, and we found the evaluation metrics to be unstable when the sample size is small, we did not apply these metrics to the generated videos. Instead, we included some sample generated examples. These example videos were generated using the default prompts provided by OpenSora.

Q10. Did you consider looking at the impact on VRAM? If there is no impact, stating so in the limitations would be appropriate.

When AST or WA-RS is applied, the hidden states from the previous timestep are stored. The shape of stored hidden states of a layer is [batchsize*2, num_of_head, seqlength, hidden_size]. For example, with the DiT XL/2 512 model, at a batchsize of 1, the VRAM usage of the original model is 11652MB. To store the FP16 hidden states, an additional 2016MB of VRAM is needed, which is 17.3% of the original model's VRAM usage. We will include a VRAM analysis in the limitations section.

Q11. If increasing the step count improves the result in Fig 9, then this suggests that there is an optimal tradeoff between attention FLOPs and step count (iso-sampling FLOPS).

In the rebuttal PDF, we have added a FLOPs-IS Pareto front plot (Figure 4.a). Interestingly, we observed that different compression ratios require different optimal settings (steps and DiTFastAttn threshold). We will include this plot and the related discussion in the appendix of the future version of the paper.

Q12. Reporting relative performance as percentages is difficult to read. It would be clearer to report them as ratios

Thanks. We've changed the format of relative performance from percentage to ratio.

Fig 5 & CLIP issues

Figure 5 is now changed accordingly (check global rebuttal and Figure 3 in the pdf). For CLIP, we use the package torchmetric to calculate it so the scale is different (0-100 vs 0-1). It is now changed to 0-1.

Thanks for your additional questions. Here is our reply:

Q1. Evaluation metrics

The experiment settings are different in the following ways:

• 1. The number of timesteps. To compare DiT with ADM and LDM, DiT uses a relatively high number of diffusion timesteps (250). For most image generation applications, 20-50 steps would be reasonable. Other methods, such as SD-XL and PixArt series, also evaluate on 20-50 steps. That is why we use 20-50 steps in our paper.
• 2. Evaluation software. We use pytorch-fid & torchmetric to calculate FID and IS scores, while DiT uses the TensorFlow evaluation suite from the ADM paper. This resulted in some differences in the FID and IS scores.
• 3. We use a cfg_scale of 4, which is the default setting for both DiT official code and diffusers DiT pipeline. In the DiT paper, they use a cfg_scale of 1.5 to match the settings used for ADM and LDM.

We have modified the code to reproduce the results. However, it is not feasible to resample the 50K images under the timestep=250 setting during the rebuttal period (this would take approximately 55 GPU hours). We will provide a partial set of the results soon.

For PixArt-Sigma, they used a curated set of 30,000 images instead of COCO to calculate FID. Since the dataset is not publicly released, we cannot reproduce those results. We will clearly explain the experiment settings in the appendix.

Q2. LPIPS implementation

Thank you for the suggestion. We have implemented LPIPS as an additional evaluation metric (by decoding the hidden states into RGB space). However, we found that LPIPS has problems for our use case:

• 1. Insensitivity to Value Changes: We observed that LPIPS is quite insensitive to value changes in our experiments. Only small value changes were observed when switching between different methods, and LPIPS always suggested using sharing across timesteps when the threshold was set to 0.005 or smaller.
• 2. Computational Overhead: Using LPIPS as a metric takes significantly more time. Therefore, we were unable to provide the LPIPS results within the rebuttal period due to the computational requirements.

Q3. Number of calibration images

We conducted an ablation study on the number of samples generated during the plan search, and found that increasing the number of samples had little effect on the resulting compression plan.

From the results, the IS and FID did not change significantly as we varied the number of samples. Nevertheless, users can easily set 'n' to a larger number if desired, as it is a configurable parameter in our code.

Q4. downscaled method like HiDiffusion be applicable?

We have carefully reviewed the HiDiffusion approach, which is designed for UNet-based diffusion models. We have considered how this method may be applicable to our framework:

• 1. HiDiffusion proposed using local attention to replace global attention in the top layers of the UNet, as global attention in the upper blocks is computationally dominant. However, in DiT, the amount of self-attention computation is equal across each layer. Therefore, we do not need to focus solely on specific layers, and instead use a search-based method to automatically identify which layers can be replaced.

2. HiDiffusion use the Modified Shifted Window Attention, where the window area is set differently across diffusion timesteps. This approach may be applicable to our method as well. However, as mentioned in the Global rebuttal (Q4), the inference speed of these attention mechanisms could be slower. It is also an interesting topic to research.

Q5. Pareto front question

Thanks for your suggestion. We have followed your advice and added marker shapes. Here is our observation from the plot:

• The point representing compression at 30 steps (yellow marker at (2.7 TFLOPs, 208 IS)) is superior to using no compression at 20 steps (blue marker at (2.9 TFLOPs, 207 IS)). This is because the compressed 30-step model has lower computational cost (TFLOPs) while achieving a higher Inception Score (IS).
• For a high FLOPs budget (> 3.5T), it is better to use moderate compression with 50 steps.
• For a medium FLOPs budget (2-3.5T), it is preferable to use slight compression with 20 steps or moderate compression with 30 steps.
• For a low FLOPs budget (< 2T), the optimal choice would be to use high compression with either 20 or 30 steps.

Thank you for your patience. Here we present the results for DiT after changing our experimental settings to align with the DiT paper.

From this table, we observe that when the threshold is increased from 0 to 0.10, the IS drops from 210 to 180, and the FID decreases from 3.16 to 3.09, then increases to 4.52. The behavior of the FID in different settings follows the pattern reported in the previous discussion[1]. We observe a higher reduction in FLOPs and Attention FLOPs in this setting. We believe the current results demonstrate the effectiveness of our method across different settings.

So far, we provide explanations, revision and examples according to each concerns. We are open to more discussion and will to provide more clarification if needed. As we nearly reached the end of the discussion period, we kindly ask you to reconsider the score.

[1] Jayasumana S, Ramalingam S, Veit A, et al. Rethinking fid: Towards a better evaluation metric for image generation[C]// CVPR. 2024: 9307-9315.

We appreciate your valuable feedback. In the revised version, we will include the LPIPS results and additional evaluation metrics, as well as more uncurated qualitative examples, as you suggested.

Regarding your interpretation of the updated evaluations (Q1), there seems to be some misunderstanding. Our findings indicate that even at a threshold of 0.025, the method was able to significantly reduce the attention FLOPs by 46%, while maintaining a satisfactory generation effect. At a threshold of 0.05, the attention FLOPs were reduced by 63% without significant degradation. The results show that as the number of steps increases, the threshold required to achieve the same compression ratio decreases. For example, with 250 steps, a threshold of 0.025 can achieve the same compression ratio as a threshold of 0.125 with 20 steps. This means that for larger timesteps, a smaller threshold is needed to reach the same compression ratio.

Therefore, we believe our results demonstrate that in the experiment setting of 250 steps, a small threshold can indeed achieve good compression performance. We will adjust the threshold values and include more results in the range of 0 to 0.05 to better illustrate this point.

We kindly ask you to reconsider the score if we have addressed all of your concern.

The results in Table 1 and Table 2 are strange. Lower IS and lower FID? I think that when the quality of the image decreases, the IS will be lower but FID will be higher.

Our findings indicate that even at a threshold of 0.025, the method was able to significantly reduce the attention FLOPs by 46%, while maintaining a satisfactory generation effect. At a threshold of 0.05, the attention FLOPs were reduced by 63% without significant degradation. The results show that as the number of steps increases, the threshold required to achieve the same compression ratio decreases. For example, with 250 steps, a threshold of 0.025 can achieve the same compression ratio as a threshold of 0.125 with 20 steps. This means that for larger timesteps, a smaller threshold is needed to reach the same compression ratio.

Therefore, we believe our results demonstrate that in the experiment setting of 250 steps, a small threshold can indeed achieve good compression performance. We will adjust the threshold values and include more results in the range of 0 to 0.05 to better illustrate this point.