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