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 studies the latency reduction of Vision Transformer model. The proposed pruning method w.r.t importance score is trained with the full pre-trained model using knowledge distillation, with latency aware regularization. In addition, the author designed a new architecture NViT with a parameter redistribution. The experiments evaluated the proposed methods with respect to the accuracy, FLOP reduction, and parameter reduction.