Chasing Sparsity in Vision Transformers: An End-to-End Exploration

Abstract: Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, Zhangyang Wang

Reviews and author responses:

This work proposes sparse ViT algorithms, named as SViTE and its variants S^2ViTE, and SViTE+, to explore high-quality sparse patterns in both ViT’s architecture and input token embeddings. The proposed methods can jointly optimize model parameters and explore connectivity throughout training. Experiments are conducted on ImageNet to validate the effectiveness of the proposed methods and find that SViTE+ can mine crucial connections and input tokens.  This work aims to explore the sparse vision transformers by unstructured and structure pruning methods. The proposed SViTE and its variants S^2ViTE, and SViTE+ achieve impressive results on baselines. 1. The reproduced result on DeiT-base is much lower than the reported results (80.98 v.s. 81.80). What's the potential reason? DeiT has opensource the training scripts and pretrained weights. If considering the reported results of DeiT as the baseline, SViTE and its variants do not achieve such impressive improvements as stated in the paper.
2. The proposed method is only validated on the DeiT but not validated on more advanced image transformer architecture such as TNT[1], T2T[2], PVT[3], etc. It would be better to validate the effects of the proposed methods (at least S^2ViTE) on more architectures.   


[1] Han, Kai, et al. "Transformer in transformer." arXiv preprint arXiv:2103.00112 (2021). 
[2] Yuan, Li, et al. "Tokens-to-token vit: Training vision transformers from scratch on imagenet." arXiv:2101.11986 (2021). 
[3] Wang, Wenhai, et al. "Pyramid vision transformer: A versatile backbone for dense prediction without convolutions." arXiv preprint arXiv:2102.12122 (2021).

 Dear Reviewer EjXo,
 
Many thanks for reviewer EjXo's positive reassessment, and the acknowledgment of our efforts in the rebuttal. We are glad that our response has solved the concern. The extra experiments and analyses will be included in our revision.
 
We agree that the attention map is related to the token sparsity if the attention weight is sparse. And we also politely point out that the token sparsity itself will not lead to attention map sparsity.
 
We are again very thankful for your time and support!
 
Best wishes,
 
Authors


 Dear Reviewer tKZe,
 
We thank reviewer tKZe’s time for the review, and we really hope to have a further discussion with reviewer tKZe to see if our response solves the concerns.
 
We would sincerely appreciate it if reviewer tKZe could reply to the most important points in our rebuttal. For example, as we pointed out that all stated accuracy gains and drops in our paper are already based on the reported accuracy of DeiT. Meanwhile, as you suggested, we have further validated the effectiveness of our proposal on extra transformer backbones (e.g., TNT).
 
We genuinely hope reviewer tKZe could kindly check our response. Thank you!
 
Best wishes,
 
Authors


The author proposes three methods for achieving sparse vision transformers to reduce the number of parameters and FLOP. By using a dynamic structured sparsity injecting method, the authors are able to achieve promising speedup without loss of accuracy. For an example, the author claims to achieve 14.7% run time reduction with 0.42% accuracy improvement. However, the training settings of the author's experiments and the DeiT baseline are different. Thus, there are some ambiguities exits.  Strength:

1. The paper is well organized and easy to follow.
2. The sparsity in vision transformers is an interesting problem to explore as it could be potentially used to apply ViTs on large input images which otherwise will cause a huge cost of the computation memory.
3. The token selector is an interesting propose to dynamically select the top-k tokens that are more important to a given input sample.

Weakness:

1. The proposed methods are trained with longer epochs than the baseline models (600 VS 300). It has been proven that training with longer epochs will further increase the performance of the ViTs with proper augmentation. Thus, it is not clear if the proposed method is able to achieve better accuracy with less parameters and computations or the method is actually applied on a stronger baseline and the the accuracy is actually reduced from the stringer baseline.

2. The idea of selecting Top-k tokens is interesting but the current analysis is not enough to support its efficiency. There is no ablations on the impacts of the selector. For example, as shown in Fig. 6 in [1], the sparsity of the attention map is increasing from bottom layers to the top layers. As the attention map sparsity is related to the token sparsity, it may be more convincing to show if following a decaying sparsity could better preserve the model accuracy and the proposed method could learned a layer-wise sparsity that is following such kind of observations.

3. Based on the results in Tab. 4 in the paper, it seems that the accuracy drop is more significant on smaller models. Considering that the proposed methods are all trained with 2x longer epochs, it is not clear what is the performance drop on large models given a well trained baseline.

Overall, I agree that injecting sparsity into ViTs is an interesting and meaningful topic to explore. However, based on the current experiments results and ablation analysis, I am not clear the impacts of the proposed method on the accuracy.  I am willing to increase my score if the author could address my concerns sufficiently.

[1] Zhou, D., Kang, B., Jin, X., Yang, L., Lian, X., Jiang, Z., ... & Feng, J. (2021). Deepvit: Towards deeper vision transformer. arXiv preprint arXiv:2103.11886. Yes

 Dear Authors,

Thanks for your response. I agree with the fact that the total training computation cost is comparable with 600 epochs, comparing to the original DeiT model. 

However, I did not agree with the comments that the attention map is not related to the token sparsity. In fact, the attention weights are multiplied with the input tokens to generate the output token. If the attention weights are sparse, the corresponding input tokens will be masked out and will not contribute to the output token which means those tokens can be removed. Do correct me if I am wrong. But the extra experiments on the layer-wise sparsity are supporting the above reasoning. So, in general, I believe the proposed methods is working to find out the reasonable sparsity. Thus, I will increase my score.

 We are very glad and appreciated that you had a positive initial impression, and we provide pointwise responses to your concerns below.

**[Cons1. Redundant tokens]** Thanks for your comments. We agree that visually many tokens might be redundant for making predictions. However, how many tokens can be abandoned should be highly dependent on the intrinsic complexity of input images and objects. Although our token selector performs input-dependent selection, it is still based on a predefined dropping ratio. Thus, a high token dropping rate (i.e., 30%) for all images might be too aggressive for those complex images and leads to overall performance degradation. Our future work is thus to explore input-adaptive sparsity ratios.

**[Cons2. SViTE+ is not compatible with structural sparsity?]** Thanks for the question. It is mainly because both structural pruning and token dropping will aggressively shrink the intermediate feature dimensions. Such aggressive reduction of ViT width might incur accuracy degradations. Based on this point, we agree that SViTE+ is not very compatible with structural sparsity. However, although the performance of S^2ViTE+ is not that impressive compared to its unstructured version (i.e., SViTE+), we think the accuracy degradation is still reasonable (e.g., 1% drop when pruning 10% tokens with 40% structural sparsity).

**[Cons3. Confusing presentation in the abstract, writing inconsistency, and typos.]** We tremendously appreciate your revision and typo comments and promise to thoroughly address them in the accepted version. 
- We will *(i) reorganize the confusing statements in the abstract to improve the readability*, and *(ii) fix the writing inconsistency issue to avoid misleading highlights*. Specifically, we change the original description to “...... Extensive results on ImageNet with diverse ViT backbones validate the effectiveness of our proposals which obtain significantly reduced computational cost and almost unimpaired generalization. Perhaps most surprisingly, we find that the proposed sparse (co-)training can sometimes improve the ViT accuracy rather than compromising it, making sparsity a tantalizing “free lunch”. For example, our sparsified DeiT-Small at (5%, 50%) sparsity for (data, architecture), improves 0.28% top-1 accuracy, and meanwhile enjoys 49.32% FLOPs and 4.40% running time savings …...”.
- Thanks so much for pointing it out, and sorry for the annoying typos. We will perform thorough proofreading and correct all of them in our final version.


 Thanks for rating our paper as interesting and meaningful. We provide pointwise responses to your concerns about our experiment settings below.

**[Cons1. (Major Con.) Different training settings and ambiguities?]**

*<Why do we choose 600 training epochs?>* We respectfully argue that our setting is fair and reasonable.  Choosing 600 training epochs for sparse training is to maintain similar **training FLOPs** compared to dense ViT training. Specifically, if training a dense DeiT model for 300 epochs needs 1x FLOPs, training SViTE-DeiT at 40% sparsity for 600 epochs needs ~0.95x FLOPs. In summary, we compare our SViTE (600 epochs) and DeiT baselines (300 epochs) based on a similar training budget. **Such comparison fashion is widely adopted in sparse training literature like [1,2] (see Table 2’s caption in [1] and Figure 2 in [2] for details).** In our final version, we will include the above discussions. Meanwhile, note that the reported running time is per epoch saving (i.e., total running time / total epoch), which would not be affected by the number of training epochs.

*<Baseline models with longer epochs.>* Actually, the performance of DeiT training without distillation saturates after 300\~400 epochs, as stated in [3]. We also conduct longer epoch (600 epochs) training for DeiT-small and -base models as you suggested. Our results collected in the table below align with the original DeiT paper [3]. It suggests that our proposed SViTE is still able to achieve better accuracy with fewer parameters and fewer training&inference computations. Specifically, at 40% structured sparsity, our sparsified DeiT-Base can achieve 0.21% accuracy gain, at **~51% training FLOPs**, 33.13% inference FLOPs, and 24.70% running time savings, compared to its dense counterpart with 600 training epochs.

|Metrics|Deit-small 300 epochs|Deit-small 600 epochs|Deit-base 300 epochs|Deit-base 600 epochs|
|:-:|:-:|:-:|:-:|:-:|
|Accuracy|79.90|80.02 (+0.12)|81.80|82.01 (+0.21)|

*<Confusion in Table 4.>* Based on the above longer epoch results, the confusion about table 4 should be addressed: the accuracy drop is more significant on smaller models compared to larger models as also indicated in line 299. A possible reason is that huge ViTs with larger capacity are more likely to have superfluous self-attentions or redundant parameters, and therefore are more amenable to sparsification compared to smaller ViTs. 

**[Cons2. The Top-K token selector: the current analysis is not enough to support its efficiency and no ablation?]** We respectfully argue that the ablation in Table 6 indicates the impact of the token selector on the accuracy, running time, and FLOPs saving. Table 6 shows that with an appropriate data sparsity from the token selector, SViTE+ models can gain significantly reduced running time, FLOPs, and even improved accuracy. It supports the efficiency and effectiveness of our proposed token selector. 

**[Cons3. The Top-K token selector: the sparsity of the attention map.]** In Figure 4, we actually have provided the attention map analyses. We observe that in the bottom layers, the attention probabilities are more centered at several heads; while in the top layers, the attention probabilities are more uniformly distributed. This kind of tendency is well preserved by our sparse ViT (SViTE) from Dense ViTs.

Note that in our case, the token sparsity isn’t necessarily correlated with the attention map sparsity. It is because our proposed token selector gains extra efficiency by reducing tokens and shrinking intermediate feature dimensions, which would not generate element-wise sparsity in attention maps. Meanwhile, as you suggested, we conduct extra ablation studies and visualizations about the layerwise sparsity of attention maps. Dense DeiT-Small and SViTE+-Small with 10% data sparsity and 50% model sparsity are adopted for experiments. We calculate the percentage of elements in attention maps whose magnitude is smaller than $10^{-4}$. We observe from the following table that the bottom layers’ attention maps of SViTE+ are denser than the ones in dense ViT, while it is opposite for the top layers. We will add the citation of deepViT and the above discussions in our revision.

|Method|Layer 1|Layer 2|Layer 3|Layer 4|Layer 5|Layer 6|Layer 7|Layer 8|Layer 9|Layer 10|Layer 11|Layer 12|
|:-|:-:|:-:|:-:|:-:|:-:|:-|:-:|:-:|:-:|:-:|:-:|:-|
|Dense DeiT-Small|28.44|44.70|37.04|13.87|13.62|13.67|12.98|5.99|5.08|3.46|2.74|15.05|
|SViTE+ -Small|24.37|33.24|28.05|19.51|9.03|13.79|13.24|8.88|6.78|5.65|2.88|35.22|

[1] Do We Actually Need Dense Over-Parameterization? In-Time Over-Parameterization in Sparse Training?

[2] Rigging the Lottery: Making All Tickets Winners

[3] Training data-efficient image transformers & distillation through attention


 **[Cons1. Lower performance than the reported results and less impressive improvements?]** We respectfully disagree. In Table 2/3/4/5, we have included both reproduced and reported results. **Most importantly, we have already used the reported results of DeiT as our baseline for comparison**. The stated accuracy gain and drop in our paper are based on the reported accuracy of DeiT. For example, in abstract, introduction and experiment sections, the stated accuracy gain of S^2ViTE-base 0.42% is calculated from 82.22% (Ours) - 81.80% (Reported); the stated accuracy gain of SViTE+-small 0.28% is also calculated from 80.18% (Ours) - 79.9% (Reported). 

Also, the most impressive point of our results is that we achieve substantially reduced computational cost, almost unimpaired generalization, and sometimes even improved accuracy. For instance, at 40% structured sparsity, our sparsified DeiT-Base can achieve 0.42% accuracy gain (**compared to the reported results**), at 33.13% FLOPs and 24.70% running time savings, compared to its dense counterpart. The impressiveness of our results is also highly acknowledged by **Reviewer 6KF1** as “... are significant in many cases” and by **Reviewer ZXSe** as “this paper establishes so far the most promising trade-off between ViT accuracy and efficiency”. 

For the lower reproduced results, the potential reasons could be the difference in computational facilities, environments, etc. We run multiple experiments during rebuttal periods, the achieved accuracies range from 80.70% to 81.36% for DeiT-base. To mitigate the confusion, we will only present the reported results in our final version.

**[Cons2. Experiments on other backbones?]** It is worth mentioning that our proposed frameworks (SViTE, S^2ViTE, SViTE+) are independent of the backbone architectures, and can be easily plugged in other vision transformer models. As suggested, we implemented both SViTE and S^2ViTE on TNT-S [1]. SViTE-TNS-S gains 0.13 accuracy improvements (Ours: 81.63 v.s. TNT-S: 81.50) and 37.54% FLOPs savings at 40% unstructured sparsity; S^2ViTE-TNS-S obtains 32.96% FLOPs and 23.71% running time reductions at 40% structured sparsity with almost unimpaired accuracy (Ours: 81.34 v.s. TNT-S:81.50). Thanks for the constructive suggestions. We will include more ViT backbone results (e.g., both T2T [2] and PVT [3]) together with extra citations and discussions in our revision.

[1] Han, Kai, et al. "Transformer in transformer." arXiv preprint arXiv:2103.00112 (2021). 

[2] Yuan, Li, et al. "Tokens-to-token vit: Training vision transformers from scratch on imagenet." arXiv:2101.11986 (2021). 

[3] Wang, Wenhai, et al. "Pyramid vision transformer: A versatile backbone for dense prediction without convolutions." arXiv preprint arXiv:2102.12122 (2021).

 We are very glad you had a positive initial impression, and we provide pointwise responses for your concerns below.

**[Cons1. Novelty incremental?]** We respectfully disagree. Our paper’s technical novelty mainly lies in the following two aspects. 

*<Studies of Structured Sparsity>* First, the prior related work ([24], [25]) restricted discussion to unstructured sparsity. We not only for the first time investigate dynamic sparsity in vision transformers (ViTs), but also explore more hardware-friendly structured sparsity by customizing a first-order importance approximation to guide the prune-and-grow of self-attention heads in ViTs (“unlike CNNs”, this novelty is also acknowledged by **Reviewer ZXSe**). 

*<Data and Model Sparsity Co-exploration>* Second, as appreciated by **Reviewer EjXo** (“the token selector is an interesting propose”), we propose a novel unified framework, i.e., SViTE+, to jointly explore critical sparse topology and important token embeddings, leading to improved efficiency and accuracy. 

**[Cons2. Random patch selection patterns?]** Thanks for pointing it out. First, we think the reason is that we only remove the most insignificant 10% tokens in Figure 5. When we increase the input data sparsity, e.g., from 10% to 40%, more irrelevant background patches are dropped, and a clearer selection pattern can be observed with a bit more accuracy drop. We will include these high data sparsity visualizations in our final version. 

Second, we argue that the selection patterns in Figure 5 are also not that random, since the patches corresponding to the main objects are typically not dropped.

**[Cons3. Visualization is not clear.]** Thanks for pointing it out. For better visualiablity, we will include zoom-in versions for visualization figures in our final version. In Figure 4, the black squares are the query pixels at the center of each attention map visualization.

Vision Transformers (ViT) have attracted much attention recently; however their large model sizes and high training cost pose practical problems for computer vision tasks. This paper introduces sparse ViT exploration algorithms ($SViTE$, $S^2ViTE$, $SViTE+$) to explore sparse patterns in the ViT architectures. The goal is to alleviate the training memory bottleneck and achieving high inference efficiency. The authors conducted comprehensive experiments on ImageNet dataset to validate their approach.  **Strengths:**

- The paper is clear, well-written and well organized.
- The authors included different illustrations and descriptions of the proposed algorithms. In addition to including several visualizations such as the sparse connectivity patterns and self-attention heatmaps.
- The experimental results show the efficiency of the proposed methods. The gain in the classification accuracy and the runtime savings are significant in many cases.
- The proposed methods could open a way for training and utilizing Vision Transformer models in an efficient way for complex computer vision tasks especially for low resources devices and real-time applications. 


**Weaknesses:** 

- The novelty of the proposed methods seem to be highly incremental. It seems the methods/ideas presented here were borrowed from the prior related work ([24], [25]); the difference seems to be applying those methods on Vision Transformers. 
- In the visualization Figure 5, the dropped/selection patches patterns seem to be random. There is high similarity between many of the background patches and other patches that weren't dropped. 
- Some visualizations are not clear enough. For example, in Figure 4, I couldn't notice the black squares mentioned in the caption.

 There is no potential negative societal impact. 

This work pioneered to study of various sparse training and pruning techniques to the recently proposed vision transformers (ViT). It shows the effectiveness of pruning/sparse training/data selection in such settings. This is a  interesting work for both ViT and sparse learning communities.   Positive 
+	This paper presents a particularly comprehensive study for bringing sparsity into ViTs. Sparsity itself is addressed from various angles, e.g., data sparsity and sparse training connectivity (structured and unstructured). 
+	The authors studied a structured version of sparse training. Unlike CNNs where structural sparsity is often enforced by channel prunin, the authors leveraged a BERT-like idea to defined the importance of self-attention heads, by first-order loss approximation. 
+	The authors experimented with various ViT and DeitT models on ImageNet. Their results are very useful and interesting. Sparsely trained VIT can outperform the dense baseline and other state-of-the-art methods to obtain sparsity. 
+	Overall, this paper establishes so far the most promising trade-off between ViT accuracy and efficiency.  Interestingly, with similar parameter counts, a large sparse ViT consistently outperforms the corresponding smaller dense ViT.
+	The work suggests sparsity conveys not only higher efficiency for transformer, but even sometimes improved accuracy. That potentially will lend a better understanding of the optimization difficulty in training ViTs, which recently received much attention.
+	Figures, algorithms, pseudocodes, and summary-of-takeaway windows help understand this paper.

Negative
-	I am not impressed by SViTE+ results. Look at Table 6, only 5- 10% of tokens can be removed without compromising accuracy. 30% removed would drop accuracy to unacceptable. This is surprising to me, since apparently many tokens would be redundant for making predictions on an image. The author might have to explain why. 
-	Also, SViTE+ seems not compatible with structural sparsity.
-	The abstract highlights two results, both improving test accuracy while saving parameters and running time. Whereas those are good, the authors organized their statements in a quite confusing way, e.g., I was only expecting some counterintuitive results after “Perhaps most surprisingly,..”, but the authors laid another similar-flavor result before that one. That caused a negative reading experience (especially it is in the abstract). 
-	Also, since most results still mildly sacrificed some accuracy, it is inappropriate or misleading to highlight only two rare results with occasional accuracy gains. This looks like an accidental writing inconsistency issue but must be fixed.
-	The paper overall has a quite clear and logical outline. But it, unfortunately, suffers from a number of annoying typos such as “ scarifying” (many times!) and “ typologies”. More careful proofreading would have helped. 
 Efficient training and inference can save ML costs, leading to more economic ML deployment in real-world applications. 

Meta-review: Four experts reviewed this paper and gave ratings 6, 6, 6, and 7, respectively. The reviewers expressed concerns about novelty, writing, and some details, but they generally appreciated the study of sparsity in Transformers. Some reviewers especially liked the results. AC felt that the concerns can be addressed in a light revision of the paper after reading the rebuttal and the post-rebuttal discussions. Hence, the decision is to recommend the paper for acceptance. The authors are encouraged to make necessary changes in the camera-ready to address the reviewers' questions to the best of their ability. We congratulate the authors on the acceptance of their paper!


Decision: accept