NViT: Vision Transformer Compression and Parameter Redistribution

1 INTRODUCTION . Self-attention based transformer models demonstrate high model capacity , easy scalability , and superior ability in capturing long-range dependency ( Vaswani et al. , 2017 ; Devlin et al. , 2018 ; Radford et al. , 2018 ; Jiao et al. , 2019 ; Brown et al. , 2020 ) . They have thus been widely applied to natural language processing ( NLP ) tasks , and recently received growing attention for computer vision tasks . Vision Transformer , i.e. , the ViT ( Dosovitskiy et al. , 2020 ) , shows that embedding image patches into tokens and passing them through a sequence of transformer blocks can lead to higher accuracy compared to state-of-the-art CNN models . DEIT , recent work by Touvron et al . ( 2021 ) , further presents a data-efficient training method such that acceptable accuracy can be achieved without extensive pretraining . Offering competitive performance to CNNs under similar training regimes , transformers now point to the appealing perspective of solving both NLP and vision tasks with the same architecture ( Zheng et al. , 2021 ; Kim et al. , 2021 ; Jiang et al. , 2021 ) . Unlike CNNs built with convolutional layers that are mainly parameterized by few dimensions like the kernel size and the number of filters , the ViT has multiple distinct components , i.e. , QKV projection , multi-head attention , multi-layer perceptron , etc . ( Vaswani et al. , 2017 ) , each defined by independent dimensions . As a result , the dimensionality of each component in each ViT block needs to be carefully designed to achieve a decent trade-off between efficiency and accuracy . However , this is typically not the case for state-of-the-art models . Models such as ViT ( Dosovitskiy et al. , 2020 ) and DEIT ( Touvron et al. , 2021 ) mainly inherit the design heuristics from NLP tasks , e.g. , use MLP expansion ratio 4 , fix QKV per head , all the blocks having the same dimensions , etc. , which may not be optimal for computer vision ( Chen et al. , 2021a ) , causing significant redundancy in the base model and a worse efficiency-accuracy trade-off upon scaling , as we show extensively in our experiments . This work targets efficient ViTs by exploring latency-aware global structural pruning , leveraging the insights to redistribute parameters for enhanced accuracy-efficiency trade-off . Our approach , as visualized in Figure 1 , starts from analyzing the blocks in the computation graph of ViT to identify all the dimensions that can be independently controlled . We apply global structural pruning over all the components in all blocks . This offers complete flexibility to explore their combinations towards an optimal architecture in a complicated design space . Our global pruning utilizes an importance score based on the first-order Taylor expansion of the pruning loss , offering comparability among all prunable components from all layers . Furthermore , we incorporate the estimated latency reduction of each neuron into its importance score . This guides the final pruned architecture to be faster on target devices , as we show in experiments . The pruned models are distilled utilizing the information of the ground truth labels , a pretrained CNN teacher akin to DEIT ( Touvron et al. , 2021 ) , and the original full model . On the ImageNet-1K benchmark ( Russakovsky et al. , 2015 ) , structural pruning enables a nearly lossless 5.14× parameter reduction , 2.57× FLOPs reduction and 1.86× speed up on V100 GPU over the DEIT-Base model . An 1 % and 1.7 % accuracy gain is observed over DEIT-Small and DEIT-Tiny models when we compress the base model to a similar latency . Using structural pruning for architectural guidance , we further make an important observation that the popular uniform distribution of parameters across all layers is , in fact , not optimal . A simple redistribution of parameters can already provide stronger architectural alternatives , as we show in our experiments for both pretraining and downstream tasks . To this end , we present a new parameter distribution rule to scale ViT architectures , enabling a breed of models named as NViT . When scaling to similar FLOPs and latency , NViT architectures achieve 0.1 % , 0.2 % and 1.1 % accuracy gains over the DEIT-Base , Small , and Tiny models respectively when trained from scratch on ImageNet-1K . Our main contributions are as follows : • Provide a systematic analysis on the prunable components in the ViT model . We identify the ability to perform structural pruning on the embedding dimension , number of heads , MLP hidden dimension , QK dimension and V dimension of each head separately ; • Propose a latency-aware , importance-based criteria that enables hardware-friendly global structural pruning of all the components , achieving a nearly lossless 1.9× speedup ; • Present a new architecture scaling rule that enables NViT , a new family of efficient vision transformer architectures that redistributes the dimensions of DEIT models to outperform them under similar FLOPs and latency . This is the first work showing the potential of discovering novel scalable architectures by pruning vision transformers ; • Demonstrate that the high performance achieved by the pruned models and NViT models transfer effectively to downstream tasks . 2 RELATED WORK . 2.1 VISION TRANSFORMER MODELS . Inspired by the success of transformer models in NLP tasks , recent research proposes to use transformer models on computer vision tasks . The inspiring vision transformer ( ViT ) ( Dosovitskiy et al. , 2020 ) demonstrates the possibility of performing high-accuracy image classification with transformer architecture only , yet also finds that learning attention between image patches is a challenging task that requires large training set and model size . This stimulates recent works to ease extensive pretraining and improve the efficiency-accuracy tradeoff . One noticeable approach DEIT ( Touvron et al. , 2021 ) provides carefully designed training schemes and data augmentation techniques to train ViT from scratch on ImageNet only . Another line of work renovates ViT transformer blocks to better capture image features , such as changing input tokenization ( Yuan et al. , 2021 ; Graham et al. , 2021 ) , using hierarchical architecture ( Liu et al. , 2021 ; Wang et al. , 2021 ; Graham et al. , 2021 ) , upgrading positional encoding ( Chu et al. , 2021 ) , and performing localized attention ( Liu et al. , 2021 ; Han et al. , 2021 ) , bringing in further accuracy improvements under similar computation cost . In this work we focus on the original ViT architecture ( Dosovitskiy et al. , 2020 ) amid its widespread usage , as illustrated in the top of Figure 1 . ViT model first divides the input image into patches that are tokenized to embedding dimension E through a linear projection . Image tokens , together with an independently initialized class token , form an input x ∈ RN×E . Input tokens pass through transformer blocks before classification is made from the class token output of the last block . In its simplest form , a ViT block includes a multi-head self attention ( MSA ) and a multi-layer perceptron ( MLP ) module . The MSA module first linearly transforms the N × E tokens into queries q ∈ RN× ( QK×H ) , keys k ∈ RN× ( QK×H ) , and values v ∈ RN× ( V×H ) . The q , k and v are then split into H heads . Each head performs the self-attention operation defined in Equation ( 1 ) in parallel : Attn ( qh , kh , vh ) = softmax ( qhk T h√ dh ) vh , ( 1 ) where 1√ dh is the scaling factor , and qh ∈ RN×QK , kh ∈ RN×QK and vh ∈ RN×V are the query , key and value of head h. The output of all the heads are then concatenated prior to a fully-connected ( FC ) linear projection back to the original dimension of RN×E . Note that though previous works set dimension QK = V in designing the model architecture ( Dosovitskiy et al. , 2020 ; Touvron et al. , 2021 ; Chen et al. , 2021a ) , setting them differently will not go against the shape rule of matrix multiplication . The MLP module includes two FC layers with a hidden dimension of M . The output of the last FC layer preserves token dimension at RN×E . Built upon the original ViT , DEIT models ( Touvron et al. , 2021 ) further exploit a distillation token , which learns from the output label of a CNN teacher during the training process to incorporate some inductive bias of the CNN model , and significantly improves the DEIT accuracy . Our work uses the DEIT model architecture as a starting point , where we explore the potential of compressing the model and better allocating of the dimensions of different blocks for enhanced efficiency-accuracy tradeoff . Our method is also applicable to other vision transformer architectures for further efficiency improvements , which will be explored in future works . 2.2 VIT COMPRESSION TECHNIQUES . Despite rapid progress of ViTs for vision tasks , ViT blocks still heavily inherit the architecture of transformers for NLP tasks ( Vaswani et al. , 2017 ) , e.g. , the dimension of each block is determined without much optimization , and all the blocks in the stacked model architecture share the same dimension . As we will show later , such design may not be optimal for computer vision tasks . To improve model efficiency , very recent works successfully use structural pruning techniques on vision transformer models , with trainable gate variables ( Zhu et al. , 2021 ) or Taylor importance score ( Chen et al. , 2021b ) . Both methods show the potential of compressing ViT models , yet only consider part of the prunable architecture like the number of heads or the MLP hidden dimension , use manually designed sparsity distribution , and do not take run time latency into account , thus may not lead to optimal compressed models . Our method resolves these issues through a latency-aware global structural pruning of all prunable components across all layers in a jointly manner . Besides pruning , AutoFormer ( Chen et al. , 2021a ) uses a neural architecture search ( NAS ) approach to search for efficient ViT models . AutoFormer only explores a small number of dimension choices , due to the constraint on the supernet training cost ; while our method continuously explores the entire design space of ViT model with a single training process , leading to the finding of more efficient architectures with less training cost . Most importantly , our pruning scheme successfully leads to a scalable design rule , where we can easily design efficient vision transformer models at different scales . This has never been attempted in previous works . There has been a series of research on improving the efficiency of attention-based models on NLP tasks , specifically pruning the number of heads in the MSA block . This includes utilizing stochastic binary gate ( Louizos et al. , 2017 ; Voita et al. , 2019 ) , pruning with sensitivity score ( Michel et al. , 2019 ) and exploring lottery ticket hypothesis ( Frankle & Carbin , 2018 ; Behnke & Heafield , 2020 ) . Our work aims on pruning all parameters including heads , and therefore , has a larger pruning space . Other great ways of improving the efficiency of transformers include weight sharing across transformer blocks ( Lan et al. , 2019 ) , dynamically controlling the attention span of each token ( Chen et al. , 2019 ; Tambe et al. , 2020 ) , and allowing the model to output the result in an earlier transformer block ( Zhou et al. , 2020 ; Schwartz et al. , 2020 ) . These techniques are orthogonal to our pruningbased method and have remained unexplored on vision models . Applying them on ViT models and combining with our pruning method would be an interesting and fruitful future direction .

Summary: This paper applies latency-aware global structural pruning to vision transformers (ViTs), which results in redistribution of the model parameters and better a speed-accuracy tradeoff. Compared to Deit-B, the pruned vision transformer model (NVP) is 1.85x faster with almost no performance loss. Based on the insights discovered in the pruning process, the paper also presents the novel vision transformer (NViT) architectures. NViT outperforms Deit on ImageNet, CIFAR, and iNat benchmarks with similar running times.