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

Vision transformers ( ViTs ) have recently received explosive popularity , but their enormous model sizes and training costs remain daunting . Conventional posttraining pruning often incurs higher training budgets . In contrast , this paper aims to trim down both the training memory overhead and the inference complexity , without sacrificing the achievable accuracy . We carry out the first-of-its-kind comprehensive exploration , on taking a unified approach of integrating sparsity in ViTs “ from end to end ” . Specifically , instead of training full ViTs , we dynamically extract and train sparse subnetworks , while sticking to a fixed small parameter budget . Our approach jointly optimizes model parameters and explores connectivity throughout training , ending up with one sparse network as the final output . The approach is seamlessly extended from unstructured to structured sparsity , the latter by considering to guide the prune-and-grow of self-attention heads inside ViTs . We further co-explore data and architecture sparsity for additional efficiency gains by plugging in a novel learnable token selector to adaptively determine the currently most vital patches . 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 . Our codes are available at https : //github.com/VITA-Group/SViTE . 1 Introduction . Recent years have seen substantial efforts devoted to scaling deep networks to enormous sizes . Parameter counts are frequently measured in billions rather than millions , with the time and financial outlay necessary to train these models growing in concert . The trend undoubtedly continues with the recent forefront of transformers [ 1–3 ] for computer vision tasks . By leveraging self-attention , reducing weight sharing such as convolutions , and feeding massive training data , vision transformers have established many new state-of-the-art ( SOTA ) records in image classification [ 1 , 2 ] , object detection [ 4–7 ] , image enhancement [ 8 , 9 ] , and image generation [ 10–12 ] . Existing vision transformers and variants , despite the impressive empirical performance , have in general suffered from gigantic parameter-counts , heavy run-time memory usages , and tedious training . That naturally calls for the next step research of slimming their inference and training , without compromising the performance . Model compression and efficient learning are no strangers to deep learning researchers , although their exploration in the emerging vision transformer field remains scarce [ 13 ] . Among the large variety of compression means [ 14 ] , sparsity has been one of the central themes since the beginning [ 15 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Conventional approaches first train dense networks , and then prune a large portion of parameters in the trained networks to zero . Those methods significantly reduce the inference complexity . However , the price is to cost even more significant computational resources and memory footprints at training , since they commonly require ( multiple rounds of ) re-training to restore the accuracy loss [ 15–17 ] . That price becomes particularly prohibitive for vision transformers , whose vanilla one-pass training is already much more tedious , slow , and unstable compared to training standard convolutional networks . An emerging subfield has explored the prospect of directly training smaller , sparse subnetworks in place of the full networks without sacrificing performance . The key idea is to reuse the sparsity pattern found through pruning and train a sparse network from scratch . The seminal work of lottery ticket hypothesis ( LTH ) [ 18 ] demonstrated that standard dense networks contain sparse matching subnetworks ( sometimes called “ winning tickets ” ) capable of training in isolation to full accuracy . In other words , we could have trained smaller networks from the start if only we had known which subnetworks to choose . Unfortunately , LTH requires to empirically find these intriguing subnetworks by an iterative pruning procedure [ 18–27 ] , which still can not get rid of the expensiveness of posttraining pruning . In view of that , follow-up works reveal that sparsity patterns might emerge at the initialization [ 28 , 29 ] , the early stage of training [ 30 , 31 ] , or in dynamic forms throughout training [ 32–34 ] by updating model parameters and architecture typologies simultaneously . These efforts shed light on the appealing prospect of “ end to end ” efficiency from training to inference , by involving sparsity throughout the full learning lifecycle . This paper presents the first-of-its-kind comprehensive exploration of integrating sparsity in vision transformers ( ViTs ) “ from end to end ” . With ( dynamic ) sparsity as the unified tool , we can improve the inference efficiency from both model and data perspectives , while also saving training memory costs . Our innovative efforts are unfolded along with the following three thrusts : • From Dense to ( Dynamic ) Sparse : Our primary quest is to find sparse ViTs without sacrificing the achievable accuracy , and meanwhile trimming down the training memory overhead . To meet this challenging demand , we draw inspirations from the latest sparse training works [ 34 , 35 ] that dynamically extract and train sparse subnetworks instead of training the full models . Sticking to a fixed small parameter budget , our technique jointly optimizes model parameters and explores connectivity throughout the entire training process . We term our first basic approach as Sparse Vision Transformer Exploration ( SViTE ) . • From Unstructured to Structured : Most sparse training works [ 32 , 33 , 36–39 , 38 , 34 , 40 , 41 , 35 ] restricted discussion to unstructured sparsity . To attain structured sparsity which is more hardware-friendly , unlike classical channel pruning available for convolutional networks , we customize a first-order importance approximation [ 16 , 42 ] to guide the pruneand-grow of self-attention heads inside ViTs . This seamlessly extends SViTE to its second variant of Structured Sparse Vision Transformer Exploration ( S2ViTE ) . • From Model to Data : We further conduct a unified co-exploration towards joint data and architecture sparsity . That is by plugging in a novel learnable token selector to determine the most vital patch embeddings in the current input sample . The resultant framework of Sparse Vision Transformer Co-Exploration ( SViTE+ ) remains to be end-to-end trainable and can gain additional efficiency . Extensive experiments are conducted on ImageNet with DeiT-Tiny/Small/Base . Results of substantial computation savings and nearly undamaged accuracies consistently endorse our proposals ’ effectiveness . Perhaps most impressively , we find that the sparse ( co- ) training can even improve the ViT accuracy rather than compromising it , making sparsity a tantalizing “ free lunch ” . For example , applying SViTE+ on DeiT-Small produces superior compressed ViTs at 50 % model sparsity plus 5 % data sparsity , saving 49.32 % FLOPs and 4.40 % running time , while attaining a surprising improvement of 0.28 % accuracy ; even when the data sparsity increases to 10 % ( the model sparsity unchanged ) , there is still no accuracy degradation , meanwhile saving 52.38 % FLOPs and 7.63 % running time . 2 Related Work . Vision Transformer . Transformer [ 43 ] stems from natural language processing ( NLP ) applications . The Vision Transformer ( ViT ) [ 1 ] pioneered to leverage a pure transformer , to encode an image by splitting it into a sequence of patches , projecting them into token embeddings , and feeding them to transformer encoders . With sufficient training data , ViT is able to outperform convolution neural networks on various image classification benchmarks [ 1 , 44 ] . Many ViT variants have been proposed since then . For example , DeiT [ 2 ] and T2T-ViT [ 45 ] are proposed to enhance ViT ’ s training data efficiency , by leveraging teacher-student and better crafted architectures respectively . In addition to image classification , ViT has attracted wide attention in diverse computer vision tasks , including object detection [ 4–7 ] , segmentation [ 46 , 47 ] , enhancement [ 8 , 9 ] , image generation [ 10–12 ] , video understanding [ 48 , 49 ] , vision-language [ 50–57 ] and 3D point cloud [ 58 ] . Despite the impressive empirical performance , ViTs are generally heavy to train , and the trained models remain massive . That naturally motivates the study to reduce ViT inference and training costs , by considering model compression means . Model compression has been well studied in both computer vision and NLP applications [ 59–61 , 42 , 62 , 21 ] . Two concurrent works [ 13 , 63 ] made initial attempts towards ViT post-training compression by pruning the intermediate features and tokens respectively , but did not jointly consider weight pruning nor efficient training . Another loosely related field is the study of efficient attention mechanisms [ 64 , 10 , 52 , 65–75 ] . They mainly reduce the calculation complexity for self-attention modules via various approximations such as low-rank decomposition . Our proposed techniques represent an orthogonal direction and can be potentially combined with them , which we leave as future work . Another latest concurrent work [ 76 ] introduced an interpretable module to dynamically and gracefully drop the redundant patches , gaining not only inference efficiency but also interpretability . Being a unique and orthogonal effort from ours , their method did not consider the training efficiency yet . Pruning and Sparse Training . Pruning is well-known to effectively reduce deep network inference costs [ 77 , 15 ] . It can be roughly categorized into two groups : ( i ) unstructured pruning by removing insignificant weight elements per certain criterion , such as weight magnitude [ 78 , 15 ] , gradient [ 16 ] and hessian [ 79 ] ; ( ii ) structured pruning [ 80–82 ] by remove model sub-structures , e.g. , channels [ 80 , 81 ] and attention heads [ 42 ] , which are often more aligned with hardware efficiency . All above require training the full dense model first , usually for several train-prune-retrain rounds . The recent surge of sparse training seeks to adaptively identify high-quality sparse subnetworks and train only them . Starting from scratch , those methods learn to optimize the model weights together with sparse connectivity simultaneously . [ 32 , 33 ] first introduced the Sparse Evolutionary Training ( SET ) technique [ 32 ] , reaching superior performance compared to training with fixed sparse connectivity [ 83 , 36 ] . [ 37–39 ] leverages “ weight reallocation '' to improve performance of obtained sparse subnetworks . Furthermore , gradient information from the backward pass is utilized to guide the update of the dynamic sparse connectivity [ 38 , 34 ] , which produces substantial performance gains . The latest investigations [ 40 , 41 , 35 ] demonstrate that more exhaustive exploration in the connectivity space plays a crucial role in the quality of found sparse subnetworks . Current sparse training methods mostly focus on convolutional networks . Most of them discuss unstructured sparsity , except a handful [ 84 , 30 ] considering training convolutional networks with structured sparsity . 3 Methodology . Our SViTE method ( and its variants S2ViTE and SViTE+ ) is inspired from state-of-the-art sparse training approaches [ 34 , 35 ] in CNNs . This section presents the sparse exploration of ViT architectures , then shows the detailed procedure of input token selection for extra efficiency gains .

Summary: 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.