Year: 2021

Abstract:
Vision transformers (ViTs) have gained popularity recently. Even without customized image operators such as convolutions, ViTs can yield competitive performance when properly trained on massive data. However, the computational overhead of ViTs remains prohibitive, due to stacking multi-head self-attention modules and else. Compared to the vast literature and prevailing success in compressing convolutional neural networks, the study of Vision Transformer compression has also just emerged, and existing works focused on one or two aspects of compression. This paper proposes a unified ViT compression framework that seamlessly assembles three effective techniques: pruning, layer skipping, and knowledge distillation. We formulate a budget-constrained, end-to-end optimization framework, targeting jointly learning model weights, layer-wise pruning ratios/masks, and skip configurations, under a distillation loss. The optimization problem is then solved using the primal-dual algorithm. Experiments are conducted with several DeiT backbones on the ImageNet dataset, and our approach consistently outperforms recent competitors. For example, DeiT-Tiny can be trimmed down to 50% of the original FLOPs almost without losing accuracy. Codes will be released.

heading: 1 INTRODUCTION
text: 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).

heading: 2 RELATED WORK

heading: 2.1 VISION TRANSFORMER
text: 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).

heading: 2.2 MODEL COMPRESSION
text: 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.

heading: 3 METHODOLOGY

heading: 3.1 PRELIMINARY
text: 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.

heading: 3.2 RESOURCE-CONSTRAINED END-TO-END VIT COMPRESSION
text: We target a principled constrained optimization framework jointly optimizing all weights and compression hyperparameters. We consider two strategies to compress ViT: (1) structural pruning of linear projection modules in a ViT block; and (2) adjusting the skipping patterns across different ViT blocks, including skipping/dropping an entire block. The second point, to our best knowledge, is the first time to be ever considered in ViT compression. The full framework is illustrated in Figure 1.
Alternatively, the two strategies could be considered as enforcing mixed-level group sparsity: the head dimension level, the head number level, and the block number level. The rationale is: when pruned to same ratios, finer-grained sparsity (i.e., pruning in smaller groups) is unfriendly to latency, while coarser-grained sparsity (i.e., pruning in larger groups) is unfriendly to accuracy. The mixedlevel group sparsity can hence more flexibly trade-off between latency and accuracy.
The knowledge distillation is further incorporated into the objective to guide the end-to-end process. Below we walk through each component one-by-one, before presenting the unified optimization.
Pruning within a Block To compress each linear projection module on width, we first denote s(l,3) as the number of rows to be pruned for weights W (l,3). Compression for weights W (l,1) is more complicated as it is decided by two degrees of freedom. As the input tensor to be multiplied with W (l,1) is the direct output of attention heads. Hence, we denote s(l,1) to be the attention head numbers that need to be pruned, and r(l,i) the number of output neurons to be pruned for each attention head i. Figure 2 illustrates the two sparsity levels: the head dimension level as controlled by r(l,i), and the head number level as controlled by s(l,3). They give more flexibility to transformer compression by selecting the optimal architecture in a more delicate and multi-grained way. Let us emphasize again that those variables above are not manually picked, but rather optimized globally.
Skipping Manipulation across Blocks The next component that we plan to jointly optimize is the skip connectivity pattern, across different ViT blocks. For vanilla ViTs, skip connection plays a crucial role in ensuring gradient flow and avoiding feature collapse (Dong et al., 2021). Raghu et al. (2021) suggests that skip connections in ViT are even more influential than in ResNets, having strong effects on performance and representation similarity, However, few existing studies have systematically studied how to strategically adjust skip connections provided a ViT architecture, in order to optimize the accuracy and efficiency trade-off.
It is known that more skip connections in general helps CNNs’ accuracy without increasing the parameter volume (Huang et al., 2017); and ViT seems to favor the same trend (Raghu et al., 2021). Moreover, adding a new skip connection over an existing block would allow us to “skip" it during inference, or “drop" it. This represents more aggressive “layer-wise" pruning, compared to the aforementioned element-wise or conventional structural pruning
While directly dropping blocks might look aggressive at the first glance, there are two unique enabling factors in ViTs that facilitate so. Firstly, unlike in CNNs, ViTs have uniform internal structures: by default, the input/output sizes of Self-Attention (SA) module, MLP module and thus the entire block, are all identical. That allows to painlessly drop any of those components and directly reconcatenating the others, without causing any feature dimension incompatibility. Secondly, Phang et al. (2021) observed that the top few blocks of fine-tuned transformers can be discarded without hurting performance, even with no further tuning. That is because deeper layer features tend to have high cross-layer similarities (while being dissimilar to early layer features), making additional blocks add little to the discriminative power. Zhou et al. (2021) also reported that in deep ViTs, the attention maps gradually become similar and even nearly identical after certain layers. Such cross-block “feature collapse" justifies our motivation to drop blocks.
In view of those, we introduce skip manipulation as a unique compression means to ViT for the first time. Specifically, for each transformer block with the skip connection, we denote gt(g,1) and gt(g,2) as two binary gating variables, to decide whether to skip this block or not (g represent the g-th block): they apparently should always sum up to 1.
The Constraints: We next formulate weight sparsity constraints for the purpose of pruning. As discussed in Sec 3.1, the target of the proposed method is to prune the head number and head dimension simultaneously, which can actually be modeled as a two-level group sparsity problem when choosing {W (l,1)}1≤l≤L as proxy compression targets. Specifically, the input dimension of {W (l,1)}1≤l≤L is equivalent to the sum of the dimensions of all heads. Then we can put a two-level group sparsity regularization on {W (l,1)}1≤l≤L input dimension to compress head number and head dimension at same time, as shown in 1a and 1b. r(l,i) corresponding to the pruned size of ith head. s(l,1) means the pruned number of heads. For {W (l,3)}1≤l≤L, it is our compression target, we just perform standard one level group sparsity regularization on its input dimension, as shown in 1c.∑
j
I (∥∥∥W (l,1)gij ,. ∥∥∥2
2 = 0
) ≥ r(l,i), (1a)
∑ i I (∥∥∥W (l,1)gi,. ∥∥∥2 2 = 0 ) ≥ s(l,1), (1b)
∑ i I (∥∥∥W (l,3)i,. ∥∥∥2 2 = 0 ) ≥ s(l,3), (1c)
∀ l = 1, 2, ..., L, ∀ i = 1, 2, ...,H (1d) where Wi,. denotes the i-th row of W ; Wgi,. denotes the i-th grouped row matrix of W ; and Wgij ,. the j-th row of Wgi,.. In other words, in Eqn. 1b, row matrices are grouped by attention heads. Hence
among the original H heads in total at block l, at least s(l,1) heads should be discarded. Similarly, Eqn. 1c demands that at least s(l,3) input neurons should be pruned at this linear projection module. Furthermore, Eqn. 1a requests that in the i-th attention head at block l, r(l, i) of the output units should be set to zeros.
Our method is formulated as a resource constrained compression framework, given a target resource budget, it will compress the model until the budget is reached. Given a backbone architecture, the FLOPs is the function of s, r and gt, denoted asRFlops(s, r, gt). We have the constraint as:
RFlops(s, r, gt) ≤ Rbudget, (2)
where s = {s(l,1), s(l,3)}1≤l≤L and r = {r(l,i)}1≤l≤L,1≤i≤H . Rbudget is the resource budget. We present detailed flops computation equations in terms of s, r and gt in Appendix.
Those inequalities can be further rewritten into equation forms to facilitate optimization. As an example, we follow (Tono et al., 2017) to reformulate Eqn. 1b as:∑
i
I ( ‖W.,gi‖ 2 2 = 0 ) ≥ s⇔ ‖W·,g‖2s,2 = 0. (3)
where ‖W·,g‖2s,2 denotes the Frobenius norm of the sub-matrix of W consisting of s groups of W with smallest group norms. Eqn. 1a and Eqn. 1c have same conversions.
The Objective: The target objective could be written as (λ is a coefficient):
min W ,gt
L(W , gt) = `(W , gt) + λ`distill(W ,Wt), (4)
where Wt denotes the weights from teacher model, i.e., the uncompressed transformer model, and `distill(∗, ∗) is defined as the knowledge distillation loss from the original teacher. We choose the simper `2 norm as its implementation, as its performance was found to be comparable to the more traditional K-L divergence (Kim et al., 2021).
The Final Unified Formulation: Summarizing all above, we arrive at our unified optimization as a mini-max formulation, by leveraging the primal-dual method (Buchbinder & Naor, 2009):
min W ,s,r,gt max p,y,z≥0 Lpruning = min W ,s,r,gt max p,y,z≥0 L(W , gt) + z (RFlops(s, r, gt)−Rbudget)︸ ︷︷ ︸ resource loss + L∑ l=1 ( y(l,1) ∥∥∥W (l,1).,g ∥∥∥2 ds(l,1)e,2 + y(l,3) ∥∥∥W (l,3).,. ∥∥∥2 ds(l,3)e,2 ) + L∑ l=1 H∑ i=1 p(l,i) ∥∥∥W (l,1).,gi. ∥∥∥2dr(l,i)e,2︸ ︷︷ ︸
sparsity loss: S(y,s,p,r,gt,W )
(5)
Eventually, we use the solution of the above mini-max problem, i.e. s and r to determine the compression ratio of each layer. Here, we select the pruned groups of the parameters by directly ranking rows by their norms, and remove those with the smallest magnitudes.

heading: 3.3 SOLVING THE UNIFIED OPTIMIZATION
text: The general updating policy follows the idea of primal-dual algorithm. The full algorithm is outlined in Algorithm 1 in Appendix. We introduce the solutions to each sub-problem in Appendix as well.

heading: 4 EXPERIMENTAL RESULTS
text: Datasets and Benchmarks We conduct experiments for image classification on ImageNet (Krizhevsky et al., 2012). We implement our compression method UVC on DeiT (Touvron et al., 2020), which has basically the identical architecture compared with ViT (Dosovitskiy et al., 2020) except for an extra distillation token. Experiment has been conducted on DeiT-Tiny/Base models; for DeiT-Tiny, we also try it without/with the distillation token. We measure all resource consumptions (including the UVC resource constraints) in terms of inference FLOPs.
Training Settings The whole process of our method consists of two steps.
• Step 1: UVC training. We firstly conduct the primal-dual algorithm to the pretrained DeiT model to produce the compressed model under the given resource budget.
• Step 2: Post training. When we have the sparse model, we finetune it for another round of training to regain its accuracy loss during compression.
In the two steps mentioned above, we strictly follow the training settings of DeiT (Touvron et al., 2020). For distillation, we select the pretrained DeiT model to teach the pruned DeiT model.
Numerically, the learning rate for parameter z is always changing during the primal-dual algorithm process. Thurs, we propose to use a dynamic learning rate for the parameter z that controls the budget constraint. We use a four-step schedule of {50, 60, 70, 80, 90} in practice.
Baseline methods We adopt several latest compression methods specifically developed for ViTs (all proposed only in the past four months), which fall under two categories:
• Category 1: Input Patch Reduction, including (i) PatchSlimming (Tang et al., 2021): which identify the effective patches in the last layer and then use them to guide the patch selection process of previous layers; and (ii) IA-RED2 (Pan et al., 2021a): which dynamically and hierarchical drops visual tokens at different levels using a learned policy network.
• Category 2: Model Weight Pruning, including (iii) Vision Transformer Pruning (VTP) (Zhu et al., 2021): which trains transformers with sparsity regularization to let important dimensions emerge. (iv) SViTE (Chen et al., 2021b): which jointly optimizes model parameters and explores sparse connectivity throughout training, ending up with one final sparse network. SViTE belongs to the most competitive accuracy-efficiency trade-off achieved so far, for ViT pruning. We choose its structured variant to be fair with UVC.

heading: 4.1 MAIN RESULTS
text: The full comparison result are listed in Tab. 1. Firstly, we notice that most existing methods cannot save beyond 50% FLOPs without sacrificing too much accuracy. In comparison, UVC can easily go with larger compression rates (up to ≥ 60% FLOPS saving) without compromising as much, showing stronger promise for resource-constrained ViT applications. For example, when compressing DeiT-Tiny (with distillation token), UVC can trim the model down to a compelling ≥50% of the original FLOPs while losing only 0.3% accuracy.
Secondly, we observe UVC to significantly outperform the latest SViTE (Chen et al., 2021b) at DeiT-Tiny, obtaining less accuracy drops (0.9/1.6, versus 2.08), at much higher FLOPs savings (0.7/0.6, versus 0.99). At the larger DeiT-Base, while UVC can compress its FLOPs to only 45% of the original, we find that SViTE cannot be stably trained at such high sparsity (the authors reported all results with ≤ 40% FLOPs saving for DeiT-Base). UVC also yield comparable accuracy with the other ViT model pruning approach, VTP (Zhu et al., 2021), with more than 10% FLOPs savings w.r.t. the original backbone.
Thirdly, UVC also performs clearly better than one latest patch reduction approach, IA-RED2 (Pan et al., 2021a), at DeiT-Base. Compared to the other patch reduction method, PatchSlimming Tang et al. (2021), we observe UVC to generally save more FLOPs, yet also sacrificing more accuracies. Moreover, as we explained in Section 3.1, those input token reduction methods represent an orthogonal direction to the model weight sparsification way that UVC is pursuing. UVC can also be seamlessly extended to include token reduction into the joint optimization - a future work that we would pursue.

heading: 4.2 ABLATION STUDY
text: As UVC highlights the integration of three different compression techniques into one joint optimization framework, it would be natural to question whether each moving part is necessary for the pipeline, and how much portion they each contribute to the final result. We conduct this ablation study, and the results are presented in Tab. 2.
Skip Manipulation only. We first present the result when we only conduct skip manipulation and knowledge distillation in our framework, but not pruning within a block. All other protocols remain unchanged. This two-in-a-way method is actually reduced into LayerDropping (Lin et al., 2018). We have the following findings: Firstly, implementing only skip connection manipulation in the framework will incur high instability. During optimization, the objective value fluctuates heavily and barely converge to a stable solution. Secondly, only applying skip manipulation will be damaging the accuracy remarkably, e.g., by nearly 4% drop on DeiT-Tiny. That is expected, as only manipulating the model architecture with such a coarse granularity (keeping or skipping a whole transformer block) is very inflexible and prone to collapse.
Pruning within a Block only. We then conduct experiments without using the gating control for skip connections. That leads to only integrating neuron-level and attention-head-level pruning, with knowledge distillation. It turns out to deliver much better accuracy than only skip manipulation, presumably owing to its finer-grained operation. However, it still has a margin behind our joint UVC method, as the later also benefits from removing block-level redundancy a priori, which is recently observed to widely exist in fine-tuned transformers Phang et al. (2021). Overall, the ablation results endorse the effectiveness of jointly optimizing those components altogether.
Apply compression methods sequentially 1 It is worth noting that our method is a joint optimization process that can simultaneously optimize for skip configuration manipulation and pruning. We mark Prune->Skip as the process of first pruning then manipulating skip configuration, while Skip>Prune vice versa. Note that Skip->Prune 71.49%->49.39% denotes that first use skip configuration to enforce the model FLOPs to be 71.49% of the original FLOPs, then apply pruning to further compress the model to 49.39% of the original FLOPs. Without loss of generality, we approximate both the pruning procedure and skipping procedure to compress 70% of the previous model, which will approximately result in a model with 50% of the original FLOPs. Based on this, we implement:
• (Skip->Prune 71.49%->49.39%) It first compresses the model to 71% by skipping then applies pruning on the original model and results in the worse performance at 66.84%.
• (Prune->Skip 69.96%->54.88%) It applies pruning to 70% FLOPs then skipping follows.
We observe that applying pruning first can stabilize the choice of skip configuration and result in a top-1 accuracy of 69.68%. Both of the results show that UVC finds the best trade-off between skip and pruning.
Furthermore, we design a better setting (Prune->Skip 63.25%->50.73%) to simulate the compression ratio found by UVC, as is observed in Figure 3b in the main text, DeiT-Tiny discarded 2 blocks in total,
1The highlighted parts are newly added from the rebuttal session.
Method DeiT-TinyAcc. (%) FLOPs remained(%)
Uncompressed baseline 72.2 100 Only skip manipulation 68.72 51.85 Only pruning within a block 70.52 50.69 Without Knowledge Distillation 69.34 51.23
Skip->Prune 71.49%->49.39% 66.84 49.39 Prune->Skip 69.96%->54.88% 69.68 54.88 Prune->Skip 63.25%->50.73% 70.02 50.73
UVC 71.3 49.23
Attention Head
Attention Head
which results in 84% of compression ratio. Thus, we first prune the DeiT-Tiny to 63% and apply skip to indulge a 50% FLOPs’ model. The obtained result of the setting (Prune->Skip 63.25%->50.73%) is much better than previous ones with sub-optimal trade-off, but still has a 1% performance gap with the result provided by UVC. We give the credit to jointly optimize the architecture and model weights to avoid local-minima convergence for architecture.

heading: 4.3 VISUALIZATION
text: To reveal more intuitions how UVC optimize the architecture, we present two groups of visualizations:
• Sparse connectivity patterns: We first visualize the sparse connectivity patterns in Fig. 3 (a), i.e, how many attention heads are preserved, and how they (sparsely) distribute over all blocks. For the smaller DeiT-Tiny model, UVC tend to more evenly drop attention heads across layers; meanwhile for the larger DeiT-Base, UVC clearly prefers to drop more attention heads in the early layers.
• Skip connection patternsL We next draw the result of skip connection patterns Fig. 3 (b), showing what layers are learned to be directly skipped under unified optimization. We observe UVC’s obvious tendency in dropping later layers, which coincides nicely with the observations by (Phang et al., 2021; Zhou et al., 2021).

heading: 5 CONCLUSION
text: In this paper, we propose UVC, a unified ViT compression framework that seamlessly assemble pruning, layer skipping, and knowledge distillation as one. A budget-constrained optimization framework is formulated to jointly learn model weights, layer-wise pruning ratios/masks, and skip configurations, under a distillation loss. Experiments demonstrate that UVC can aggressively trim down the prohibit computational costs in an end-to-end way. Our future work will extend UVC to incorporating weight quantization as part of the end-to-end optimization as well.

heading: A UPDATING POLICY
text: Updating Weights W Different from other pruning methods in CNN that uses fixed pretrained weights to select the pruned channels/groups with certain pre-defined metrics (Lin et al., 2020; He et al., 2019), here our subproblem could be considered as following a dynamic pruning criteria that will be updated along. Specifically we solve the following subproblem:
Proxη1S(y,s,p,r,W )(W̄ ) = arg minW 1
2 ∥∥W − W̄∥∥2 + η1S (y, s,p, r,W ) , (6) where W̄ = W t − η1∇̂W ` (W t). The solution admits a bi-level projection (Yang et al., 2016):
W (l,1)∗.,gij = W̄ (l,1).,gij , if ∥∥∥W̄ (l,1).,gij ∥∥∥2 2 ≥ ∥∥∥W (l,1).,g ileast-dr(l,i)e ∥∥∥2 2 ,
1 1+2η1p(l,i) W̄i, otherwise,
W (l,1)∗.,gi = W̄ (l,1).,gi , if ∥∥∥W̄ (l,1).,gi ∥∥∥2 2 ≥ ∥∥∥W (l,1).,g least-ds(l,1)e ∥∥∥2 2 ,
1 1+2η1y(l,1) W̄.,gi , otherwise,
W (l,3)∗ .,i = W̄ (l,3).,i , if ∥∥∥W̄ (l,3).,i ∥∥∥2 2 ≥ ∥∥∥W (l,3).,least-ds(l,3)e∥∥∥22,
1 1+2η1y(l,3) W̄.,i, otherwise,
where least-j denotes the index of the (group) columns of W that have j-th least (group) norm.
Updating gt (gt(l,0), gt(l,1)) are used to generate a binomial categorical distribution to decide whether pass through block l or directly skip it. As the two variables are discrete, we apply the renowned Gumbel-Softmax (GSM) trick (Jang et al., 2016) to obtain differentiable and polarized sampling. For gtl = (gt(l,0), gt(l,1)), given i.i.d Gumbel noise g drawn from Gumbel(0, 1) distribution, a soft categorical sample can be drawn by
Gl = GSM(gtl) = Softmax((log(gtl) + g)/τ) ∈ R2, (7)
where G(l,1) refers to the continuous possibility to preserve current block l while G(l,0) to drop it. Directly applying the chain rule on L(W , gt) w.r.t gt can now calculate ∇̃gtL(W t, gtt). Also, gt participates in the calculation of FLOPs, by deciding whether to pass certain blocks. Since passing or skipping a block is a dynamic choice during training, we estimate the FLOPs of the l-th blockRFlops(s, r, gt) with skip gating by using its expectation. To be specific,
RFlopsl(sl, rl, gtl) = E[RFlopsl(sl, rl, gtl)|sl, rl] (8a) = G(l,0)RFlopsl(Identity) +G
(l,1)RFlopsl(sl, rl) (8b) = G(l,1)RFlopsl(sl, rl) (8c)
Hence, updating policy for gt is formulated as: gtt+1 = gtt − η4 ( ∇̃gtL ( W , gtt ) + ∇̃gtz ( RFlops ( s, r, gtt ) −Rbudget )) (9)
= gtt − η4 ( ∇̃gtL ( W , gtt ) + ∇̃gtzRFlops(s, r)GSM(gtt) ) (10)
Updating s and r Similar to the updating policy of gt, one gradient term w.r.t. s and r are ∇̃sz (RFlops (s, r, gt)−Rbudget), ∇̃rz (RFlops (s, r, gt)−Rbudget) respectively. The other gradient term is calculated on the unified formulation Eqn. 5. Refer to that, s and r are floating-point numbers during the optimization process. In pratice, ceiling functions are operated
on them to determine the integer number that should be pruned for each layer. However, the ceiling function d.e is non-differentiable. To solve this problem, we implement Straight-through estimator(STE) (Bengio et al., 2013) to provide a proxy of gradient when performing the backward pass. We set ∂̃dse
∂̃s = 1.
As for ‖W.,g‖2s,2 term in the sparsity loss, we use ‖W.,g‖ 2 s+1,2 − ‖W.,g‖ 2 s,2 as the proxy of partial derivative of ‖W.,g‖2s,2 with respect to s:
∂̃ ‖W.,g‖2s,2 ∂̃s = ∥∥W.,gleast-min{Dim(W ),s+1}∥∥22 , (11)
where Dim(W ) is the number of column groups of W . The other two terms in the sparsity loss can be processed similarly.

heading: B MAIN ALGORITHM
text: The general updating policy follows the idea of primal-dual algorithm. The full algorithm is outlined in Algorithm 1.
Algorithm 1: Gradient-based algorithm to solve problem (5) for Unified ViT Compression. Input: Resource budget Rbudget, learning rates η1, η2, η3, η4, η5, η6, number of total iterations τ . Result: Transformer pruned weights W ∗.
1 Initialize t = 1,W 1 ; // random or a pre-trained dense model 2 for t← 1 to τ do 3 W t+1 = Proxη1S(yt,st,pt,rt,W t) ( W t − η1∇̂WL ( W t, gt )) ; // Proximal-SGD
4 st+1 = st − η2 ( ∇̃sS ( yt, st,pt, rt, gtt,W t+1 ) + ∇̃szt ( RFlops ( st, rt, gtt ) −Rbudget )) ;
// Gradient (STE) Descent 5 rt+1 = rt − η3 ( ∇̃rS ( yt, st+1,pt, rt, gtt,W t+1 ) + ∇̃rzt ( RFlops ( st+1, rt, gtt ) −Rbudget )) ;
// Gradient (STE) Descent 6 gtt+1 = gtt − η4 ( ∇̃gtL ( W t+1, gtt ) + ∇̃gtzt ( RFlops ( st+1, rt+1, gtt ) −Rbudget )) ;
// Gradient Descent 7 zt+1 = zt + η7 ( RFlops ( st+1, rt+1, gtt+1 ) −Rbudget ) ; // Gradient Ascent
8 y(l,1) t+1 = y(l,1) t + η5 (∥∥∥W (l,1).,g t+1∥∥∥2⌈ s(l,1) t+1 ⌉ ,2 ) ; // Gradient Ascent
9 y(l,3) t+1 = y(l,3) t + η5 (∥∥∥W (l,3).,. t+1∥∥∥2⌈ s(l,3) t+1 ⌉ ,2 ) , ∀ l = 1, · · · , L
p(l,i) t+1 = p(l,i) t + η6 (∥∥∥W (l,1).,gi. t+1∥∥∥2⌈ r(l,i) t+1 ⌉ ,2 ) , ∀ i = 1, · · · , H, ∀ l = 1, · · · , L
10 W ∗ = W

heading: C EXTRA RESULT
text: Here we provide the experimental results on the backbone Deit-Small in Table 3. Here we introduce three new baselines that was also compared in Tang et al. (2021). In the future, we will also provide the results on T2T-ViT Yuan et al. (2021) backbone and the performance after integrating token selection.

Review summaries:
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.