Aligned Weight Regularizers for Pruning Pretrained Neural Networks (2022)

Abstract: Pruning aims to reduce the number of parameters while maintaining performance close to the original network. This work proposes a novel self-distillation based pruning strategy, whereby the representational similarity between the pruned and unpruned versions of the same network is maximized. Unlike previous approaches that treat distillation and pruning separately, we use distillation to inform the pruning criteria, without requiring a separate student network as in knowledge distillation. We show that the proposed cross-correlation objective for self-distilled pruning implicitly encourages sparse solutions, naturally complementing magnitude-based pruning criteria. Experiments on the GLUE and XGLUE benchmarks show that self-distilled pruning increases mono- and cross-lingual language model performance. Self-distilled pruned models also outperform smaller Transformers with an equal number of parameters and are competitive against (6 times) larger distilled networks. We also observe that self-distillation (1) maximizes class separability, (2) increases the signal-to-noise ratio, and (3) converges faster after pruning steps, providing further insights into why self-distilled pruning improves generalization.

## introduction deep neural networks (dnns) have grown increasingly large in the recent years. this has led to models requiring more storage requirements, more resources for training and inference (e.g., gpus and tpus), longer compute times and larger carbon footprints. this is largely due to the rise of masked self-supervised learning (ssl) which trains dnns (e.g., transformers in nlp) on a large collection of samples that do not have task labels but instead use a subset of the inputs as labels. given the aforementioned challenges, it has become more difficult for machine learning practitioners to use these ssl pretrained models for fine-tuning on downstream tasks. while training tricks such as effective batch sizes, gradient accumulation and dynamic learning rate schedules @xcite have improved the efficiency of fine-tuning dnns under resource constraints, it can still come at a cost, e.g. gradient accumulation leads to less updates. pruning @xcite is a type of model compression method @xcite that aims to address these shortcomings by zeroing out a subset of weights in the dnn, while maintaining performance close to the original model. retraining is often carried out directly after each pruning step to recover from pruning induced performance drops. this process is referred to as iterative pruning. although, iterative pruning has been extensively studied in the ssl setting @xcite @xcite and the transfer learning setting @xcite @xcite , little is known about pruning dnns in the zeroshot setting foot_0 where a model is required to make predictions on a set of samples from classes that are unobserved during training. one salient example is pretrained cross-lingual language models (xlms) @xcite whereby the model is trained with a masked/translation language model (mlm/tlm) objective to predict tokens for a large set of different languages whereby the objective forces the xlm model to learn similar representations for different languages. after cross-lingual pretraining, the model is further fine-tuned to a downstream task in one language (e.g., english) and then evaluated on different languages in the zero-shot setting (e.g., spanish, french, chinese, etc.). in this context, applying current pruning methods can damage the xlm cross-lingual alignment that has been learned during pretraining. ideally, we would aim to prune xlms in such a way that avoids this alignment distortion which effects the zero-shot performance of pruned xlms. additionally, overfitting to the language used for fine-tuning becomes more of an issue due to the progressive reduction in parameters throughout iterative pruning as the remaining weights are relatively large, moving away from an "aligned" xlm state. this is an important problem to address as the application of large pretrained models in the zero shot-setting for natural language and other modalities (e.g images and audio) is of practical importance e.g., using xlms in production for multiple languages by only requiring annotations in a single language for fine-tuning, making predictions on unseen classes at test time from pretrained visual representations @xcite using only semantic descriptions (i.e., label similarity to known classes) or zero-shot predictions in pretrained multimodal models such as clip @xcite . hence, this work addresses the alignment distortion pruning problem by introducing alignreg, a class of weight regularizers for magnitude-based pruning that force pruned models to have parameters that point in a similar direction or have a similar distribution to the parameters of the original pretrained network. to our knowledge, this is the first study on how iteratively pruned models perform in the zero-shot setting and how the solution differs from solutions found in the non-zero shot setting. we believe our findings have a strong practical implication as well-established pruning criteria may not be suitable given the observed discrepancy between zero-shot performance and the typically reported non-zero shot performance. moreover, our proposed weight regularizer improves overall pruning generalization in zero-shot cross-lingual transfer. below, we summarize our contributions. ## related work below we describe regularization-based pruning, other non-magnitude based pruning and how masked language modeling (mlm) implicitly learns to align cross-lingual representations. regularization-based pruning. pruning can be achieved by using a weight regularizer that encourages network sparsity. three well-established regularizers are l 0 @xcite , l 1 regularization @xcite and the commonly used l 2 regularization for weight sparsity @xcite . @xcite have proposed an l 2 regularizer that increases in influence throughout retraining and shows the increasing regularization improves pruning performance. for structured pruning where whole blocks of weights are removed, group-wise brain damage @xcite and ssl @xcite propose to use group lasso @xcite to learn structured solutions. importance-based pruning. magnitude-based pruning (mbp) relies on the assumption that weight or gradient magnitudes have correlation with its importance to the overall output of the network. mozer and smolensky instead use a learnable gating mechanism that approximates layer importance, finding that weight magnitudes reflect importance statistics. to measure weight importance as the difference in loss between pruned and unpruned network, lecun et al. approximate this difference with a taylor series up to the second order. this involves the product of the gradient and weight magnitude in the 1st term and an approximation of the hessian and the square of the weight magnitude for the second term. however, computing the hessian and even its approximations @xcite @xcite @xcite ) can significantly slow down retraining. in our work, we avoid the requirement of computing the hessian or approximations thereof, as it is not scalable for models such as xlm-r @xcite . @xcite have extended mbp to block approximations to avoid pruning lowest weight magnitudes that may be connected to weights in adjacent layers that have high weight magnitude. lee et al. have provided a method to automatically choose the sparsity of layers by using the rescaled version of weight magnitude to incorporates the modellevel distortion incurred by pruning. ## methodology in this section, we describe how our proposed alignreg weight regularizers can improve pruning performance in both supervised learning and zero-shot pruning settings. we focus on two regularizers, namely, a neuron correlation-based regularizer (cosine-mbp) and frobenius layer-norm regularizer (frobenius-mbp). let d := {x i , y i } d @xmath0:= (x 1 , . . . , x n ) and x @xmath1.g., @xmath2. for structured prediction (e.g., ner and pos), y @xmath3, y @xmath4. let θ = (θ 1 , . . . , θ l ) be the parameters of a pretrained network f with l layers, where θ l refers to the parameters, including weight matrix w l and bias b l , at layer l. let f θ be a network with parameters θ consisting of weights @xmath5. here, wl := w l m l where m is the pruned mask. for mbp @xcite we remove weights of w l , ∀@xmath6. note that this is a layer-wise process and requires the pruned weights to be masked with m l which has 0 entries corresponding to weights to be pruned and 1 entries for unpruned weights w l . global mbp can also be used whereby the weights {w l } l @xmath7. unlike layer-wise mbp, the percentage of weights removed in each layer can vary for global-mbp. typically, weight regularization is used with mbp to encourage weight sparsity. thus the objective for iterative pruning can be expressed as, ## experimental setup datasets. we perform experiments on multilingual tasks from the xglue benchmark @xcite with pretrained xlm-r base . this covers pairwise classification (qam, qadsm, wpr, xnli), sentence classification (nc) and structured prediction (ner and pos) tasks. ## empirical results we now discuss results on the xglue tasks. news classification (nc) figure 1 shows the results on news classification where a category for news article is predicted and evaluated in 5 languages and trained and iteratively pruned on english text. firstly, we observe the trend in iterative pruning performance degradation is somewhat volatile. from preliminary experiments we found news classification to require only 3 epochs to converge for standard fine-tuning on xlm-roberta base . we find that this task is relatively "similar" to the pretraining task and therefore able to easier recover from pruning steps. overall, both cosine-mbp and frobenius-mbp consistently lead to the best zero-shot test performance across both pruning steps and languages. ## conclusion in this paper, we analysed iterative pruning in the zero-shot setting where a pretrained masked language model uses self-supervised learning on text from various languages but can only use a single language for downstream task fine-tuning. we find that some languages degrade in iterative pruning performance faster than others for some tasks (ner and xnli) and propose a weight regularizer that biases the iteratively pruned model towards learning weight distributions close to the crosslingually aligned pretrained state. this improves over well-established weight regularization methods for magnitude-based pruning in both the standard supervised learning setting and the zero-shot setting.