Unified Visual Transformer Compression
Authors: Shixing Yu
Tianlong Chen
Jiayi Shen
Huan Yuan
Jianchao Tan
Sen Yang
Ji Liu
Zhangyang Wang
Venue: ICLR 2022 Poster (2022)
Keywords: Vision Transformer
Model Compression
Pruning
Layer Skipping
Distillation

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 ViT variants, e.g. DeiT and T2T-ViT 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 are available online:~\url{https://github.com/VITA-Group/UVC}.

Published as a conference paper at ICLR 2022
UNIFIED VISUAL TRANSFORMER COMPRESSION
Shixing Yu' Tianlong Chen "eiayi Shen', Huan Yuan", Jianchao Tan",
Sen Yang", Ji Liu", Zhangyang Wang'
University of Texas at Austin, "Texas A&M University, KKwai Inc.
(shixingyu, tianlong.. chen, atlaswang) Butoxas edd, aslylya-617@taam (yuanhuan9412, conyang.nlpt, si.liu.uwisse wttp://wwws..mm....oommm
agmall.com, jianchaotam@rualsoou.com
ABSTRACT
Vision transformers (ViTs) have gained popularity recently. Even without cus- tomized image operators such as convolutions, ViTs can yield competitive perfor- mance when properly trained massive data. However, the computational over- head of ViTs remains prohibitive, due to stacking multi-head self-attention modules and else. Compared 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 sion. This paper proposes a unified ViT compression framework that seamlessly assembles three effective techniques: pruning, layer skipping, and knowledge dis- tillation. We formulate a budget-constrained. end-to-end optimization framework, targeting jointly learning model weights, layer-wise pruning ratios/masks, and skip configurations, under distillation loss. The optimization problem is then solved using the primal-dual algorithm. Experiments are conducted with several ViT variants, e.g. DeiT and T2T- VIT 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 are available online: ht.p:s: / thub. wwwpehwww......oo.....e......tee cOm/VITA VAVVVAAAA-GGUUUVAA
two aspects of compres-
INTRODUCTION
Convolution neural networks (CNNs) (LeCun et al. 1989 Krizhevsky al. 2012 He 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-basse 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 (Parmrr et al. [2018] Child al. 2019 Chen et al. 2020a). The Vision Transformer (ViT) architecture (Dosovitskiy et al. 2020). and its variants, have been demonstrated to achieve comparable superior results series of image understanding tasks compared 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 more resource- -intensivv 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 large wealth of successful CNN compression techniques (Liu et al. 2017 Li et al. [2066) He et al., (017],2099) is not immediately clear whether they are as effective as for ViTs. One further open question is how to best integrate their power for VIT compression, one often needs to jointly exploit multiple
compression means for CNNs MMishra & Marr 2018 Yang et al. 20206 Zhha et all 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.
"Equal Contribution.

---

Published as a conference paper at ICLR 2022
2020 and dynamic width/depth inference (Hou et al. 2020). Lately, earlier works compressing
ViTs have also drawn ideas from those similar aspects: examples include weight/attention pruning
(Zhu et al., [022 Chen 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 either compare 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 impede more performance gais from ViT compression.
This paper aims to establish the first all-in-one compression framework that organically integrates three different compression strategies: (structured) pruning, block skipping, and knowledge dis- tillation. Rather than ad-hoc composition, we propose a unified vision transformer compression (UVC) framework, which seamlessly integrates the three effective compression techniques and jointly optimizes towards the task utility goal under the budget constraints. UVC is mathematically formulated as 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. simultane- ously 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 popular variants of ViT, including several DeiT backbones and TiT on ImageNet, and our proposal consistently performs better than or comparably with existing methods. For example, UVC on DeiT-Tiny (with/without distillation tokens) yields around 50% FLOPs reduction, with little performance degradation (only 0.3%/0. loss compared to the uncompressed baseline).
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 the 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 CNNs (Dosovitskiy al., 2020). The main point of Vision Transformer is that they encode the images by partitioning them into sequences of patches, projecting them into token embeddings, and feeding them to transformer encoders (Dosovitskiy et al. 2020). ViT outperforms convolutional nets given sufficient training data on various image classification benchmarks.
Since then, ViT has been developed to various different variants first data efficiency towards train- ing, like DeiT (Touvron et al. 2020) and T2T-ViT (Yuan et al. (021) 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 al. 2020 Dai et al.) Zhu et al. 2020). semantic segmentation (Wang et al. 2021al b), image enhancement (Chen 2021a] Yang et al. (2020)), image generation (Jiang et al. 2021). video understanding (Bertasius et al.
2021
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 2021 Chen et al. 222 lb: Shen et al. 2021) by zero out parameters in structured group
Yu et al.
manner.
Instructured pruning can be magnitude- based (Han et al., [2015a] b). hessian- (LeCun
et al.
1990: Dong al. 2017), and on. They result in irregular sparsity, causing sparse matrix
2

---

Published as a conference paper at ICLR 2022
operations that are hard to accelerate hardware (Buluc & Gilbert, 2008 Gale 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. (016) proposes weights as the metric. [[in et al. (2020) proposes to use channel rank as sensitivity metric while (He et al. 2019) uses the geometric median of the convolutional filters as pruning criteria. Particularly for transformer- beaded models, the basic structures that many consider pruning with include blocks, attention heads, and/or fully-com collecconde matrix rows (Chen et al. 2020b 202Ib). For example, (Michel 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.
channel-wise summation
Knowledge Distillation. Knowledge distillation (KD) is special technique that does not explicitly compress the model from any dimension of the network. KD lets a student model leverage "'soft" labels coming from a teacher network (Hinton al., 2015) to boost the performance of student model. This can be regarded as form of compression from the teacher model into smaller student. The soft labels from the teacher are 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 crucial role in transformers (Raghu et al. 2021), by tackling the vanishing gradient problem (Vaswani et al., (017), 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 a 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.] (220)) proposes to randomly replace the original modules with their designed compact substitutes to train the compact modules mimic the behavior of the original modules. (Zhang & He| 2020) designs Switchable-Tranee 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., (020) first trains 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.
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 transformer blocks. In each block of the ViT, there are two constituents, namely the Multi-head Self Delf-Attention (MSA) module and the MLP module. Uniformly, each MSA for transformer block has H attention heads originally.
In the Multi-head Self-Attention module, WhO, (!) w(O are the weights of the three linear K projection matrix in block that uses the block input X to calculate attention matrices: Q', K", V The weights of the projection module that follows self-attention calculation are denoted as W((.1), represent the first linear projection module in block l. The MLP module consists of two linear projection modules (1.2) and (1..3)
Compression Targets The main parameters that can be potentially compressed in ViT block are w ((())
K' w-( and w(t.1), w((.2), (:...), 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 reduce other dimensions such 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, (Wu.)))sssL 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(.3) Jigggt as our
we innovate to use (W(.)1) Jisig
3

---

Published as a conference paper at ICLR 2022
Made:
UVC
maseng gooddl
Comprassed
Transformer
Knosdedge
thask hights
0bc Pepao
Aasevable
The overall framework of UVC, that integrates three compression strategies: (1) Pruning within a block: In transformer block, we targeting on pruning Self-Attention head numbers (s(1.1)), neuron numbers within a Self-Attention head (r"))) and the hidden size of MLP module (s(1.)) as well. (2) Skipping manipulation across blocks: When gt. dominates, directly skip block 7 and send X into block 1 + 1; Otherwise, pass X into block without skip connection. (3) Knowledge distillation: The original uncompressed model is used to provide soft labels for knowledge distillation.
Figure 1:
pruning targets, since these linear layers do not have dimension alignment issues with other parts, they can be freely pruned, while the output of (W(1.2) should match with the dimension of block input. We not prune W w w inside each head, since Q', K V 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 al. 2021). The dimension between X and the output of the linear projection module should be aligned. For this sake, W(..2) is excluded from our compression. Eventually, the weights to be compressed in our subsequent framework are (W(l.1), W(l.3) IIsSs.
3.2 RESOURCE-CONSTRAIIE CONSTRAINED END-TO-END VIT COMPRESSION
We target a principled constrained optimization framework jointly optimizing all weights and com- pression hyperparameters. We consider two strategies to compress VIT: (1) structural pruning of linear projection modules in 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
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 enforced with the same pruning ratios, models that are under finer-grained sparsity (i.e., pruning in smaller groups) are unfriendly to latency, while models that are under coarser- coarser-grained sparsity (i.e., pruning larger groups) is unfriendly to accuracy. The mixed-level 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.
To compress each linear projection module on width, we first denote (1.33) as the number of columns to be pruned for weights W(((3), Compression for weights (((,) is more complicated as it is decided by two degrees of freedom. As the input tensor to be multiplied with ((..)) is the direct output of attention heads. Hence, within each layer I, we denote s(t.1) be the attention head numbers that need to be pruned, and p((,)) 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 p((.1), and the head number level as controlled by (t.1). They give more flexibility to transformer compression by selecting the optimal architecture in multi calti-graiied way. We emphasize again that those variables above are not manually picked, but rather optimized globally.
Pruning within Block
Skipping Manipulation across Blocks The next component to jointly optimize is the skip connec- tivity patterm, across different VIT blocks. For vanilla ViTs, skip connection plays crucial role in
4

---

Published as a conference paper at ICLR 2022
ensuring gradient flow and avoiding feature collapse (Dong et al. 2021 Raghu et al. (0021) 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.
In general, it is known that more skip connections help CNNs' accuracy without increasing the parameter volume (Huang et al., 2017); and ViT seems favor the same trend (Raghu et al. (022)) 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-wiss or conventional structural pruning
While directly dropping blocks might look aggressive the first glance, there are two unique enabling factors in ViT's 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 re- concatenating the others, without causing any feature dimension incompatibility. Secondly, Phang et al. (202T) 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-layet layer similarities (while being dissimilar to early layer features), making additional blocks add little to the discriminative power. /hho et al. (202T) 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 ViT for the first time. Specifically, for each transformer block with the skip connection, we denote gt((..) and (t...) as two binary gating variables (summing up to 1), to decide whether to skip this block or not.
The Constraints: We next formulate weight sparsity constraints for the purpose of pruning. As discussed in See 3.1 the target of the proposed method is prune the head number and head dimension simultaneously, which can actually be modeled as a two-level group sparsity problem when choosing hiigh L as proxy compression targets. Specifically, the input dimension of (W(.1)) J sig.t. is equivalent to the sum of the dimensions of all heads. Then we can put a two-level group sparsity regularization on (W'(1.))
(((..)))eee | input dimension to compress head number and head dimension at the same time, as shown in la and 1b (1.+)) corresponding to the pruned size of in. head. ((.11) means the pruned number of heads. For { (W(1)3) sisL.. it is our compression target, we just perform standard one level group sparsity regularization on its input dimension, as shown Ic
Σ
1
w Sur
(1.1)
2
=
0
2
(1.))
(1a)
2
Σ
2
=0
> s1.11),
(lb)
12
Σ
[ (t.3
0
2
(1.3)
(1c)
V1 1,2, L, Vi 1,2 .
(1d)
where W., denotes the i-th column of W; W denotes the i-th grouped column matrix of W, i... the i-th head; and W W the j-th column of W 9:... which is the j-th column of the i-th head. In other words, in Eqn. Tb column matrices are grouped by attention heads. Hence among the original H heads in total at block I, at least g g(.1)) heads should be discarded. Similarly, Eqn. demands that at least (1.3) input neurons should be pruned at this linear projection module. Furthermore, Eqn. la requests that in the i-th attention head at block 1, i) of the output units should be set to zeros.
Our method is formulated as 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 and gt, denoted as [ r, gt). We have the constraint as:
Reisps (s, rgtt) S Rhudget;
5

---

min Published as a conference paper at ICLR 2022
Inbed
200
All
JR
ongratstepe 58%
poired 18'9
Figure 2: The two sparsity levels for pruning within block: the head dimension level as controlled by y(l.)) and the head number level as controlled by g(l.1), When reaching the same pruning ratio, neuron-level sparsity will not remove any head, which is usually unfriendly to latency; while head level sparsity will only remove attention heads, which usually unfriendly to accuracy.
where (s(111,, (((.3)) hssgs. and (r((.)) (p(.,)))ssss hisss. Reudge is the resource budget. We present detailed flops computation equations in terms of 8, r and gt in Appendix.
Those inequalities can be further rewritten into equation forms facilitate optimization. As example, we follow (Tono et al. (017) to reformulate Eqn. ab
"n di 11 = o) 2 8 + W gl =0
(3)
where || W. denotes the Frobenius norm of the sub-matria W consisting of groups of W with smallest group norms. Eqn. Ta and Eqn. Tcl have same conversions.
The Objective: The target objective could be written (A a coefficient):
L(W. .t) f(W.gt gt) + Xlasstu((W Wt),
W Wgtt
where W, denotes the weights from teacher model, i.e., the uncompressed transformer model, and ldistsu (+, *) is defined as the knowledge distillation loss. We choose the simper &g norm as its performance found to be comparable to the more traditional K-L
its implementation. divergence (Kim et al. 2021).
The Final Unified Formulation:
Summarizing all above, we arrive our unified optimization as mini-max formulation, by leveraging the primal-dual method (Buchbinder & Naor 2009):
min max Craning W. 200.220
min max L(W, gt) + (Hriop (s, r, r,gt) Rmadger parntppp..22
resource loss
L
H
w
Iramm sparsity
(1.3)) (3))
Σ
(..) .9t
(5)
1=
8 g.p...gt.W W)
Eventually, we use the solution of the above mini-max problem, i.e. and r to determine the compression ratio of each layer, and gt to determine the selection of skip configuration. Here, select the pruned groups of the parameters by directly ranking columns by their norms, and remove those with the smallest magnitudes.
The general updating policy follows the idea of primal-dual algorithm. The full algorithm is outlined in Algorithm Tin Appendix. We introduce the solutions each sub-problem Appendix as well.
EXPERIMENTAL RESULTS
Datasets and Benchmarks We conduct experiments for image classification on Ima- geNet (Krizhevsky al. 2012). We implement UVC on DeiT (Touvron al., (2020), which has basically the identical architecture compared with ViT (Dosovitskiy et al. (2020) except for an extra distillation token; and variant of -2T-ViT Yuan et al. 2021). with special token embedding

---

Published as a conference paper at ICLR 2022
layer that aggregates neighboring tokens into token and backbone with deep-narrow structure. Experiment has been conducted on Deil-Tiny/Small/Bus and T2T-ViT-14 models. For DeiT-Tiny, we also try it without/with the distillation token. We measure all resource consumptions (including the UVC resource constraints) terms of inference FLOPPs.
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 mainly follow the training settings of DeiT Touvron et al. 2020 except for a relatively smaller learning rate which benefits finetuning of converged models.
For distillation, we select the pretrained uncompressed model to teach the pruned model.
Numerically, the learning rate for parameter z always changing during the primal-dual algorithm process. Thurs, we propose to use dynamic learning rate for the parameter that controls the budget constraint. We use a four-step schedule of (1,5, 5, 9, 13, 17] practice.
Baseline methods We adopt several latest compression methods which fall under two categories:
Category 1: Input Patch Reduction, including (1) POWER (Goyal et al. 2020 which accelerates language model BERT inference by eliminating word-vector in a progressive way. (fi) HVT (Pan et al. 2021bb: which reduce the token sequence dimension by using max-pooling hierarchically. (27)) PatchSlimming TTang et al. 2021): which identifies the effective patches in the last layer and then use them guide the patch selection process of previous layers; and (iv) IA-RED² (Pan et al. 2021a): which dynamically and hierarchical drops visual tokens at different levels using learned policy network.
Category 2: Model Weight Pruning, including (v) SCOP (Tang et al.
2020): which is
originally a channel pruning method used on CNN, we follow (Tang et al. 2021) and implement it ViT. (vi) Vision Transformer Pruning (VTP) (Zhu al.
(021): which
trains transformers with sparsity regularization to let important dimensions emerge. (vii) SVITE (Chen al. 2021b): which joinhy optimizes model parameters and explores sparse connectivity throughout training, ending up with one final sparse network. SViTE belongs the most competitive accuracy-efficieet
trade-off achieved so far, for VIT pruning. We choose its structured variant to be fair with UVC.
4.1
MAIN RESULTS
The main results are listed in Tab. 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 2 60% FLOPS saving) without compromising as much. For example, when compressing DeiT-Tiny (with distillation token), UVC can trim the model down to compelling 50% of the original FLOPs while losing only 0.3% accuracy (Tab. 3 Especially, compared with the latest SVITE (Chen et al. 2021b) that can only save up to around 30% FLOPs, we observe UVC to significantly outperform ig at DeiT-Tiny/Smalll at less accuracy drops (0.9/0.4, versus 2. 1/0.6) with much more aggressive FLOPs savings (50.7 30.7%% versus 23.7% 163). For the larger DeiT-Base, while UVC can save 45% of its FLOPs, fin that SVITE cannot be stably trained at such high sparsity. UVC also yield comparable accuracy with the other ViT model pruning approach, VTP (Zhu et al. 202T), with more than 10% FLOPs savings w.r.t. the original backbone.
Secondly, UVC obtains strong results compared with Patch Reduction based methods. UVC performs
clearly better than IA-RED (Pan et al. 202 at DeiT-Base, and outperforms HVT Paan et al. 2021b) on both DeiT- TinySSall models, which are latest and strong competitors. Compared to Patch Slimming 'Tang et al. (2021). UVC at DeiT-Small reaches 79.44% top-1 accuracy, while the compression ratio measured by FLOPs is comparable. On other models, we observe UVC to generally save more FLOPs, yet also sacrificing more accuracies. Moreover, as explained in
7

---

Published as a conference paper at ICLR 2022
Table Comparison of the vision transformers compressed by UVC with different benchmarks on ImageNet. FLOPs remained denotes the remained ratio of FLOPs to the full-model FLOPs.
Model
Method
Top-1 Acc. (%)
FLOPs(G)
FLOPs remained
Baseline
72.2
1.3
100
SVITE
70.12 (-2.08)
0.99
76.31
PatchSlimming
72.0 (-0.2)
0.7
53.8
DeiT-Tiny
UVC
71.8 (-0.4)
0.69
53.1
HVT
69.7 (-2.5)
0.64
49.23
UVC
71.3 (-0.9)
0.64
49.23
UVC
70.6 (-1.6)
0.51
39.12
Baseline
79.8
4.6
100
SVITE
79.22 (-0.58)
3.14
68.36
PoWER
78.3 (-1.5)
2.7
58.7
DeiT-Small
UVC
79.44 (-0.36)
2.65
57.61
SCOP
77.5 (-2.3)
2.6
56.4
PatchSlimming
79.4 (-0.4)
2.6
56.5
HVT
78.0 (-1.8)
2.40
52.2
UVC
78.82 (-0.98)
2.32
50.41
Baseline
81.8
17.6
100
IA-RED
80.3 (-1.5)
11.8
67.04
SVITE
82.22 (+0.42)
11.87
66.87
DeiT-Base
VTP
80.7 (-1.1)
10.0
56.8
PatchSlimming
81.5 (-0.3)
9.8
55.7
UVC
80.57 (-1.23)
8.0
45.50
Baseline
81.5
4.8
100
PoWER
79.9 (-1.6)
3.5
72.9
T2T-VTT-14
UVC
80.4 (-1.1)
2.90
60.4
UVC
79.6 (-1.9)
2.47
51.5
UVC
78.9 (-2.6)
2.11
44.0
Table 2: Ablation study the modules implemented in UVC. The first part present single technique ablation. The second part presents the result to sequentially apply Skip configuration and Pruning.
DeiT-Tiny
Method
Acc. (%)
FLOPs remained %)
Uncompressed baseline
72.2
100
Only skip manipulation
68.72
51.85
Only pruning within block
70.52
50.69
Without Knowledge Distillation
69.34
51.23
Skip 71.49% +49.39%
66.84
49.39
Prune->Skip 69.96% -55,8%5548899
69.68
54.88
Prune Prune->Skip 63.25% +50.73%
70.02
50.73
UVC
71.3
49.23
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 future work that we would pursue.
Thirdly, we test UVC on compressing TIT-VIT Yuaan et al. 2021). In the last row of Tab. UVC with 44% FLOPPs achieves only 2.6% accuracy drop. Meanwhlle, UVC achieves a 1.9% accuracy drop with 51.5% of the original FLOPs. With 60.4% FLOPs, UVC only suffers from 1.1% 1% accuracy drop, already outperforming PoWER with 72.9% FLOPs.
8

---

Published as a conference paper at ICLR 2022
4.2 ABLATION STUDY
As UVC highlights the integration of three different compression techniques into one joint opti- mization framework, it would be natural to question whether each moving part is necessary for the pipeline, and how they each contribute the final result. We conduct this ablation study, and the results are presented in Tab.
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 actually reduced into LayerDropping (Lin al. 2018). We have the following findings: Firstly, implementing only skip connection manipulation will incur high instability. During optimization, the objective value fluctuates heavily due to the large architecture changes (adding or removing one whole block) and barely converges to a stable solution. Secondly, only applying skip manipulation will be damaging the accuracy remarkably, e.g., by nearly 4% drop on Deiy--Tiny. That is expected, only manipulating the model architecture with such coarse granularity (keeping or skipping whole transformer block) 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-leve pruning, with knowledge distillation. It turns out to deliver much better accuracy than only skip manipulation, presumably owing to its finer-grained operation. However, still has margin behind our joint UVC method, as the latter also benefits from removing block-level redundancy a priori, which is recently observed to widely exist in fine-tuned transformers [Paagg et al. (202T). Overall, the ablation results endorse the effectiveness of jointly optimizing those components altogether.
Applying Individual Compression Methods Sequentially Our method is joint optimization process for skip configuration manipulation and pruning, but is this joint optimization necessary?
To address this curiosity, we mark Prune the process of pruning first, manipulating skip configuration follows, while Skip->Prune vice versa. Note that Skip->rune 71 denotes the pipeline that use skip configuration first enforce the model FLOPs be 71.49% of the original FLOPs, then apply pruning to further compress the model to 49.39%. 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 model with 50% of the original FLOPs. Based on this, we implement:(i) Skip Skip->Prune Prune 71.49% 449.39% 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.8888 % applics pruning to 70% FLOPs then skipping follows. We observe that applying pruning first can stabilize the choice of skip configuration and result in top-1 accuracy of 69.68%. Both results endorse the superior trade-off found by UVC.
939%
Furthermore, design a better setting (Prune->Skip 63. ratio found by UVC. Details can be found in Appendix. The obtained result is much better than previous ones, but still at a 1% performance gap from the result provided by UVC. We assign the extra credit to jointly optimizing the architecture and model weights.
simulate the compression
5 CONCLUSION
In this paper, we propose UVC, a unified VIT compression framework that seamlessly assembles pruning, layer skipping, and knowledge distillation as one. A budget-constrained optimization framework is formulated for joint learning. 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.
REFERENCES
Haoli Bai, Wei Zhang, Lu Hou, Lifeng Shang, Jing Jin, Xin Jiang Qun Liu, Michael Lyu, and Irwin King. Binarybert: Pushing the limit of bert quantization. arXiv preprint arXiv:2012. /5701, 2020.

---

Published as a conference paper at ICLR 2022
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through
stochastic neurons for conditional computation. arXiv preprint arXiv:/308.3433: 2013.
Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space-timn simee-timeaee attention all you need for video understanding? arXiv preprint arXiv:2102.05009 2021.
Niv Buchbinder and Joseph Naor. The design of competitive online algorithms via a primal-dual approach. Now Publishers Inc, 2009.
Aydin Buluc and John R Gilbert. Challenges and advances in parallel sparse matrix-matrix mul- tiplication. In 2008 37th International Conference Parallel Processing, pp. 503-510. IEEE, 2008.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision, pp. 213-229, Springer, 2020.
Hanting Chen, Yunhe Wang, Tianyu Guo, Chang Xu, Yiping Deng, Zhenhua Liu, Siwei Ma, Chunjing Xu, Chao Xu, and Wen Gao. Pre-trained image processing transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 12299- 12310, 2021a.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David I Luan, and Ilya Sutskever. Generative pretraining from pixels. In International Conference on Machine Learning, pp. 1691 1703. PMLR, 2020a.
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis for pre-trained bert networks. Advances in neural information processing systems, 33: 33:15834 15846, 2020b.
Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, and Zhangyang Wang. Chasing sparsity in vision transformers: An end-to-end exploration. arXiv preprint AXii:21006.0553. 202lb.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 190,,09994,,009,, 2019.
Zhigang Dai. Bolun Cai, Yugeng Lin, and Junying Chen. Up-detr: Unsupervised pre-training for object detection with transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1601-1610, 2021.
Jacob Devlin, Ming-Wes Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04880, 2018.
Xin Dong. Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer optimal brain surgeon. In Advances in Neural Information Processing Systems, pp. 4857-4867, 2017.
Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention not all you need: Pure attention loses rank doubly exponentially with depth. arXiv preprint arXiv:2/03.03404, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 /929, 2020.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint arXiv: 1909. 1556, 2019.
Trevor Gale, Erich Elsen, and Sara Hooker. The state of sparsity in deep neural networks. arXiv preprint 1rXiv:1902.09574 2019.
Prakhar Ganesh, Yao Chen, Xin Lou, Mohammad Ali Khan, Yin Yang, Deming Chen, Marianne Winslett, Hassan Sajaa, and Preslav Nakov. Compressing large-scale transformer based models: A case study bert. Xiv preprint arXiv:2002. 1985, 2020.
10

---

Published as a conference paper at ICLR 2022
Mitchell A Gordon, Kevin Duh, and Nicholas Andrews. Compressing bert: Studying the effects of weight pruning on transfer learning. arXiv preprint arXiv:2002.08307, 2020.
Saurabh Goyal, Anamitra Roy Choudhury, Saurabh Raje, Venkatesan Chakaravarthy, Yogish Sab- harwal, and Ashish Verma. Power-bert: Accelerating bert inference via progressive word-vector elimination. In International Conference on Machine Learning, pp. 3690-3699. PMLR, 2020.
Fu-Ming Guo, Sijia Liu, Finlay S Mungall, Xue Lin, and Yanzhi Wang. Reweighted proximal pruning for large-scale language representation. arXiv preprint arXiv: 1909, /2486, 2019.
Yi Guo, Huan Yuan, Jianchao Tan, Zhangyang Wang, Sen Yang, and Ji Liu. Gdp: Stabilized neural network pruning via gates with differentiable polarization. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pp. 5239-5250, 2021.
Kai Han, Yunhe Wang. Hanting Chen, Xinghao Chen, Jianyuan Guo. Zhenhua Liu, Yehui Tang. An Xiao, Chunjing Xu, Yixing Xu, Zhaohui Yang, Yiman Zhang, and Dacheng Tao. A survey on vision transformer, 2021.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint 15000:15100000449 2015a. 5a.
Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network, In Advances neural information processing systems, pp. 1135-1143, 2015b.
Kaiming Hc. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Yang He, Ping Liu, Ziwei Wang, Zhilan Hu, and Yi Yang. Filter pruning via geometric median for deep convolutional neural networks acceleration. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 4340 4349, 2019.
Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. In The IEEE International Conference on Computer Vision (ICCV), Oct 2017.
Yihui He, Ji Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. Ame: Automl for model compression and acceleration on mobile devices. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 784-800) 2018.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.02533 2015.
Lu Hou, Zhiqi Huang, Lifeng Shang. Xin Jiang. Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. arXiv preprint arXiv:2004.04037 2020.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Weinberger. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4700-4708, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:161 1.01 144, 2016.
Ding Jia, Kai Han, Yunhe Wang, Yehui Tang, Jianyuan Guo, Chao Zhang, and Dacheng Tao. Efficient vision transformers via fine- ane-grained manifold distillation. arXiv preprint arXiv:2107.0177, 2021.
Yifan Jiang, Shiyu Chang, and Zhangyang Wang. Transgan: Two pure transformers make one strong gan, and that can scale up. Advances in Neural Information Processing Systems, 34, 2021.
Tachyeon Kim, Jaehoon Oh, Nak Yil Kim, Sangwook Cho, and Se- Ye-Young Yun. Comparing kullback-leibler divergence and mean squared error loss knowledge distillation. arXiv preprint 100.v220550091/ 2021.
11

---

Published as a conference paper at ICLR 2022
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep con- volutional neural networks. Advances in neural information processing systems, 25:1097-1105, 2012.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv: 1942, 2019.
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard E Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied handwritten zip code recognition. Neural computation, 1(4):541-55) 1989.
Yann LeCun, John S Denker, and Sara Solla. Optimal brain damage. In Advances neural information processing systems, pp. 598-605, 1990.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shot network pruning 2018.
based on connection sensitivity. arXiv preprint
Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. arXiv preprint 160iv: 08710, 2016.
Mingbao Lin, Rongrong Ji, Yan Wang, Yichen Zhang, Baochang Zhang, Yonghong Tian, and Ling Shao. Hrank: Filter pruning using high-rank feature map. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1529-1538, 2020.
Shaohui Lin, Rongrong Ji, Yuchao Li, Yongjaan Wu, Feiyue Huang, and Baochang Zhang. Accel- erating convolutional networks via global & dynamic filter pruning. In IJCAI, pp. 2425-2432 2018.
Zhuang Liu, Jianguo Li, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang. Learning efficient convolutional networks through network slimming. In ICCV, 2017.
Jia Luo, Jianxin Wu, and Weiyao Lin. Thinet: A filter level pruning method for deep neural network compression. In Proceedings of the IEEE international conference on computer vision, pp. 5058-5066, 2017.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? arXiv preprint arXiv: 1905 2019.
Asit Mishra and Debbie Marr. Apprentice: Using knowledge distillation techniques to improve low- now-precision network accuracy. In International Conference Learning Representations, 2018.
Bowen Pan, Yifan Jiang, Rameswar Panda, Zhangyang Wang, Rogerio Feris, and Aude Oliva. Ia-red Interpretabilit--awee redundancy reduction for vision transformers. arXiv preprint arXiv:2106 12620, 2021a.
Zizheng Pan, Bohan Zhuang, Jing Liu, Haoyu He, and Jianfei Cai. Scalable vision transformers with hierarchical pooling. In Proceedings of the IEEE/CV International Conference Computer Vision, pp. 377-386, 2021b.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, pp. 4055-4064. PMLR, 2018.
Jason Phang, Haokun Liu, and Samuel Bowman. Fine-tuned transformers show clusters of similar representations across layers. arXiv preprint w:Xiv:2109.0844000 2021.
Maithra Raghu, Thomas Unterthiner, Simon Korbblihh, Chiyuan Zhang, and Alexey Dosovitskiy. Do vision transformers see like convolutional neural networks? arXiv preprint 10X.v:210..088/1 2021.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv: 1910.0/100 2019.
12

---

Published as a conference paper at ICLR 2022
Jiayi Shen, Haotao Wang. Shupeng Gui, Jianchao Tan, Zhangyang Wang. and Ji Liu. Umec: Unified model and embedding compression for efficient recommendation systems. In International Conference Learning Representations, 2021. /BPp:/pp///pppeeer.e....t for/ommmmm
Yehui Tang, Yunhe Wang, Yixing Xu, Dacheng Tao, Chunjing Xu, Chao Xu, and Chang Xu. Scop: Scientific control for reliable neural network pruning. arXiv preprint arXiv:2010. 2020.
Yehui Tang, Kai Han, Yunhe Wang, Chang Xu, Jianyuan Guo, Chao Xu, and Dacheng Tao. Patch slimming for efficient vision transformers. arXiv preprint arXiv:2106.02860 2021.
Katsuya Tono, Akiko Takeda, and Jun-ya Gotoh. Efficient de algorithm for constrained sparse
optimization. arXiv preprint arXiv: 0ee./7111...999 2017.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. arXiv preprint arXiv:2012. 12877, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint 1exXiv::706.3333 2017.
Huiyu Wang, Yukun Zhu, Hartwig Adam, Alan Yuille, and Liang-Chich Chen. Max-dceplab: End-to- end panoptic segmentation with mask transformers. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 5463 4474-,474, 2021a.
Yuqing Wang. Zhaoliang Xu, Xinlong Wang, Chunhua Shen, Baoshan Cheng, Hao Shen, and Huaxia Xia. End-to-end video instance segmentation with transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8741-8750, 2021b.
Longhui Wei, An Xiao, Lingxi Xie, Xiaopeng Zhang, Xin Chen, and Tian. Cireumventing outliers of autoaugment with knowledge distillation. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part III 16, pp. 608 608-625. Springer, 2020.
Xia Xiao, Zigeng Wang, and Sanguthevar Rajasekaran. Autoprune: Automatic network pruning by regularizing auxiliary parameters. In Advances Neural Information Processing Systems, pp. 13681-13691,200
Canwen Xu, Wangchunshu Zhou, Tao Ge, Furu Wei, and Ming Zhou, Bert-of-theseus: Compressing bert by progressive module replacing. arXiv preprint arXiv:2002.02925, 2020.
Fuzhi Yang, Huan Yang, Jianlong Fu, Hongtao Lu, and Baining Guo. Learning texture transformer network for image super-resolution In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 5791-5800, 2020a.
Haichuan Yang, Yijun Huang, Lam Tran, Ji Liu, and Shuai Huang. On benefits of selection diversity via bilevel exclusive sparsity. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 5945-5954, 2016.
Haichuan Yang, Shupeng Gui, Yuhao Zhu, and Ji Liu. Automatic neural network compression by joint learning: A constrained optimization- based approach. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2178 21 2188, 2020b.
Ruichi Yu, Ang Li, Chun-Fu Chen, Jui-Hsin Lai, Vlad I Morariu, Xintong Han, Mingfei Gao, Ching- Yung Lin, and Larry S Davis. Nisp: Pruning networks using neuron importance score propagation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 9194-9203, 2018
Shixing Yu. Zhewei Yao, Amir Gholami, Zhen Dong, Michael W Mahoney, and Kurt Keutzer. Hessian-aware pruning and optimal neural implant. arXiv preprint 101.::2:0000094000 2021.
Li Yuan, Francis EH Tay, Guilin Li. Tao Wang, and Jiashi Feng. Revisiting knowledge distillation via label smoothing regularization. In Proceedings of the IEEÉ/CVF Conference on Computer Vision and Pattern Recognition, pp. 3903-3911, 2020.
13

---

Published as a conference paper at ICLR 2022
Li Yuan, Yunpeng Chen, Tao Wang. Weihao Yu, Yujun Shi, Zihang Jiang, Francis EH Tay, Jiashi Feng. and Shuicheng Yan. Tokens-to-tokers vit: Training vision transformers from scratch on imagenet. 2021.
Minjia Zhang and Yuxiong He. Accelerating training of transforrerr-bsee language models with progressive layer dropping. arXiv preprint arXiv:2010.1333 13366, 2020.
Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. Ternarybert: Distillation-awarl aware ultra-low bit bert. arXiv preprint arXiv:2009. 12811 2020.
Hengshuang Zhao, Li Jiang. Jiaya Jia, Philip Torr, and Vladlen Koltun. Point transformer. arXiv preprint arXiv:20/2.09/100 2222002..9114.. 2020a.
Yang Zhao, Xiaohan Chen, Yue Wang, Chaojian Li, Haoran You, Yonggan Fu, Yuan Xie, Zhangyang Wang, and Yingyan Lin. Smartexchange: Trading higher-cost memory storage/access for lower-cost computation. 2020 ACM/IEEE 47th Annual International Symposium on Computer Architecture (ISCA), pp. 954-967. IEEE, 2020b.
Minghang Zheng, Peng Gao, Xiaogang Wang, Hongsheng Li, and Hao Dong. End-to-end object detection with adaptive clustering transformer. arXiv preprint XXi::2011.09311 2020.
Daquan Zhou, Bingyi Kang, Xiaojie Jin, Linjie Yang. Xiaochen Lian, Zihang Jiang, Qibin Hou, and Jiashi Feng. Deepvit: Towards deeper vision transformer, 2021.
Mingjian Zhu, Kai Han, Yehui Tang, and Yunhe Wang. Visual transformer pruning. arXiv preprint arX4v:2104.08500, 2021.
Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai. Deformable detr: Deformable transformers for end-to-end object detection. arXiv preprint arXiv:2010.04159, 2020.
14

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 UPDATING POLICY
Updating Weights W Different from other pruning methods in CNN that use fixed pre-trained weights to select the pruned channels/groups with certain pre-defined metrics (Lin et al. [020 He al.) 2019), here subproblem could be considered following dynamic pruning criteria that will be updated along. Specifically, we solve the following subproblem:
1 pr S(y..... w) (W) arg minw 2
w
w/l + 1 (U,8.p.r,
(6)
where W = W - -m nwwl (W*). The solution admits a bi-level projection
Yang et al.
2016
W(11)
if (111)
2
(1.1) deat (o(11)] 2
W(..).
9963
W otherwise,
1+2gp(
(-ll.1)
if ((.1) (1.1) w
2
((il))
otherwise, (2011
:9:
W
(tay
w1l13)
W (1.3)
(1.3)
w((.)),
2
lesst [x(1.3)
W
otherwise,
1+2myita
where least-j denotes the index of the (group) columns of W that have j-th least (group) norm.
Updating gt gt' (gt(t.) (t(..)) are used to generate a binomial categorical distribution to decide whether pass through block directly skip it. As the variables are discrete, apply the renowned Gumbel SorellSooftt Softmax (GSM) trick (Jang et al.) (016)) to obtain differentiable and polarized sampling. For gt' (.t( gt(()),, ), given i.i.d Gumbel noise drawn from Gumbel 1) distribution, soft categorical sample can be drawn by
G² GSM(gtt) SofImar + g)/r) E ERR
where ((..) refers to the continuous possibility preserve current block 1 while ((1,0) drop it. Directly applying the chain rule L(W, .gt) w.r.t gt can now calculate ggtC(W', "gt.()
Also, gt participates in the calculation of FLOOPs, by deciding whether to pass certain blocks. Since passing skipping block is dynamic choice during training, estimate the FLOPS of the 1-th block Reoops (,, r, gt) with skip gating by using its expectation. To be specific,
Rioops, (8), rt, gti)
= Ri (st, Tl, (1,gt)) 81, ri]
(8a)
G(I.)) Reiops (Identity) + ((..1) Ryiops (81,
(8b)
= G(I.1) Relops, (st, rı)
(8c)
Hence, updating policy for gt is formulated as:
gl = gt' 774 (VgeL (W. Wtt)) + & gez (Rewps (s,r.tt") Rbuugea)
(9)
= gt' 7/4 (v (WWtth + v (B,,
(10)
15

---

Published as a conference paper at ICLR 2022
Updating and Similar to the updating policy of gt, one gradient term w.r.t. and vaz (Relops (s, r.t) Roudgat) v (Rriops (s, r. .t) Reudger) respectively.
The other gradient term calculated on the unified formulation Eqn. 5 Refer to that, s and r are floating-point numbers during the optimization process. In practise, ceiling functions are operated on them to determine the integer number that should be pruned for each layer. However, the ceiling function [.] is non-differentiable. To solve this problem, implement thrrightthhrugh (STi) (Bengio et al. 2013) to provide a proxy of the gradient when performing the backward
0[s] =1
pass. We set
Ōs
As for || [ .g 22 term in the sparsity loss, we use | 11
[ ll the proxy of partial
derivative of [W.. [W.gll with respect to
Ô W w. W
(D ).++1) 11
(1I)
Is
where Dim( W) is the number column groups of W. The other two terms in the sparsity loss can be processed similarly.
A.2 MAIN ALGORITHM
The general updating policy follows the idea of primal-dual algorithm. The full algorithm is outlined in Algorithm
Algorithm 1: Gradient- based algorithm to solve problem (5) for Unified ViT Compression.
Input: Resource budget Roudges. learning rates 7/, Ija, 77,, Tl. Y,, number of total iterations T. Transformer pruned weights W /1 random
Result:
Initialize t W
pre-trained dense model
for t 1
W Prox,, S(y pt.... ( n WWw.
11 Proximal-SGD
s s41 s 1/2
V.S (y', s', ,'.r''r' gt', w Va2' (Hrups (s', rt gt') -Rrutgs):
// Gradient (STE) Descent
r11 r' 75
(v,S (y', st+1, p¹ .rr'ggt" w + 2 (Briops (s'-1, rr' rgggt) Reutga)
1/ Gradient (STE) Descent
ge gt 174 oe (We+1 gt') + ge (Rreops (s'-1, nlt1 gt)) Reudga)
// Gradient Descent 141 = 3¹ + 1r (Rrogs
(s +1
gt
Reudea)
// Gradient Ascent
g(a1)*+1 = -000 + n5
W :9
11
// Gradient Ascent
((a) y n7
(1l
1A 1,...,
p p100ttt = p(.)) + m/8
W (1.1ccc
Vi= H, 1,.,
W W
A.3 VISUALIZATION
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-Ting model, UVC tends 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.
16

---

Published as a conference paper at ICLR 2022
Dopped
1-2
1-3
1-4
1-1
1-4
1n7
1-1
10:30
Iell
(a) attention head
(b) skip connection
Figure 3: Visualizations of: (a) sparse attention head patterns for DeiT-Tiny (left) and DeiT-Base
(right); (b) skip connection patterns for DeiT-Tiny (top) and DeiT-Base (bottom).
Table 3: Result with distillation token of DeiT
Model
Method
Top-1 Acc. (%)
FLOPs(G)
FLOPs remained(%)
DeiT-Tiny
Baseline
74.4
1.3
100
(dist token)
UVC
74.1(-0.3)
0.66
50.58
Skip connection patterns: We next draw the result of skip connection patterns Fig. (b), showing what layers learned be directly skipped under unified optimization. We observe UVC's obvious tendency to drop later layers, which coincides nicely with the observations by (Phang
2021
Zhou al. 2021).
As is observed in Fig. DeiT-Tiny discarded 2 blocks in total, which results in 84% of compression ratio. Thus, in Sec. 4.2, to design an optimal architecture for sequentially applying compression methods, first prune the DeiT-Tiny to 63% and apply skipping to indulge 50% FLOPs' model to mimick the behavior of UVC.
A.4 EXTRA RESULTS
In this section, we present the result of DeiT with distillation token extract the best performance vision transformer can reach under UVC framework. Results are shown in 3
17


Reviews:
[{"review_id": "h0cmFmURDLp", "reviewer": "Reviewer_K4Gq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors’ efforts to formulate and solve ViT compression as an end-to-end budget-constrained optimization is appreciated. Directly skipping layers is a fresh idea as a compression means and seems particularly to fit ViT due to its homogenous block structure and common feature collapse. This makes a brand-new angle to compress ViTs. The formulation and algorithm are correct to my best knowledge. UVC only requires specifying an overall resource budget and can automatically optimize the composition of different techniques. It is easier to use than ad-hoc alternatives while being very flexible. UVC experimentally outperformed a few very recent baselines (e.g, surpassing SViTE of NeurIPS’21), accompanied by sufficient empirical dive-in with ablation studies and visualization. \n\nThe novelty is okay but not so great. Combining two or three compression means is well studied in CNN compression. For UVC’s specific combination, the distillation loss is a rather straightforward plug-in, and isn’t the same technically interesting as the other two in joint optimization.\nThe performance isn’t clearly better than Patch Slimming, mostly due to the current mismatched FLOPs. If the authors cannot show better than it in a fair setting, then I suggest the authors to demonstrate they can be combined so UVC’s empirical benefits can still stand.  \nA major caveat in experiments is that the authors never reported their due ablation experiments. A clearly missing baseline is: what if you apply three compression methods sequentially, such as first pruning then skipping layer then distilling (or in whatever other reasonable orders)? This missing is crucial as no existing experiment can directly support joint optimization is essential\nThere is no DeiT-small in table and I don’t understand why. The compression ratios chosen to display are very “sparse” (kinda understandable, since training VITs is heavy). How will UVC perform if increasing compression ratios further?\n(Optional) it would further strengthen the paper if more strong ViT variants can be reported such as Swin Transformer or T2T.\n\n---\nThe authors have well addressed my major concerns. I am glad to increase my score to 6. ", "summary_of_the_paper": "This is an interesting work trying to assemble three effective techniques for pruning, layer skipping, and knowledge distillation. ", "summary_of_the_review": "While this joint optimization is a reasonably novel effort in the (somehow niche) ViT compression field, I cannot endorse this paper further given several critical experiments are missing. I’m willing to revisit my rating if the authors can provide the above questioned results. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "X2E-cC8DT3Z", "reviewer": "Reviewer_iRBq", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n\n* Model compression is well studied in the fields of CNNs in CV and Transformer models in NLP. But such techniques still need to be adopted and verified in visual Transformers.\n\n* The paragraphs in this paper are clear and easy to understand.\n\n* Experiments in ImageNet with DeiT models prove the effectiveness of the proposed method. Compared with baseline models and recent Transformer compression methods, UVC can obtain larger compression rates, resulting in higher accuracies.\n\nWeaknesses\n\n* The novelty of the paper is quite limited. It is like a combination of existing methods, such as Slimming, BigNAS.\n\n* The whole process of UVC is consists of compression training and post-training. I suspect that the comparison between different methods is not aligned, and post-training leads better performance.\n\n* I suggest that UCV can be verified on more Transformer architectures, such as Swin/T2T/CvT/CeiT, showing the generalization ability.\n\n* All tables are unmarked, which are hard to follow.\n\n", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, UVC has made a good attempt at model compression for visual Transformer models. But the novelty is limited. This method needs to provide a more rigorous comparison, and at the same time, it needs to be verified on more visual Transformer architectures to prove its effectiveness and generalization.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Hl6BSGjEO68", "reviewer": "Reviewer_UdwX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Motivation and Overall Structure:\nWhile a few prior arts already considered compressing ViTs, this paper stands out with its elegant formulation as a resource-constrained optimization problem, that is solved from end to end. It looks more synergistic and less ad-hoc. \n\nI also like the rigorous optimization in this paper. Different from previous ad-hoc methods that need tune more hyperparameters, UVC can compress with one global budget control, and jointly coordinate different compression forms automatically under the hood.\n\nModel and Algorithm: \nBesides the common distillation loss, one of the authors’ main new ideas is to enforce mixed-level group sparsity: the head dimension level, the head number level, and the block number level. The authors point out the rationale as: 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 mixed-level group sparsity can hence more flexibly trade-off between latency and accuracy. \nParticularly, including layer-skipping as a compression means is novel, and to my best knowledge not considered by previous joint compression works yet. The authors also present good rationale why directly skipping layers in ViT might be acceptable: due to the observed layer collapse in many trained ViTs. \nFor the algorithm part: while I am familiar with primal-dual or ADMM, I got lost when reading the ~2 pages of optimization derivation in section 3. Specifically, while laying out all details, it is difficult to follow what main algorithm novelty the authors intend to claim: is just this formulation idea novel? Or the overall primal-dual algorithm being novel? Or the solution to some subproblem novel?  \n\nExperiment and Analysis: \nOne biggest issue I have had - I am not convinced by the comparison with the current SOTA method of patch slimming. It shows UVC can compress to higher ratios, but meanwhile suffer from notably more accuracy losses. If I just look at DeiT-Tiny PatchSlimming (0.7 FLOPs) versus UCV (0.6 FLOPs), I’m under the impression that UVC drops accuracy more quickly. \n\nWhile the authors argued UVC can also be seamlessly extended to include token reduction, leaving this vaguely as “future work” cannot directly support the current work’s merit over prior arts. Could the authors report a more aligned apple-to-apple comparison with PatchSlimming?\n\nMore questions regarding experiments: Why not reporting DeiT-small? And why not more baselines for DeiT-Tiny with token? Also, the choice of target FLOPs looks a bit arbitrary to me, and different methods are not fully aligned. Could the authors clarify how they’re chosen, or whether they’re cherry picked?\n\nWriting:\nThe paper’s writing is overall polished. BUT the algorithm part (page 5-6, Algorithm 1) is unnecessarily lengthy, I think the readability would be better if many here were moved to supplementary. Also, all tables in the paper are not numbered – very weird and hard to refer to.\n", "summary_of_the_paper": "ViT compression is a new sub-field just being studied this year. With a few previous methods exploring pruning or data-dropping separately, this paper is the first to integrate multiple compression methods as one joint optimization.  ", "summary_of_the_review": "This seems to be an overall well-executed paper, with reasonable amounts of novelty and potential. However, the authors need clarify on a few raised questions, on which my final rating will depend. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper receives positive reviews. The authors provide additional results and justifications during the rebuttal phase. All reviewers find this paper interesting and the contributions are sufficient for this conference. The area chair agrees with the reviewers and recommends it be accepted for presentation.
Decision: Accept (Poster)