Attention Map Guided Transformer Pruning for Edge Device

The aim of person re-identification is to retrieve interested persons across camera views. In recent years, re-ID has made great progress [5, 3, 6], and applied in various scenarios, such as video surveillance, autonomous driving, and activity analysis [12, 13]. After observing the device applied to crowded public places, Zhuo et al. [24] propose the occluded person Re-ID topic. Different from holistic person Re-ID, the query set is constructed with occluded pedestrian images, while the training set and gallery set are still holistic. Recent two mainstream methods deal with occlusion problems with the assistance of pose estimation [25, 26] and human parsing [27, 28]. With the popularity of ViT, some works also introduce transformer into occluded person Re-ID models [29, 8, 7, 30]. Li et al. [8] design a transformer encoder-decoder architecture for diverse part discovery to solve the occluded problem. TransReid [30] takes full advantage of ViT’s powerful feature extraction capability by aggregating globally enhanced local information. Considering the simplicity of structure and outstanding performance, we choose TransReid as the baseline model for our pruning method.

Recent ViT pruning work can be divided into two categories: reducing number of tokens and removing parameters. Tang et al. [17] develop a top-down patch slimming approach that removes redundant patches based on the reconstruction error of the pre-trained model. Xu et al. [20] propose the structure preserving token selection and slow-fast updating strategies to fully utilize the complete spatial structure. These methods can effectively save the computational cost, but they cannot reduce the model size. Zhu et al. [31] introduce control coefficients and remove neurons with smaller ones to reduce the number of embedding dimensions, hence the model parameters can be compressed. However, the unstructured pruning and zero masking operation are unfriendly for hardware efficiency which can provide limited acceleration and compression on edge devices. Our structured head pruning with dimension align operation is designed for friendly edge device deployment. Our work aims to prune ViTs based on attention maps, which can afford guidance for retaining discriminative tokens for occluded Re-ID. Besides, we propose the entropy of the attention map as the prune criterion for head pruning rather than the commonly used Taylor criterion.

The ViT model [1] splits input image to $\mathrm{N}$ patches and appends a class token, then feeds them to encoder transformer blocks which is kept same as the vanilla transformer [32]. The multi-head self-attention (MSA) and multi-layer perceptron (MLP) are the main components which cost most of the computation in the transformer block. We denote $\mathrm{Z}_{l-1}\in\mathbb{R}^{\mathrm{N\times d}}$ as the input of the $l$-th layer, then the attention calculation for head $\mathrm{h}$ can be formulated as:

where $\mathrm{d}$ is the head embedding dimension. $\mathrm{Q}_{l}^{\mathrm{h}}=\mathrm{W}_{l}^{\mathrm{q,h}}\mathrm{Z}_{l-1}$, $\mathrm{K}_{l}^{\mathrm{{h}}}=\mathrm{W}_{l}^{\mathrm{k,h}}\mathrm{Z}_{l-1}$, and $\mathrm{V}_{l}^{\mathrm{h}}=\mathrm{W}_{l}^{\mathrm{v,h}}\mathrm{Z}_{l-1}$ are the query, key, and value of the $\mathrm{h}$-th head in the $l$-th layer, respectively. Thus, the formula for the MSA module can be written as:

$\mathrm{H}$ is the number of heads. For $\mathrm{N}$ patches which are sent to the MSA module, we can calculate the complexity and parameters with Eq. (1) and Eq. (2):

and

The computation complexity is composed of three parts. The projection of input features $\mathrm{Z_{in}}$ to $\mathrm{Q,K,V}$, the projection of attention features $\mathrm{Attn}$ to output features $Z_{out}$, and the attention computation $(\mathrm{Q\times K^{T}})\times\mathrm{V}$. The computational cost of the three parts are $3\mathrm{NDHd}$, $\mathrm{NDHd}$, and $\mathrm{2N^{2}Hd}$, respectively. $\mathrm{D}$ is the embedding dimension, and $\mathrm{D=H\cdot d}$ is true when the model is not pruned. We leverage this form to present the change in our pruning directly. We factor the formula to get the final expression. The computation of softmax and division are neglected. Parameters are the sum of four projection matrices $\mathrm{W^{q},W^{k},W^{v},W^{o}}$.

For our pruning methods, the number of attention heads and key-value tokens is pruned to ${\mathrm{H}}^{\prime}$ and ${\mathrm{N}}^{\prime}$, respectively. Then the Eq. (3) and Eq.(4) can be rewritten as:

and

Eq. (5) shows the computational cost of our pruned model. The formula indicates that with $\eta\%$ heads or tokens pruned, $\eta\%$ computation will be reduced in the MSA module. If both are pruned (heads and tokens), $\eta^{2}\%$ of the complexity of the MSA module will be reduced. From Eq. (6), the parameters of the MSA module are proportional to head numbers, while pruning tokens can’t reduce the parameters.

Similar to the channel pruning of the convolution layer, head pruning is a structural parameter pruning mechanism. Thus, we can employ traditional channel pruning methods to prune the attention heads. Our initial idea is to leverage classical methods like L1 Norm [33] and Taylor Pruning [14, 15], but don’t succeed in achieving better results. We consider the reason for our failure is that we haven’t taken into consideration of the transformer’s special characteristics.

The computation of the attention map is the core stage of the transformer block. To this end, we visualize the changes of the attention maps during the training from scratch process of the ViT model. The visualization results are present in Fig. 3. All the attention maps showed are the mean result of one epoch training. We accumulate the attention maps and calculate the average for the common information learned by attention maps. By observing Fig. 3, we can notice that in the initial stage of training, the model hasn’t learned any useful information, so the attention maps are disordered and have a large entropy (e.g., epoch 1 in Fig. 3). After a certain training epoch, the model learns some basic information and begins to present a certain pattern. However, attention to important information hasn’t been learned, which means uniform attention and the attention maps still have large uncertainty. Therefore, the entropy is only reduced a little (e.g., epoch 50 in Fig. 3). When the model finally converges, the entropy of all layers has reduced, which indicates the model has learned the information of the dataset. Important tokens are paid great attention, while unimportant tokens are paid fairly low attention. The shallow layers learn local attention and only focus on a few tokens. Hence, the entropy is significantly reduced (e.g., layer 0 and layer 2 at epoch 300 in Fig. 3). In contrast, the deep layers tend to learn global attention, which needs to pay attention to more tokens. As a result, deep layers have higher entropy (e.g., layer 4 and layer 6 at epoch 300 in Fig. 3).

Inspired by this, we can have the following conclusions. With the increase of useful information learned by the attention head, the model will pay more attention to some tokens, which makes the attention map more deterministic. On the contrary, when the attention head learns less information, it will have a unified focus on the overall situation, which results in great uncertainty. Information entropy can well measure the uncertainty of the attention map. In other words, the attention map with a large entropy value indicates that this head can only provide limited useful information and make little contribution to the overall situation, so it can be pruned.

For attention map $\mathrm{A}^{\mathrm{h},l}$ of layer $l$ and head $\mathrm{h}$, $\mathrm{A}_{i,j}^{\mathrm{h},l}$ denotes the $i$-th row and $j$-th column of attention map. In Eq. (1), a softmax operation is conducted along the key dimension after attention computation. It can be considered the similarity of $j$-th key token to the $i$-th query token in the original paper [32]. We can also treat $\mathrm{A}_{i}^{\mathrm{h},l}$ as a probability distribution for key tokens to the $i$-th query token mathematically. Hence the information entropy of attention map $\mathrm{A}^{\mathrm{h},l}$ can be obtained by:

According to the principle of maximum entropy [34, 35, 36], a uniform distribution of probability has maximum entropy. Therefore, we can prune the attention heads with a big entropy.

The similarity calculation between query and key tokens is the main memory usage in a transformer block as it is proportional to the power of token length. To be specific, each key token needs to calculate similarity with all query tokens to obtain the attention map. However, not all tokens have a powerful contribution to the learning of global attention. In fact, a sparse token sequence may even lead to a better performance [37] as images contain large regions of low-level texture and uninformative background. In the occlusion case, tokens with background and occlusion can be removed without impairing the model accuracy.

By observing Fig. 3, we can notice that the ViT model has learned a certain pattern on the attention map at epoch 300. The blue column of the visualized attention map means a small similarity between the corresponding key token to global query tokens. Small global similarity tokens typically make fewer contributions to the whole model and are less discriminative, thus we can prune these tokens which correspond to the blue columns. However, the simple summation of the similarity may neglect the global effect of the loss. Therefore, we propose to calculate the gradients of the attention map and perform element-wise multiplication for similarities and gradients.

Our token pruning strategy is inspired by Taylor Pruning [14, 15]. The difference lies in that we compute the importance of the corresponding elements of key tokens on the attention map instead of simply calculating the importance of key tokens. For Taylor Pruning, the importance $\mathcal{I}$ of the $i$-th key token $k_{i}$ can be formulated as:

where $\frac{\delta\mathcal{L}}{\delta k_{i,\mathrm{d}}}$ is the gradient of loss function $\mathcal{L}$ with respect to the $\mathrm{d}$-th dimension of $k_{i}$. We compute the sum of the importance of key tokens corresponding to the column features on the attention map to replace the importance of key tokens. So our importance $\mathcal{I}$ of the $i$-th key token $k_{i}$ can be obtained by rewritten Eq. (8) as:

$\mathrm{A}_{i,j}^{\mathrm{h}}$ denotes the $i$-th row and $j$-th column of the $\mathrm{h}$-th attention map. It’s noticeable that we ignore the first token in calculating importance as the first token is the class token which shouldn’t be pruned in the classification task. In addition, different from head pruning that we can remove the parameters to change numbers of head, token pruning has no relation to model structure. As shown in Fig. 2, we add a index layer for pruning the key and value tokens.

Instead of the handcraft structure or the searched result in a large searching space, we employ the global pruning scheme to autonomously form the final structure of the pruned model. Compared with hand designing and heuristic researching, global pruning needs less prior experience and searching time. In addition, the pruning message can be fully utilized for finding the most suitable structure for the current task. Our global pruning scheme is simple to imply. First, we calculate the importance score of all the heads/tokens in the model. Then, we arrange them in descending order, removing the heads/tokens with the lowest score according to the pruning rate. When pruning the model in the global range, the difference between layers should be considered. As proposed in [17], deeper layers are less important for the redundant global attention. For better-allocated layer structures, we introduce a weight coefficient $\lambda$ for the calculation of importance score $\mathcal{I}$ in layer $l$. The formula is as follows:

We also add a soft distillation loss for better accuracy recovery when fine-tuning. The final loss function can be formulated as:

where $p$ is the output softmax probability of pruned model, $q$ is the corresponding value of original model, and $y$ is the true label. $\mathcal{L}_{CE}$ and $\mathcal{L}_{KL}$ denote the cross-entropy loss and the KL-divergence loss, respectively. $\alpha$ balances the weight of the teacher’s knowledge.