# Human Not in the Loop: Objective Sample Difficulty Measures for Curriculum
Learning

###### Abstract

Curriculum learning is a learning method that trains models in a meaningful
order from easier to harder samples. A key here is to devise automatic and
objective difficulty measures of samples. In the medical domain, previous work
applied domain knowledge from human experts to qualitatively assess
classification difficulty of medical images to guide curriculum learning,
which requires extra annotation efforts, relies on subjective human
experience, and may introduce bias. In this work, we propose a new automated
curriculum learning technique using the variance of gradients (VoG) to compute
an objective difficulty measure of samples and evaluated its effects on elbow
fracture classification from X-ray images. Specifically, we used VoG as a
metric to rank each sample in terms of the classification difficulty, where
high VoG scores indicate more difficult cases for classification, to guide the
curriculum training process We compared the proposed technique to a baseline
(without curriculum learning), a previous method that used human annotations
on classification difficulty, and anti-curriculum learning. Our experiment
results showed comparable and higher performance for the binary and multi-
class bone fracture classification tasks.

Index Terms—  Curriculum learning, Variance of Gradient, Elbow fracture,
Medical imaging, Classification

##  1 Introduction

Curriculum learning (CL) has been proposed as a training strategy in deep
learning to expedite learning and improve the model performance. In this
technique, samples are presented in the order of from “easy” to “hard” to the
model, instead of randomly selecting samples [1]. This is motivated by the
learning process of humans and animals, who start learning with an easy
‘concept’ before gradually increasing the difficulties of the task. It has
been shown that starting training with “easy” samples enables machine learning
models to converge faster and shows a significant improvement in the
generalization of the models [1]. Due to those benefits, CL has been
implemented in various machine learning tasks such as image classification
[2], object detection [3], multi-task learning [4], etc. The existing CL
methods can be divided into predefined CL or automated CL, in terms of
assigning the difficulty by humans or algorithms [5].

Previous studies in CL for medical image analysis are mostly based on
predefined CL. Amelia et al. proposed strategies to assign a degree of
difficulty based on human knowledge such as medical decision trees and
inconsistencies between the annotations of multiple clinical experts [6]. Jun
et al. proposed a strategy based on clinically known knowledge from human
experts who qualitatively assigned the difficulty level of diagnosis for each
elbow fracture subtype [7]. Although both strategies improved the model
performance by using the medical knowledge guided curriculum learning, they
need extra efforts from human experts to annotate the difficulty level, which
is subjective, qualitative, less reproducible, and may introduce bias.

Table 1: Distribution of the normal and different elbow fracture subtypes in our dataset Type | normal |  | Ulnar
---
Fracture
| Radial
---
Fracture
| Humeral
---
Fracture
| Dislocation
---
Fracture
| Complex
---
Fracture
| Coronoid
---
Fracture
| Total of
---
fracture
Total
Train | 800 | 88 | 340 | 84 | 11 | 42 | 27 | 592 | 1,392
Test | 400 | 10 | 44 | 9 | 2 | 4 | 4 | 73 | 473
Total | 1,200 | 98 | 384 | 93 | 13 | 46 | 31 | 665 | 1,865

In this work, we developed a novel and automated CL strategy using variance of
gradients (VoG) as a scoring criterion for classification difficulty. Here,
samples with higher VoG scores indicate more difficult to classify, since
images with a high VoG score are more likely to have clutter and complex
patterns, making it more difficult for a deep learning model to classify the
images [8]. We implemented a learning process using VoG ranking through
permutations of the training epochs by sampling without replacement at the
beginning, and then an algorithm was used to update the probability, until it
reached a random shuffle status. We performed experiments on an elbow X-ray
dataset of 1865 patients, by conducting binary classification (normal vs.
fracture) and multi-class classification (normal vs. $3$ individual fracture
subtypes). We compared the proposed technique to a baseline (without
curriculum learning), a previous method that used human annotations on
classification difficulty, and anti-curriculum learning. Results show that our
method outperforms the compared methods and because of its automatic nature,
it is objective, reproducible, and independent of human input. Our
contributions can be summarized as follows:

  * •

We adapted the VoG method from computer vision of natural image classification
to the curriculum learning context and showed it is effective for medical
imaging applications.

  * •

We showed in elbow fracture classification tasks that the VoG-based curriculum
learning can either achieve comparable or higher performance than several
compared methods, and our method is more efficient because it is an automated
process and independent of human input.

  * •

We showed that the VoG-based objective scores and the human radiologists’
subjective scores measure different aspects of the difficulty of the
classification samples, while they both can guide curriculum learning.

The rest of the paper is organized as follows. Section 2 describes the methods
of our study. Section 3 presents experiment designs and datasets. Results are
shown in Section 4. Section 5 discusses conclusions and future work.

##  2 Methods

To design a curriculum paradigm in deep curriculum learning, two major
components need to be determined. First, how to measure the training data and
assign or derive a metric of sample difficulty. Second, how to present the
measured samples in the training procedure [5]. In the following, we describe
the difficulty measurer and training scheduler to address the two components.

###  2.1 Difficulty Measurer

The input image $X$ can be divided into individual pixels, referred to as
$x_{i}$, with $i$ representing the index ranging from $1$ to $N$, where $N$ is
the total number of pixels in the image. For the given image, the gradient of
the activation $A$ for each pixel is computed where $l$ denotes as the pre-
softmax layer of the network and $p$ is the index of either the true or
predicted class probability. Let $S$ be a matrix that represents the gradient
of $A^{l}_{p}$ concerning individual pixels $x_{i}$, and $S$ is calculated by
Equation (1). The gradient matrix $S$ will have the same size as an image [8].

$S=\frac{\partial A^{l}_{p}}{\partial x_{i}}.$ (1)

In order to better learn features during training [9], we take the average
over the color channels and the generated matrix $T$ records the gradient of
each pixel $x_{i}$. In the next step, it shows how the variance of gradients
are computed. We generate the above gradient matrix $T$ for all individual
checkpoints in a given set of $K$ checkpoints to get gradients at different
steps, i.e., ${T_{1},...,T_{K}}$. We then calculate the mean gradient $\mu$ by
by taking the average of the $K$ gradient matrices. Note, $\mu$ is the mean
across different checkpoints and is of the same size as the gradient matrix
$T$. Then the variance of gradients across each pixel is calculated as
following [8]:

$\mu=\frac{1}{K}\sum_{i=1}^{K}T_{i}.$ (2)
$VoG_{p}=\sqrt{\frac{1}{K}}\sum_{i=1}^{K}(T_{i}-\mu)^{2}.$ (3)

Once we have variance of gradients for each pixel, we then average the pixel-
wise variance of gradients to compute the VoG score for the given input images
and obtain the ranking of training samples based on VoG score.

Table 2: Model performance on binary (normal vs. fracture) classification.
RS-CL: using radiologist-assigned difficulty score. VoG-CL: using VoG score.  | Accuracy | Recall | AUC | F1 Score
---|---|---|---|---
| Mean | Std | Mean | Std | Mean | Std | Mean | Std
Baseline | 0.7497 | 0.0164 | 0.7517 | 0.0084 | 0.8407 | 0.0095 | 0.7208 | 0.0052
RS-CL | 0.7850 | 0.0106 | 0.7877 | 0.0076 | 0.8642 | 0.0109 | 0.7595 | 0.0103
VoG-CL | 0.7858 | 0.0041 | 0.7885 | 0.0060 | 0.8652 | 0.0099 | 0.7605 | 0.0087
Anti-VoG-CL | 0.7610 | 0.0201 | 0.7570 | 0.0116 | 0.8356 | 0.0148 | 0.7206 | 0.0148
Table 3: Model performance on multi-class classification. | Accuracy | Balanced Accuracy | AUC
---|---|---|---
| Mean | Std | Mean | Std | Mean | Std
Baseline | 0.7175 | 0.0277 | 0.5841 | 0.0213 | 0.8369 | 0.0097
RS-CL | 0.7322 | 0.0200 | 0.5613 | 0.0364 | 0.8496 | 0.0098
VoG-CL | 0.7337 | 0.0170 | 0.6294 | 0.0625 | 0.8570 | 0.0107
Anti-VoG-CL | 0.7227 | 0.0279 | 0.5775 | 0.0374 | 0.8417 | 0.0095
| Binary Accuracy | Binary AUC | F1 Score
| Mean | Std | Mean | Std | Mean | Std
Baseline | 0.7678 | 0.0368 | 0.8325 | 0.0154 | 0.7139 | 0.0269
RS-CL | 0.7715 | 0.0152 | 0.8386 | 0.0156 | 0.7246 | 0.0167
VoG-CL | 0.7853 | 0.0265 | 0.8527 | 0.0104 | 0.7358 | 0.0203
Anti-VoG-CL | 0.7620 | 0.0199 | 0.8377 | 0.0132 | 0.7174 | 0.0211

###  2.2 Training Schedular

Once the difficulty score is computed, a training schedular is then applied
during the training based on the VoG difficulty score. Our curriculum learning
method is equipped with a sampling-without-replacement strategy to permute and
reorder the training data at the beginning of each epoch. Each training sample
$s_{i}$ has a probability, $p_{i}$, and its value is updated at the beginning
of each training epoch in order to implement curriculum learning. Each sample
is first assigned with the initial value $p_{i,(1)}$ at the first epoch at
initial and then updated based on the updating algorithm. Finally, after
certain epochs, the data are randomly shuffled and each sample shares the same
probability.

For the initialization of $p_{i}$, the $p_{i}$ of all images is calculated
based on their VoG ranking scores. These scores reflect the level of
difficulty in learning the features of the samples during the training
process. To distinguish the importance of ‘easy’ and ‘hard’ samples, easier
samples with lower VoG scores are assigned with a higher probability according
to according to Equation (4):

$p_{i,(1)}=\frac{s_{i}}{\sum_{j=1}^{N}s_{j}},$ (4)

where $s_{k}$ represents the rank of the sample $x_{k}$ in the sorted VoG
scores of all training samples from highest to lowest, and $N$ denotes the
number of images in the training set.

After the first epoch, the sampling probability of all the training samples is
updated and we denote the updated sample probability for each $x_{i}$ at epoch
$e$ as $p_{i,(e)}$, and $\lambda_{i}$ (Equation (5)) is a scalar for updating
sampling probability of $xi$. $p_{i,(e)}$ can be calculated by the sampling
probability of the previous epoch $p_{i,(e-1)}$ multiple $\lambda_{i}$. This
ensures $p_{i}$ for every image to exponentially and smoothly transit to
$p_{i,(final)}=1/N$. $L$ is the number of epochs after which $p_{i}$ equals to
$p_{i,(final)}$; for the rest of the training process and after $L$ epoch, the
sampling probability of each sample will not be updated. $L$ is treated as a
hyperparameter for determining when the training reaches a stage of randomly
shuffling the dataset. This way, the model can gradually reduce the emphases
we initially given on the easier cases and then reach a stage of random
selection of samples in the entire training set after epoch $L$ [7].

$\lambda_{i}=\sqrt[L]{\frac{1/N}{p_{i,(1)}}}.$ (5)

##  3 Experiment

(a) Binary classification

(b) Multi-classification

Fig. 1: Comparison between VoG-based and radiologist-assigned difficult scores
at the class level for each fracture subtype.

Dataset. The dataset used for evaluation of our proposed method contained a
total of 1865 different X-ray elbow images retrospectively collected at our
institution. There are 665 fracture-positive and 1200 non-fracture (normal)
cases. All fracture-positive images were categorized by an experienced
radiologist into six different fracture subtypes, that is, (a) Ulnar fracture,
(b) Radial fracture, (c) Humeral fracture, (d) Dislocation, (e) Complex
fracture and (f) Coronoid process fracture.

We defined two different experimental tasks in this study, including binary
and multi-class classification. In the binary classification task, we merged
all six fracture subtypes in one class as a fracture-positive cohort versus
non-fracture (normal) cohort. We used a stratified technique based on
different fracture subtypes to split the entire dataset into training and test
sets. Table 1 displays the data distribution in the training and test set for
each elbow fracture subtype. During the test process, we randomly split the
400 normal cases into 4 subsets, each containing one hundred images. We then
evaluated our methods on an unseen subset of 100 normal cases and all 73
fracture cases to form a balanced test. We repeated the test 4 times, each
time selecting a different set of normal cases, and the average performance
results were reported.

In the multi-class classification task, images with Dislocation, Complex, and
Coronoid fracture were excluded due to their very small sample sizes. Similar
to the binary classification task, the test was repeated 4 times, and average
results were reported.

Metrics. In the binary classification, accuracy, recall, Area Under the ROC
Curve (AUC), and F1 score were used to evaluate models’ performance. In the
multi-task classification task, we calculated accuracy, balanced accuracy,
AUC, binary task accuracy, binary task AUC and F1 score. Balanced accuracy was
calculated by averaging the ratios between the number of true positives and
that of total samples in each class, which reduces the effect caused by class
imbalance. Binary task accuracy and binary task AUC were computed by
transforming the prediction to ’normal’ and ’fracture’. If the ground truth
and the predicted labels are ’normal’, they are set to 0, and set to 1
otherwise.

Implementation Details. We employed the VGG16 convolutional neural network
architecture [10] to implement the classification models. We compared our
proposed training strategy to three methods on both the binary and multi-class
classification tasks: 1) Baseline model: it is based on randomly shuffling the
training data for each epoch. 2) Radiologists Score (RS)-CL: This is a
clinical knowledge-guided curriculum learning, which uses the difficulty score
assigned by an experienced radiologist to each fracture subtype. We used the
scores given in a prevous work [7]. 3) Anti-curriculum learning: This method
reverses the ranking scores computed by VoG, and starts the training from the
“difficult” samples. The model architecture and curriculums were implemented
by PyTorch and ran on an NVIDIA TESLA V100 GPU from a local supercomputing
center. All the hyperparameter tuning and data augmentation were the same for
all experiments. The experiment results are summarized using the mean and
standard deviation across the 5 runs. All models were run for 30 epochs using
the cross-entropy loss. VoG was computed using checkpoints from the 26th,
28th, and 30th epochs, which were experimentally determined.

In addition, in order to gain insights on the VoG scores, we also compared the
VoG difficult scores derived automatically from the samples to the difficulty
scores given by the radiologist at the class level. Based on the VoG scores
during training, the ranking of each training sample can be calculated at the
instance level. To calculate the class-level difficulty score based on the
instance level VoG scores, each training sample is assigned a difficult score
$D_{i}$, as follows:

$D_{i}=\frac{(N-r_{i})}{N}*100,$ (6)

where $r_{i}$ is the ranking of each sample and $N$ represents the number of
the training samples. Then we computed the average of the difficulty scores
across the training sample in each class to generate a class-level VoG
difficult score, for each of the fracture subtype. We plot histograms to
compare the distributions of VoG scores and radiologist-assigned scores to see
how they may be related at the class level.

##  4 Results

The main results are shown in Tables 2 and 3. As can be seen, in the binary
classification task, VoG-based CL achieved comparable performance to the RS-CL
method (Table 2), with a mean accuracy of 0.786 vs. 0.785 (RS-CL) and an AUC
of 0.865 vs. 0.864 (RS-CL). Both the VoG-CL and RS-CL outperformed the
baseline method and anti-curriculum method, as expected. Note that even though
the performance is very close between VoG-CL and RS-CL, the proposed VoG
method is superior as it is an automated curriculum and independent of human
input. As shown in Table 3, our proposed VoG-based CL outperformed all the
compared methods (including RS-CL) in the multi-class classification task,
showing a mean balanced accuracy of 0.629 and AUC of 0.857. This experiment
shows advantages of VoG-CL over RS-CL in a more complex classification task.
In addition, the anti-curriculum settings showed inferior results compared to
the curriculum settings, as expected in both classification tasks.

Fig. 1 shows the comparisons between VoG-based class difficulty score and
radiologist-assigned scores. In the binary classification task, class “normal”
had the lowest difficulty score while class Complex fracture had the highest
score, which was in line with the radiologist’s difficulty scores. For the
multi-class classification task, no obvious patterns were found. One
consideration here is the smaller sample sizes for class Ulnar fracture and
Humeral fracture, which may have an influence to the VoG scores. Overall, from
what we see in Fig. 1, there are no obvious matching between the two different
types of scores, indicating they measure the samples in different aspects,
while they both can guide curriculum learning, and achieved similar
performance in the binary classification task. Further experiments are needed
to evaluate this finding.

##  5 Conclusion

In this work, we proposed the VoG-based method to automatically measure sample
difficulty for curriculum learning. In the elbow fracture subtype
classification tasks, our method showed comparable or outperforming effects.
Our method can save the efforts from human annotation on the difficulty score,
prevent the subjectivity from human input and thus potential bias. While we
focused on one specific application of fracture classification, our method can
be easily extended to other diseases or medical imaging classification
problems. One limitation of our study is the small sample sizes for several
subtype fractures and the data are not well balanced. In future work, we will
perform more evaluation using larger and other datasets as well as different
imaging classification tasks.

##  6 Compliance with Ethical Standards

This retrospective single-center study was approved by Institutional Review
Board (IRB) at our institution and was conducted compliant to Health Insurance
Portability and Accountability Act (HIPAA). All medical data were de-
identified before used for machine learning modeling.

Acknowledgements. This project was supported in part by a National Institutes
of Health (NIH) grant (1R01CA218405), the grant 1R01EB032896 as part of the
National Science Foundation (NSF)/NIH Smart Health and Biomedical Research in
the Era of Artificial Intelligence and Advanced Data Science Program, a NSF
grant (CICI: SIVD: #2115082), and the University of Pittsburgh Momentum Funds
(a scaling grant) for the Pittsburgh Center for AI Innovation in Medical
Imaging. This work used the computational support by the Extreme Science and
Engineering Discovery Environment (XSEDE), which is supported by NSF grant
number ACI-1548562. Specifically, it uses the Bridges system, which is
supported by NSF award number ACI-1445606, at the Pittsburgh Supercomputing
Center.

## References

  * [1] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston,  “Curriculum learning,”  in Proceedings of the 26th annual international conference on machine learning, 2009, pp. 41–48.
  * [2] William Lotter, Greg Sorensen, and David Cox,  “A multi-scale cnn and curriculum learning strategy for mammogram classification,”  in Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support, pp. 169–177. Springer, 2017.
  * [3] Dingwen Zhang, Junwei Han, Long Zhao, and Deyu Meng,  “Leveraging prior-knowledge for weakly supervised object detection under a collaborative self-paced curriculum learning framework,”  International Journal of Computer Vision, vol. 127, no. 4, pp. 363–380, 2019.
  * [4] Nikolaos Sarafianos, Theodoros Giannakopoulos, Christophoros Nikou, and Ioannis A Kakadiaris,  “Curriculum learning of visual attribute clusters for multi-task classification,”  Pattern Recognition, vol. 80, pp. 94–108, 2018.
  * [5] Xin Wang, Yudong Chen, and Wenwu Zhu,  “A survey on curriculum learning,”  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021\.
  * [6] Amelia Jiménez-Sánchez, Diana Mateus, Sonja Kirchhoff, Chlodwig Kirchhoff, Peter Biberthaler, Nassir Navab, Miguel A González Ballester, and Gemma Piella,  “Medical-based deep curriculum learning for improved fracture classification,”  in International Conference on Medical Image Computing and Computer-Assisted Intervention. Springer, 2019, pp. 694–702.
  * [7] Jun Luo, Gene Kitamura, Emine Doganay, Dooman Arefan, and Shandong Wu,  “Medical knowledge-guided deep curriculum learning for elbow fracture diagnosis from x-ray images,”  in Medical Imaging 2021: Computer-Aided Diagnosis. SPIE, 2021, vol. 11597, pp. 247–252.
  * [8] Chirag Agarwal, Daniel D’souza, and Sara Hooker,  “Estimating example difficulty using variance of gradients,”  in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022, pp. 10368–10378.
  * [9] Avanti Shrikumar, Peyton Greenside, and Anshul Kundaje,  “Learning important features through propagating activation differences,”  in International conference on machine learning. PMLR, 2017, pp. 3145–3153.
  * [10] Karen Simonyan and Andrew Zisserman,  “Very deep convolutional networks for large-scale image recognition,”  arXiv preprint arXiv:1409.1556, 2014.


# Automatic Channel Pruning for Multi-Head Attention

Eunho Lee, Youngbae Hwang
Department of Intelligent Systems and Robotics
Chungbuk National University
Cheongju-si, 28644, Republic of Korea
{ehlee<EMAIL_ADDRESS>
Corresponding author

###### Abstract

Despite the strong performance of Transformers, their quadratic computation
complexity presents challenges in applying them to vision tasks. Automatic
pruning is one of effective methods for reducing computation complexity
without heuristic approaches. However, directly applying it to multi-head
attention is not straightforward due to channel misalignment. In this paper,
we propose an automatic channel pruning method to take into account the multi-
head attention mechanism. First, we incorporate channel similarity-based
weights into the pruning indicator to preserve more informative channels in
each head. Then, we adjust pruning indicator to enforce removal of channels in
equal proportions across all heads, preventing the channel misalignment. We
also add a reweight module to compensate for information loss resulting from
channel removal, and an effective initialization step for pruning indicator
based on difference of attention between original structure and each channel.
Our proposed method can be used to not only original attention, but also
linear attention, which is more efficient as linear complexity with respect to
the number of tokens. On ImageNet-1K, applying our pruning method to the
FLattenTransformer, which includes both attention mechanisms, shows
outperformed accuracy for several MACs compared with previous state-of-the-art
efficient models and pruned methods. Code will be available soon.

##  1 Introduction

Transformer has achieved remarkable success in various computer vision tasks
based on attention mechanisms that effectively capture long-range
dependencies. The attention module generates an attention map by utilizing the
query $Q\in\mathbb{R}^{N\times C}$ and key $K\in\mathbb{R}^{N\times C}$ to
extract the relationships between tokens, and then projects the value
$V\in\mathbb{R}^{N\times C}$ to obtain a feature map with global information.
It requires a computational cost of $\Omega(N^{2}C)$, which is quadratic with
respect to the number of tokens $N$. Despite providing excellent performance,
the quadratic complexity with respect to $N$ poses significant challenges for
deployment on mobile and edge devices.

Recent research attempts to mitigate this issue by designing efficient
transformers. Some approaches propose network architectures that limit the
number of tokens [36; 25; 14]. By reducing the number of tokens, they cut down
computation cost while trying to maintain the performance. However, this
results in limitations of the receptive field, constraining the capture of
global dependencies. Another approach proposes a new attention mechanism that
can replace original attention [18; 30; 12]. Linear attention is a method that
maintains the ability to capture long-range dependencies while reducing
computational complexity linearly. They first approximate Softmax function by
replacing it with a simple activation or a tailored function. By changing the
computation order from $(Q\cdot K^{T})\cdot V$111Original attention mechanism
to $Q\cdot(K^{T}\cdot V)$222Linear attention mechanism, they obtain an
attention mechanism with computational complexity of $\Omega(NC^{2})$, which
is quadratic with respect to channels instead of tokens.

Another strategy to improve the efficiency of the network is network pruning.
It aims to achieve a lightweight network by removing redundant parts from the
existing network while minimizing performance degradation. Human knowledge-
based analysis is conducted on each element such as gradient, Hessian
distribution, and filter properties to measure the extent of redundancy [10;
43; 41]. These methods require delicate tuning tailored to the network for
optimal performance. To tackle this, some approaches find the optimal network
structure during the learning process. Automatic pruning involves a hyper-
network or pruning indicator with learnable parameters to the target network
[40; 22; 21]. It finds the optimal configuration for each layer within a
constrained budget through the learning process to achieve the best
performance. This eliminates the need for non-trivial hand-crafted designs
tailored to specific networks.

Figure 1:  Problems arising when multi-head is not considered. (a) Multi-head
attention is applied to each head, forming different representation subspaces.
(b) If multi-head is not considered, reconfiguration leads to channel
misalignment problems where representation subspaces get mixed, or subspaces
are completely removed. It causes significantly reduction of the expression
capacity. (c) Our method resolves these issues through an automatic pruning
method by considering multi-head.

Despite the strength of automatic pruning, applying it directly to Transformer
is not trivial. The original attention mechanism is structured with multiple
heads, as shown in Fig. 1 (a). It enables the model to collectively attend to
information from various representation subspaces at different positions [34].
Fig. 1(b) illustrates the issues when applying automatic pruning without
considering multi-head attention. If all channels of a specific head are
removed, the corresponding representation subspace is eliminated, leading to
limitations in the features that can be extracted. As a result, there is a
significant degradation in the expression capacity compared to the original
network. Pruning without considering multi-head attention leads to mixed
channels from different representation subspaces, resulting in the extraction
of features that are entirely different from those of the original network. We
refer to this issue as the _channel misalignment problem._

In this paper, we propose an automatic pruning method that takes into account
the multi-head attention mechanism. As in previous approaches [40; 21; 19], we
apply a pruning indicator to estimate and remove redundant elements through
learning. During the training of the pruning indicator, we incorporate weights
based on head-wise similarity to preserve more informative factors for each
head. This approach ensures that the representation subspaces remain similar
to the original attention mechanism, minimizing the loss of expression
capacity. To prevent the channel misalignment problem that occurs during the
reconfiguration process, as shown in Fig. 1(b), we propose a pruning indicator
adjustment process, which consists of head-wise ranking and rank-based
average. It resolves the channel misalignment problem by equalizing the
pruning proportions for each head, as shown in Fig. 1(c). Our method utilizes
similarity between each channel when training the pruning indicator, allowing
removed channels to be represented by the combination of remaining channels.
Consequently, we compensate for information loss by adding a reweight module
to adjust the scale of the remaining channels. We present a method to
initialize the pruning indicator for linear attention to prevent excessive
channel removal.

Our method can effectively extend beyond original attention to include linear
attention, without requiring tailored heuristics. We apply it to
FlattenTransformer [12], a model integrating both attention mechanisms, to
discover the optimal model at the target compression rate. Through comparisons
with various manually designed models and pruned models on the ImageNet-1K
benchmark [7], our approach demonstrates improved efficiency and superior
accuracy for various MACs. It is worthy note that compared with NViT [41], a
model obtained through structure pruning, our approach demonstrates
approximately 0.3% better performance at equivalent computational cost. Even
when compressing Flatten-Swin-T, our method exhibits superior performance by
up to 0.25%, despite being more efficient.

##  2 Related Works

###  2.1 Efficient Vision Transformer

Transformer [34] exhibits exceptional performance in the NLP field due to
their ability to effectively capture long-range dependencies. The Vision
Transformer (ViT) has successfully adapted these models for image
classification, achieving outstanding performance [8]. However, the quadratic
computational complexity of the original attention module in Transformers
poses challenges for various vision applications. Several methods address this
concern by limiting the number of tokens. Pyramid Vision Transformer (PvT)
[36] progressively limits the number of tokens using spatial reduction
attention, which controls the feature map size in the patch embedding.
Deformable Attention Transformer (DAT) [39] selects the positions of keys and
values in a data-dependent manner, considering only attentive tokens.
SwinTransformer [25] divides the input into windows and performs original
attention within these limited regions. Neighborhood Attention Transformer
(NAT) [14] localizes original attention by considering only the nearest
neighboring pixels related to the query. Other researches improve efficiency
by integrating convolution operations into transformer models. Convolutional
vision Transformer (CvT) [38] uses convolutional projections instead of linear
projections to control efficiency. CMT [11] proposes a hybrid network of
transformers and convolutions, achieving a better trade-off between accuracy
and efficiency.

Another approach to achieve efficient vision transformers is to approximate
the original attention with linear complexity operations. EdgeNeXt [28]
applies transposed attention along the channel dimension instead of the
spatial dimension, achieving linear complexity with respect to tokens.
Reformer [18] replaces the dot-product operation with locality-sensitive
hashing, reducing the computation to $\Omega(n\log n)$. LinFormer [35]
approximates the original attention matrix using low-rank matrix
factorization, achieving linear complexity. SwiftFormer [30] demonstrates that
key-value interaction can be replaced with a linear layer without performance
degradation, effectively reducing computation. A different line of research
involves changing the order of original attention computations to achieve a
complexity of $\Omega(NC^{2})$. This requires effectively replacing Softmax.
CosFormer [29] replaces Softmax with ReLU activation and cosine-based distance
re-weighing. SOftmax-Free Transformer (SOFT) [27] uses a Gaussian kernel
function to replace Softmax and achieves linear complexity through low-rank
decomposition. Castling-ViT [42] extracts spectral similarity between Q and K
using a linear angular kernel. FLattenTransformer [12] reduces computation to
linear complexity with focused linear attention, effectively maintaining the
expressiveness of original attention.

###  2.2 Network Pruning

Pruning involves identifying and removing redundant components to reduce the
network size. Pruning has proven highly effective in original Convolutional
Neural Networks (CNNs) [24; 16; 10; 6]. Building on the strong performance of
CNNs, there has been researches aimed at enhancing the efficiency of Vision
Transformers through pruning. One approach, token pruning, involves finding
and removing unnecessary tokens, effectively reducing the computational load
of the attention mechanism [1; 37; 31]. Another approach is to reduce the size
of the Transformer model itself. SViTE [4] determines redundant components via
sparse training and prunes accordingly. Unified Visual Transformer Compression
(UVC) [43] reduces network size by combining pruning methods with various
compression techniques. Novel ViT (NViT) [41] enables global structural
pruning based on a Hessian-based structural pruning criterion. It uses only
their heuristic criterion when pruning multi-heads. In contrast, the proposed
method takes into account the representation space of each head, minimizing
the information loss of the original model’s attention mechanism.

To address these issues, there are automatic pruning methods that identify
redundant components through training. AutoPrune [40] determines the necessity
of weights based on trainable auxiliary parameters. MetaPruning [26] involves
training a meta-network, PruningNet, to find the optimal structure for a given
target network. DHP [22] applies differentiable hypernetworks to determine the
configuration of each channel in the backbone network. Instead of using a
hypernetwork, another approach involves a pruning module to each component,
learning the importance and removing components [21; 19]. VTP [45] applies
automatic pruning to Transformers by using trainable gate variables to
identify and remove unnecessary components through learning. Automatic pruning
can identify the optimal network structure without heuristics, but directly
applying them to Transformers is not trivial. Unlike the traditional CNN
structure, an automatic pruning method should take into account the multi-head
structure in Transformers without losing expression capacity. In this paper,
we propose a multi-head-aware automatic pruning method, which can be applied
to both original attention and linear attention mechanisms.

##  3 Method

###  3.1 Multi-head Automatic Pruning

Figure 2:  Multi-head Automatic Pruning process. To consider multi-head, we
first undergo the (a) Computing pruning indicator process. We incorporate
similarity-based weights into pruning indicators, enabling the pruning
indicator to consider saliency channels in each head. (b) Through Pruning
indicator adjustment, we share rank-wise pruning indicators for each head. It
ensures equal channel removal across all heads, preventing the channel
misalignment.

In this section, we propose an automatic pruning method for handling multi-
head models. It consists of two processes, computing pruning indicator and
adjusting pruning indicator. Fig. 2(a) illustrates a step of computing pruning
indicator, where the importance of each channel is estimated to assign scores.
As in previous automatic pruning methods, a pruning indicator consisting of
learnable parameters $m\in\mathbb{R}^{C_{out}}$ is employed to determine
importance during training. We assign weights to the pruning indicator based
on head-wise similarity, considering more informative elements of the
representation subspace. First, the projection matrix
$P\in\mathbb{R}^{C_{in}\times C_{out}}$ is divided into
$P^{1},P^{2},...,P^{h}\in\mathbb{R}^{C_{in}\times C_{h}}$ based on the number
of heads $h$, where $C_{in},C_{out}$ and $C_{h}$ are the number of input,
output and head channels, respectively. To compute the importance, we first
compute the cosine similarity of the channels for each head.

$\centering S^{k}=\frac{P^{k}\cdot(P^{k})^{T}}{\mathinner{\\!\left\lVert
P^{k}\right\rVert}\mathinner{\\!\left\lVert
P^{k}\right\rVert}}=\begin{pmatrix}s^{k}_{1,1}&\cdots&s^{k}_{1,C_{h}}\\\
\vdots&\ddots&\vdots\\\ s^{k}_{C_{h},1}&\cdots&s^{k}_{C_{h},C_{h}}\\\
\end{pmatrix}\@add@centering$ (1)

where $S^{k}$ is the similarity matrix obtained from the cosine similarity of
$P^{k}$, where element $s^{k}_{i,j}$ represents the relationship between the
$i$-th channel and the $j$-th channel in head $k$. To incorporate this into
the pruning indicator $m$, we calculate the similarity-based weights
$w\in\mathbb{R}^{C_{out}}$ as follows.

$\begin{split}w^{k}_{i}=&1+\lim_{p\to\infty}(\sum_{n=1,n\neq
i}^{C_{h}}|s^{k}_{i,n}|^{p})^{\frac{1}{p}},\\\
w=&\text{Concat}[w^{1},w^{2},\cdots,w^{h}].\end{split}$ (2)

For each row, we estimate similarity by applying the Chebyshev norm
($L_{\infty}$) to the remaining elements excluding itself. Through this
process, $w_{i}^{k}$ represents the similarity between the $i$-th channel and
the other channels in head $k$. If $w_{i}^{k}$ has a high score for a specific
channel, it indicates the presence of other channels with high similarity,
implying that it can be replaced by those channels. By combining
$w_{1}^{k},w_{2}^{k},\dots,w_{C_{h}}^{k}\in w^{k}$ using the concatenation
operation Concat$(\cdot)$, we obtain the similarity-based weight $w$. These
values are weighted in the pruning indicator as follows.

$m^{*}=w\odot m.$ (3)

where $\odot$ represents the element-wise product. Fig. 2(b) illustrates the
method for adjusting pruning indicators to ensure equal distribution of pruned
channels across all heads. First, the rank for each head is computed according
to $m^{*}$. The score is calculated by taking the average of $m^{*}$ with the
same rank in each head. This average is then reassigned to the head positions
with the same rank, enabling the adjustment of pruning indicators to share the
same importance score across all heads. Consequently, redundant channels are
masked out by passing through the following indicator function:

$\mathbb{I}(m^{*})=\begin{cases}1,&\text{if }m^{*}<\tau\\\
0,&\text{otherwise.}\end{cases}$ (4)

Since the same rank across all heads holds the same pruning indicator, the
removal of channels corresponding to each rank is simultaneously determined.
This allows training to ensure that pruning occurs at the same proportion
across all heads, thereby preventing channel misalignment.

###  3.2 Reweight Module

The removal of channels results in information loss, requiring a method to
compensate for this loss to minimize performance degradation. The proposed
method learns indicators by considering the similarity of channels for each
head, and subsequently removes channels based on these values. The channels
being removed can be represented by the weighted sum of the remaining
channels, due to their high similarity with other channels.

For the reweighting method, a simple channel attention mechanism inspired by
SENet [17] is employed. The input token $F\in\mathbb{R}^{N\times C}$ is
compressed into a feature map $F_{C}\in\mathbb{R}^{1\times C}$ by taking the
mean. The compressed feature map is encoded into weights to compensate for
each channel through a linear layer and $Tanh$ activation function. This is
multiplied with the query, allowing the passage of information from the
removed channels to the remaining channels. By applying this module,
compensation can be provided for the removed channels, minimizing performance
degradation.

###  3.3 Pruning Indicator Initialization for Linear Attention

Unlike original attention in Transformer, linear attention with a
computational complexity of $\Omega(NC^{2})$, significantly improves in
efficiency as channels are removed. Even though the channels of linear
attention are likely to be severely removed during the pruning process, it is
non-trivial to properly determine an initial constant value to prevent it. As
a result, sophisticated initialization of the pruning indicator is required.
We introduce a data-driven method to solve the problem. When the embedding
vector is received as input, it is projected into query
$Q\in\mathbb{R}^{N\times C}$, key $K\in\mathbb{R}^{N\times C}$, and value
$V\in\mathbb{R}^{N\times C}$. The relationship between each token $Q{K}^{T}$
can be expressed as a linear combination of each channel as shown below.

$QK^{T}V=(Q_{1}K^{T}_{1}+Q_{2}K^{T}_{2}+\cdots+Q_{C}K^{T}_{C})V.$ (5)

The importance of each channel $Q_{i}K_{i}^{T}V$ can be determined by
calculating the matrix distance from $QK^{T}V$ in Eq. 5. By projecting
$Q_{i}K_{i}^{T}$ onto $V$, we compare matrices in $\mathbb{R}^{N\times C}$
instead of $\mathbb{R}^{N\times N}$, significantly reducing computational
complexity when $N\gg C$. We use the difference in singular values to measure
the matrix distance. Let $\Sigma^{j}$ and $\Sigma^{j}_{i}$ be the singular
values of the $QK^{T}V$ and the $Q_{i}K_{i}^{T}V$ in $j$-th image of training
database, respectively. The importance score of channel i is given by:

$T^{j}_{i}=|\Sigma^{j}_{i}-\Sigma^{j}_{i}|.$ (6)

The pruning indicator is initialized by accumulating the distance differences
for each image and normalizing them between 0 and 1. Using this
initialization, we can retain a sufficient number of channels, ensuring the
attention mechanism operates effectively.

###  3.4 Training Automatic Pruning Method

In the $l$-th layer of the Transformer, the pruning module ${m^{*}}^{l}$ is
applied to the original model’s weight $o^{l}\in\mathbb{R}^{C_{in}\times
C_{out}}$ as follows:

$\tilde{o}_{i}^{l}=(\mathbb{I}({m^{*}}^{l}_{i})\odot o_{i}^{l}).$ (7)

where $o^{l}_{i}$ and $\tilde{o}_{i}^{l}$ denote the weights of the $i$-th
channel in the original and pruned layers, respectively. $\mathbb{I}(\cdot)$
is the indicator function in Eq. 4. Since $\mathbb{I}(\cdot)$ is a non-
differentiable binary operation, the straight-through estimator (STE) is
applied in back-propagation, allowing the pruning indicator to be learned by
directly passing the gradient from $\mathbb{I}(m_{i}^{l})$ to
${m^{*}}_{i}^{l}$.

The loss for the automatic pruning method is as follows:

$\mathcal{L}=\mathcal{L}_{CE}(f(x),y)+\mathcal{L}_{MAC}(M_{prune},M_{target})$
(8)

where $L_{CE}$ is the cross-entropy loss, comparing the model output $f(x)$
for input $x$ with the true label $y$. MAC aware loss $\mathcal{L}_{MAC}$ is
the Euclidean distance between the current pruned model’s MAC, $M_{prune}$ and
the target MAC, $M_{target}$. To compute the current MAC, the following
formula is used:

$M_{prune}=\sum_{l\in
P}^{l}N\times\tilde{C}_{in}\times\tilde{C}_{out}+\sum_{l\in
LA}^{l}(N\times\tilde{C}_{q}\times\tilde{C}_{k}+N\times\tilde{C}_{k}\times\tilde{C}_{v})+\sum_{l\in
OA}^{l}(N^{2}\times\tilde{C}_{q}+N^{2}\times\tilde{C}_{v})$ (9)

where $P$ denotes the projection matrix layer, and
$\tilde{C}_{in},\tilde{C}_{out},$ and $N$ represent the numbers of remaining
input and output channels and tokens, respectively. Unlike traditional CNNs,
Transformers have additional computational costs due to the attention
mechanism. $LA$ and $OA$ represent linear attention and original attention,
respectively, with $\tilde{C}_{q},\tilde{C}_{k},$ and $\tilde{C}_{v}$
indicating the remaining numbers of query, key, and value channels.
Considering the reconfiguration process, the pruning indicator $m^{*}$ for
query and key projection is shared to ensure they are pruned at the same
proportion.

##  4 Experiments

To verify the proposed method, we apply our Automatic channel Pruning for
Multi-head Attention (APMA) to the FLattenTransformer [12]. This network is
composed of original attention blocks from the SwinTransformer [25] and linear
attention blocks. Unlike previous pruning methods, the proposed approach can
be applied to both types of attentions without specific heuristics. APMA-B
compresses the FLatten-Swin-S model to 6.2G, while APMA-S and APMA-T compress
the FLatten-Swin-T model to 4.2G and 1.3G, respectively. While other models
follow the hybrid structure of FLattenTransformer, the smallest model, APMA-T,
consists solely of linear attention to achieve a higher compression ratio. We
evaluate the performance of the compressed models on the classification task
using the ImageNet-1K dataset [7].

To compress the network, we perform search and refine process. The search
process involves determining the optimal network through an automatic pruning
module over 30 epochs. We use the AdamW optimizer, starting with a learning
rate of 5e-4 and decaying to 5e-6, with a weight decay set to 1e-6. The batch
size is 1024, and the training is conducted on 8 RTX A6000 GPUs. The refine
process aims to recover any information loss during the search phase. Except
for setting the weight decay, the experimental setup remains the same as in
the search process. APMA-B and APMA-S apply a weight decay of 0.05, while
APMA-T uses a weight decay of 1e-6 to minimize additional sparsity. In all
experiments, the threshold $\tau$ for the indicator function is set to 0.5.

###  4.1 Comparison with other methods

| Type | Acc. (%) $\uparrow$ | MAC (G) $\downarrow$ |  | Type | Acc. (%) $\uparrow$ | MAC (G) $\downarrow$
---|---|---|---|---|---|---|---
SViTE-B-40 [4] | P | 82.92 | 11.7 | T2T-ViT-14 [44] | M | 81.7 | 6.1
EViT-DeiT-B [23] | P | 82.1 | 11.6 | CaiT-XS-24 [33] | M | 82.0 | 5.4
AutoFormer-B [3] | P | 82.90 | 11 | AS-ViT-S [5] | M | 81.2 | 5.3
ConViT-S+ [9] | M | 82.90 | 10 | TNT-S [13] | M | 81.5 | 5.2
VTP-DeiT-B [45] | P | 80.7 | 10.0 | DeiT-S [32] | M | 81.2 | 4.6
CaiT-S-24 [33] | M | 83.5 | 9.4 | CvT-13 [38] | M | 81.6 | 4.5
Swin-S [25] | M | 83.00 | 8.7 | Swin-T [25] | M | 81.3 | 4.5
FLatten-Swin-S [12] | M | 83.50 | 8.7 | FLatten-Swin-T [12] | M | 82.1 | 4.5
SPViT [15] | P | 82.40 | 8.4 | GLiT-S [2] | M | 80.5 | 4.4
CaiT-XS-36 [33] | M | 82.90 | 8.1 | NViT-S [41] | P | 82.19 | 4.2
UVC-DeiT-B [43] | P | 80.6 | 8.0 | APMA-S | P | 82.35 | 4.2
STEP-Swin-S [20] | P | 79.6 | 6.3 | AutoFormer-T [3] | P | 75.70 | 1.3
NViT-H [41] | P | 82.95 | 6.2 | NViT-T [41] | P | 76.21 | 1.3
APMA-B | P | 83.19 | 6.2 | APMA-T | P | 76.53 | 1.3
Table 1:  ImageNet-1K results of various efficient models and our proposed method. Our approach demonstrates higher performance at lower MAC than both manually designed networks (M) and efficient networks using pruning methods (P).  | Acc. (%) $\uparrow$ | MAC (G) $\downarrow$ | Throughput $\uparrow$ | Speed up
---|---|---|---|---
FLatten-Swin-S [12] | 83.5 | 8.7 | 492 | x1
APMA-B | 83.19 | 6.2 | 663 | x 1.4
FLatten-Swin-T | 82.1 | 4.5 | 805 | x 1.6
APMA-S | 82.35 | 4.2 | 921 | x 1.9
APMA-T | 76.53 | 1.3 | 1221 | x 2.5
Table 2:  Comparison of Throughput between FlattenTransformer and our proposed
method. All measurements are conducted under the same computational
environment on RTX A6000. Applying our proposed method demonstrates
effectively improved inference speed. Note that APMA-S achieves better
throughput compared to Flatten-Swin-T, while also showing better performance.

In Table 1, we compare the networks compressed using our proposed method with
other manually designed models (M) and pruned networks (P). Compared to
manually designed models (M), our method demonstrates higher performance with
lower computational cost. For instance, when comparing APMA-B with Swin-S, our
method achieves improved performance with approximately 30% less computation.
Notably, APMA-S is more efficient than both Swin-T and DeiT-S, showing a
significant performance improvement of about 1.15%.

We compare our compressed networks with other pruned networks (P). Our methods
are shown to be the most efficient while maintaining high performance,
surpassing the previous state-of-the-art NViT on all MACs. APMA-B
significantly reduces MACs while achieving superior accuracy compared to other
pruning methods. Particularly, the model composed solely of linear attention,
APMA-T, demonstrates a 0.32% higher performance at the same computational cost
compared to other compressed networks. Other compressed networks are based on
original attention with a computational complexity of $\Omega(N^{2}C)$. In
contrast, our proposed method compresses linear attention-based networks with
a computational complexity of $\Omega(NC^{2})$. Consequently, even with less
channel pruning, we achieve a significantly more efficient reduction in
computational cost, offering greater expressive capacity at the same
computational cost.

We also compare the performance of FLattenTransformer with our proposed method
applied. APMA-B compresses the FLatten-Swin-S model, achieving approximately
30% less computational cost with only a 0.31% performance drop. However, as
shown in Table 2, there is a substantial throughput improvement of around 40%.
APMA-S demonstrates a 0.25% performance improvement over the FLatten-Swin-T
model, despite having lower computational cost and faster computation speed.
This indicates the effectiveness of our method in finding the optimal model.
APMA-T, the fastest network, shows a 2.5x speed improvement over the FLatten-
Swin-S model.

###  4.2 Ablation Study

Figure 3:  Top-3 singular value norms for each attention layer of the pruned model. When using Similarity-based weights, it demonstrates larger singular value norms. It indicates that when employing the proposed method, salient channels can be effectively retained.  | Acc. (%) $\uparrow$ | MAC (G) $\downarrow$
---|---|---
Original | 76.53 | 1.3
(-) Reweight module | 76.36 (-0.17) | 1.3
(-) Similarity-based weight | 75.90 (-0.63) | 1.3
(-) Multi-head automatic pruning | 72.78 (-3.75) | 1.3
(-) Pruning indicator initialize | 71.04 (-5.49) | 1.3
Table 3:  Evaluation of the effectiveness of proposed methods. Significant
performance degradation is observed when each module and method is removed.

In this section, we estimates the effectiveness of each module and method
proposed by removing them individually and comparing the performance. All
experiments are conducted on APMA-T, and the experimental setup is
consistently configured. The Similarity-based weight and Multi-head automatic
pruning are proposed in Sec. 3.1, Reweight module in Sec. 3.2, and Pruning
indicator initialization in Sec. 3.3.

#### Similarity-based Weight

To demonstrate the effectiveness of the similarity-based weight method, we
examine the singular value norms of each attention layer in the model
immediately after the search process, as shown in Fig. 3. Due to the varying
number of channels in the models with and without the similarity-based weight,
we calculate the norm of the top 3 singular values for each head. Fig. 3 shows
that the singular value norms are higher across all layers when the
similarity-based weight is used. This indicates that the model possesses
eigenvectors with higher importance, which means it can inherently extract
more distinct features. Consequently, as shown in Table 3, the performance
difference between the model using the proposed similarity-based weight and
the model without it is a significant 0.63%.

Figure 4:  The effect of the reweight module. The red block represents the
query token, while the blue blocks depict the relationship between this query
token and other tokens. It demonstrates that the reweight module can
compensate for information loss caused by pruning, directing attention to more
relevant tokens.

#### Reweight Module

We apply a reweight module to compensate for the information loss that occurs
when channels are removed. As shown in Table 3, when the module is removed in
the proposed automatic pruning method, there is a decrease in accuracy of 0.17
without any change in MAC. Fig. 4 illustrates the attention for tokens with or
without applying the module. When applying the reweight module, the
relationship with relevant object tokens strengthens, resembling the original
model, while the relationship with background tokens weakens. The result
demonstrates that the reweight module compensates for information loss,
allowing relevant tokens to receive better attention.

Figure 5:  Pruned model structure after reconfiguration. (a) Without applying
multi-head pruning, the pruning ratio for each head is inconsistent, leading
to a channel misalignment problem. (b) When pruning indicator initialization
for linear attention is not performed, the attention mechanism does not
function properly. (c) Our proposed method ensures consistent pruning ratios
across heads to resolve the channel misalignment problem and demonstrates
effective pruning of each module in appropriate proportions. Zoom in for
details.

#### Multi-head Automatic Pruning

In contrast to conventional automatic pruning methods, our proposed approach
can remove channels considering the multi-head structure. Without considering
multi-heads, the accuracy drops by 3.75, indicating a significant performance
degradation, as shown in Table 3. In Fig. 5(a), the channels for query, key,
and value vary across each head. During the reconstruction process in the
network, where redundant channels are actually removed for real acceleration,
channel misalignment occurs across heads. On the other hand, applying our
proposed method, as shown in Fig. 5(c), guarantees that channels are removed
at the same proportion across heads.In Fig. 5(a), the almost complete removal
of channels in specific heads leads to restricted representation space and
severe performance degradation. In contrast, our proposed method ensures an
equal distribution of channels per head, thereby preserving the network’s
expression capacity.

#### Pruning Indicator Initialize for Linear Attention

Fig. 5(b) shows the model structure obtained from the search process when the
pruning indicator for linear attention is not initialized. It illustrates that
channels of the projection matrix for the query and key are severely pruned,
whereas the feed-forward network is not pruned. This imbalance leads to
improper operation of the attention mechanism. In contrast, Fig. 5(c)
demonstrates how our method alleviates such issues. Table 3 presents a
significant performance degradation of 5.49% when the pruning indicator is not
initialized. This shows the effectiveness of our proposed initialization
method.

##  5 Conclusion

In this paper, we introduce an automatic pruning method tailored for multi-
head attention mechanisms. Leveraging similarity-based weights and a pruning
indicator adjustment process, our approach handles multi-head attention,
considering the representation subspace effectively. Moreover, the integration
of a reweight module compensates for the information loss incurred by pruned
channels. By initializing the pruning indicator for linear attention in a
data-driven manner, our method facilitates the discovery of optimal models.
Comparative analysis against manually designed models and various pruning
techniques demonstrates the superior efficiency and performance trade-off of
our proposed approach. In this paper, we only focus on experiments on the
classification task. Future work will extend our method to various tasks such
as object detection and semantic segmentation.

## References

  * [1] Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman.  Token merging: Your vit but faster.  arXiv preprint arXiv:2210.09461, 2022.
  * [2] Boyu Chen, Peixia Li, Chuming Li, Baopu Li, Lei Bai, Chen Lin, Ming Sun, Junjie Yan, and Wanli Ouyang.  Glit: Neural architecture search for global and local image transformer.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 12–21, 2021.
  * [3] Minghao Chen, Houwen Peng, Jianlong Fu, and Haibin Ling.  Autoformer: Searching transformers for visual recognition.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 12270–12280, 2021.
  * [4] Tianlong Chen, Yu Cheng, Zhe Gan, Lu Yuan, Lei Zhang, and Zhangyang Wang.  Chasing sparsity in vision transformers: An end-to-end exploration.  Advances in Neural Information Processing Systems, 34:19974–19988, 2021.
  * [5] Wuyang Chen, Wei Huang, Xianzhi Du, Xiaodan Song, Zhangyang Wang, and Denny Zhou.  Auto-scaling vision transformers without training.  In International Conference on Learning Representations, 2022.
  * [6] Yanqi Chen, Zhengyu Ma, Wei Fang, Xiawu Zheng, Zhaofei Yu, and Yonghong Tian.  A unified framework for soft threshold pruning.  In The Eleventh International Conference on Learning Representations, 2023.
  * [7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248–255, 2009.
  * [8] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In International Conference on Learning Representations, 2021.
  * [9] Stéphane d’Ascoli, Hugo Touvron, Matthew L Leavitt, Ari S Morcos, Giulio Biroli, and Levent Sagun.  Convit: Improving vision transformers with soft convolutional inductive biases.  In International conference on machine learning, pages 2286–2296. PMLR, 2021.
  * [10] Shangqian Gao, Zeyu Zhang, Yanfu Zhang, Feihu Huang, and Heng Huang.  Structural alignment for network pruning through partial regularization.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 17402–17412, 2023.
  * [11] Jianyuan Guo, Kai Han, Han Wu, Yehui Tang, Xinghao Chen, Yunhe Wang, and Chang Xu.  Cmt: Convolutional neural networks meet vision transformers.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12175–12185, 2022.
  * [12] Dongchen Han, Xuran Pan, Yizeng Han, Shiji Song, and Gao Huang.  Flatten transformer: Vision transformer using focused linear attention.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 5961–5971, 2023.
  * [13] Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang.  Transformer in transformer.  Advances in neural information processing systems, 34:15908–15919, 2021.
  * [14] Ali Hassani, Steven Walton, Jiachen Li, Shen Li, and Humphrey Shi.  Neighborhood attention transformer.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6185–6194, 2023.
  * [15] Haoyu He, Jianfei Cai, Jing Liu, Zizheng Pan, Jing Zhang, Dacheng Tao, and Bohan Zhuang.  Pruning self-attentions into convolutional layers in single path.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
  * [16] Zejiang Hou, Minghai Qin, Fei Sun, Xiaolong Ma, Kun Yuan, Yi Xu, Yen-Kuang Chen, Rong Jin, Yuan Xie, and Sun-Yuan Kung.  Chex: Channel exploration for cnn model compression.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12287–12298, 2022.
  * [17] Jie Hu, Li Shen, and Gang Sun.  Squeeze-and-excitation networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7132–7141, 2018.
  * [18] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya.  Reformer: The efficient transformer.  In International Conference on Learning Representations, 2020.
  * [19] Donghyeon Lee, Eunho Lee, and Youngbae Hwang.  Pruning from scratch via shared pruning module and nuclear norm-based regularization.  In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 1393–1402, 2024.
  * [20] Jiaoda Li, Ryan Cotterell, and Mrinmaya Sachan.  Differentiable subset pruning of transformer heads.  Transactions of the Association for Computational Linguistics, 9:1442–1459, 2021.
  * [21] Yanyu Li, Pu Zhao, Geng Yuan, Xue Lin, Yanzhi Wang, and Xin Chen.  Pruning-as-search: Efficient neural architecture search via channel pruning and structural reparameterization.  In Thirty-First International Joint Conference on Artificial Intelligence, 2022.
  * [22] Yawei Li, Shuhang Gu, Kai Zhang, Luc Van Gool, and Radu Timofte.  Dhp: Differentiable meta pruning via hypernetworks.  In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part VIII 16, pages 608–624. Springer, 2020.
  * [23] Youwei Liang, Chongjian Ge, Zhan Tong, Yibing Song, Jue Wang, and Pengtao Xie.  Not all patches are what you need: Expediting vision transformers via token reorganizations.  In International Conference on Learning Representations, 2022.
  * [24] 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, pages 1529–1538, 2020.
  * [25] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 10012–10022, 2021.
  * [26] Zechun Liu, Haoyuan Mu, Xiangyu Zhang, Zichao Guo, Xin Yang, Kwang-Ting Cheng, and Jian Sun.  Metapruning: Meta learning for automatic neural network channel pruning.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 3296–3305, 2019.
  * [27] Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang.  Soft: Softmax-free transformer with linear complexity.  Advances in Neural Information Processing Systems, 34:21297–21309, 2021.
  * [28] Muhammad Maaz, Abdelrahman Shaker, Hisham Cholakkal, Salman Khan, Syed Waqas Zamir, Rao Muhammad Anwer, and Fahad Shahbaz Khan.  Edgenext: efficiently amalgamated cnn-transformer architecture for mobile vision applications.  In European Conference on Computer Vision, pages 3–20. Springer, 2022.
  * [29] Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong.  cosformer: Rethinking softmax in attention.  In International Conference on Learning Representations, 2022.
  * [30] Abdelrahman Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, Ming-Hsuan Yang, and Fahad Shahbaz Khan.  Swiftformer: Efficient additive attention for transformer-based real-time mobile vision applications.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 17425–17436, 2023.
  * [31] Quan Tang, Bowen Zhang, Jiajun Liu, Fagui Liu, and Yifan Liu.  Dynamic token pruning in plain vision transformers for semantic segmentation.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 777–786, 2023.
  * [32] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  In International conference on machine learning, pages 10347–10357. PMLR, 2021.
  * [33] Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou.  Going deeper with image transformers.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 32–42, 2021.
  * [34] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [35] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity.  arXiv preprint arXiv:2006.04768, 2020.
  * [36] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao.  Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 568–578, 2021.
  * [37] Siyuan Wei, Tianzhu Ye, Shen Zhang, Yao Tang, and Jiajun Liang.  Joint token pruning and squeezing towards more aggressive compression of vision transformers.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2092–2101, 2023.
  * [38] Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang.  Cvt: Introducing convolutions to vision transformers.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 22–31, 2021.
  * [39] Zhuofan Xia, Xuran Pan, Shiji Song, Li Erran Li, and Gao Huang.  Vision transformer with deformable attention.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4794–4803, 2022.
  * [40] Xia Xiao, Zigeng Wang, and Sanguthevar Rajasekaran.  Autoprune: Automatic network pruning by regularizing auxiliary parameters.  Advances in neural information processing systems, 32, 2019.
  * [41] Huanrui Yang, Hongxu Yin, Maying Shen, Pavlo Molchanov, Hai Li, and Jan Kautz.  Global vision transformer pruning with hessian-aware saliency.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18547–18557, 2023.
  * [42] Haoran You, Yunyang Xiong, Xiaoliang Dai, Bichen Wu, Peizhao Zhang, Haoqi Fan, Peter Vajda, and Yingyan Celine Lin.  Castling-vit: Compressing self-attention via switching towards linear-angular attention at vision transformer inference.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14431–14442, 2023.
  * [43] Shixing Yu, Tianlong Chen, Jiayi Shen, Huan Yuan, Jianchao Tan, Sen Yang, Ji Liu, and Zhangyang Wang.  Unified visual transformer compression.  In International Conference on Learning Representations, 2022.
  * [44] Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan.  Tokens-to-token vit: Training vision transformers from scratch on imagenet.  In Proceedings of the IEEE/CVF international conference on computer vision, pages 558–567, 2021.
  * [45] Mingjian Zhu, Yehui Tang, and Kai Han.  Vision transformer pruning.  arXiv preprint arXiv:2104.08500, 2021.