Vision Transformer Compression with Structured Pruning and Low Rank Approximation

Hinton et al., (2015) proposed to train a smaller neural network by first training a much bigger model and then training the smaller one using soft probabilities of the previous one. Since then, this technique has been successfully applied in multiple applications (Gou et al.,, 2021). However, we leave out this type of compression technique since it involves training the teacher model first and then running it again to train a student model, which may require significant computational resources.

A simple way to compress a model is to use low precision for weights and activations. With 8-bit quantization, model size can be reduced by a factor of 4 with negligible accuracy loss (Krishnamoorthi,, 2018; Nagel et al.,, 2021). Multiple efforts try to further reduce the precision, train a model with low precision etc (Gupta et al.,, 2015). However, it is difficult to lower the precision beyond 16-bit without significant accuracy degradation and specialized hardware (Wang et al.,, 2018). It has become a trend to use half precision for mixed precision training to avoid the accuracy drop and make use of the available hardware (Micikevicius et al.,, 2018).

Over-parameterization of neural networks is considered to be one of the reasons behind its superior performance. Many works have also established that a lot of unimportant weights can be removed to compress the models without sacrificing accuracy (Han et al.,, 2015; Blalock et al.,, 2020). There are several criteria to decide the importance score (LeCun et al.,, 1989; Hassibi and Stork,, 1992; Molchanov et al.,, 2019). A common and convenient way is to zero out the weights with small magnitude. This is a type of unstructured pruning. However, simply zeroing out does not improve the model size and speed unless there is specialized hardware. Another method, called structured pruning (Li et al.,, 2017; Wen et al.,, 2016), avoids this issue by dropping a group of weights such as entire convolution filter or attention head. Therefore, we consider structured pruning in this work and combine it with other compression techniques (Han et al.,, 2016).

Decomposition of weight matrices into low rank matrices is yet another technique to reduce the model size without affecting its structure. In case of convolution, factorizing channel and spatial convolution weights is used in place of traditional convolutional operation due to lesser number of parameters. Feedforward weight matrices can be factorized into two lower rank matrices to reduce the number of parameters (Sainath et al.,, 2013). However, it is not necessary that weight matrices are low rank, which we confirm for vision transformer from our experiments. The low rank structure can still be used if input feature is believed to lie in smaller dimensions (Chen,, 2021).

With the success of BERT, several works have studied how to make transformer models smaller. DistilBERT (Sanh et al.,, 2019) is a widely used compressed version of BERT models. Another line of work tries to make transformer models more efficient since computational complexity of transformer model grows quadratically in the number of input tokens. This is challenging for natural language processing, specially speech processing. Wang et al., (2020) show that attention matrix is typically low rank, and therefore, can benefit from low rank approximation.

Most of the works in computer visoin domain focus on removing uninformative patch tokens to speed up the training and inference steps (Wang et al.,, 2021). Vision transformer pruning (VTP) (Zhu et al.,, 2021) removes unimportant dimensions (columns or rows) of matrices in a transformer block. Inspired by several works combining multiple compression strategies (Han et al.,, 2016; Mao et al.,, 2020), we build upon VTP and first show the shortcomings of this technique. Then we propose to use low rank approximation in combination with VTP to mitigate some of the issues. Our work is orthogonal to token pruning, quantization and knowledge distillation and can be combined with these approaches.

An input image $x\in\mathbb{R}^{H\times W\times C}$ is divided and reshaped to have $P\times P$ patches, each having $P^{2}C$ pixels. The reshaped input $x_{P}$ has shape $\mathbb{R}^{N\times P^{2}C}$ where $N=HW/P^{2}$. Each patch is linearly projected into a $d$ dimensional space using a learnable embedding matrix. Following BERT, a learnable embedding corresponding to special token $[class]$ is prepended to $x_{P}$ making total sequence length $N+1$. This 1-d sequence of $d$ dimensional token can be fed to standard transformer model. Dosovitskiy et al., (2021) use learnable 1-d positional encoding which is added to the token embeddings. Structure of a transformer block is as shown in figure 1(a). We do not show GELU non-linearity which is applied after each of the top two linear layers in the block.

Output embedding corresponding to the $[class]$ token after the final transformer block serves as the input image representation. A classification head is attached on top of this to train the model. The head may contain additional hidden layer.

Each transformer block has almost all of the parameters in weight matrices of linear layer and query, key, value of self-attention layers. Vision transformer pruning Zhu et al., (2021) aims to drop unimportant columns or rows from a matrix. To achieve this, dimension pruning layers are inserted between two layers containing weight matrices as shown in figure 1(b). The pruning layer is a binary mask with dimension equal to the input dimension of next layer (or output dimension of previous layer) and is multiplied with the output of previous layer. Therefore, if $i^{th}$ entry in the mask is 0, then it is equivalent to dropping $i^{th}$ column from the next layer and $i^{th}$ row from the previous layer. If the previous layer happens to be a self-attention layer, then $i^{th}$ row is dropped from each of the key, query and value matrices.

Since binary mask is non-differentiable, the above method relaxes mask to be a continuous value during training. More precisely, all mask values in a dimension pruning layer is initialized to 1 indicating that initially no column or row is dropped at the beginning. The following loss is optimized to learn the weights and mask

$\displaystyle L_{CE}+\lambda\sum_{l=1}^{L}|m_{l}|_{L1}$

where $L_{CE}$ is the regular cross entropy loss for classification task and $|m_{l}|_{L1}$ is the L1 norm of the mask or dimenison pruning layer in $l^{th}$ transformer block. Intuitively, the L1 loss promotes sparsity making some of the mask values close to 0, which results in lower contribution of corresponding dimensions in the overall activation. Therefore, we can drop the entire column to obtain less parameters. During inference, we again discretize the mask as follows

$\displaystyle m_{l}=m_{l}>\tau$

where $\tau$ is the threshold determined from desired pruning rate.

Figure 2(a) shows the effect of applying VTP to vision transformer. We observe that mainly matrices from feedforward modules (top two linear layers in a transformer block) are pruned whereas attention weight matrices and projection layers are majorly unpruned. We discuss the reasons in more detail in section 6. However, this motivates us to compress self-attention block (query, key, value matrices in scaled dot-product layer and the following linear layer) in combination with above pruning approach for extra compression. We apply low rank approximation (LRA) for this purpose.

One issue with low rank approximation is that new rank $k$ must be much smaller than original rank $d$ since reducing the rank from $d$ to $k=d/2$ does not lead to lesser parameters if the original matrix has shape $d\times d$. We partly avoid this problem by applying LRA only to the attention block, where we share one of the low rank matrices among key, query and value matrices. It can be implemented by projecting $d-$dimensional input of a transformer block to a lower dimensional space of size $k$ by inserting a linear layer or projection matrix as shown by blue layer in figure 2(b)(b). This changes input dimension of query, key and value matrices from $d$ to $k$. We can now compute attention weights in either $d-$dimensional space or a lower $k-$dimensional space. We refer to the former as LRA($d\times k$) and the later as LRA($k\times k$). LRA($d\times k$) does not affect the size of linear layer after self-attention layer. However, LRA($d\times k$) reduces the input dimension of this layer also to $k$.

We apply LRA only to the attention block and not the feedforward block. This is because by design, transformer feedforward block is a factorized matrix where hidden dimension in between them is kept higher. We experiment with only LRA applied to vision transformer. We also apply VTP discussed earlier in the feedforward block (as shown in dashed green layer in figure 2(b)(b)) with LRA in attention block. This gives our hybrid compression framework.

We conduct hyper-parameter tuning for vision transformer for which results are shown in table 1. We experiment with patch size $P$ of 2, 4 and 8. Results are shown in the first 3 rows. Patch size of 8 performs the worst and patch size of 2 performs the best. However, we choose patch size of 4 because of better training efficiency since transformer computational complexity grows quadratically with number of input tokens.

Then we varied the attention dimension and hidden dimension of feedforward blocks. Attention dimension of 512 and feedforward hidden dimension of 1024 worked the best as given in $6^{th}$ row. Finally, we changed the number of transformer blocks. We did not find significant accuracy difference between 6 and 9 layers. Therefore, we chose 6 transformer blocks in our network for better computation time. All the models compute self-attention with 8 heads (512/8 = 64 dimension for each head). Our final choice of hyper-parameter is shown in the second last row of table 1. It achieves 72.61% accuracy with 13M parameters on CIFAR-10 test set.

For pruning using VTP, we again train vision transformer with the final choice of hyperparameters. We insert the dimension pruning layers in each transformer block as shown in figure 1(a)(b). We have only one hyper-parameter $\lambda$, which measures the strength of the sparsity or L1 loss. We experimented with values of $10^{-2},10^{-3}\text{ and }10^{-4}$. We found that higher values of $\lambda$ significantly degraded the model performance but $10^{-4}$ worked well with a very small drop in model accuracy. This value is also the same as that in the original work although for ImageNet experiments. We show the results for VTP with different pruning rate in table 2. Prune rate of $\alpha$ means we want to zero out $\alpha$ fraction of all the mask values combined across all the pruning layers inserted in our model. Note that we have to re-train only one model with pruning layers as pruning rate is specified only at the inference time.

We can see from the second last row that we can compress the model by 36% with around 12% increase in error. However, the error goes up significantly if we compress the model further. We make two observations in this table. First, prune rate of 0.1 has higher error than prune rate of 0.3 which is counter-intuitive. We discuss this in detail in section 6. Second, pruning rate is not exactly the same as fractions of weights pruned because pruning rate is specified for fraction of dimensions to be pruned and not weights. We also observed that prune rate greater than 0.6 gave error perhaps because all dimensions of some weight matrix was getting dropped. This is indeed possible with VTP whereas low rank approximation in our method systematically compresses all attention blocks.

We chose $k=d/4$ for low rank compression. The results for LRA are shown in table 3. We note that we train a new model for all the LRA experiments (with and without pruning). We observe that only applying LRA is better than VTP for similar compression. For example, LRA $(d\times k)$ compresses the model by 24% and has 5.8% increase in error whereas pruning has 10% increase in error for similar compression. Similarly, LRA $(k\times k)$ has 39% compression with 8% increase in error whereas pruning has 11.5% increase in error. When we use pruning in feedforward block and LRA in attention block (hybrid compression), we achieve 50% compression with 14% increase in error whereas pruning by itself has 20.5% increase in error for only 44% compression, which is about 2% absolute accuracy gap.