Adaptive Sparse ViT: Towards Learnable Adaptive Token Pruning by Fully Exploiting Self-Attention

1 Introduction

Recently, Vision Transformer (ViT) has made remarkable progress on image classification [Dosovitskiy et al., 2020; Touvron et al., 2021a; Liu et al., 2021], object detection [Carion et al., 2020; Zhu et al., 2020], semantic segmentation [Zheng et al., 2021; Xie et al., 2021], and other vision tasks. However, as the model complexity is quadratic to the number of tokens, ViT suffers from expensive computational costs, which limits its application and deployment.

Not all image patches are helpful for the final prediction. For instance, the large number of image tokens in the background region do not contribute to the recognition and can be pruned during the inference process, greatly accelerating the model runtime without significant impact on performance. Token pruning has attracted the interests of many researchers. We classify the existing methods according to whether additional calculations are introduced to evaluate the token score. Evo-ViT [Xu et al., 2022] utilizes class attention [Xu et al., 2022] to estimate token score and develop a novel slow-fast token evolution approach to improve the throughput of ViT. EViT [Liang et al., 2022] employ a similar method to measure token importance and fuse discarded tokens. Both of these approaches require manually specifying the pruning ratio for each stage, and perform the same pruning policy for different input instances, which may result in simple samples being under-pruned or complex samples being over-pruned in the beginning stages, as illustrated in Figure 2 top. Another type of work identifies token importance via extra measures. DynamicViT [Rao et al., 2021] prune tokens in a fixed ratio by inserting lightweight predictors to predict token scores. IARED2 [Pan et al., 2021] introduce a multi-head interpreter and employed reinforcement learning to generate pruning scheme for each token. A-ViT [Yin et al., 2021] compute halting scores for all tokens to adaptively discard unimportant tokens. The latter two achieve sample-adaptive pruning at the cost of additional computation of significance scores for all tokens.

Fixed Pruning Ratio

In this work, we propose an adaptive token sparse framework for ViT acceleration, named AS-ViT, which fully exploits Multi-Head Self-Attention (MHSA) to estimate token importance scores, and uses a minimal cost, only three learnable thresholds, to perform the corresponding token pruning policy for a specific input instance as shown in Figure 2 bottom. Specifically, we first propose the attention head importance weighted class attention score. It uses the intermediate results of MHSA to calculate token-level head importance, which is then multiplied as a weighting factor on the class attention [Xu et al., 2022] score to better identify informative tokens. Then, we insert three learnable thresholds in ViT hierarchically. Only tokens with a score greater than the current threshold will be kept, and the pruned token will not participate in the later computations. Lastly, we propose a budget-aware loss to optimize the thresholds to achieve a trade-off between accuracy and computational effort across the dataset. During testing, the learnable threshold is fixed and we just compare the threshold and individual token scores to actually discard uninformative tokens for hardware acceleration. We investigate the intrinsic features of MHSA and reuse its computational results to evaluate the informativeness of the token. Threshold-based comparison avoids sorting all tokens. We use the above method to minimize the cost of token pruning.

We conduct extensive token pruning experiments for the widely used vision transformer backbones, DeiT [Touvron et al., 2021a] and LV-ViT [Jiang et al., 2021] on ImageNet. For instance, our method improves 1.5x throughput with only

0.2% decrease in accuracy while reducing the 35% GFLOPs of the DeiT-small model. For other model and pruning rates, our method also achieves a better accuracy and complexity balance compared with previous approaches.

2 RelatedWork

Vision Transformer. Transformer [Vaswani et al., 2017], developed in NLP, has been successfully applied to various vision tasks and tends to replace CNN [He et al., 2016; Tan and Le, 2019] gradually. ViT [Dosovitskiy et al., 2020] illustrates that transformer lacks inductive bias and requires large-scale dataset pre-training to achieve an approximate performance with the state-of-the-art CNN. DeiT [Touvron et al., 2021a] eliminates the above problem with well-tuned training parameters and the introduction of distillation token. LV-ViT [Jiang et al., 2021] explores a variety of techniques for training vision transformer, significantly improving the performance of ViT. PVT [Wang et al., 2021] constructs a hierarchical ViT similar to CNN and proposes spatial-reduction attention to reduce the computation complexity. Swin Transformer [Liu et al., 2021] proposes a window based self-attention that makes the complexity linear with respect to token number and becomes a generic visual backbone. DGT [Liu et al., 2022a] propose the dynamic group attention to accelerate inference.

Static ViT Pruning. Analogous to weights pruning in CNN [Han et al., 2015; Li et al., 2016; Molchanov et al., 2016; Liu et al., 2017; Luo et al., 2017; Frankle and Carbin, 2018], there are many works [Zhu et al., 2021; Yang et al., 2021; Chen et al., 2021; Yu et al., 2022] for ViT parameters compression. VTP [Zhu et al., 2021] prunes parameters in MHSA and FFN indiscriminately through L1 regularized sparse training. NViT [Yang et al., 2021] establishes the global weights importance via performing Taylor expansion to the loss, then conducts structured pruning and parameter reassignment based on dimensional trends. SViTE [Chen et al., 2021] fully explores the sparsity of ViT, compressing transformer with structured pruning, unstructured sparsity and token pruning. In this paper, we focus on token compression, which is orthogonal to static ViT pruning, and we can further improve the compression rate, benefit fromweights pruning.

Dynamic ViT Pruning. Thanks to the transformer's parallel computing mechanism, pruning image tokens can bring real acceleration without the need of additional support. Both DynamicViT [Rao et al., 2021] and IA-RED2 [Pan et al., 2021] dynamically keep informative tokens by inserting prediction modules. PS-ViT [Tang et al., 2021] belongs to static token pruning, which statistically obtains the importance distribution of tokens across the dataset and prunes them top-down. EViT [Liang et al., 2022] and Evo-ViT [Xu et al., 2022] use the class attention score to distinguish how informative a token is. A-ViT [Yin et al., 2021] adaptively calculates the halting score for each token. Our method strives to implement instance-wise token pruning without additional computational costs.

3 Method

3.1 Preliminary

Vision Transformer [Dosovitskiy et al., 2020] provides a new paradigm for image recognition. ViT first splits the image into N × N non-overlapping patches and embeds them into a D dimensional feature space, and then adds a class token before the image tokens for final classification. Considering the position relationship, all tokens are added with a learnable position encoding and then fed into a stacked transformer block. We summarize the operations inside the block by using the following two equations:

where MHSA stands for Multi-Head Self-Attention [Vaswani et al., 2017], FFN is feed-forward neural network, and LN stands for layer normalization.

MHSA represents features in multiple subspaces using the same parameters. The input is first projected into three matrices Query Q ∈ R(N+1)×D , Key K ∈ R(N+1)×D , and Value V ∈ R(N+1)×D , respectively, and then sliced into H attention heads to perform the parallel operations:

where Dh = D H is the feature dimension of the single head output. In particular, the attention scores of class token xcls and other tokens can be written as:

which is used to reflect which tokens are contributing to the classification. Next, MHSA concatenates the Context of all attention heads together and project them through a matrix.

3.2 Class Attention ScoreWeighted by Attention Head Importance

Popular metrics for evaluating the importance of a token include its similarity score to the class token [Liang et al., 2022; Xu et al., 2022] and the attention it receives from other tokens [Goyal et al., 2020; Kim et al., 2021]. However, we observed that both approaches ignore the diversity of attention heads, and the scores of different heads are treated equally, in other words, the final score is their average value.

Considering that different tokens may receive diverse attention in multiple attention heads, we propose a metric to estimate the token-level head importance, and it simply relies on the intermediate results of MHSA. Taking the i-th input token xi as an example, firstly, we define the Context of the h-th head of the l-th layer as Context(h,l) (xi) ∈ RDh , then we calculate the l2-norm along the last dimension as the importance H(h,l) (xi) of the h-th head like Equation (5):

Our motivation for estimating head importance derives from CNN pruning, of which l2-norm [He et al., 2018] is a commonly used metric for filter significance rating. In addition, our metrics are token-level, different from the previous head pruning work [Michel et al., 2019]. Next, we take the proportion of the importance of the h-th attention head to the sum of all head as its weighting factor like Equation (6):

We visualize the token-level head importance in the Figure 3 and brighter areas indicate that the current attention head is more important for this token. Taking the second image as an example, the foreground tokens focus more on head 1, 4 and 5, while the background tokens clearly favor head 2 and 3. It is clear that the diversity of individual heads should be taken into account when scoring tokens.

The vanilla class attention score A(h,l) (xcls,i) of the i-th token is chosen as the basic evaluation metric, and we can estimate token score S l (xi) more accurately through head importance weighting as Equation (7). The flow of head importance weighted class attention score is shown in Figure 4. Compared to existing adaptive token pruning works, our approach takes into account the attention head diversity and has no reliance on any extra scoring mechanisms.

3.3 Adaptive Token Pruning Based on Learnable Thresholds

In order to achieve sample-adaptive token pruning while minimizing the operation cost. We introduce three stage-wise learnable thresholds to control token sparse behavior, following LTP [Kim et al., 2021]. Figure 4 illustrates the overall framework of our proposed method. We usually divide the successive transformer blocks into four stages and insert a learnable threshold θ before the second, third and fourth stages, respectively. Adaptive Sparse Module is a comparator based on thresholds and token scores. We keep tokens with scores greater than the threshold, as in Equation (8):

where M n (xi) is a binary mask to indicate whether the i-th token is pruned or not, and S l (xi) denotes token score of the layer before the n-th stage. The token scores corresponding to different input images usually present different distributions, so using a threshold to truncate the distribution can yield a specific pruning strategy for each instance.

Considering images in the same mini-batch have different pruning configurations, for the sake of parallel training, we cannot simply discard uninformative tokens during the training process. M n (xi) can be used to explicitly cut the connection between uninformative tokens and useful tokens. There are two masking methods, attention mask [Rao et al., 2021] and activation mask [Kim et al., 2021], the former acting on the attention matrix and the latter on the output of the FFN.

However, it is not easy to optimize the learnable thresholds. The pruning mask comes from comparison, blocking the gradient back propagation, making the threshold untrained. We transform the hard mask into a soft differentiable mask using the sigmoid function:

To approximate the hard mask, we employ a temperature parameter T , where the Sigmoid function behaves close to the step function at a sufficiently high temperature. The soft mask is differentiable, and by using the gradient straight through estimator (STE), we are able to optimize the learnable threshold as normal.

3.4 Budget-Aware Training

We achieve token pruning by constraining the target computation across the dataset. Compared to DynamicViT [Rao et al., 2021], which manually sets the token sparsity ratio at each stage to indirectly control the complexity, our method automatically achieves a good trade-off between accuracy and complexity given only a budget. Specifically, we propose a budget-aware loss (LFLOPs). Given a mini-batch inputs x with size of B, we can obtain their average FLOPs and then make the actual computational cost close to the target budget using MAE loss as Equation (10):

where FLOPs is a function to calculate the actual operations for different inputs under the effect of all thresholds Θ, and t is the expected complexity. The intention of designing the budget-aware loss is consistent with our sample-adaptive approach. The learnable thresholds are optimized to the appropriate range in a data-driven manner under the budget constraint, without imposing other artificial restrictions.

Our training objective include other two parts. The first part is the regular cross-entropy loss (LCE) as following equation:

where y denotes the ground truth labels and ˆ y the Softmax output. To further improve the performance, we consider transferring the knowledge of the full model to the compressed network during training. Let the distribution ˆ yt denote the prediction of the teacher network, and we use KL divergence to minimize the gap between the output of the student network and the teacher network. For vision transformers like LV-ViT [Jiang et al., 2021] with an additional linear layer to integrate all image tokens, it also needs to be aligned with the teacher network. Therefore, distillation loss (Ldistill) can be written as:

The overall training objective is a combination of the above three components:

where λ is used to control the loss balance, and we set λ1 = 2, λ2 = 0.5 in our experiments.

Inference. Learnable thresholds are fixed after the training. Given an input image to do inference, we only need the intermediate result of MHSA to get the token score effortlessly. The pruning process is also simple enough that our method only needs one comparison to know which tokens are to be kept, saving the topK computation cost compared to the previous work [Rao et al., 2021; Xu et al., 2022; Liang et al., 2022].

4 Experiments

4.1 Implementation Details

Our experiments are conducted on the ImageNet-1K [Deng et al., 2009] classification dataset, with token pruning performed on the popular DeiT [Touvron et al., 2021a] and LV-ViT [Jiang et al., 2021] models. We finetune the pre-trained model by 30 epoch to obtain the compressed network, and most of the training settings stay the same as the originals.

4.2 Main Results

Performance Comparisons With Existing Pruning Methods. We test the Top-1 accuracy, throughput and hardware latency of three models, DeiT-S, LVViT-S and LVViT-M, under different pruning budgets and compare them with previous token pruning work. The experimental results on accuracy are illustrated in Table 1. Compared to previous work, our proposed adaptive sparse ViT achieves state-of-the-art performance with similar complexities. For all models, the Top-1 accuracy degradation of our pruned models is controlled within 0.2% when the computation decreases by 30%∼35%. When the compression rate of DeiT-small is further increased to 50%, the advantage of sample-adaptive token sparsity becomes more obvious over the fixed pruning rate approaches [Rao et al., 2021; Liang et al., 2022], as they may be forced to discard some important tokens. Moreover, our method is far better than the sample-adaptive [Pan et al., 2021; Yin et al., 2021], owing to the more accurate token scoring mechanism. The reinforcement learning strategy employed by IA-RED2 [Pan et al., 2021] is difficult to converge.

Efficiency Comparisons With Existing Pruning Methods.

The experimental results on efficiency are in Table 1. Note that some previous methods did not release code and trained weights, so we cannot compare and report them. Compared to previous work, AS-ViT achieves the best results in terms of accuracy, throughput and hardware latency. In contrast to methods like DynamicViT [Rao et al., 2021] that use extra modules to calculate token scores, AS-ViT relies only on the intermediate results of MHSA to evaluate tokens, thus saving expensive computations and achieving higher throughput metrics. The threshold-based comparator employed by Adaptive Sparsity Module saves the computational cost of ranking all tokens, which makes our latency metrics significantly better than previous work.

Comparisons With Other Models. We compare the complexity and accuracy of our adaptive sparse LV-ViT (abbreviated as AS-LV-ViT) with other sota models on ImageNet in Figure 1. Our AS-LV-ViT shows quite competitive performance under different complexities, with far higher throughput than other CNN [Tan and Le, 2019; Liu et al., 2022b] and ViT [Chu et al., 2021; Wang et al., 2021; Xu et al., 2021; Wu et al., 2021; Touvron et al., 2021b; Jiang et al., 2021; Yuan et al., 2021] while still providing advanced accuracy. In addition, by simply adjusting the budget, our approach achieves a better accuracy-complexity trade-off compared automatically, avoiding the tedium of manual design.

4.3 Ablation Analysis

Effectiveness of Each Sub-Module. In Table 2, we study the effect of each module in detail. ASM represents the Adaptive Sparsity Module, and HS is attention head importance weighted class attention score. In the ablation experiments, we use fixed ratio topK module instead of ASM and vanilla class attention score to replace HS. It is obvious that the Adaptive Sparsity module significantly improves the model performance compared with the fixed ratio module by 0.2%, which fully illustrates the necessity and effectiveness of sample adaptive token pruning. With ASM, the attention head weighted class attention score improves the precision by 0.07% compared to the original metric without a noticeable reduction in throughput. In addition, we can observe that HS needs to be used with ASM to have better results.

Effectiveness of Training Techniques. We apply masking strategy to achieve parallel training and knowledge distillation to stabilize the training process. The respective experimental results are listed in Table 3. We conduct experiments with DeiT-Small. Obviously, attention_mask is far better than acti-vation_mask, which can shield the uninformative tokens from interacting with other tokens. And by transferring the knowledge of full model to the compressed model, we can further improve the accuracy.

Comparison With Raw Attention Baseline of Pre-Trained Model. In this part, we consider the attention of the pretrained model as a token evaluation metric and discard unimportant tokens in a fixed proportion. In addition, we apply head importance weighted class attention score (HS) directly on the raw attention baseline to fully demonstrate the effectiveness of our method. The data in Table 4 present the accuracy of each method at different keeping rates. The attention map of the pre-trained model is no longer reliable and the accuracy decreases significantly as the compression ratio gradually increases. While our method can still achieve the similar performance of the original model by fine-tuning. Furthermore, our head importance weighted token score can produce positive results without training.

Batch Inference. Our approach achieves both sampleadaptive and batch-adaptive token pruning. The fact that ASViT is well suited for single-image computation scenarios does not mean that it cannot be used for parallel inference. Since we use the average complexity of mini-batch to calculate the budget-aware loss, we can also do parallel inference from the mean value of the number of kept tokens within a batch. And the experimental results in the Table 5 show that this does not cause significant accuracy degradation.

Pruning Location. Referring to previous works [Rao et al., 2021; Liang et al., 2022], we adopts a progressive token sparse strategy. The position and number of our Adaptive Sparsity Module are need to consider. We experiment with several insertion configurations to demonstrate that the current strategy is accuracy-latency optimal. The results are given in Table 6. The current method has higher accuracy compared to the [3,6,9] pruning strategy. This may be attributed to the early layer's class attention score not being stable enough, causing some information tokens to be discarded incorrectly. When increasing to 4 modules, there is no significant change in throughput and latency. We further insert thresholds from the 3rd to 11th layers, and the training becomes unstable and brings severe accuracy degradation, which we speculate is attributed to the difficulty of optimizing multiple thresholds in the limited fine-tuning. In addition, frequent reorganization of tokens causes non-negligible time consumption.

Performance on Large Model and Different Resolutions.

To fully demonstrate the effectiveness of our method, we perform token sparse on the large baseline model DeiT-Base at different input resolutions. As illustrated in the Table 7, AS-ViT performs better than previous work under the same resolution and complexity. When using larger resolutions, our method is significantly better than IA-RED2 [Pan et al., 2021]. Furthermore, at a resolution of 384x384, the accuracy degradation is smaller compared to 224x224, suggesting that there is greater redundancy at higher resolutions, which can improve efficiency through token pruning.

Visualization. To analyze the pruning behavior of our adaptive token sparse method on different images, we count the amount of kept tokens of the AS-DeiT-S model in each stage for ImageNet validation dataset images and plot their distribution in Figure 5. The number of each stage basically shows a Gaussian distribution, peaking at a certain value. This suggests that a fixed proportion of pruning is also relatively reasonable, since the easy and difficult samples are only a minority. However, our adaptive token sparse method can give the corresponding pruning configurations for samples with different recognition difficulties, which is reflected in the distribution extending to both sides. Further, we select three representative images for visualization of token pruning results in Figure 6. For easy samples, AS-ViT discards plenty of useless tokens in the early stage to save computational cost, While for complicated and hard to recognize instances, the model tends to discard more tokens in later stages with higher confidence.

5 Conclusion

In this work, we propose a sample-adaptive token pruning method, which effortlessly evaluates token importance via fully exploiting the MHSA mechanism, and then introduces learnable thresholds to accomplish pruning. Our proposed budget-aware loss can effectively constrain the thresholds so that the pruned model reaches the complexity budget. Compared to previous work, our approach does not require additional sub-networks to compute token scores and also performs specific pruning strategies for different samples just by comparing with thresholds. Experimental results on various models show that our AS-ViT greatly improves the throughput and achieves lower latency without significantly affecting the accuracy. In the future, we will transfer token pruning to downstream tasks and combining it with static parameters pruning for further acceleration.

Contribution Statement

Xiangcheng Liu and Tianyi Wu contribute equally to this work. This work was done by Xiangcheng Liu as an intern at the Institute of Deep Learning, Baidu Research. The corresponding author is Guodong Guo.