COMCAT: Towards Efficient Compression and Customization of Attention-Based Vision Models

TL;DR

COMCAT tackles the efficiency bottleneck of attention-based vision models by introducing head-level low-rank MHA compression with automatic rank selection, outperforming pruning while reducing parameters. The approach reformulates multi-head attention to exploit low-rank structure at the head level and uses a differentiable NAS-like search to choose ranks under a cost constraint. It demonstrates improved accuracy with DeiT-small/base on ImageNet and enables faster, storage-efficient customization of text-to-image diffusion models. The findings offer practical deployment benefits for resource-constrained vision systems and for rapid diffusion-model personalization.

Abstract

Attention-based vision models, such as Vision Transformer (ViT) and its variants, have shown promising performance in various computer vision tasks. However, these emerging architectures suffer from large model sizes and high computational costs, calling for efficient model compression solutions. To date, pruning ViTs has been well studied, while other compression strategies that have been widely applied in CNN compression, e.g., model factorization, is little explored in the context of ViT compression. This paper explores an efficient method for compressing vision transformers to enrich the toolset for obtaining compact attention-based vision models. Based on the new insight on the multi-head attention layer, we develop a highly efficient ViT compression solution, which outperforms the state-of-the-art pruning methods. For compressing DeiT-small and DeiT-base models on ImageNet, our proposed approach can achieve 0.45% and 0.76% higher top-1 accuracy even with fewer parameters. Our finding can also be applied to improve the customization efficiency of text-to-image diffusion models, with much faster training (up to $2.6\times$ speedup) and lower extra storage cost (up to $1927.5\times$ reduction) than the existing works.

Figures (11)

• Figure 1: The heatmap of the cumulative singular values after performing SVD for each weight matrices in the MHA layers of the pre-trained DeiT-base model. It is seen that some matrices in some layers exhibit weaker low-rankness than others, implying that directly factorizing individual weight matrix may not be efficient.
• Figure 2: The number of parameters when the ratio of the cumulative singular values reaches 90% with factorizing $W^Q,W^K,W^V,W^O$, and the corresponding $W^{QK}, W^{VO}$ of the pre-trained DeiT-base. It is seen that exploring head-level low-rankness is more parameter efficient than matrix-level.
• Figure 3: (Left) Standard MHA layer. (Middle) Direct factorization of the individual weight matrices in MHA layer. (Right) Our proposed exploring head-level low-rankness of MHA layer.
• Figure 4: Compressing vision transformer using low-rank MHA layers and automatic rank selection.
• Figure 5: Customizing text-to-image diffusion model using low-rank MHA mechanism. Latent image feature $\mathbf{x_{new}}$ corresponds to $X_Q$ and text feature $\mathbf{c_{new}}$ corresponds to $X_K$ and $X_V$.