Towards Lossless Head Pruning through Automatic Peer Distillation for Language Models

1 Introduction

Transformer-based language models [Devlin et al., 2018; Yang et al., 2019] have been proven to be highly effective in learning universal language representations and applicable to downstream tasks with slight fine-tuning, as opposed to other machine learning methods [Hochreiter and Schmidhuber, 1997; He et al., 2016; Zhou et al., 2020]. These models tend to suffer from high computational cost and memory usage due to their large model size. To downsize the transformer models, head pruning, a sub-category of weight pruning technique has been widely investigated [Michel et al., 2019; Kovaleva et al., 2019; Voita et al., 2019; Lee et al., 2020].

∗Z. Wang is now affiliated with Walmart Global Tech, Sunnyvale, CA

Head pruning usually requires iterative but independent sparsification and retraining operations, which usually results in a significant drop in model accuracy. Specifically, during sparsification, head masks are updated to select redundant heads for pruning. During retraining, the retained attention heads are updated to recover the previous model accuracy [Michel et al., 2019; Voita et al., 2019]. We observe that such a workflow usually leads to a significant drop in model accuracy. The reason behind the drop in model accuracy is the "knowledge" loss from the pruned attention heads.

In this paper, we propose a counter-traditional pruning strategy, peer distillation, and reveal that the knowledge learned in advance but discarded from the pruned attention heads plays a crucial role in maintaining model accuracy. The pruned attention heads contain useful grammatical and semantic information to maintain model performance and provide valuable guidance for the training of the retained heads. We also develop an automatic framework to locate the to-be-pruned attention heads in each layer, freeing the timeconsuming human labor in tuning hyperparameters. Figure 1 illustrates the "peer distillation" process, where we recycle the knowledge from pruned heads to retained heads. To our best knowledge, peer distillation is the first attempt to recycle pruned knowledge and mitigate accuracy degradation. Benefiting from the recycled knowledge, the accuracy of the model is better maintained at higher model sparsity. Furthermore, peer distillation accelerates pruning convergence by pruning more attention heads while maintaining model accuracy.

Our contributions are as follows:

• We provide an automatic head pruning approach that achieves the best accuracy-sparsity tradeoff by reformulating the head pruning optimization. Specifically, we define a straightforward gating function for head pruning and propose a softplus STE function to address the mixed integer back propagation. Additionally, we provide head pruning visualizations to show the variable choice of model substructures during pruning.
• We propose the peer distillation strategy to recycle the discarded knowledge from the pruned heads. Specifically, we distill the pruned heads into the retained heads for each transformer layer. Since the pruned head maintains important "knowledge" to preserve the model accuracy, this strategy helps the retained head absorb this "knowledge" and thus be more capable of maintaining model accuracy.

Experimental results on nine GLUE benchmark tasks [Wang et al., 2018] show that we achieve high compression rates with zero or minor accuracy degradation. The pruned models outperform the original ones with only 41.71% heads left on average. In extreme cases, our proposed method can prune 99.3% (143 of the 144) heads without accuracy degradation on WNLI dataset. On average, we could prune 79.74% heads with only 0.82% accuracy degradation. Compared with other differentiable head pruning, e.g. [Xiao et al., 2019], peer distillation increases the head sparsity by 22.17% with the same model accuracy on MRPC dataset. The method provides an efficient tool to analyze and reduce the redundancy of multi-heads and is suitable for other attention-based models.

2 RelatedWorks and Background

Related Works. The redundancy of the multi-head mechanism has been discovered and investigated. [Michel et al., 2019; Kovaleva et al., 2019] first discusses duplicated attention head patterns. [Raganato et al., 2020] presents that models with fixed single attention head for each layer would nicely preserve model accuracy. [An et al., 2020] analyzes head redundancy from a Bayesian perspective and explains the causes of such redundancy.

Different attention head pruning algorithms are developed. [Michel et al., 2019] prunes attention head greedily based on predefined head importance metric but the pruned heads can never be recovered during training. [Kovaleva et al., 2019] shows the attention head redundancy and manually disables attention heads to improve model performance. [Voita et al., 2019] employs Gumbel softmax to relax the head pruning problem to be a differentiable subnetwork searching problem but more experiments and discussion are expected to prove its effectiveness. [Lee et al., 2020] applies deep Q-learning to automatically prune attention heads but the search time can be comparatively long. More recently, a self-supervised metapruning framework (SMP) [Zhang et al., 2021] is designed by combining head importance scoring and Gumbel softmax pruning through representation distance minimization.

Challenges. Accuracy degradation is inevitable during and after head pruning. While performing head pruning, since the pruned heads will no longer participate in the training process, the output of the model will change significantly, resulting in an irreversible accuracy loss. During pruning, the retained heads have to recover the model accuracy without any auxiliary information, which is extremely difficult, especially when the pruning granularity (single weight, row, column, head, etc.) is large.

Multi-head Attention. Self-attention plays an important role in Transformer-based language models. In Transformer layers, multiple attention heads work in parallel. The selfattention is calculated based on Query (Q), Key (K), and Value (V ) matrices and the attention score matrix is calculated as follows:

where A is the attention score matrix and Dk is the dimension of matrix K.

The multi-head attention mechanism uses different matrics of (Q, K, V ) to learn different representation subspaces. After concatenating the derived attention heads, a feed-forward layer is utilized to project the concatenation:

where X denotes the input of the i-th attention head, W Q i , W K i , and W V i are attention matrices, W O is projection matrix, and Hi denotes attention head.

Transformer-based model relies on multiple encoder layers. For each layer, the output is formulated as follows:

where O is the output of the transformer layer, Norm(·) is the layer normalization function, ResConnect(·) denotes the residual connection [Devlin et al., 2018].

3 Peer Distillation for Automatic Head Pruning

We propose the peer distillation head pruning framework for automatic head selection and model performance preservation as shown in Fig. 2. First, differentiable head pruning is performed to select heads by automatically updating binary gates. Then, the discarded "knowledge" related to the pruned heads is recycled and peer distillation is performed to distill the discarded "knowledge" into the derived compact model. It is worth noting that the pruning and peer distillation operations are performed in the same training loop.

Differentiable Head Pruning. In this section, we propose a differentiable method for head pruning. Unlike pruning methods with hard constraints [Han et al., 2015; Boyd et al., 2011; Li et al., 2016; Li et al., 2020], our method obtains model sparsity by updating the gate parameters. This leads to two important benefits: first, we do not have to set expected sparsity for each layer and can obtain model sparsity automatically; second, without hard constraints, the pruning process prunes the model more "smoothly" and nicely remains model accuracy. This makes the retraining process not a necessity and leads to faster training convergence.

In order to achieve sparse attention heads, we define the pruning loss consisting of model accuracy loss and model sparsity loss. The head pruning problem can be formulated by the following steps: first, we introduce the attention head mask M, in which M is composed of lists of binary variables representing the status of their corresponding heads and can be defined as:

where i and j denote the index of attention head and layer, respectively, and Mij denotes the pruning status of the head.

Second, inspired by the early works on neural network quantization and pruning [Hubara et al., 2016; Xiao et al., 2019], we employ learnable discrete functions called straight through estimators (STEs) g to describe the head mask M and thus M can be formulated as:

where W ′ ij is the auxiliary parameter to control the open and close of the binary gates, the binary head mask Mij is represented as a step function g with a continuous auxiliary parameter W ′ ij .

Third, we formulate the head pruning problem as:

where Lp is the pruning loss to update auxiliary parameters W ′ , LA is the model accuracy loss with model weight W, and µ is the penalty factor.

However, due to the binary nature of g(W ′ ) and the continuous weights W values, the problem described in Eq. 7 is a mixed integer programming problem, which brings difficulties in optimizing it directly using back-propagation.

To update the sparse head structure, we introduce coarse gradients [Hubara et al., 2016] to make the binarized function g differentiable. Coarse gradients provide a good approximation for updating parameter W ′ through back-propagation and could ensure that the update direction of W ′ ij gradient reflects the accuracy and sparsity objectives of themodel [Xiao et al., 2019].

Different coarse gradients have been practiced and discussed in literature [Srinivas et al., 2017; Yin et al., 2019]. We use Softplus STE in [Xiao et al., 2019] and the auxiliary parameter W ′ can be updated as:

where lrg is the learning rate of the optimizer to update W ′ . Knowledge Recycling. To connect the separate sparsification and training stages, and preserve the "knowledge" of the pruned heads, we introduce knowledge recycling and peer distillation. After differentiable pruning, we obtain the binary head mask M and thus a compact model with fewer attention heads. Previous research [Hu et al., 2018; Voita et al., 2019; Lee et al., 2020; Wang et al., 2020; Huang et al., 2022] just discarded the pruned heads and trained the remaining model parameters directly. However, this will introduce significant accuracy loss. By introducing knowledge recycling, we provide the opportunity to maintain and make full use of the existing knowledge in the entire model to boost model performance.

We select head attention score matrix, A, and layer output, O, as representations of head knowledge. A contains the abstraction of the input and the relationship extraction between input tokens. As an interface with other model layers, the layer output, O, integrates the effective information of all heads and directly reflects the impact of head pruning.

To update the model weight W and model structure (represented by auxiliary parameters W ′ ), two back-propagations are performed. For each iteration (one batch of input), the knowledge recycling process is performed by four steps:

Step1: We do model forward to derive the complete knowledge of the dense model according to Eq. 1, 4, and 5:

Oc = Norm(ResConnect(MH))

where M is the current head mask, Ac and Oc are the complete head attention score matrix and layer output before updating M, respectively. Initially, M is an all-ones matrix.

Step2: By leveraging differentiable head pruning in Eq. 8, we update M and derive the structure of the pruned model according to Eq. 6 and 5.

Step3: For the derived pruned model with updated M, we do model forward to derive the retained knowledge (Ar , Or ), where Ar and Or are the retained attention score matrix and layer output after updatingM, respectively.

Step4: To preserve the model performance, we do peer distillation (described in the following subsection).

Peer Distillation. Traditional Knowledge Distillation (KD) method distills "knowledge" from the original model (teacher) to another model (student). For transformer-based models, different heads in the same layer work in parallel to provide abstract representations of input. The pruned heads could provide additional learned model information, which is, however, discarded during the training of the compact model.

The proposed Peer Distillation (PD) method distills the "knowledge" from the pruned model structure to the retained model structure in the same layer. Specifically, the retained heads learn to fit the matrices of head attention score before pruning (teacher) and the objective is defined as:

where N is the number of transformer layer of the model (e.g. 12 for BERTBASE model), S is the number of attention heads of a transformer layer(e.g. 12 for BERTBASE

Algorithm 1 Peer distillation head pruning procedure

Input: pretrained tranformer-based model TM, head pruning optimizer OPH, weight optimizer OPW

Output: updated model weightWand head masksM

• 1: for every layer Lj in TM do
• 2: for every head Hi in Lj do
• 3: Create and initialize auxiliary parameters W′ij
• 4: end for
• 5: end for

• 6: for every training iteration do
• 7: Domodel forward to derive head attention socreAc and layer output Oc in Eq. 10 and 11
• 8: Calculate head pruning loss Lp = in Eq. 7
• 9: UpdateW′ through back propagation by minimizing Lp using OPH
• 10: Detach Ac and Oc from the computation graph to avoid gradient calculation
• 11: Update head masksM in Eq. 6
• 12: Domodel forward to derive head attention socreAr and layer output Or in Eq. 10 and 11 with updated head masks M
• 13: Calculate the mixed loss Lm in Eq. 14
• 14: UpdateW through back propagation by minimizing Lm using OPW
• 15: end for

model), MSE(·, ·) is the mean squared error loss function, Acji and Arji are the i-th attention score matrices (formulated in Eq. 10) in the j-th layer before and after differentiable pruning, respectively.

We also distill the layer outputs from the unpruned to retained heads to better maintain the original model knowledge. The objective is defined as:

where Ock and Ork are the outputs (formulated in Eq. 11) of the k-th layer before and after differentiable pruning, respectively.

Finally, we formulate the parameter update problem of the derived compact model as follows:

where Lm is the mixed loss to update model weights W and λ is the penalty factor. PDloss1 and PDloss2 are peer distillation losses with arguments W (including W Q, W K , W V and W O ) according to Eq. 2 and 3. To notice, while calculating Eq. 14, all model parameters, W, will be updated.

By minimizing the mixed loss, Lm, the retained model parameters are updated to preserve the model accuracy. The previously discarded model information of the pruned heads plays as a "teacher" and transfers its "knowledge" to the "students" (i.e. retained heads).

Algorithm 1 illustrates the peer distillation head pruning procedure. We conduct head pruning and peer distillation in the same training loop and update the model weight and structure in parallel. To notice, we detach the derived discarded knowledge Ac and Oc ("teacher" in the peer distillation method) from DNN graph to ensure that only weights of the compact model are updated.

4 Evaluation

Datasets.We test our method on GLUE benchmark [Wang et al., 2018] and report the performance of the unpruned and pruned models following the conventions by using accuracy for SST-2, QNLI, MNLI, QQP, RTE and WNLI; Matthews Correlation Coefficient (MCC) for CoLA, F1 scores for MRPC, and Spearman for STS-B. Our pre-trained model is the BERTBASE [Devlin et al., 2018] model. The model consists of 12 attention layers and 12 heads for each layer. We define the head sparsity as:

Implementation Details. We follow the default finetuning steps for 9 tasks according to Huggingface [Wolf et al., 2019] and obtain baseline models after training for 4 epochs. Then, differentiable pruning is executed for the models.

Baselines. To validate the effectiveness of our proposed method, we introduce four baselines in our experiments. For method Random, we report the results of 50% head sparsity from [Zhang et al., 2021]. In [Michel et al., 2019], the Head Importance Score for Pruning (HISP) is proposed by ranking the head importance and removing the heads with lower importance score. In our test, we calculate the head importance and prune 50% heads with lower importance scores. Method L0 Norm represents the Gumbel softmax based pruning method proposed in [Voita et al., 2019]. And the Single-Shot Meta-Pruner (SMP) is proposed by [Zhang et al., 2021] in which head importance and Gumbel softmax based pruning are nicely combined.

Experimental Results. We show our comparison results in Table 1. For fairness, we compare our head pruning model accuracy with state-of-the-art algorithms with fixed global head sparsity of 50%. Compared with Random method, our proposed method enjoys better model accuracy. HISP [Michel et al., 2019] calculates head importance for pruning and suffers a significant accuracy drop since the importance is not estimated directly according to the final model performance. Our method outperforms L0 Norm approach [Voita et al., 2019] with a large margin in all existing 7 GLUE benchmark tasks. SMP [Zhang et al., 2021] improves Gumbel softmax based approach [Voita et al., 2019] by combining head importance scoring and self-supervision, but the discrepancy between model training and model testing prohibits its further advances. Comparing with state-of-the-art head pruning methods, peer distillation takes the lead in 7 of the 8 tasks.

To fully investigate the performance of peer distillation, we test different sparsity levels. As shown in Fig. 3, our method can achieve more than 80% sparsity with only 1.03% accuracy drop and outperforms all existing head pruning methods. Additionally, for 8 of the 9 tasks, our pruned models outperform the original (unpruned) models as shown in Table 2, which is consistent with [Kovaleva et al., 2019]'s study. More specifically, we could achieve 1.20% accuracy increase while pruning more than 58% heads on average. Surprisingly, the pruned model on CoLA dataset achieves 6.99% accuracy increase after pruning 54.17% heads and the pruned model on WNLI dataset has the same accuracy as the original one after pruning 99.3% heads (only one head left). In extreme cases (within 1% accuracy drop), our proposed method could prune 79.74% heads with 0.82% accuracy drop on average as shown in Table 3.

5 Discussion of Differentiable Head Pruning and Peer Distillation

Differentiable Head Pruning. Different from the quantified importance scores of heads [Michel et al., 2019] for pruning, in our method, we use learnable gate parameters to determine the retention of the heads. Fig. 4 shows the update process of the auxiliary parameter W ′ and head masking gate status g(W ′ ) of the 144 attention heads jointly trained with the model weights W by following three stages:

• In the initialization stage, the auxiliary parameters are ini-

(a) W ′ before pruning head

(b) Mask before pruning

(c) W ′ during pruning

(d) Mask during pruning tialized by following a truncated normal distribution with values greater than 0 and gates are open (Fig. 4a and 4b).

• In the intermediate stage, the gate auxiliary parameters of less important heads are receiving negative gradients. Through training, some auxiliary parameters are below zero (denoted in cold colored boxes in Fig. 4c, which leads to pruning of corresponding attention heads in Fig. 4d.
• In the final stage, in Fig. 4e, more auxiliary parameters drop below zero. The optimization converges when the balance is achieved between the model accuracy component and sparsity component in loss function Eq. 7.

Peer Distillation Performance. To verify the effectiveness of the peer distillation, we conduct experiments of differentiable head pruning with and without peer distillation as shown in Fig. 5. Without peer distillation (baseline as the starred blue curve in Fig. 5b), the head difference (PD loss) changes slightly (±4.32%) during pruning and reveals that the retained heads converge to a different value instead of similar ones to the pruned heads. This proves our assumption that the retained compact model updates independently without absorbing useful knowledge from the pruned heads. With the help of peer distillation (λ = 0.35), the PD loss is reduced by 68.55% and the retained heads become much close to the pruned ones. Compared with differentiable head pruning without peer distillation, the head sparsity is increased by 22.17% as shown in Fig. 5c while the evaluation model accuracy (F1 score) are the same (Fig. 5a). The result verifies the effectiveness of the peer distillation to preserve model performance with a shrunk model structure.

• (a) Model Performance (mcc for CoLA)

6 Ablation Study

In this section, we perform ablation study over several hyperparameters when doing differentiable pruning and peer distillation with BERT model and check their effects.

Differentiable Pruning Learning Rate. To solve the optimization problem in Eq. 7, different optimizers are utilized to update weight and gate parameters. Specially, we use the

(a) Model Performance (F1 for MRPC)

default initial learning rate (3e-5) to update model weights. For auxiliary parameters updating, a larger initial learning rate, lrg, enhances the ability of the gate optimizer to adjust gate parameters and leads to higher sparsity. In our test, while increasing the initial lrg, we could increase the sparsity from 44% to 76% with only 4% performance (mcc for CoLA dataset) drop as shown in Fig. 6. Interestingly, we observe the obvious compete between accuracy and sparsity, since different optimizers (weight and gate optimizer) tend to reduce the loss function in different directions.

Peer Distillation Loss Penalty Factor. While selecting different knowledge loss penalty factors, λ, from 0.2 to 0.5, we observe the similar compact model performance (values changed within ±5.62% in F1 score and ±0.01 in final training mixed loss) in Fig. 7. However, while increasing λ, significant head sparsity increase (Fig. 8b) is achieved due to the reduced PD loss (Fig. 8a). To notice, the penalty factor should not be chosen too large since it might affect the model accuracy loss in Eq. 14 and leads to model accuracy degradation. In experiments, we choose λ = 0.35 for the final evaluation.

7 Conclusion

In this work, we propose a novel peer distillation head pruning method. We leverage differentiable pruning to automatically select heads and derive a compact model structure. Then knowledge recycling and peer distillation are proposed to make use of the discarded knowledge represented by the pruned heads. The retained heads exploit the recycled knowledge from the pruned heads and ensure the maintenance of the learned knowledge and thus model accuracy. With the help of recycled knowledge from the pruned heads, higher head sparsity is achieved with a slight loss in accuracy. Our results outperform the state-of-the-art head pruning results and validate the effectiveness of our method.

Acknowledgements

We would like to acknowledge the USDA-NIFA Agriculture and Food Research Initiative (Award No.: 2022-67023-36399 to CD).