Know What You Don’t Need:Single-Shot Meta-Pruning for Attention Heads

Zhengyan Zhang1,2,3, Fanchao Qi1,2,3, Zhiyuan Liu1,2,3†, Qun Liu4, Maosong Sun1,2,3

1Department of Computer Science and Technology, Tsinghua University, Beijing, China2Institute for Artificial Intelligence, Tsinghua University, Beijing, China

3State Key Lab on Intelligent Technology and Systems, Tsinghua University, Beijing, China4Huawei Noah’s Ark Lab

{zy-z19, qfc17}@mails.tsinghua.edu.cn

Abstract

Deep pre-trained Transformer models haveachieved state-of-the-art results over a varietyof natural language processing (NLP) tasks.By learning rich language knowledge withmillions of parameters, these models are usu-ally overparameterized and significantly in-crease the computational overhead in applica-tions. It is intuitive to address this issue bymodel compression. In this work, we proposea method, called Single-Shot Meta-Pruning,to compress deep pre-trained Transformersbefore fine-tuning. Specifically, we focuson pruning unnecessary attention heads adap-tively for different downstream tasks. To mea-sure the informativeness of attention heads, wetrain our Single-Shot Meta-Pruner (SMP) witha meta-learning paradigm aiming to maintainthe distribution of text representations afterpruning. Compared with existing compressionmethods for pre-trained models, our methodcan reduce the overhead of both fine-tuningand inference. Experimental results show thatour pruner can selectively prune 50% of at-tention heads with little impact on the perfor-mance on downstream tasks and even providebetter text representations. The source codewill be released in the future.

1 Introduction

Pre-trained language models (PLMs), such asBERT (Devlin et al., 2019), XLNet (Yang et al.,2019) and RoBERTa (Liu et al., 2019), haveachieved state-of-the-art results across a varietyof natural language processing (NLP) tasks. Tofully utilize large-scale unsupervised data duringpre-training, PLMs are becoming larger and larger.For example, GPT-3 (Brown et al., 2020) has 175billion parameters. With the growing number of

† Corresponding author: Z.Liu (liuzy@tsinghua.edu.cn)

Pre-trained Model

Fine-tuned Model

InefficientFine-tuning

Efficient Fine-tuning

PrunedPre-trained

Model

PrunedFine-tuned

Model

CompressedFine-tuned

Model

Inference Efficient

Efficient Inefficient

Our Single-Shot Pruning

Previous PLMCompression

Figure 1: An illustration of the three-step paradigm(gray blocks) adopted by most of PLMs and twotypes of model compression methods. Compared toprevious work, our method makes both fine-tuningand inference efficient.

model parameters, the computational overhead, in-cluding memory and time, becomes tremendouslyheavy, which severely limits the application ofPLMs to downstream NLP tasks. Therefore, modelcompression for PLMs is increasingly important.

Most PLMs are Transformer-based (Radfordet al., 2018; Devlin et al., 2019), and they are uti-lized following a three-step paradigm: pre-training,fine-tuning and inference, as illustrated in Figure1. Quite a few methods have been proposed tocompress pre-trained Transformers during or afterfine-tuning to reduce the computational overhead inthe inference phase (Tang et al., 2019; Turc et al.,2019; Jiao et al., 2019; McCarley, 2019; Wanget al., 2019b; Fan et al., 2020), while little workattempts to perform model compression before fine-tuning.

In fact, compressing pre-trained Transformersbefore the fine-tuning phase is more significant. Forone thing, computational overhead of pre-trainedTransformers during fine-tuning is usually heav-ier than that during inference, because of extra

arX

iv:2

011.

0377

0v1

[cs

.CL

] 7

Nov

202

0

computations for gradient descent. Furthermore,compressing pre-trained Transformers before fine-tuning can reduce the overhead during both fine-tuning and inference, which is more helpful.

In this paper, we make the first attempt to con-duct model compression for deep pre-trained Trans-formers before fine-tuning. According to previouswork (Kovaleva et al., 2019; Michel et al., 2019),deep pre-trained Transformers are overparameter-ized, and only part of attention heads are actuallyuseful for downstream tasks. Therefore, we pro-pose to prune the unnecessary attention heads ofpre-trained Transformers to reduce the overhead.

We train a pruner to measure the importance ofattention heads and identify the unnecessary ones.We assume the unnecessary attention heads cannotprovide useful information, and pruning them willhave little effect on the distribution of the text repre-sentations learned by the pre-trained Transformers.Therefore, we design a self-supervised objectivefunction for the pruner, which trains the model tomaintain the distribution of representations afterpruning. In addition, to make our pruner more gen-eral, we adopt a meta-learning paradigm to trainit. To provide diverse task distributions, we sampledata from multiple corpora to form the training setof meta-learning.

Our pruning strategy is single-shot, which meansit can compress the pre-trained Transformers oncebefore fine-tuning rather than using an iterativeoptimization procedure (McCarley, 2019; Voitaet al., 2019). We name our model Single-ShotMeta-Pruner (SMP). In the experiments, we applySMP to the representative pre-trained TransformerBERT, and conduct evaluations on GLUE (Wanget al., 2019a) and the semantic relatedness tasks ofSentEval (Conneau and Kiela, 2018). Experimen-tal results show that SMP can prune as many as50% of attention heads of BERT without sacrific-ing much performance on GLUE, and even bringperformance improvement on the semantic related-ness tasks of SentEval. In addition, SMP is alsocomparable to, if not slightly better than, the base-line method which conducts model compressionafter fine-tuning. Moreover, we find the patternsof unnecessary heads learned by SMP are transfer-able, which means SMP could work with differentTransformer models and downstream tasks.

2 Related Work

To compress pre-trained Transformers, there aretwo mainstream approaches, namely knowledgedistillation and parameter pruning.

(1) Knowledge distillation (Sanh et al., 2019;Chen et al., 2020a; Sun et al., 2020) treats the orig-inal large model as a teacher to teach a lightweightstudent network. Sun et al. (2019) design the stu-dent networks to learn from multiple intermediatelayers of the teacher model. Jiao et al. (2019) pro-pose to learn from both teachers’ hidden states andattention matrices.

(2) Parameter pruning aims to remove unnec-essary parts of networks, such as weight magni-tude pruning (McCarley, 2019; Li et al., 2020) andlayer pruning (Fan et al., 2020; Sajjad et al., 2020).Given a complete BERT after fine-tuning, Michelet al. (2019) propose to prune attention heads ac-cording to the change of loss function when slightlyperturbing the attention matrices. They argue thatthe loss function is situated in a local minimumafter fine-tuning and sensitive to the change of im-portant attention heads. Compared to this work, ourSMP meta-learns the criterion and prunes PLMsbefore fine-tuning.

In addition to these two approaches, researchersalso explore other methods, such as weight factor-ization (Wang et al., 2019b), weight sharing (Lanet al., 2020), and parameter quantization (Zafriret al., 2019). Most of current compression stud-ies focus on reducing the overhead of inference.There are also some researches trying to directlycompress the models during pre-training (Gordonet al., 2020; Sanh et al., 2019), but this kind of com-pression has severe impact on the performance ofdownstream tasks. In this work, our SMP aims toreduce the overhead of both fine-tuning and infer-ence and better maintain the performance of PLMs.

From the more general perspective of pruningneural networks, our SMP prunes models beforetraining (fine-tuning), which is different from prun-ing after and during training. Pruning after train-ing aims to identify unnecessary parts in a fullytrained model based on weight magnitude (Hanet al., 2015) or effects on the loss (LeCun et al.,1990). Pruning during training (Louizos et al.,2018; Voita et al., 2019) attempts to combine prun-ing and training procedures together. These meth-ods require approximately the same computationaloverhead as training a full network.

Single-shot pruning (Lee et al., 2019b,a;

}Sampled Data

Representation Space

Representation Space

Relative Distance Distribution for

Relative Distance Distribution for

Representation Space

Objective Function�KL Divergence

Update

PLM

SMP

SMP

SMP

1 2 3 4

3

2

1

Laye

r

Head Input: Know what you don’t need

1 2

3

2

1

Laye

r

Head

Reconstruct

Attention Matrix

0.2

0.0

1.0

Knowwhatyou

don’tneed

PruneImportantHead

UnimportantHead

Procedure of Pruning

AppropriateO\Pruned PLM

Inappropriately Pruned PLM

Relative Distance Distribution for

1

3

4

5

2

Figure 2: An example of training our SMP and pruning a 3-layer 4-head Transformer by SMP. (1) Giventhe sampled data from the training corpora, a full pre-trained Transformer represents these instancesin its representation space. (2) Based on the attention behaviors when encoding these instances, ourSMP identifies the unnecessary heads and prunes the model. (3) In the beginning, the representationdistribution of the pruned model is much different from that of the full model. (4) We use the relativedistance distribution to parameterize the representation space and compute the KL divergence betweendistributions. (5) After optimization, the SMP can provide a good pruned model, which maintains thedistribution.

Dettmers and Zettlemoyer, 2019), which prunesnetworks before training, is more efficient thantraditional pruning, which leads to lower compu-tational overhead. Most existing studies of single-shot pruning focus on the weight pruning of ran-domly initialized networks by pre-defined criteria,but the models pruned by weight pruning are diffi-cult to accelerate (Han et al., 2015). In this work,we focus on directly pruning the structures (atten-tion heads) in Transformers, which makes prunedmodels easy to accelerate. We also consider howto maintain the knowledge in pre-trained models,which is different from pruning randomly initial-ized networks.

3 Method

In this section, we elaborate on our SMP model aswell as its objective function and training method.Figure 2 illustrates the overall framework and work-flow of SMP.

The goal of SMP is to find and prune the unim-portant attention heads in pre-trained Transformers.To this end, SMP calculates the importance scoreof each attention head. Then the attention headswith low importance scores are pruned to obtain apruned Transformer.

To train SMP, we design a self-supervised ob-jective function, which aims to keep the output ofthe PLMs not changing a lot after pruning. Specif-ically, we propose to preserve the distribution oftext representations.

Furthermore, we adopt the meta-learningparadigm in training to make SMP general andbe able to apply to almost all sentence-level tasks.

3.1 Score CalculationTransformer is composed of a stack of identicallayers, and each layer has two sub-layers: a multi-head self-attention network and a point-wise feed-forward network. For the multi-head self-attentionnetworks, each attention head yields an attentionmatrix. For example, as illustrated in the bottomleft of Figure 2, a 3-layer 4-head Transformer has3× 4 = 12 attention heads in total. To obtain theimportance score of an attention head, we first com-pute the importance scores of its attention matricesfor all instances, and then average them out as thefinal result.

We formulate importance score calculation ofan attention matrix as an image classification task,whose input is an attention matrix and output isits importance score. We adopt a convolutionalneural network (CNN) as the encoder of attention

matrices, which is widely used in image processing.SMP concatenates a sigmoid non-linear function tothe matrix encoder to output a score ranging from0 to 1.

Considering the difference between single-sentence and sentence-pair downstream tasks, SMPactually outputs a two-dimensional vector compris-ing two scores ssing and spair, which are designedfor single-sentence and sentence-pair tasks respec-tively. Formally, the importance score of an atten-tion matrix is calculated by

[ssing, spair] = σ(CNN(Matt)), (1)

where Matt represents an attention matrix.After calculating the importance score of each

attention head in the full pre-trained TransformerT , we can prune the unimportant heads and obtaina pruned Transformer T :

T = SMP(T ). (2)

3.2 Self-supervised Objective Function

Considering pre-trained Transformers essentiallyencode input instances into vector representations,it is reasonable to assume that the unimportantheads have little effect on the distribution of the rep-resentations of a set of sampled instances. In otherwords, the representation distribution of sampledinstances should be maintained after pruning thoseunimportant attention heads. For example, giventhree instances {xi, xj , xk}, we compute their rep-resentations before and after pruning by

hi = T (xi), hi = T (xi).

An appropriately pruned model should make surethat if hi is closer to hj than hk, hi should also becloser to hj than hk.

Based on this assumption, we design the trainingobjective function for SMP. We parameterize therepresentation distribution of a set of sampled in-stances using the relative distance distribution. Therelative distance distribution for an instance recordsthe normalized distances between the instance andother instances.

Given a set of instances {x1, . . . , xN}, the rela-tive distance distribution for an instance xn is anN-dimensional normalized vector rn, whose i-thentry is the relative distance between xn and xi:

rni =eDist(hn,hi)∑Nj=1 e

Dist(hn,hj), (3)

where Dist is the function measuring the distancebetween two representations. In this work, we sim-

ply use cosine distance.To quantify the variation of relative distance dis-

tribution after pruning, we use the Kullback-Leibler(KL) divergence (Kullback and Leibler, 1951). TheKL divergence between the relative distance dis-tributions associated with the original and prunedTransformers are

DKL(rn||rn) = −

N∑i=1

rni ln(rnirni

), (4)

where rn and rn denote relative distance distri-butions associated with the original and prunedTransformers respectively.

Our SMP intends to maintain the representationdistribution after pruning, which means makingDKL(r

n||rn) as small as possible for all instances.Therefore, the objective function of SMP is

LSMP =N∑

n=1

DKL(rn||rn). (5)

3.3 Model Training via Meta-learning

To train our SMP, we design a meta-learning pro-cess. We show a simple example of this trainingparadigm in Figure 2.

At the beginning of each episode, we sample kinstances from the training data to construct a minidataset, which is a set of sentence pairs or singlesentences.

During pruning, we first compute the impor-tance score for each head according to the type ofthe mini dataset (sentence-pair or single-sentence),as in Equation (1). Then we apply Gumbel-softmax (Jang et al., 2016) to the importance scoresof all heads, which is a common reparameterizationmethod and can transform the importance scoresto discrete 0 or 1. We multiply the outputs of anattention head by its discrete importance score of0 or 1, and the unimportant heads, whose impor-tance scores are 0, will be pruned. Meanwhile,Gumbel-softmax can make sure that the pruningoperation is differentiable and we can conduct back-propagation for SMP.

After pruning the pre-trained Transformer fordifferent synthetic mini datasets, SMP is trainedto adapt to different corpora and master the meta-knowledge about pruning attention heads.

4 Experiments

In this section, we evaluate our SMP onGLUE (Wang et al., 2019a) and SentEval (Con-

neau and Kiela, 2018). The pre-trained Transform-ers used here are BERTBASE and BERTLARGE.1

4.1 Experiment Setup

SMP architecture. We set the size of input atten-tion matrices to (128×128), which can cover mostdownstream tasks. Considering some tasks, suchas question answering with some input sequenceslonger than 128, we downsize the attention matri-ces. Our SMP is composed of five CNN layers. Thedimension of the output of the first layer is 8, andthe following layers’ dimensions are twice as largeas the former layers. As a result, the dimension ofthe output representation for the attention matrix is8× 24 = 128. To compute the matrix scores spairand ssimp, we feed the output representation to afull-connected layer.

Pruning. In this work, we follow the head prun-ing paradigm of (Michel et al., 2019) and prunethe same number of heads for each layer. We setthe pruning ratio to 50%, which can significantlyimprove the computation efficiency and effectivelymaintain the original performance according to ourexperiments. We also report the influence of prun-ing ratio in this section.

Training data for meta-learning. We selectseven datasets from GLUE (Warstadt et al., 2018;Socher et al., 2013; Dolan and Brockett, 2005;Williams et al., 2018; Rajpurkar et al., 2016; Daganet al., 2006) as the training data. The statistics areshown in Table 1. In particular, we split the originaltraining data of these datasets into the training andvalidation part in the ratio of 9:1.

Training details of meta-learning. We use twokinds of sequence-level representations of BERT:[CLS] token representation for sentence-pair dataand mean pooling on the sequence outputs forsingle sentence data. For each training episode,we set the number of sampled instances k to 60,which lets SMP make full use of the memory ofthe GPUs. The model is updated every 8 episodes.The quicker update cycle will lead to an unstabletraining and the slower update cycle will bring ex-tra training time. We choose Stochastic GradientDescent as the optimizing algorithm and the bestlearning rate on the validation set is picked from{1, 2, 5} × 10−2. Based on the observation in theexperiments, we set the total number of episodes to48, 000, which is enough for the full convergenceof SMP. We choose the checkpoint with the low-

1https://github.com/google-research/bert

Dataset MNLI QQP SST-2 CoLA STS-B MRPC RTE

Type Pair Pair Sing Sing Pair Pair Pair

Size 392k 363k 67k 8.5k 5.7k 3.5k 2.5k

Table 1: Statistics of the corpora used to train SMP.Sing refers to the single sentence task. Pair refers tothe sentence-pair task.

est loss on the validation set as the final model.We train two SMP models based on BERTBASEand BERTLARGE respectively. SMP was trained onfour 16-GB V100 GPUs for approximate 6 hoursusing BERTBASE and 18 hours using BERTLARGE.

Baselines. To validate the effectiveness of SMP,we introduce four baselines in our experiments.

(1) Fine-tune (None). We fine-tune a completeBERT on downstream tasks, which can provide anoracle result without pruning.

(2) Random. We randomly select the same num-ber of heads to prune as SMP in each layer beforefine-tuning. We repeat the random experiments fivetimes and report the mean of model performances.Since the number of head combinations is verylarge, random experiments only give a rough esti-mation of performances. For example, BERTBASEhas 144 attention heads, and there are C72

144 ≥ 1042

combinations for the pruning ratio of 50%.(3) L0 Norm. Following Voita et al. (2019), we

multiply the output of each head with a scalar gateand introduce an L0 regularization loss to thesegates. Using this method, we can search the opti-mal value of each gate by gradient descent.

(4) HISP and HISP-retrain. Besides L0 Norm,we adopt the attention head pruning method intro-duced by Michel et al. (2019), called Head Impor-tance Score for Pruning (HISP). The original al-gorithm directly evaluates the model performanceafter pruning. According to previous studies ongeneral neural network pruning (Han et al., 2015),retraining after pruning can further promote theperformance of pruned models. Hence, we intro-duce HISP-retrain, which retrains pruned modelsfor better performance. In our experiments, weretrain the pruned model given by HISP for 3 addi-tional epochs as HISP-retrain. Since HISP prunesmodels after fine-tuning, pre-trained Transformerscould better learn from fine-tuning due to the largermodel capacity.

4.2 GLUE

The GLUE benchmark (Wang et al., 2019a) is usedto validate the effectiveness of SMP on the gen-

Model Pruning Method MNLI-(m/mm) QQP SST-2 CoLA STS-B MRPC RTE Average

None 83.85/83.82 90.96 92.43 57.84 88.71 85.78 65.70 81.14

L0 Norm 79.70/79.83 85.82 91.74 52.10 88.30 77.45 62.45 77.17HISP 81.69/81.90 86.88 91.85 54.84 88.46 81.12 65.34 79.01

BERTBASE HISP-retrain 83.56/83.73 91.03 92.20 53.24 88.58 85.04 66.78 80.52

Random 82.43/82.63 90.34 91.83 52.37 87.83 80.88 65.77 79.26SMP 83.36/83.75 90.96 92.31 57.26 88.49 85.04 67.87 81.13

None 87.87/87.62 91.49 93.69 63.89 90.99 88.72 85.92 86.27

L0 Norm 85.93/85.83 90.26 93.46 56.02 90.33 86.51 81.94 83.79HISP 85.44/85.67 85.17 93.11 62.54 89.65 87.50 81.58 83.83

BERTLARGE HISP-retrain 87.42/87.26 91.55 93.46 60.09 90.14 89.70 83.03 85.33

Random 86.36/86.43 91.26 92.45 58.78 90.31 86.85 80.35 84.10SMP 86.86/86.96 91.46 93.34 63.57 90.95 89.70 82.31 85.64

Table 2: Results on seven tasks in GLUE (%). HISP-retrain is an inference-oriented pruning method, whichonly reduces the overhead of inference. SMP is our method, which reduces the overhead of both fine-tuning andinference. MNLI contains two validation sets and provides two results. We underline the truly best results (fromthe original models) and boldface the best results among the pruned models.

eral fine-tuning tasks. We compare four methodsmentioned above on seven downstream tasks inthe GLUE benchmark. We exclude two tasks inGLUE, namely the Winograd Schema Challengeand QNLI. The former is excluded due to the smallsize of the dataset while the latter is excluded forthe experiment on model transferability. The fine-tuning experiments follow the hyperparameters re-ported in the original study (Devlin et al., 2019)except the number of epochs. The random baselineand SMP adopt the same hyper-parameters used infine-tuning a complete BERT. For small datasetscontaining less than 10,000 instances, we set thenumber of epochs to 10. For the others, we keepthe original number unchanged (3 epochs).

We report the results on the validation, ratherthan test data, so the results differ from the originalBERT paper. From Table 2, we observe that:

(1) The average performance of random pruningis consistently worse than that of fine-tuning, whichshows the serious impact of pruning on pre-trainedTransformers. In the experiments, we find thatsome random seeds lead to a good performancewhile some random seeds significantly degrademodel performance. The variation of random prun-ing proves the assumption that there are importantattention heads, which should not be pruned beforefine-tuning, in pre-trained Transformers. It is re-lated to the lottery ticket hypothesis for pre-trainedTransformers (Chen et al., 2020b).

(2) The overall performances of L0 Norm andHISP are worse than that of random pruning, whichindicates that head pruning on a converged model

leads to serious performance degradation. Mean-while, we find that HISP-retrain significantly out-performs HISP, which reflects the importance of re-training in pruning-after-training approaches. Mostresults of HISP-retrain are better than those of ran-dom pruning and close to the result of fine-tuning,which indicates that HISP-retrain can select im-portant heads for downstream tasks and providea good pruned model. However, there still existsthe case that retraining degrades the performance.For BERTBASE, the performance of HISP-retrainon CoLA is lower than HISP by about 1.5%.

(3) SMP achieves the best results on averageamong these pruning methods, and have compara-ble performance with fine-tuning, which indicatesSMP significantly reduces the impact of pruningon downstream tasks. Besides, SMP even outper-forms the fine-tuning method in some tasks, suchas RTE for BERTBASE and MRPC for BERTLARGE.It shows that pruning unnecessary structure canalso promote the performance of downstreamtasks of Transformers. Besides, SMP works wellon both BERTBASE and BERTLARGE, which re-veals the generality of SMP.

4.3 SentEval

SentEval (Conneau and Kiela, 2018) is used to vali-date the representation ability of the models prunedby SMP. The goal of SMP is to preserve the dis-tribution of text representations after pruning formaintaining important prior knowledge learned bypre-training. Hence, we use SentEval to investigatewhether SMP maintains important prior knowledge

Model Pruning Method STS-12 STS-13 STS-14 STS-15 STS-16 STS-B SICK-R

GloVe BoW — 52.10 49.60 54.60 56.10 51.40 64.70 79.90InferSent — 59.20 58.90 69.60 71.30 71.50 75.60 88.30

BERTBASE

Full 46.87 52.77 57.15 63.46 64.50 65.49 80.57

Random 51.07 48.19 57.66 64.48 61.00 65.98 79.67SMP 57.59 63.94 64.64 69.06 66.80 70.18 82.19

BERTLARGE

Full 54.87 60.78 64.21 68.07 66.65 69.91 83.91

Random 55.47 55.04 63.85 67.70 64.47 70.42 83.18SMP 62.13 62.57 71.18 74.38 71.55 71.19 84.52

Table 3: Results on semantic relatedness tasks in SentEval. Numbers reported are Pearson correlations x100. Theresults of GloVe and InferSent are from the paper of SentEval (Conneau and Kiela, 2018). We underline the overallbest results and boldface the best results among BERT models.

Model Ratio Memory (MB) Speed (IPS)

BERTBASE0% 841 — 18.2 —

50% 538 −36.0% 24.5 +34.6%

BERTLARGE0% 2,156 — 5.3 —

50% 1,514 −29.8% 7.3 +37.7%

Table 4: Average memory overhead per instance andthe speed in instances per second (IPS) on QNLI.

after pruning. As mentioned in the experimentalsetup, there are two approaches to compute text rep-resentations from pre-trained BERT. Based on thefindings of previous work (Ma et al., 2019), meanpooling is better than [CLS] token so that we usethe representations of mean pooling in SentEvalexperiments. We use the sentence relatedness tasksin SentEval to evaluate the unsupervised represen-tation ability of pruned models because these tasksdirectly use text representations to compute the co-sine similarity without additional architecture andtraining. Sentence relatedness tasks are composedof six STS tasks (Agirre et al., 2012, 2013, 2014,2015; Cer et al., 2017) and SICK-R (Marelli et al.,2014). Since these tasks are unsupervised, we ex-clude L0 Norm and HISP, which need supervision.

We report the results on semantic relatednesstasks in Table 3. From this table, we have twoobservations: (1) Random pruning has similar per-formance to a full BERT. In some tasks, randomlypruned models are even better than full models.It indicates a full BERT cannot provide good textrepresentations for semantic relatedness althoughthere is rich language knowledge in the pre-trainedmodel. For these tasks, BERT might have manyunnecessary heads so even random pruning couldbring in minor performance improvements. (2)SMP significantly improves the performance ofthe pruned BERT, which indicates SMP makes

2 4 6 8 10The number of pruned heads in each layer

81

82

83

84

Accu

racy

(%)

fine-tune: 83.85

SMPHISP-retrain

Figure 3: Evolution of accuracy on the MultiNLI-matched when heads are pruned from BERTBASE.

full use of important prior knowledge and helpspruned models provide informative representationsfor unsupervised tasks.

4.4 Effect on Fine-tuning Efficiency

In this subsection, we investigate the effect ofpruning on fine-tuning efficiency, which is the ad-vantage of single-shot pruning. Experiments areconducted on a machine equipped with Tesla P40GPUs. As shown in Table 4, pruning 50% of themodel’s heads speeds up fine-tuning by more than34% and reduce the memory overhead per instanceby around 30%. In this case, we train more in-stances simultaneous for the pruned model dueto the less memory overhead. Besides, the timeof running SMP is nearly 1/30 of the fine-tuningtime, which is negligible. According to the re-sults, single-shot pruning on deep Transformerscan make fine-tuning more efficient.

4.5 Influence of Pruning Ratio

In this subsection, we investigate the influenceof pruning ratio. We test pruned models on theMultiNLI-matched validation set. As shown inFigure 3, we observe that: (1) Pruning a small

SMP Model MNLI-(m/mm) QQP SST-2 CoLA STS-B MRPC RTE Average

SMP-LARGE 86.86/86.96 91.46 93.34 63.57 90.95 89.70 82.31 85.64SMP-BASE 86.72/86.63 91.43 93.23 64.79 91.00 89.46 83.39 85.83

Table 5: Results of transferability on seven tasks in GLUE (%). We prune BERTLARGE using two different SMPmodels. SMP-LARGE is the model trained on BERTLARGE while SMP-BASE is the model trained on BERTBASE.

(a) 0.989 (b) 0.994 (c) 0.027 (d) 0.003

Figure 4: Detection of the implicit pruning rules learned by SMP. Given input sentence “These are issues whichfurther studies may seek to address.”, we present four attention matrices with their corresponding scores as subtitles.The left two matrices’ scores are close to 1 while the right two matrices’ scores are close to 0.

number of unnecessary attention heads promotesthe performance by nearly 1%. (2) SMP is betterwhen the number of pruned heads is smaller than6 while HISP-retrain is better in the other cases.It indicates that pruning too many parameters be-fore fine-tuning will influence the performance ondownstream tasks where the process of retrainingcould save this degradation of performance.

4.6 Transferability

In this part, we evaluate the transferability of SMP.We consider two kinds of transferability, includingthe transferability to new Transformer encodersand new datasets.

Transferability to New Transformer En-coders. We use the SMP trained on BERTBASEto prune BERTLARGE. The results are shown inTable 5. We observe that the SMP trained onBERTBASE achieves comparable results to the SMPtrained on BERTLARGE when pruning BERTLARGE.It indicates that the attention patterns learned bySMP are general in the Transformer encoders withdifferent sizes.

Transferability to New Datasets. We chooseQNLI, which is a natural language inference (NLI)dataset, as the task. QNLI is used to validatewhether SMP can transfer to a new NLI dataset.Note that we use BERTBASE as the pre-trainedTransformer. As shown in Table 6, SMP increasesthe average performance of random pruning andachieves comparable result to HISP-retrain. Itdemonstrates that SMP captures general patternsin attention matrices, which can transfer to pruning

Method Fine-tune HISP-retrain Random SMP

Acc. 91.3 91.0 89.4 90.6

Table 6: Accuracy on QNLI.

pre-trained Transformers on other tasks.

4.7 Visualization

In this subsection, we investigate the implicit ruleslearned by SMP. We compute the attention matri-ces for a given sentence, and score each attentionmatrix. In Figure 4, we show four attention matri-ces. The first matrix seems to be a lower diagonalmatrix, which refers to the attention to previouswords. This attention head implicitly captures thesequential information of the sentence. The secondmatrix shows a strong relation between “address”and “issues”, which is long-term dependency. Inthe third matrix, every element is small. Therefore,this matrix does not bring any useful information.In the last matrix, although there are extremely highvalues, all tokens attend to the same token ([CLS]),which is not informative. SMP gives high scoresto the first two matrices and low scores to the lasttwo matrices, which shows the implicit pruningrules learned by SMP are consistent with humanintuition.

5 Conclusion and Future Work

In this work, we propose Single-Shot Meta-Pruningto reduce the computational overhead of both fine-tuning and inference when using deep pre-trained

Transformers. Specifically, SMP learns the implicitrules for pruning in terms of attention matricesand adaptively prunes unnecessary attention headsbefore fine-tuning. In our experiments, we findpruning 50% of attention heads with SMP has littleimpact on the performances on downstream tasks.What’s more, pruning a few unnecessary heads canfurther improve the model performance in somecases.

There are four important directions for futureresearch: (1) Explore task-aware pruning, such astaking the labels of instances into account. (2) Jointpruning in each layer to maintain more diversityin pruned models, such as limiting the number ofthe attention heads sharing similar patterns in eachlayer. (3) Discover more unnecessary structuresin Transformer, such as point-wise feed-forwardnetworks. (4) Apply implicit pruning rules to con-straining the pre-training procedure of Transform-ers, which guides pre-training models through themore efficient use of parameters and attention.

ReferencesEneko Agirre, Carmen Banea, Claire Cardie, Daniel

Cer, Mona Diab, Aitor Gonzalez-Agirre, WeiweiGuo, Inigo Lopez-Gazpio, Montse Maritxalar, RadaMihalcea, et al. 2015. Semeval-2015 task 2: Seman-tic textual similarity, english, spanish and pilot oninterpretability. In Proceedings of SemEval.

Eneko Agirre, Carmen Banea, Claire Cardie, DanielCer, Mona Diab, Aitor Gonzalez-Agirre, WeiweiGuo, Rada Mihalcea, German Rigau, and JanyceWiebe. 2014. Semeval-2014 task 10: Multilingualsemantic textual similarity. In Proceedings of Se-mEval.

Eneko Agirre, Daniel Cer, Mona Diab, Aitor Gonzalez-Agirre, and Weiwei Guo. 2013. Semeval-2013shared task: Semantic textual similarity. In Proceed-ings of SemEval.

Eneko Agirre, Mona Diab, Daniel Cer, and AitorGonzalez-Agirre. 2012. Semeval-2012 task 6: A pi-lot on semantic textual similarity. In Proceedings ofSemEval.

Tom B. Brown, Benjamin Mann, Nick Ryder, MelanieSubbiah, Jared Kaplan, Prafulla Dhariwal, ArvindNeelakantan, Pranav Shyam, Girish Sastry, AmandaAskell, Sandhini Agarwal, Ariel Herbert-Voss,Gretchen Krueger, Tom Henighan, Rewon Child,Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu,Clemens Winter, Christopher Hesse, Mark Chen,Eric Sigler, Mateusz Litwin, Scott Gray, BenjaminChess, Jack Clark, Christopher Berner, Sam Mc-Candlish, Alec Radford, Ilya Sutskever, and Dario

Amodei. 2020. Language models are few-shot learn-ers. arXiv preprint arXiv:2005.14165.

Daniel Cer, Mona Diab, Eneko Agirre, Inigo Lopez-Gazpio, and Lucia Specia. 2017. Semeval-2017task 1: Semantic textual similarity-multilingual andcross-lingual focused evaluation. arXiv preprintarXiv:1708.00055.

Daoyuan Chen, Yaliang Li, Minghui Qiu, Zhen Wang,Bofang Li, Bolin Ding, Hongbo Deng, Jun Huang,Wei Lin, and Jingren Zhou. 2020a. AdaBERT:Task-adaptive BERT compression with differen-tiable neural architecture search. arXiv preprintarXiv:2001.04246.

Tianlong Chen, Jonathan Frankle, Shiyu Chang, SijiaLiu, Yang Zhang, Zhangyang Wang, and MichaelCarbin. 2020b. The lottery ticket hypothesisfor pre-trained bert networks. arXiv preprintarXiv:2007.12223.

Alexis Conneau and Douwe Kiela. 2018. Senteval: Anevaluation toolkit for universal sentence representa-tions. In Proceedings of LREC.

Ido Dagan, Oren Glickman, and Bernardo Magnini.2006. The PASCAL recognising textual entailmentchallenge. In Machine learning challenges.

Tim Dettmers and Luke Zettlemoyer. 2019. Sparse net-works from scratch: Faster training without losingperformance. arXiv preprint arXiv:1907.04840.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, andKristina Toutanova. 2019. BERT: Pre-training ofdeep bidirectional transformers for language under-standing. In Proceedings of NAACL-HLT.

William B Dolan and Chris Brockett. 2005. Automati-cally constructing a corpus of sentential paraphrases.In Proceedings of IWP.

Angela Fan, Edouard Grave, and Armand Joulin. 2020.Reducing transformer depth on demand with struc-tured dropout. In Proceedings of ICLR.

Mitchell A Gordon, Kevin Duh, and Nicholas Andrews.2020. Compressing BERT: Studying the effects ofweight pruning on transfer learning. arXiv preprintarXiv:2002.08307.

Song Han, Jeff Pool, John Tran, and William Dally.2015. Learning both weights and connections forefficient neural network. In Proceedings of NIPS.

Eric Jang, Shixiang Gu, and Ben Poole. 2016. Categor-ical reparameterization with gumbel-softmax. arXivpreprint arXiv:1611.01144.

Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang,Xiao Chen, Linlin Li, Fang Wang, and QunLiu. 2019. TinyBERT: Distilling bert for nat-ural language understanding. arXiv preprintarXiv:1909.10351.

Olga Kovaleva, Alexey Romanov, Anna Rogers, andAnna Rumshisky. 2019. Revealing the dark secretsof BERT. In Proceedings of EMNLP-IJCNLP.

Solomon Kullback and Richard A Leibler. 1951. Oninformation and sufficiency. The annals of mathe-matical statistics.

Zhenzhong Lan, Mingda Chen, Sebastian Goodman,Kevin Gimpel, Piyush Sharma, and Radu Soricut.2020. ALBERT: A lite bert for self-supervised learn-ing of language representations. In Proceedings ofICLR.

Yann LeCun, John S Denker, and Sara A Solla. 1990.Optimal brain damage. In Proceeding of NIPS.

Namhoon Lee, Thalaiyasingam Ajanthan, StephenGould, and Philip HS Torr. 2019a. A signal prop-agation perspective for pruning neural networks atinitialization. arXiv preprint arXiv:1906.06307.

Namhoon Lee, Thalaiyasingam Ajanthan, and PhilipTorr. 2019b. SNIP: Single-shot network pruningbased on connection sensitivity. In Proceedings ofICLR.

Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin,Kurt Keutzer, Dan Klein, and Joseph E Gonzalez.2020. Train large, then compress: Rethinking modelsize for efficient training and inference of transform-ers. arXiv preprint arXiv:2002.11794.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Man-dar Joshi, Danqi Chen, Omer Levy, Mike Lewis,Luke Zettlemoyer, and Veselin Stoyanov. 2019.RoBERTa: A robustly optimized BERT pretrainingapproach. arXiv preprint arXiv:1907.11692.

Christos Louizos, Max Welling, and Diederik P.Kingma. 2018. Learning sparse neural networksthrough L0 regularization. In Proceedings of ICLR.

Xiaofei Ma, Peng Xu, Zhiguo Wang, Ramesh Nallap-ati, and Bing Xiang. 2019. Universal text represen-tation from bert: An empirical study. arXiv preprintarXiv:1910.07973.

Marco Marelli, Luisa Bentivogli, Marco Baroni, Raf-faella Bernardi, Stefano Menini, and Roberto Zam-parelli. 2014. Semeval-2014 task 1: Evaluation ofcompositional distributional semantic models on fullsentences through semantic relatedness and textualentailment. In Proceedings of SemEval.

JS McCarley. 2019. Pruning a BERT-basedquestion answering model. arXiv preprintarXiv:1910.06360.

Paul Michel, Omer Levy, and Graham Neubig. 2019.Are sixteen heads really better than one? In Pro-ceedings of NIPS.

Alec Radford, Karthik Narasimhan, Tim Salimans, andIlya Sutskever. 2018. Improving language under-standing by generative pre-training. In Proceedingsof OpenAI Technical report.

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, andPercy Liang. 2016. SQuAD: 100,000+ questions formachine comprehension of text. In Proceedings ofEMNLP.

Hassan Sajjad, Fahim Dalvi, Nadir Durrani, andPreslav Nakov. 2020. Poor man’s BERT: Smallerand faster transformer models. In Proceedings ofACL.

Victor Sanh, Lysandre Debut, Julien Chaumond, andThomas Wolf. 2019. DistilBERT, a distilled versionof BERT: smaller, faster, cheaper and lighter. arXivpreprint arXiv:1910.01108.

Richard Socher, Alex Perelygin, Jean Wu, JasonChuang, Christopher D Manning, Andrew Ng, andChristopher Potts. 2013. Recursive deep modelsfor semantic compositionality over a sentiment tree-bank. In Proceedings of EMNLP.

Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. 2019.Patient knowledge distillation for bert model com-pression. In Proceedings of EMNLP.

Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu,Yiming Yang, and Denny Zhou. 2020. MobileBERT:a compact task-agnostic bert for resource-limited de-vices. In Proceedings of ACL.

Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, OlgaVechtomova, and Jimmy Lin. 2019. Distilling task-specific knowledge from BERT into simple neuralnetworks. arXiv preprint arXiv:1903.12136.

Iulia Turc, Ming-Wei Chang, Kenton Lee, and KristinaToutanova. 2019. Well-read students learn better:The impact of student initialization on knowledgedistillation. arXiv preprint arXiv:1908.08962.

Elena Voita, David Talbot, Fedor Moiseev, Rico Sen-nrich, and Ivan Titov. 2019. Analyzing multi-headself-attention: Specialized heads do the heavy lift-ing, the rest can be pruned. In Proceedings of ACL.

Alex Wang, Amanpreet Singh, Julian Michael, FelixHill, Omer Levy, and Samuel R. Bowman. 2019a.GLUE: A multi-task benchmark and analysis plat-form for natural language understanding. In Pro-ceedings of ICLR.

Ziheng Wang, Jeremy Wohlwend, and Tao Lei. 2019b.Structured pruning of large language models. arXivpreprint arXiv:1910.04732.

Alex Warstadt, Amanpreet Singh, and Samuel R. Bow-man. 2018. Neural network acceptability judgments.arXiv preprint 1805.12471.

Adina Williams, Nikita Nangia, and Samuel R. Bow-man. 2018. A broad-coverage challenge corpus forsentence understanding through inference. In Pro-ceedings of NAACL-HLT.

Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Car-bonell, Russ R Salakhutdinov, and Quoc V Le.2019. XLNet: Generalized autoregressive pretrain-ing for language understanding. In Proceedings ofNeurIPS.

Ofir Zafrir, Guy Boudoukh, Peter Izsak, and MosheWasserblat. 2019. Q8BERT: Quantized 8bit BERT.arXiv preprint arXiv:1910.06188.