CAIT: Triple-Win Compression towards High Accuracy, Fast Inference, and Favorable Transferability For ViTs

TL;DR

CAIT tackles the high computational cost of Vision Transformers by proposing a joint compression framework with asymmetric token merging (ATME) and consistent dynamic channel pruning (CDCP). ATME reduces tokens while preserving spatial structure, enabling robust transfer to downstream tasks, while CDCP prunes channels dynamically with head- and attention-level consistencies to maintain MHSA integrity. Across ImageNet and multiple pixel-level downstream tasks, CAIT delivers state-of-the-art accuracy-speed tradeoffs and demonstrated transferability to segmentation, detection, and medical/aerial datasets. The work offers a practical, end-to-end pruning strategy that preserves spatial information and allows fast, scalable ViT deployment on resource-constrained devices.

Abstract

Vision Transformers (ViTs) have emerged as state-of-the-art models for various vision tasks recently. However, their heavy computation costs remain daunting for resource-limited devices. To address this, researchers have dedicated themselves to compressing redundant information in ViTs for acceleration. However, existing approaches generally sparsely drop redundant image tokens by token pruning or brutally remove channels by channel pruning, leading to a sub-optimal balance between model performance and inference speed. Moreover, they struggle when transferring compressed models to downstream vision tasks that require the spatial structure of images, such as semantic segmentation. To tackle these issues, we propose CAIT, a joint \underline{c}ompression method for ViTs that achieves a harmonious blend of high \underline{a}ccuracy, fast \underline{i}nference speed, and favorable \underline{t}ransferability to downstream tasks. Specifically, we introduce an asymmetric token merging (ATME) strategy to effectively integrate neighboring tokens. It can successfully compress redundant token information while preserving the spatial structure of images. On top of it, we further design a consistent dynamic channel pruning (CDCP) strategy to dynamically prune unimportant channels in ViTs. Thanks to CDCP, insignificant channels in multi-head self-attention modules of ViTs can be pruned uniformly, significantly enhancing the model compression. Extensive experiments on multiple benchmark datasets show that our proposed method can achieve state-of-the-art performance across various ViTs.

Figures (8)

• Figure 1: Comparison of throughput and accuracy between CAIT (Ours) and other models. The top-1 accuracy is tested on ImageNet-1K and the throughput is evaluated on a NVIDIA RTX-3090 GPU with a batch size of 256.
• Figure 2: The comparisons between CAIT and others. (a) indicates that previous token pruning methods usually sparsely drop the tokens and disrupt the complete spatial structure when transferring to downstream tasks. (b) shows that ATME leverages horizontal token merging (HTM) and vertical token merging (VTM) to prune tokens while maintaining the spatial integrity, where image feature maps can be easily upsampled for downstream tasks. (c) denotes that previous channel pruning methods usually suffer from irreversible removal of important channels due to the direct pruning. (d) presents that instead of directly modifying the model structure like previous works, CDCP dynamically determines the importance of channels periodically and encourages structured sparsity network gradually during training, which can thus recover the important channels.
• Figure 3: The overview of our proposed joint compression method for ViTs. We design an asymmetric token merging (ATME) strategy with horizontal token merging (HTM) and vertical token merging (VTM) to prune tokens while preserving their spatial integrity. Consistent dynamic channel pruning (CDCP) is further introduced to enable dynamic fine-grained compression optimization for all learnable weights with minimal performance degradation.
• Figure 4: The proposed asymmetric token merging strategy prunes tokens horizontally or vertically.
• Figure 5: Left: Framework of the compactor. Right: Consistencies for pruning MHSA.