Unified Visual Transformer Compression

1 INTRODUCTION . Convolution neural networks ( CNNs ) ( LeCun et al. , 1989 ; Krizhevsky et al. , 2012 ; He et al. , 2016 ) have been the de facto architecture choice for computer vision tasks in the past decade . Their training and inference cost significant and ever-increasing computational resources . Recently , drawn by the scaling success of attention-based models ( Vaswani et al. , 2017 ) in natural language processing ( NLP ) such as BERT ( Devlin et al. , 2018 ) , various works seek to leverage the Transformer architecture to computer vision ( Parmar et al. , 2018 ; Child et al. , 2019 ; Chen et al. , 2020 ) . The Vision Transformer ( ViT ) architecture ( Dosovitskiy et al. , 2020 ) , and its variants , have been demonstrated achieves comparable or superior results on a series of image understanding tasks compared to the state of the art CNNs , especially when pretrained on datasets with sufficient model capacity ( Han et al. , 2021 ) . Despite the emerging power of ViTs , such architecture is shown to be even more resource-intensive than CNNs , making its deployment impractical under resource-limited scenarios . That is due to the absence of customized image operators such as convolution , the stack of self-attention modules that suffer from quadratic complexity with regard to the input size , among other factors . Owing to the substantial architecture differences between CNNs and ViTs , although there is a large wealth of successful CNN compression techniques ( Liu et al. , 2017 ; Li et al. , 2016 ; He et al. , 2017 ; 2019 ) , it is not immediately clear whether they are the same effective for ViTs . One further open question is how to best integrate their powers for ViT compression , as one often needs to jointly exploit multiple compression means for CNNs ( Mishra & Marr , 2018 ; Yang et al. , 2020b ; Zhao et al. , 2020b ) . On the other hand , the NLP literature has widely explored the compression of BERT ( Ganesh et al. , 2020 ) , ranging from unstructured pruning ( Gordon et al. , 2020 ; Guo et al. , 2019 ) , attention head pruning ( Michel et al. , 2019 ) and encoder unit pruning ( Fan et al. , 2019 ) ; to knowledge distillation ( Sanh et al. , 2019 ) , layer factorization ( Lan et al. , 2019 ) , quantization ( Zhang et al. , 2020 ; Bai et al. , 2020 ) and dynamic width/depth inference ( Hou et al. , 2020 ) . Lately , earlier works on compressing ViTs have also drawn ideas from those similar aspects : examples include weight/attention pruning ( Zhu et al. , 2021 ; Chen et al. , 2021b ; Pan et al. , 2021a ) , input feature ( token ) selection ( Tang et al. , 2021 ; Pan et al. , 2021a ) , and knowledge distillation ( Touvron et al. , 2020 ; Jia et al. , 2021 ) . Yet up to our best knowledge , there has been no systematic study that strives to either compare or compose ( even naively cascade ) multiple individual compression techniques for ViTs – not to mention any joint optimization like ( Mishra & Marr , 2018 ; Yang et al. , 2020b ; Zhao et al. , 2020b ) did for CNNs . We conjecture that may potentially limit the performance gain of ViT compression . This paper aims to establish the first all-in-one compression framework that organically integrate three different compression strategies : ( structured ) pruning , block skipping , and knowledge distillation . Rather than ad-hoc composition , we for the first time propose a unified vision transformer compression ( UVC ) framework , which seamlessly integrates the three effective compression and jointly optimizes towards the task utility goal under the budget constraints . UVC is mathematically formulated as a constrained optimization problem and solved using the primal-dual algorithm from end to end . Our main contributions are outlined as follows : • We present UVC that unleashes the potential of ViT compression , by jointly leveraging multiple ViT compression means for the first time . UVC only requires to specify a global resource budget , and can automatically optimize the composition of different techniques . • We formulate and solve UVC as a unified constrained optimization problem . It simultaneously learns model weights , layer-wise pruning ratios/masks , and skip configurations , under a distillation loss and an overall budget constraint . • Extensive experiments are conducted with several DeiT backbones on ImageNet , which consistently verify the effectiveness of our proposal . For example , UVC on DeiT-Tiny ( with/without distillation tokens ) yields around 50 % FLOPs reduction , with little performance degradation ( only 0.3 % /0.9 % loss compared to the baseline ) . 2 RELATED WORK . 2.1 VISION TRANSFORMER . Transformer ( Vaswani et al. , 2017 ) architecture stems from natural language processing ( NLP ) applications first , with the renowned technique utilizing Self-Attention to exploit information from sequential data . Though intuitively , transformer model seems inept to the special inductive bias of space correlation for images-oriented tasks , it has proved itself of capability on vision tasks just as good as CNN ( Dosovitskiy et al. , 2020 ) . The main point of Vision Transformer is that they encode the images by partitioning it into sequences of patches , projecting them into token embeddings , and feeding them to transformer encoders ( Dosovitskiy et al. , 2020 ) . ViT outperforms convolutional nets if given sufficient training data on various image classification benchmarks . Since then , ViT has been developed to various different variants first on data efficiency towards training like DeiT ( Touvron et al. , 2020 ) and T2T-ViT ( Yuan et al. , 2021 ) are proposed to enhance ViT ’ s training data efficiency , by leveraging teacher-student and better crafted architectures respectively . Then modifications are made to the general structure of ViT to tackle other popular downstream computer vision tasks , including object detection ( Zheng et al. , 2020 ; Carion et al. , 2020 ; Dai et al. , 2021 ; Zhu et al. , 2020 ) , semantic segmentation ( Wang et al. , 2021a ; b ) , image enhancement ( Chen et al. , 2021a ; Yang et al. , 2020a ) , image generation ( Jiang et al. , 2021 ) , video understanding ( Bertasius et al. , 2021 ) , and 3D point cloud processing ( Zhao et al. , 2020a ) . 2.2 MODEL COMPRESSION . Pruning . Pruning methods can be broadly categorized into : unstructured pruning ( Dong et al. , 2017 ; Lee et al. , 2018 ; Xiao et al. , 2019 ) by removing insignificant weight via certain criteria ; and structured pruning ( Luo et al. , 2017 ; He et al. , 2017 ; 2018 ; Yu et al. , 2018 ; Lin et al. , 2018 ; Guo et al. , 2021 ; Yu et al. , 2021 ; Chen et al. , 2021b ; Shen et al. , 2021 ) by zero out parameters in a structured group manner . Unstructured pruning can be magnitude-based ( Han et al. , 2015a ; b ) , hessian-based ( LeCun et al. , 1990 ; Dong et al. , 2017 ) , and so on . They result in irregular sparsity , causing sparse matrix operations that are hard to accelerate on hardware ( Buluc & Gilbert , 2008 ; Gale et al. , 2019 ) . This can be addressed with structured pruning where algorithms usually calculate an importance score for some group of parameters ( e.g. , convolutional channels , or matrix rows ) . Liu et al . ( 2017 ) uses the scaling factor of the batch normalization layer as the sensitivity metric . Li et al . ( 2016 ) proposes channel-wise summation over weights as the metric . Lin et al . ( 2020 ) proposes to use channel rank as sensitivity metric while ( He et al. , 2019 ) uses the geometric median of the convolutional filters . Particularly for transformer-based models , the basic structures that one many consider pruning with include blocks , attention heads , and/or fully-connected matrix rows . For example , ( Michel et al. , 2019 ) canvasses the behavior of multiple attention heads and proposes an iterative algorithm to prune redundant heads . ( Fan et al. , 2019 ) prunes entire layers to extract shallow models at inference time . Knowledge Distillation . Knowledge distillation ( KD ) is a special technique that do not explicitly compress the model from any dimension of the network . KD lets a student model leverages “ soft ” labels coming from a teacher network ( Hinton et al. , 2015 ) to boost the performance of a student model . This can be regarded as a form of compression from the teacher model into a smaller student . The soft labels from the teacher is well known to be more informative than hard labels , and leads to better student training ( Yuan et al. , 2020 ; Wei et al. , 2020 ) Skip Configuration . Skip connection plays a crucial role in transformers ( Raghu et al. , 2021 ) , by tackling the vanishing gradient problem ( Vaswani et al. , 2017 ) , or by preventing their outputs from degenerating exponentially quickly with respect to the network depth ( Dong et al. , 2021 ) . Meanwhile , transformer has an inborn advantage of uniform block structure . A basic transformer block contains a Self-Attention module and a Multi-layer Perceptron module , and its output size matches the input size . That implies the possibility to manipulate the transformer depth by directly skipping certain layers or blocks . ( Xu et al. , 2020 ) proposes to randomly replace the original modules with their designed compact substitutes to train the compact modules to mimic the behavior of the original modules . ( Zhang & He , 2020 ) design a Switchable-Transformer Blocks to progressively drop layers from the architecture . To flexibly adjust the size and latency of transformers by selecting adaptive width and depth , DynaBERT ( Hou et al. , 2020 ) first trains a width-adaptive BERT and then allows for both adaptive width and depth . LayerDrop ( Fan et al. , 2019 ) randomly drops layers at training time , while at test time , it allows for sub-network selection to any desired depth . 3 METHODOLOGY . 3.1 PRELIMINARY . Vision Transformer ( ViT ) Architecture . To unfold the unified algorithm in the following sections , here we first introduce the notations . There are totally L transformer blocks , while H denotes number of attention heads in each transformer block . In each block l of the ViT , there are two constituents , namely the self-attention module and the MLP module . In the self-attention module , W ( l ) Q , W ( l ) K , W ( l ) V are the weights of the three linear projection matrix using the block input X l to calculate Ql , Kl , V l. The weights of the projection module that follows the calculation of self-attention is denoted as W ( l,1 ) , represents the first linear projection module in block l. The MLP module consists of two linear projection modules W ( l,2 ) and W ( l,3 ) . Compression Targets The main parameters that can be potentially compressed in a ViT block are W ( l ) Q , W ( l ) K , W ( l ) V and W ( l,1 ) , W ( l,2 ) , W ( l,3 ) . Our goal is to prune the head number and head dimensions simultaneously inside each layer , associated with the layer level skipping , solved in a unified framework . Currently , we do not extend the scope to reducing other dimensions such as input patch number or token size . However , our framework can also pack these parts together easily . For head number and head dimensions pruning , instead of going into details of QKV computation , we innovate to use { W ( l,1 ) } 1≤l≤L to be the proxy pruning targets . Pruning on these linear layers is equivalent to the pruning of head number and head dimension . We also add { W ( l,3 ) } 1≤l≤L as our pruning targets , since these linear layers do not have dimension alignment issues with other parts , they can be freely pruned . We do not prune W ( l ) Q , W ( l ) K , W ( l ) V since Q l , Kl , V l should be of the same shape for computing self-attention . Besides , skip connection is recognized as another important component to enhance the performance of ViTs ( Raghu et al. , 2021 ) . The dimension between X l and the output of linear projection module should be aligned . For this sake , W ( l,2 ) is excluded from our compression . Eventually , the weights to be compressed in our subsequent framework are { W ( l,1 ) , W ( l,3 ) } 1≤l≤L .

Summary: This paper proposed a unified ViT compression framework that assembles pruning, layer skipping, and knowledge distillation as one. By adding the sparsity loss and the resources loss into the objective, ViT models can be stably trained at a high sparse ratio. By introducing the knowledge distillation loss, the compressed model can maintain performance with around 50% of the original FLOPs. Experimental results in ImageNet with DeiT models prove the effectiveness of the proposed method.