A Fast Post-Training Pruning Framework for Transformers
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Pruning is an effective way to reduce the huge inference cost of large Transformermodels.
However, prior work on model pruning requires retraining the model.
This can add high cost and complexity to model deployment, making it difﬁcult touse in many practical situations.
To address this, we propose a fast post-trainingpruning framework for Transformers that does not require any retraining.
Givena resource constraint and a sample dataset, our framework automatically prunesthe Transformer model using structured sparsity methods.
To retain high accuracywithout retraining, we introduce three novel techniques: (i) a lightweight masksearch algorithm that ﬁnds which heads and ﬁlters to prune based on the Fisherinformation; (ii) mask rearrangement that complements the search algorithm; and(iii) mask tuning that reconstructs the output activations for each layer.
We applyour method to BERT BASE and DistilBERT, and we evaluate its effectiveness on
GLUE and SQuAD benchmarks.
Our framework achieves up to 2.0 ⇥ reduction in
FLOPs and 1.56 ⇥ speedup in inference latency, while maintaining < 1% loss inaccuracy.
Importantly, our framework prunes Transformers in less than 3 minuteson a single GPU, which is over two orders of magnitude faster than existing pruningapproaches that retrain.
Our code will be publicly available at GitHub.
Introduction
In recent years, Transformer [74] has become a de facto standard model architecture in Natural
Language Processing [11, 44, 4] and it is becoming common in Computer Vision [13, 73, 46] and
Speech Recognition [2, 24, 7] as well.
However, efﬁcient deployment of Transformer architectureshas been challenging due to their large model size and high latency.
To address this, structuredpruning has become a promising technique for efﬁcient Transformer deployment.
While prior work on pruning Transformers substantially reduces inference time, it is often difﬁcult touse in practice for several reasons.
First, previous approaches require retraining the pruned modeland/or jointly learning the pruning conﬁgurations during training.
For instance, Block Movement
Pruning [36] increases the training time by 10 ⇥ for pruning.
Such additional training adds signiﬁcantcomputational overhead, given the large training cost of Transformers.
Second, previous methodsadd many moving parts to the model deployment process.
Pruning pipelines are often complex andinvolve additional hyperparameter tuning.
For instance, ROSITA [45] uses a three-stage knowledge distillation [22] with sophisticated pruning schedules.
Such techniques demand signiﬁcantengineering efforts for reproducing and debugging, which impedes their adoption in productionpipelines.
Third, these previous methods do not directly adapt to the users’ constraints.
They eitherrely on vague regularization hyperparameters or ﬁxed architectures selected independently of the user settings.
This can result in sub-optimal models not tailored for the given constraint and hardware
[Figure]
To address these limitations, we propose a fast post-training pruning framework for Transformersthat does not require any retraining of the models.
As illustrated in Figure 1, our framework takesas input a Transformer model, a sample dataset, and a FLOPs/latency constraint.
It then outputs apruned Transformer model that can be deployed immediately.
By avoiding expensive retraining, theend-to-end compression pipeline can be extremely fast and simpliﬁed, typically in a few minutes,without any user interventions that complicate the whole process.
Indeed, post-training compression has been widely studied for quantization, and it gained considerable attention in both academia and industry [3, 91, 26].
Although quantization-aware trainingmethods achieve higher compression rates in general, post-training quantization (PTQ) has oftenbeen more preferred in practice due to its retraining-free advantage.
PTQ allows quantization tohappen seamlessly when deploying models through the various frameworks such as TensorRT [53],
TFLite [17], and OpenVINO [28].
Similar to these PTQ frameworks, our framework provides anout-of-the-box tool that enables pruning of Transformers without engineering efforts.
Our contributions can be summarized as follows:
• We propose a novel post-training pruning framework for Transformers that does not require model
[Figure]
retraining.
To retain accuracy without retraining, our framework consists of three stages of: (i)the mask search process guided by the Fisher information matrix to ﬁnd which heads/ﬁlters toprune (Section 4.1); (ii) the mask rearrangement process that rearranges the pruned heads/ﬁlters bycapturing intra-layer interactions (Section 4.2); and (iii) the mask tuning process that adjusts the
mask variables to ensure that the output signal is recovered for each layer (Section 4.3).
• We extensively test our framework by applying it to BERT BASE and DistilBERT on GLUE andSQuAD tasks (Section 5.2).
With 1 % of accuracy drop, our framework reduces 30–50 % of theoriginal FLOPs (Figure 4), resulting in up to 1.56 ⇥ speedup on an NVIDIA V100 GPU (Table 3).
• We show that our method achieves comparable or even better FLOPs-accuracy trade-off than prior
structured pruning methods without retraining (Section 5.3, Figure 5).
Our end-to-end pruningﬁnishes in only 39 and 135 seconds on average for GLUE and SQuAD (Section 5.4, Table 5),which is over 100 ⇥ faster than the retraining methods.
Related Works
Efﬁcient Transformers.
In order to improve the inference speed and reduce the memory footprint of
Transformers, multiple different approaches have been proposed.
These can be broadly categorizedas follows: (i) efﬁcient architecture design [37, 34, 69, 79, 27, 85]; (ii) hardware-software codesign [18, 78, 19, 70]; (iii) knowledge distillation [60, 30, 80, 68]; (iv) quantization [90, 63, 89, 32];(v) neural architecture search [77, 86, 65, 64, 6, 88]; and (vi) pruning.
In this paper, we only focus onpruning and brieﬂy discuss the related works.
Transformers Pruning.
Pruning has been a popular choice for reducing unimportant weights in
Transformers.
Pruning can be largely categorized into unstructured and structured pruning.
Forunstructured pruning, magnitude-based pruning [16], the lottery-ticket hypothesis [15, 55, 8, 9], movement pruning [61] have been explored for Transformers.
While these methods compress themodel size, commodity hardware cannot take advantage of the unstructured sparsity for speedup.
For this reason, a number of structured pruning methods have been introduced to remove structuredsets of parameters.
For example, [49, 75] drop attention heads in multi-head attention layers.
Anotherthread of work prunes entire Transformer blocks based on a simple heuristic [59] or an adaptivetraining using a layer-wise dropout method [14].
Relatedly, [81] structurally prunes weight matricesvia low-rank factorization and l 0 regularization; [31, 42] attempt to jointly prune attention heads andﬁlters of weight matrices; and [47, 23] take a step further by dynamically determining the pruningconﬁgurations at run time.
Recent block pruning schemes chunk weight matrices into blocks andprune them based on group Lasso [41], adaptive regularization [87], and movement pruning [36].
While the structured pruning techniques have achieved high compression rates and speedups, theimprovement largely attributes to retraining of the models during or after pruning, often combinedwith knowledge distillation.
However, model retraining introduces several practical problems.
Mostnotably, retraining in general increases training time signiﬁcantly.
For instance, Block Movement
Pruning (BMP) [36], a state-of-the-art structured Transformer pruning method, requires 10 ⇥ longertraining time than normal ﬁne-tuning.
In addition, the extra hyperparameters introduced by thepruning methods add additional moving parts and lead to increased training cost.
Post-training Model Compression.
Post-training compression methods have been widely studiedin quantization.
These methods, categorized as post-training quantization (PTQ), perform quantization without any retraining, thereby minimizing the training cost and user intervention.
Multipleapproaches have been proposed for PTQ in order to mitigate the accuracy degradation without retraining, including analytic computation of optimal clipping ranges [3], outlier channel splitting [91],and adaptive rounding methods [51, 26].
Although not as much as for quantization, post-training schemes have also been explored for unstructured [25, 38] and structured pruning.
For structured pruning, Srinivas and Babu [67] propose todetect and merge similar convolution ﬁlters iteratively, and Neuron Merging [33] compensates for theinformation loss due to pruning by merging similar neurons based on cosine distance.
While thesemethods are effective for CNNs, their applicability is limited to models with simple architectures,and they rely on the characteristics of ReLU nonlinearity [25, 33].
Hence, the prior methods cannotbe extended to the general Transformer architectures.
Overview
3.1 Background
[Figure]
Transformer Architecture.
In this paper, we focus on the pruning of encoder-based Transformer [74]models, especially the BERT [11] architecture family.
BERT is a stack of homogeneous Transformerencoder blocks, each of which consists of a multi-head attention (MHA) layer followed by a pointwise Feed-Forward Network (FFN) layer.
Speciﬁcally, an MHA layer consists of H independentlyparameterized attention heads:
[Equation]
where Att is a dot product attention head, LN is layer normalization, and x is the input sequence.
Theoutput of the MHA layer is then fed into the FFN layer, which consists of N ﬁlters:
[Equation]
where W (1) , W (2) , b (1) and b (2) are the FFN parameters, and  is the activation function, typically
GELU [21].
Note that ( H , N ) is (12, 3072) for BERT BASE , and (16, 4096) for BERT LARGE .
We alsodenote L as the number of Transformer layers.
Granularity of Pruning.
Our framework considers the structured pruning of both heads in MHA andﬁlters in FFN layers.
We do not prune the embedding and the ﬁnal classiﬁer, as computation of thoselayers takes a negligible portion of the total inference latency.
Since our pruning framework alwaysproduces a smaller dense architecture, the model can be readily accelerated without the need forspecialized hardware logic, which is often required to gain latency speedup for unstructured sparsity.
d   :
1+,-    : 1// (b) Mask Search (c) Mask Rearrangement       : 1+,- : 1//    : 1+,-    : 1// (a) Initial Mask (d) Mask Tuning    : 1+,- : 1// 1.
Sum × H d
Attn W (!)
× N
N MHA Mul
Mul : 1+,-
W (')
d x x $%& x (() d : 1//
Notations.
For mathematical simplicity, we introduce mask variables associated with the outputs ofheads and ﬁlters:
[Equation]
where m MHA l 2 R H and m FFN l 2 R N are the mask variables for MHA and FFN in the l -th layer,respectively, and m MHA l,i and m FFN l,i are their i -th elements.
Furthermore,  denotes Hadamard product.
Originally, the mask variables are initialized to 1.
Zeroing out a mask variable is equivalent to pruninga head/ﬁlter associated with it.
That is, setting m MHA l,i and m MHA l,i as zero is equivalent to pruning thei -th head and ﬁlter, respectively.
Overall, there are HL head mask variables and NL ﬁlter mask variables, summing up to ( H + N ) Lnumber of total mask variables in a Transformer model.
To simplify notations, we additionally deﬁnem MHA 2 R HL , m FFN 2 R NL , and m 2 R
( H + N ) L as the ﬂattened vectors of the head, ﬁlter, and totalmask variables, respectively, across all layers.
In what follows, we discuss how to ﬁnd the optimalsparse masks under a given cost constraint and how to adjust their values to recover accuracy.
3.2 Framework Overview
Figure 1(b) and Figure 2 illustrate the overview of our framework.
Inputs.
Our framework has 3 inputs: a Transformer model; a sample dataset; and a resource constraint.
The input Transformer model should contain weights ﬁne-tuned for a downstream task.
The sampledataset is a small partition of the training dataset (typically 1–2K examples) of the downstream task.
The resource constraint can be given either as the number of ﬂoating point operations (FLOPs) or asan actual latency on target hardware.
In the later case, we further assume that a latency lookup tablefor the target hardware is provided.
Compression Pipeline.
As illustrated in Figure 2, our framework consists of 3 stages: Fisherbased mask search; Fisher-based mask rearrangement; and mask tuning.
During the Fisher-basedmask search stage (Section 4.1), we search for a binary mask applied to the heads and ﬁlters byincorporating the Fisher information of the mask variables.
Intuitively, the mask variables withrelatively higher Fisher information are considered more important, and they should be less likelyto be pruned [39, 50, 43].
As ﬁnding the optimal mask that minimizes the Fisher information lossis intractable using the full Fisher matrix, we propose a lightweight search algorithm that ﬁnds theoptimal mask under reasonable approximations.
Second, in the Fisher-based mask rearrangementstage (Section 4.2), the framework modiﬁes the searched mask patterns in a layer-wise manner tobetter take into account the intra-layer interactions of the mask variables.
Lastly, in the mask tuningstage (Section 4.3), the framework tunes the nonzero mask variables to recover the accuracy drop byreconstructing the layer-wise output signal.
Methodology
We pose the pruning problem as ﬁnding a binary mask to zero out a particular head or ﬁlter.
Afterﬁnding the initial set of mask variables, we then allow adaptation/tuning of the remaining maskvariables to avoid changes in the output norm of a layer.
This process is done without any retraining.
Note that the number of the mask variables is much less than the number of the parameters in a
Transformer (e.g., 37K vs. 110M in case of BERT BASE ).
This allows the framework to use only a small number of examples without overﬁtting to the sample dataset, and thus to be extremely faster thanthe retraining-based pruning methods which typically use the entire dataset.
As the framework keepsthe model “as is” and only decides the mask variables, we henceforth regard the model parameters asconstants and consider the mask variables as the only parameters for our pruning problem.
Problem Formulation.
We formulate Transformer pruning as a constrained optimization problemon the mask m:
[Equation]
where L denotes the loss function, Cost is the FLOPs/latency of the architecture pruned by the mask,and C is the given FLOPs/latency constraint.
Unfortunately, such a problem is generally intractableas Cost is usually a function of l 0 -norm of the mask m, which is non-differentiable.
Thus, in whatfollows, we introduce several assumptions and approximations to simplify the problem.
We start by approximating the loss function using the second-order Taylor expansion around theinitial mask 1 :
[Equation]
where g = E [ @@ m L ( 1 )] and H =
E [ @ 2 @ m 2 L ( 1 )] .
Eq.
3 is deduced from an assumption that the modelhas converged to a local minima, where the gradient term is close to 0 [39].
As L ( 1 ) is a constant, wecan rewrite the optimization objective as follows:
[Equation] [Equation]
Eq.
4 shows that the optimal mask is determined by the Hessian of the loss with respect to the maskvariables.
Since forming the exact Hessian matrix explicitly is infeasible, we approximate the Hessian
H with the empirical Fisher information matrix I of the mask variables:
[Equation]
where D is the sample dataset and ( x, y ) is a tuple of an input example and its label.
4.1 Fisher-Based Mask Search
Diagonal Approximation of Fisher Information Matrix.
It is intractable to solve the optimizationobjective in Eq.
4 using the full empirical Fisher information matrix I .
Thus, we ﬁrst make a simpleassumption that I is diagonal .
This further simpliﬁes Eq.
4 as follows:
[Equation]
Since we restrict the possible mask values to either 0 or 1, the following can be derived from Eq. 6:
[Equation]
Z
[Equation]
We can interpret the diagonal element of I as the importance score of the head/ﬁlter associated witheach mask variable, and Eq.
7 as a process of minimizing the total importance scores of the prunedheads and ﬁlters.
Such an importance score has also been introduced in [72, 50] to guide pruning.
Solving FLOPs-constrained Problem.
We need to solve Eq.
7 given a cost constraint.
For a giventarget FLOP cost, denoted by C, we can formulate the binary mask search problem as follows:
[Equation]
[Equation]
where F head 2 R and F ﬁlter 2 R are the FLOPs for computing a head and a ﬁlter, respectively.
Notethat the number of FLOPs of a head/ﬁlter is constant across all layers.
While such an optimization
Input: FLOPs constraint C , diagonal Fisher information matrix I 1: for n = 0 to HL do
=
HI = heads 4: f =
F c
5: k
= NL f .
FI =
2 least important ﬁlters 7:
S [ n
=
P i
[ FI I ii
R n
= (
, FI ) 8: end for 9: n ⇤ = arg min n S [ n ]
10: HI ⇤ , FI ⇤ =
R n ⇤ ]
12: m MHA
Output:
) problem can be generally solved by a knapsack algorithm [1, 62], the following observations allow asimpler and faster solution: (1) having more heads and ﬁlters unpruned always optimizes Eq.
8 sincethe diagonal elements of I are non-negative; and (2) if a certain number of heads needs to be pruned,they should be the ones with the lowest importance scores because each head accounts for the sameamount of FLOPs.
The same statement also holds for pruning ﬁlters.
These lead to our mask searchalgorithm described in Algorithm 1.
Algorithm 1 partitions the solution space by the number of remaining heads in the pruned architecture( n in line 1).
For each n , by observation (1), the number of remaining neurons should be as line 4.
Then by observation (2) the heads/ﬁlters with the lowest important scores are selected to be pruned.
S [ n ] is the evaluation of Eq.
8 with this mask.
When the loop terminates, the output is the mask withthe smallest S [ n ] .
In Section A.1, we prove that the output mask m ⇤ of Algorithm 1 is optimal.
Thatis, any other mask m satisfying the given FLOPs constraint will have a higher loss:
[Equation]
Solving Latency-constrained Problem.
If the cost constraint is given in terms of an actual latencyon target hardware, we have a new cost constraint formula in the optimization target in Eq. 8:
[Equation]
where the function LAT indicates the latency of a MHA/FFN layer after pruning.
We assume that alatency lookup table on the target hardware is provided so that evaluating LAT takes negligible time.
Unfortunately, the latency constraint makes the problem more challenging such that it cannot besolved by directly applying Algorithm 1.
This is because LAT is not linear to the number of remainingheads or ﬁlters after pruning [56], as shown in Figure 3 (Left).
We can interpret this as follows: (1)with a sufﬁcient number of heads/ﬁlters in a layer, the hardware resources such as parallel corescan be fully utilized, resulting in latency roughly proportional to the number of heads/ﬁlters; and(2) otherwise, the hardware is underutilized and a constant overhead dominates the latency [35, 48].
Thus, pruning more heads/ﬁlters below a certain threshold does not translate into actual speedup.
[Figure]
[Equation]
[Equation]
= )
Based on the above analysis, we approximate LAT as a piece-wise linearfunction as in Figure 3 (Right) such that
LAT ( m l ) is 0 if || m l ||
= 0 , c
if 0 <|| m l || 0  T , and a ( || m l ||
T )
+ c if|| m l || 0 > T , where c 2 R is the con- 217 stant overhead, T 2 N is the threshold 218 number of heads/ﬁlters that the latencystarts to be linear, and a 2 R is theslope of the linear part.
This can be easily obtained by ﬁtting the actual latencyin the lookup table with the minimummean squared error.
This LAT approximation allows us to extend Algorithm 1 to solving the problems with latencyconstraints.
The core idea is to consider separately the constant part of LAT and the linear part of
LAT; after handling the constant part, we can apply the Algorithm 1 to the linear part.
The detailedmodiﬁcation to Algorithm 1 is described in Section A.2.
4.2 Fisher-based Mask Rearrangement
Block Diagonal Approximation of Fisher Information Matrix.
Although it simpliﬁes the problem,the diagonal assumption in Section 4.1 alone might not ﬁnd the best solution, as it does not take intoaccount the interactions between different mask variables.
We can better capture the interactionsby using a block diagonal approximation to the Fisher operator, where a block corresponds to a
MHA layer or a FFN layer.
However, the block diagonal approximation results in an intractableoptimization problem over the binary mask.
To alleviate this, we use the results from the previousstep to warm start the optimization problem with the block-diagonal approximation.
That is, giventhe mask m ⇤ obtained in Section 4.1, we constrain ||
m l || 0 to be equal to || m
⇤ l || 0 for all layers l .
Given the two assumptions, that (i) there is no interaction between the mask variables in differentlayers (i.e., the block diagonal approximation), and (ii) the number of pruned heads/ﬁlters are predetermined for each layer (i.e., warm-start), Eq.
4 breaks down to a set of layer-wise optimizationproblems, as follows based on the derivation in Section A.3:
[Equation]
(11) where I l is the l -th diagonal block of I .
This can be approximately solved by greedily exploring thepossible binary masks for m l since || m l || 0 is ﬁxed.
Because this process does not change the number 243 of heads/ﬁlters in each layer, the searched mask leads to the same FLOPs/latency as the one obtainedin Section 4.1.
In effect, this process rearranges the binary mask variables of each layer to ﬁnd abetter arrangement for pruning locations.
4.3 Mask Tuning
In the previous two stages, the possible mask values are restricted to either 0 or 1 in order to simplifythe search process.
In this stage, we further relax this restriction.
The nonzero variables in the maskˆ m from Section 4.2 are tuned to any real values such that the pruned model recovers its accuracy.
Layer-wise Reconstruction via Linear Least Squares.
We tune the mask variables toward minimizing the layer-wise reconstruction error , similarly to [20].
For each layer, we reconstruct the outputactivation of the original model with the remaining heads/ﬁlters in the pruned model.
This can beformally written as follows:
[Equation]
where layer is the either MHA or FFN, and x and x 0 are the inputs to the layer of the prunedmodel and the original model, respectively.
Note that this stage does not incur any change in model
FLOPs/latency, as we only tune the non-zero valued mask variables.
We show in Section A.4 that
Eq.
12 can be reduced to a linear least squares problem of arg min m l || Am l  b || 22 , where the matrix
A denotes a collection of the output activations of the model pruned by the binary mask and thevector b is the output activation of the original model.
Because the size of the matrix A can be large, our framework uses the LSMR solver in CuPy [52] tosolve the linear least squares problem.
For the regularization hyperparameter (i.e., damp ) of LSMR,we ﬁx its value to 1.
Then, to increase stability, we restrict the acceptable range of the tuned maskvariables to [-10, 10].
When the solver ﬁnds a layer mask that exceeds this range, we discard the maskfor that layer and stop mask tuning.
While the use of LSMR solver involves two hyperparameters(i.e.
, damp and the acceptable range), we empirically ﬁnd that these need not be tuned for differenttasks and models.
In all of our experiments, we used the ﬁxed hyperparameter values.
Evaluation
5.1 Experimental Setup
Our framework is implemented on top of PyTorch [54] and HuggingFace Transformers [84] library.
We evaluate the effectiveness of our approach using BERT BASE [11] and DistilBERT [60] on
GLUE [76] and SQuAD [58, 57] benchmarks.
We use 2K examples from the training sets for pruning,and we evaluate the resulting models on the development sets.
All of the results are averaged over theruns with 10 different seeds.
More details on the experimental setup can be found in Section A.5.
Note that these results can be achieved in only 39 and 135 seconds for GLUE and SQuAD benchmarks respectively, on a single GPU system, as described in Table 5 (in Section A.11).
5.2 Performance Evaluation
FLOPs.
Figure 4 shows the accuracy of BERT BASE and DistilBERT with different FLOPs constraints on GLUE and SQuAD datasets.
As can be seen in the plots, with only 1% of accuracydrop, BERT BASE achieves 60–70% of the original FLOPs for all tasks.
DistilBERT also shows asimilar pattern and shows 50% FLOPs reduction (in STS-B and MRPC) even though it is already acompressed architecture.
Results with larger sample datasets are provided in Section A.6.
Latency.
We further measure the latency on real hardware by pruning BERT BASE with latencyconstraints and deploying the resulting models on an NVIDIA V100 GPU.
Table 3 (in Section A.7)lists the latency speedup with maximum accuracy drop of 1% for GLUE and SQuAD datasets.
Withbatch size of 256, we achieve speedup of 1.47 ⇥ on average and up to 1.56 ⇥ .
5.3 Comparison with the Prior Methods
FLOPs and Accuracy Comparison.
Here, we compare our method with the prior structured pruningmethods for Transformers including Flop [81], SLIP [71], Sajjad et al.
[59], DynaBERT [23],
EBERT [47], and Block Movement Pruning (BMP) [36] by the FLOPs-accuracy trade-off of BERT BASE on GLUE tasks.
We use the results without knowledge distillation and data augmentationreported in each paper.
Since the baseline accuracy differs slightly from paper to paper, we comparethe amount of the accuracy drop from the baseline instead of the absolute accuracy.
The results areplotted as Figure 5.
We include the comparison details and full table in Section A.8.
Interestingly, our method exhibits comparable or sometimes better results than the prior methodswithout any model retraining and with substantially lower pruning costs.
This empirically demonstrates that retraining and a complex pruning pipeline are not necessary for moderate level of pruningof Transformers.
For high sparsity, we ﬁnd that our framework with retraining works comparably toor better than the prior methods at the same pruning cost (See Section A.9).
[Table]
[Table]
[Table] [Table] [Figure] [Table]
[Table]
Retraining Cost.
We select DynaBERT [23], EBERT [47]
[Table]
and BMP [36] that achieve comparablygood accuracy in Figure 5, and we systematically analyze their end-to-end retraining costs on the MNLI dataset.
As shown in Table 1, these methods require 5  17 hours of retraining.
On the other hand, our method ﬁnishes in less than a minute, which is 2  3 orders of magnitude faster.
We alsohighlight that this training latency only accounts for a single hyperparameter, and the entire costshould be multiplied by the size of the hyperparameter space.
While the prior methods rely on anumber of hyperparameters, ours introduce only two hyperparameters (in Section 4.3) which we ﬁxfor all experiments.
See Section A.10 for more details.
Unlike other methods (including ours) that take as input the ﬁne-tuned models, Sajjad et al.
[59] startsfrom the pre-trained model, heuristically drops the top layers, and then ﬁne-tunes on downstreamtasks.
Therefore, one can regard it as a zero-cost pruning method.
However, as can be seen in
Figure 5, the resulting models suffer from large accuracy degradation in most cases.
5.4 Discussion
Ablation Studies.
Table 2 lists an ablation of mask rearrangement (Section 4.2) and tuning (Section 4.3) stages for pruned BERT BASE with 60% of FLOPs.
We ﬁnd that both stages help recover thebaseline accuracy, and that mask tuning is in particular critical, recovering up to 2.88% accuracy.
Tofurther investigate the importance of mask search and rearrangement, we compare the retraining-freeperformance of the binary masks obtained by our method and other pruning criteria: weight magnitude and gradient-based method used in DynaBERT.
We uniformly pruned the layers using the twocriteria with different width multipliers.
Figure 6 shows that the two methods signiﬁcantly degradethe accuracy under non-trivial sparsity.
Even with mask tuning, the accuracy is not fully recovered.
The results demonstrate that our mask search and re-arrangement are necessary to get optimal binarymasks, and that mask tuning is only effective when the mask mostly preserves the accuracy.
Time Breakdown.
We break down our pruning pipeline into 4 parts—gradient computation, masksearch, rearrangement, and tuning—and we measure the latency for each stage as Table 5 (Section A.11).
For GLUE and SQuAD tasks, our framework ﬁnishes in 39 and 135 seconds on average.
Conclusion
In this work, we have proposed a novel post-training pruning framework for Transformers that doesnot require model retraining.
The core of our framework is to the three-stage decomposition of thepruning process.
It uses a fast Fisher-based mask search algorithm to decide which heads/ﬁlters toprune, rearranges the pruned heads/ﬁlters, and tunes the mask variables to recover the output signalfor each layer.
We empirically evaluate our framework using BERT BASE and DistilBERT, whereour pruning method achieves up to 50% FLOPs reduction with only 1% accuracy degradation on
GLUE and SQuAD datasets.
This results in up to 1.56 ⇥ latency speedup on an NVIDIA V100 GPU.
End-to-end pruning only needs 39 and 135 seconds for GLUE and SQuAD, which is 2  3 ordersof magnitude faster than the prior methods.
Overall, our method shows comparable or even bettercompression performance, as compared to the prior retraning-based methods.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes](c)
Did you discuss any potential negative societal impacts of your work?
[No](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
Yes
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] The notationsare deﬁned in Section 3.1 and our problem is formally described in Section 4.(b)
Did you include complete proofs of all theoretical results?
[Yes] Section A.1 containsthe proof of the optimiality of Algorithm 1.
Section A.3 and Section A.4 describe thederivation processes omitted in the paper.3.
If you ran experiments...(
a ) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Yes ] Our code isincluded in the supplementary material
b
) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] Our method has only two hyperparameters which were ﬁxed inall of our experiments (See Section 4.3).
For details on the datasets, see Section 5.1.
Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?
[No] Because our method is training-free, the results are very
stable across seeds.
For simplicity, we omitted the error
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] For all experiments, we used an
AWS p3.2xlarge instance which has 1 NVIDIA V100 GPU.
See Section 5.1 for details.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We cited the
HuggingFace transformers library, from which we got the pre-trained Transformer
weights.
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]