Layer-wise Model Pruning based on Mutual Information

Introduction

In spite of impressive results of neural networks, the huge model size has hindered their applications in cases where computation and memory resources are limited.For example, the GPT-3 model (Brown et al., 2020) has 175B parameters in total, with 96 layers and 96 attention heads (Vaswani et al., 2017) per layer. As a result, training and using existing huge models not only requires rich hardware resources, but also consumes high environmental costs (Strubell et al., 2019).

Model pruning, reduces model sizes by dropping a fraction of the model parameters, to reduce computation intensity and memory footprint of large models at the lowest cost of accuracy on end tasks (Joulin et al., 2016; Ganesh et al., 2020; Gordon et al., 2020). Among pruning techniques, weight based pruning is a widely-used group of methods. It focuses on removing weights according to their importance under different specific criteria, e.g., the magnitude (Han et al., 2015b, a), first-order derivative (Lee et al., 2018; Sanh et al., 2020) and second-order derivative information (LeCun et al., 1990; Hassibi and Stork, 1993), and it has been successfully applied to a large variety of model architectures (Guo et al., 2016; Gale et al., 2019; Molchanov et al., 2019) and downstream tasks (McCarley, 2019; Gordon et al., 2020).

While weight-based methods have been successfully applied to a wide range of neural models for model pruning, they come with the following shortcomings: (1) weights in matrices are pruned irregularly, which lead to irregular memory access, resulting in runtime inefficiency; (2) weight matrices are pruned independently, and this neglect of global supervision from training signals at the top layer and ignorance of information propagation between consecutive layers may result in sub-optimality of pruned networks.

In this paper, inspired by mutual information (MI) based feature selection (Kuncheva, 2007) in SVMs and logistic regression, we propose MI based layer-wise pruning, to address the aforementioned drawbacks of weight-based pruning methods in NLP. For each layer of a multi-layer neural network, neurons with higher values of MI with respect to the preserved neurons in the upper layer are preserved. Starting from the top softmax layer, layer-wise pruning proceeds until reaching the bottom input word embedding layer in a top-down fashion. Once the preserved neurons in each layer are selected, the redundant dimensions along with the corresponding rows and columns of the weight matrices can be pruned or squeezed, inducing model sparsity at different levels.

The proposed pruning strategy naturally addresses the aforementioned two shortcomings of weight-based methods: (1) it avoids irregular memory access since it squeezes the pruned representations and matrices into their smaller but dense counterparts. This enables significantly faster computations than weight-based pruning methods at the same sparsity level; (2) rather than viewing each weight matrix separately based on their own weight values, the proposed method operates from a more global perspective based on training signals at the top layer, and prunes each layer by propagating the effect of global training signals through consecutive layers in a top-down fashion. This leads to better performances at the same sparsity level.

We conduct extensive experiments on both generative tasks (MT) and discriminative tasks (question answering) in NLP to examine the effectiveness of the proposed strategy. We show that compared to weight-based pruning methods including magnitude pruning (Han et al., 2015b), movement pruning (Sanh et al., 2020) and $L_{0}$ pruning (Louizos et al., 2017), the proposed method yields greater speedup along with better performances for the same sparsity levels on generative NLP tasks of WMT’14 En$\rightarrow$Fr and WMT’14 En$\rightarrow$De, and discriminative NLP tasks of SQuAD v1.1 (Rajpurkar et al., 2016), MNIL Williams et al. (2017) and SST-5 Socher et al. (2013). In addition, we also show that the proposed method serves the feature selection purposes, where we observe significant performance boosts when fixing preserved neurons and relearning the pruned ones, leading to a state-of-the-art performance of 43.9 BLEU score for En$\rightarrow$Fr translation in setups without back-translation or external data.

Related Work

Model pruning refers to reducing the model size by dropping a fraction of the model parameters, which dates back to early works of Optimal Brain Damage (PBD) (LeCun et al., 1990) and Optimal Brain Surgeon (OBS) (Hassibi and Stork, 1993). One major branch of neural model pruning methods is magnitude pruning (Han et al., 2015b; See et al., 2016; Narang et al., 2017; Molchanov et al., 2019; Gale et al., 2019; Frankle et al., 2020), which prunes model parameters measured by their importance scores. Han et al. (2015b) removed all parameters with weight values below a threshold, and then retrained the remaining sparse network. Guo et al. (2016) proposed dynamic network surgery, allowing for model connection recovery from incorrect pruning decisions made in previous iterations. Michael H. Zhu (2018) adopted a gradual pruning schedule, in which the sparsity level increases from an initial sparsity value to a specified final sparsity value during training. Other methods for neural model pruning include $L_{0}$ regularization pruning (Louizos et al., 2017), variational dropout pruning (Kingma et al., 2015; Molchanov et al., 2017; Gomez et al., 2019) and movement pruning (Sanh et al., 2020), etc. Recent works have proposed a line of techniques to prune and produce sparsity in a structured way (Anwar et al., 2017; Zhou et al., 2016; Hu et al., 2016; Liu et al., 2019b), which aims at pruning full convolutional filters or whole layers. Methods for structured pruning mainly include group Lasso (Alvarez and Salzmann, 2016; Wen et al., 2016; He et al., 2017), sparsity regularization (Li et al., 2016; Liu et al., 2017; Huang and Wang, 2018; Gordon et al., 2018) and automatic network searching (He et al., 2018; Yu and Huang, 2019; Dong and Yang, 2019; Ding et al., 2019).

Pruning Transformer based models has been of growing interest (Guo et al., 2019; Chen et al., 2020; Li et al., 2020). Fan et al. (2019) proposed LayerDrop to reduce Transformer depth. Michel et al. (2019) proposed to use head importance score to prune BERT attention heads. Attention heads can also be pruned by using $L_{0}$ regularization (Voita et al., 2019) and cascade pruning (Wang et al., 2021). Wang et al. (2020) combined $L_{0}$ regularization with matrix factorization to prune BERT. Gordon et al. (2020) proposed that BERT can be pruned once during pre-training rather than separately for each task without sacrificing performance.

2 Mutual Information Feature Selection

Feature selection is the process of selecting a proper subset of features for better model performances (Kira and Rendell, 1992; Guyon and Elisseeff, 2003; Chandrashekar and Sahin, 2014; Bolón-Canedo et al., 2016; Cai et al., 2018). A widely used method for feature selection is Mutual Information Based Feature Selection (Vergara and Estévez, 2014; Liu et al., 2009; Beraha et al., 2019), which selects features that minimize the redundancy and maximize the relevance w.r.t. the target variable. Various approaches including minimum-Redundancy-Maximum-Relevance (mRMR) (Estévez et al., 2009; Brown et al., 2012; Bennasar et al., 2015) are proposed to accurately select features.

Model

Given a set of inputs $\mathcal{M}=\{(X,Y)\}$, where each input is a word sequence $X=\{x_{1},...,x_{t},...,x_{N_{x}}\}$ and $N_{x}$ denotes the length of the input, our goal is to predict the label(s) for $X$, denoted by $Y$.

The output from the last layer $h_{t}^{L}$ is fed to the final softmax layer for predictions. To prune a neural network model, let $m^{l}\in\{0,1\}^{D\times 1}$ denote the mask for representation dimensions at layer $l$. The number of $1$s in $m^{l}$ is a pre-defined hyper-parameter, denoted by $K$, controlling the sparsity of the network. $M^{l}\in\{0,1\}^{D\times N}$ makes $N$ copies of $m^{l}$, making the dimensionality of the mask the same as that of layer representations for $X$. Let $u^{l}$ denote the set of indexes for preserved dimensions, where $m^{l}[j\;\text{for}\;j\;\text{in}\;u^{l}]=1$. Eq.(1) can be rewritten as:

where $\otimes$ is the Hadamard product. We need special attentions for the uppermost softmax layer. No dimension should be pruned for this layer since each dimension corresponds to an output label. $m^{\text{softmax}}=^{|\mathcal{Y}|}$, where $|\mathcal{Y}|$ denotes the size of the output label set.

2 Layer-wise Pruning

The key point of layer-wise pruning is to construct correlations between dimensions in two consecutive layers $l-1$ and $l$. Then based on the correlations, we can prune the network in a top-down fashion: with respect to output labels in the final softmax layer, we select the top $K$ correlated dimensions in the $L^{\text{th}}$ layer based on the correlation measure, zeroing out the rest. Let $I(A,B)$ denote correlation between two set of dimensions:

Next, we go to the $(L-1)^{\text{th}}$ layer, preserving dimensions in the $(L-1)^{\text{th}}$ layer that are most correlated with preserved dimensions in the $L^{\text{th}}$ layer

This process proceeds until the bottom input embedding layer. An illustration of the proposed layer-wise pruning method is show in Figure 1. Algorithm 1 describes the pruning process.

3 Mutual Information between Dimensions

Here, we describe quantitative ways to compute correlation scores $I(A,B)$ between dimensions in layer $l-1$ and layer $l$ using MI.

Mutual information (MI) is a measure between two random variables to quantify the amount of information obtained about one variable through the other variable. In our case, we wish to compute the MI between dimensions $u^{l}$ at layer $l$ and dimensions $u^{l-1}$ at layer $l-1$. Let $v_{d^{l}_{k}}$ denote the variable for the neuron value of the ${d^{l}_{k}}$-th dimension at the $l^{\text{th}}$ layer. MI between $u^{l}$ and $u^{l-1}$ is given by:

To tangibly compute Eq.(5), we make assumptions that both $v_{d^{l}_{1}},...,v_{d^{l}_{K}}$ and $v_{d^{l-1}_{1}},...,v_{d^{l-1}_{K}}$ are samples from Gaussian distributions:

where ${v}_{t,d^{l}_{{u^{l}}}}$ is a vector of length $K$, corresponding to a sub-vector within $h^{l}_{t}$ with dimension $u^{l}$. $\eta^{l-1}_{u^{l}}$, $\eta^{l,l-1}_{u^{l}}$, $\Sigma^{l-1}_{u^{l}}$, $\Sigma^{l,l-1}_{u^{l}}$ can be computed similarly.

It is worth noting that the proposed model relies on the Gaussian assumption for MI computations, and several recent efforts have been proposed to release this strong assumption, such as training independent neural nets to estimate MI Belghazi et al. (2018), using variational distributions to approximate the distribution Cheng et al. (2020); Poole et al. (2019). These workarounds to avoid the Gaussian assumption requires learning another model (an independent neural model in Belghazi et al. (2018) and variational distributions in Cheng et al. (2020)) through gradient updates, and thus cannot be adapted to the scale in our situation, where we have to estimate MI for all dimensions across all layers. The adopted Gaussian model is efficient in estimating MI values in bulk, and achieve satisfying performances. We leave how to relax this assumption to future work.

3.2 Greedy Selection

Inspired by Brown et al. (2012), further assumptions are made that the selected dimensions are independent and class-conditionally independent given unselected features, transforming Eq.(8) to the following form:

It is straightforward to see that the first part of Eq.(9), i.e., $I(u^{l},d)$ models the relevance of selected dimensions, against the redundancy compared to the dimensions already selected , manifested in the second and the third part. The model degenerates to the model of Maximum Relevancy Minimum Redundancy (mRMR) (Peng et al., 2005) when $\beta=0$.

3.3 Squeezing Weights and Features

For weight matrixes $W$ and feature $H^{l}$ involved in the matrix manipulation $WH^{l}$, we do not need to actually compute the Hadamard product in Eq2. Instead, for $H$, we squeeze all preserved dimensions to the left side and truncate the rest. For $W$, rows and columns that correspond to pruned dimensions will be erased and the remaining dimensions will be squeezed. For example, with $m^{l}=$ and $m^{l+1}=$, the third row and first column of the original matrix $W=[w_{ij}]$ can be pruned, the result of which is squeezed into a smaller matrix:

This avoids irregular memory accesses and thus can significantly speed up matrix-vector product. Figure 1 gives a tangible illustration.

4 Iterative Pruning

Instead of aggressively reducing dimensions from $D$ to $K$ in only one iteration, iterative pruning Han et al. (2015b) gradually reduces model dimensions in multiple steps: in each iteration, pruning is followed by model retraining using preserved dimensions. As we will show in experiments, this strategy achieves better performances than the single-step pruning with the same sparsity levels.

5 Retraining Pruned Dimensions

The proposed MI based pruning strategy can not only be used for reducing model size, but also for improving model performances. We can view the MI pruning model from a feature selection perspective: given fixed size of features (where we view each neural dimension as a feature), we wish that all features in each neural layer be informative and relevant. To this end, we can first remove redundant or irrelevant features, add new features, retrain the model, and repeat this process. This strategy is akin to feature selection methods in SVMs or logistic regression Kuncheva (2007).

In the neural setup, we can achieve this goal by (1) pruning irrelevant dimensions; (2) reinitializing pruned dimensions (adding new features); and (3) retraining the model. Preserved dimensions and weight matrices are fixed during model retraining, and we only update pruned dimensions. We report the performances of pruning and retraining 60% dimensions. It is worth noting that the strategy of retraining pruned dimensions does not serve as the goal of speedup and model compressing, as pruned dimensions are relearned, making the model of the same size as the model before pruning. We as view retraining pruned dimensions as a byproduct of the pruning, with the goal of improving performances.

6 Discussions

For the $Wh$ matrix multiplication in neural models, we refer to $W$ as weights, and $h$ as features. Weight-based methods (Han et al., 2015a, b) prune networks based on values of $W$, removing features with smaller weights, which are comparable to $L1$ or $L2$ regularizers for feature selection Ng (2004); Ravikumar et al. (2010). MI-based pruning method is comparable to MI based feature selection, which attaches attentions to the features by measuring feature-label correlations Kuncheva (2007); Yu et al. (2008).

Experiments

We conduct experiments on both generative and discriminative NLP tasks. For generative tasks, we conduct experiments on WMT14 En-Fr and WMT14 En-DE. The WMT14 En-Fr dataset consist of 36M and is split into 32000 word-piece vocabulary. The WMT 2014 En-DE dataset consisting of about 4.5 million sentence pairs. We use BPE (Sennrich et al., 2016b) to maintain a source-target vocabulary of 37,000. We use Transformers (Vaswani et al., 2017) as the model backbone. We use En-Fr to perform comprehensive analysis where we use four model setups: extra-large, large, base and tiny. The model statistics are shown in Table 1. It is worth noting that the large and base models are identical to models in Vaswani et al. (2017). We train different models with 16 V100 GPUs with 32G memories. We follow protocols in Vaswani et al. (2017). Adam (Kingma and Ba, 2014) is used for all models with $\beta_{1}$ = 0.9, $\beta_{2}$ = 0.98 and $\epsilon=10^{-6}$. A dropout rate of 0.1 is applied to all layers across all models, and the strategy of label smoothing (Szegedy et al., 2016) is used with smoothing value set to $0.1$.Since our goal is to test the performances of different pruning techniques in the vanilla supervised setup, no advanced MT techniques such as backtranslation (Sennrich et al., 2016a; Edunov et al., 2018), self-learning He et al. (2020); Sun et al. (2020), data noising Xie et al. (2017); Bengio et al. (2015), nearest neighbor search Khandelwal et al. (2020); Meng et al. (2021); Zheng et al. (2021) are used. We use beam search with a beam size of 20, with no penalty on length. We report BLEU scores based on multi-bleu.perl of single models (no ensemble), average floating-point operations (FLOPs), and average practical speedup.

For discriminative tasks, we followed the current trend of LM pretraining Devlin et al. (2018); Liu et al. (2019a); Jiao et al. (2019); Radford et al. (2019); Lan et al. (2019); Brown et al. (2020); Clark et al. (2020); Sun et al. (2021). We test different pruning models on the tasks of question answering (Rajpurkar et al., 2016, 2018), natural language inference Bowman et al. (2015); Williams et al. (2017) and text classification Socher et al. (2013); Tang et al. (2014); Howard and Ruder (2018); Chai et al. (2020); Lin et al. (2021). We use BERT (Devlin et al., 2018) as the backbone, and fine-tune BERT on different datasets. Adam (Kingma and Ba, 2014) is used for all models, with batch size, learning rate and the number of epochs treated as hyper-parameters to be tuned on the dev set. We compare the proposed strategy with the following weight based pruning models:

Magnitude Pruning (Han et al., 2015b): removing weights based on their absolute weight values.

Movement Pruning (Sanh et al., 2020): removing weights based on the first-order derivative.

L0 Pruning (Louizos et al., 2017): using the $L_{0}$ loss to regularize the number of non-zero weights.

MT results are shown in Tables 2 and 3. Observations can be summarized as follows: (1) When comparing with movement and magnitude pruning, at the same levels of sparsity, the proposed MI method yields greater speedup. This is due to the fact that using MI, the weight matrix $W$ can be squeezed avoiding irregular memory accesses. For magnitude and movement pruning: though $W$ is sparse, pruned dimensions in $W$ are scattered and irregular memory accesses are inevitable.

(2) The MI model yields not only speedup but also performance boosts: we find that the proposed MI pruning consistently works better, both in the low-sparsity and high-sparsity situations. This is because the mutual information strategy provides a more global feature (dimension) selection strategy based on the output label, rather than focusing on the local matrix weights in matrix manipulations. Regarding magnitude pruning and movement pruning, we find that movement pruning underperforms magnitude pruning at lower sparsity levels but works better at higher sparsity levels.

(2) Based on MI, training a big model and then pruning it to a smaller one outperforms directly training a smaller model of the same size, e.g., pruning extra-large to large yields a BLEU score of 42.4 for En-Fr, which is +0.6 higher than vanilla large (41.8). This is also the case with pruning extra-large to base and tiny, and pruning large to base and tiny. The explanations are as follows: a directly trained model contains redundant and irrelevant dimensions; for the large-training-then-pruning strategy, the model first learns a larger set of feature dimensions, and then prunes irrelevant ones. This makes the model consist of fewer irrelevant feature dimensions than the one directly trained, leading to better performances.

(3) Pruning and then retraining yields consistent performance boosts over direct training: +0.6 for extra-large (43.3 vs 43.9), +0.5 for large (41.8 vs 42.3) and +0.5 for base (37.9 vs 38.4) for En-Fr. This is because direct training introduces redundant and less relevant features; retraining pruned dimensions can help the model replace less relevant dimensions with relevant ones, obtaining a state-of-the-art performance of 43.9 BLEU score for En$\rightarrow$Fr translation in setups without back-translation or external data. Similar phenomenon are observed for En-De with +0.4 for the large model, and +0.6 for base models.

2 BERT Pruning

We carry out experiments on the pretrained model of BERT-large which contains 24 layers, 1,024 hidden units per layer, 16 heads per layer and 340M parameters in total. We select different degrees of sparsities from 0$\%$ to 90$\%$ at an interval of $10\%$. Model pruning can happen either in the pretraining stage (pretrain-prune), the fine-tune stage (finetune-prune), and both (hybrid): For hybrid, pruning happens at both stages, with the ultimate sparsity level $\gamma$ being the product of the sparsity level of two stages, $\gamma_{\text{pretrain}}$ and $\gamma_{\text{finetune}}$. We compare the performance of the three strategies on the SQuAD v1.1, MNIL and and SST-5 in Figure 2 and Figure 4. Generally, pretrain-prune works consistently better than finetune-prune with the same level of sparsity. This is because the training objective at the pretraining stage is a more general one than that at the finetuning stage, with more training data points and categories. Pruning at the finetuning stage is more prone to overfitting, leading to inferior performances. The hybrid method outperforms the pretrain-prune strategy if the sparsity levels at two stages are carefully calibrated. This is because the hybrid model can progressively prune less relevant dimensions in pretraining and then less relevant dimensions in task-specific finetuning, leading to better final performances.

For both pretrain-prune and finetune-prune, we find that the proposed MI method offers greater speedup and better performances at the same sparsity levels. Similar phenomenon are found for MNIL and SST-5. Figure 4 shows the speedup gains for different models for the pretrain-prune setup. With the same sparsity, random pruning and the proposed MI based pruning lead to the largest speedup, followed by magnitude pruning, movement pruning and $L_{0}$ pruning. This observation validates that condensed weights serve as an effective remedy for irregular memory access.

Ablation Studies

In this section, we conduct ablation studies to get a better understanding of model behaviors. We use SQuAD for analysis, where BERT-large is used.

The value of $\alpha$ and $\beta$ in Eq.(9) controls the tradeoff between selecting relevant dimensions and removing redundant dimensions. Based on the pretrain-prune strategy with sparsity level of 20$\%$, we can see from Figure 5 that the model works best when the value of $\alpha$ is set to 0.4, and then deteriorates as $\alpha$ increases when fixing $\beta=0$. With fixed value of $\alpha=0.4$, we find that the influence from $\beta$ is less significant. This shows that given the conditional independency assumption, the improvement from the class-conditionally independent assumption is marginal. We thus suggest omitting this part if computing resources are limited.

2 The Effect of Iterative Pruning

Table 4 presents results with different number of pruning iterations, where we use linear interpolation to obtain sparsity levels for different iterations. As can be seen, though more pruning iterations lead to better performances, the boost becomes marginal when iteration number exceeds 2.

Fixing the overall sparsity of 0.2, we explore the effect of $\gamma_{\text{pretrain}}$ and $\gamma_{\text{finetune}}$. When $\gamma_{\text{finetune}}=1$, it means we only perform pruning at the pretraining stage; When $\gamma_{\text{finetune}}=0.2$, it means we only perform pruning at the finetuning stage. As can be seen from Table 5, performance peaks when $\gamma_{\text{finetune}}$ is slightly lower than 1 ($\gamma_{\text{finetune}}=0.8$ , $\gamma_{\text{pretrain}}=0.25$), and then declines as we increase $\gamma_{\text{finetune}}$. This further validates that the final performance benefits more when most pruning happens at the pretraining stage.

4 Layers with Different Sparsity Values

We explore the situation where given fixed overall sparsity value, different layers can have different levels of sparsity. We additionally consider two setups, pyramid, where lower layers are denser and thus less sparse than upper layers, and inverted pyramid where upper layers are less sparse than lower layers. For pyramid, with the overall sparsity of 0.2, the lowest word embedding starts with a sparsity level of 0.1, with the sparsity of all layers forms an arithmetic sequence. inverted pyramid has the same overall sparsity value of 0.2, with the lowest word embedding starts with a sparsity level of 0.3. Results are shown in Table 6. We can observe that inverted pyramid outperforms vanilla, which outperforms pyramid. These results illustrate that to obtain better performances in model pruning with fixed overall sparsity, upper layers should be less sparse than lower layers. This is because upper layers contain more high-level and dense information about the input. Therefore, pruning upper layers does more harm to the model. Lower layers contain more noise, and thus hurt the model less when get pruned.

Conclusion and Future Work

In this paper, we propose MI based methods for model pruning in NLP. The proposed model avoids the issue of irregular memory access, leading to higher speedup with the same level of sparsity. Also, the proposed strategy prunes the model in a top-down fashion based on global training signals, and thus achieves higher accuracies. In future work, we should release the strong assumption that neuron values come from a Gaussian distribution.

References