CAT Pruning: Cluster-Aware Token Pruning For Text-to-Image Diffusion Models

The paper has two overarching weaknesses that become apparent in a multitude of smaller issues.

Presentation Key aspects of the paper are presented badly, making it hard for readers to grasp the contributions made by CAT-Pruning

• Underlying fundamentals of the method are not sufficiently explained. For example nowhere in the Abstract, Introduction or Conclusion to the authors elaborate that they prune tokens in the attention layers of the DM. In fact when asking multiple computer vision researchers what they assumed the method to be about, all of them assumed that tokens where pruned in the embeddings of the text prompt.
• The authors use MACs as a key performance metric throughout the paper without ever explaining it
• Similarly it is never explained between which samples the CLIP Score is calculated during evaluation
• a third of Page 7 is just empty, but Page 8 is almost exclusively Figures
• Tables 2 and 3 do not contain any bold numbers
• The diffusion notations are somewhat disconnected from common conventions. For example, noise (estimates) are usually reffered to as $\epsilon$, during generation diffusion steps $t$ should decrease since this is the reverse diffusion process. The authors also claim that Diffusion involves solving a reverse-time SDE, which is indeed one valid mathmatical foundation for diffusion. However, it specifically does not encompas the models actually used in the remainder of the paper, with SD3 being a rectified flow model that is specifically incompatibel with stochhastic algorithms.

Evaluation Further, there are problems with the validation of the poposed method.

• The main claim is that the output image after pruning is perceptually similar to the image generated without pruning. However, that specific aspect is never empirically evaluated. The obvious choice here would have been to simply report LPIPS distance between the pruned und unpruned image.
• Similarly the tradeoff between compute reduction and (un)pruning percentage $\alpha$ is demonstrated with 4 examples but not ablated over empirically
• In the same vain many of the design choices of the CAT algorithm are showcased with 1 or 2 qualitiative examples with limited to no empirical ablations. A more structured analysis on the importance/downstream influence of the different components in the selection algorithm would have been important
• The empriical results in Tab. 2 and 3 do not contain any confidence intervals or standard deviations
• There is no qualitative comparison with competing methods
• The authors only compare against one other baseline, although other methods exist, including ∆-DiT [1], Faster Diffusion [2], or TGATE [3] or basic methods like KV caching
• The most prominent choices for speeding up DM inference are, of course, distillation methods or consistency models. Consequently, it would be important to consider if CAT-Pruning still offers advantages when applied to distilled or consistency models.

[1] Pengtao Chen, et al. ∆-DIT: A training-free acceleration method tailored for diffusion transformers. arXiv:2406.01125 [2] Senmao Li et al. Faster diffusion: Rethinking the role of unet encoder in diffusion models. arXiv:2312.09608, 2023. [3 Wentian Zhang, et al. Cross-attention makes inference cumbersome in text-to-image diffusion models. arXiv:2404.02747, 2024.

Other

• Method appears to be limited to DiT architecture. At least no other architecture is considered

Weaknesses:

• [Minor] The authors are suggested to provide a brief overview of the diffusion model architecture (e.g., SD-3 like architecture) to which their method is applicable. This will improve readability by quite an extent.
• [Major] The paper does not compare with some of the caching only baselines (e.g., TGATE which they cite). This would be important to understand the effectiveness of combining caching with pruning techniques. Although not directly comparable, I’d also suggest the authors to provide the distillation based baselines (less number of steps in the inference) in the paper, to provide a full picture of the effectiveness of the caching + pruning family of methods.
• [Major] The authors provide the CLIP-score for generation quality — but it’d be more effective to show a more holistic evaluation of the method in terms of generation quality. For example, the authors can test on compositionality, long caption generation etc for their method in terms of generation.
• [Minor]: The authors are suggested to provide more qualitative results comparing the generation across different methods improving efficiency.

The paper has limited theoretical justification: The paper primarily relies on empirical observations and intuitions to justify the effectiveness of CAT Pruning. While the authors present Proposition 1 and provide a simplified proof in the appendix, a more rigorous theoretical analysis could strengthen the paper's contribution. A deeper theoretical understanding of the relationship between relative noise magnitude, token staleness, and spatial clustering could lead to more informed design choices and potentially improved performance. For example, exploring the convergence properties of the algorithm or deriving bounds on the error introduced by pruning would provide valuable insights.

There is a lack of comparison with other pruning techniques, and other techniques in general: The paper compares CAT Pruning only with AT-EDM, another token pruning technique. However, a comprehensive comparison with a wider range of pruning methods for diffusion models, such as those leveraging model distillation, quantization, or low-rank factorization, would provide a more complete picture of the proposed method's strengths and limitations. This would allow for a more informed assessment of the relative performance and efficiency of CAT Pruning compared to other state-of-the-art techniques. In particular the paper seems to ignore the highly influential VQVAE and VQGAN based methods for leveraging compressed latent representations of data. These methods are extremely popular ways for reducing the computational load of diffusion models. See Gu et al., (2022) for example (among many other examples), i.e.

Gu, S., Chen, D., Bao, J., Wen, F., Zhang, B., Chen, D., Yuan, L. and Guo, B., 2022. Vector quantized diffusion model for text-to-image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 10696-10706).

"CAT Pruning" focuses on optimizing token processing within a U-Net architecture. However, VQ-Diffusion demonstrates that utilizing a VQVAE to learn a compressed, discrete latent representation can lead to significant computational savings. By shifting the diffusion process to this lower-dimensional latent space, VQ-Diffusion achieves notable speed improvements. The "CAT Pruning" paper does not acknowledge or compare its approach to this architectural shift towards latent space diffusion using VQVAEs, which constitutes a notable weakness. Similarly, the paper on "High-Resolution Image Synthesis with Latent Diffusion Models" of

Rombach, R., Blattmann, A., Lorenz, D., Esser, P. and Ommer, B., 2022. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 10684-10695).

discusses the popularity and advantages of VQGANs, particularly their ability to learn compressed latent representations that can be used for high-resolution image synthesis. VQGANs, in conjunction with autoregressive models, have emerged as powerful tools for high-resolution image synthesis. These models operate on a compressed, discretized latent space learned by a VQGAN, potentially offering significant computational advantages over pixel-based approaches. The "High-Resolution Image Synthesis" paper uses VQ-regularization as one method for training the autoencoder that produces the latent space for their LDMs. They discuss how their model "can be interpreted as a VQGAN but with the quantization layer absorbed by the decoder."

Since the proposed CAT Pruning technique emphasizes reducing computational costs as a primary goal, comparing this CAT token pruning method with VQVAE and VQGAN style approaches in a much more direct manner would be very helpful. These methods inherently operate in a compressed latent space, and understanding the relationships between these kinds of methods and the proposed method would be very helpful, comparing empirically would be strongly desired and it could demonstrate the relative efficiency of CAT Pruning much more clearly.

In summary, with respect to the issue of comparing the work here to these popular methods, the paper would be strengthened by:

1. Explicitly at least acknowledging the popularity of VQVAE and VQGAN-based approaches for diffusion model acceleration.
2. Discussing the potential trade-offs between their token pruning within the U-Net architecture and the use of VQGANs and other approaches for latent space diffusion.
3. Ideally, conducting more direct experiments to compare the performance, efficiency, and image quality of CAT Pruning against a VQGAN-based diffusion model.

1. The paper lacks a thorough discussion of prior token pruning work, making it difficult to assess the proposed method’s novelty and improvements over existing methods. A more detailed overview of token pruning techniques, their limitations, and how the proposed method addresses these would clarify its contributions.
2. The study lacks comparisons with a diverse set of token pruning baselines and established training-free methods (e.g., caching), which limits a comprehensive view of the proposed method’s effectiveness relative to existing techniques.
3. The qualitative differences shown in Figure 6 between clustering and non-clustering approaches are subtle, and there is no rigorous ablation study to quantitatively assess clustering’s impact on model performance. A thorough ablation study is needed to substantiate the claimed benefits of clustering.
4. The paper relies solely on CLIP Score to assess image quality, which measures text-image alignment but not visual fidelity. Including metrics such as FID would provide a more complete evaluation of image quality and support claims of fidelity preservation.