Pruning Pre-trained Language Models with Principled Importance and Self-regularization (2023)

Abstract: Iterative pruning is one of the most effective compression methods for pre-trained language models. We discovered that finding the optimal pruning decision is an equality-constrained 0-1 Integer Linear Programming problem. The solution to this optimization problem leads to a principled importance criterion which we use to rank parameters during iterative model pruning. To mitigate the poor generalization at high sparsity levels, we propose a self-regularization scheme where model prediction is regularized by the latest checkpoint with increasing sparsity throughout pruning. Our experiments on natural language understanding, question answering, named entity recognition, and data-to-text generation with various Transformer-based PLMs show the effectiveness of the approach at various sparsity levels.

## introduction pre-trained language models (plms) @xcite have significantly advanced the state-of-the-art in various natural language processing tasks @xcite @xcite . however, these models often contain a vast amount of parameters, posing nontrivial requirements for storage and computation. due to this inefficiency, the applications of plms in resource-constrained scenarios are still limited. to resolve the above challenge, model compression @xcite @xcite has been actively studied to make plms meet the practical requirement. among them, iterative pruning methods are widely adopted at only a tiny expense of model performance when adapting plms to downstream tasks. during the course of iterative pruning, model parameters can not only be updated but also * the corresponding author. despite the empirical success, existing importance criteria for model pruning still face two major limitations: (1) they are heuristically defined and may not accurately quantify a parameter's contribution to the learning process, e.g., absolute weight value in magnitude-based pruning and gradient-weight product in sensitivity-based pruning; (2) they determine the importance of each parameter individually without considering the effect of coinstantaneous parameter updates on model performance, e.g., sensitivity is estimated by the absolute change in training error if only a single parameter is pruned and others remain unchanged. in this paper, we rethink the design of the importance criterion for model pruning from an optimization perspective. we begin by analyzing the temporal variation of any given learning objective based on a single-step gradient descent update under the iterative pruning setting. we show that finding the optimal pruning decision can be framed as solving an equality-constrained 0-1 integer linear programming (ilp) problem, where the constraint is defined by the specified sparsity. the resulting problem is a particular case of a general 0-1 knapsack problem in which the weight for each item is the same. the solution to this problem naturally leads to a principled importance criterion which we use to rank all model parameters and derive the optimal stepwise pruning decision. when a high sparsity (e.g., 80%∼90%) is pursued, the limited capacity often renders the pruned model fails to retain satisfactory performance with conventional fine-tuning. to further improve the model's generalization ability, we propose a selfregularization scheme, where the model prediction is regularized by the latest best-performing model checkpoint during pruning. we show that such a scheme eases model learning with decreasing capacity and effectively yields a tighter upper bound of expected generalization error than learning from training data alone. to validate the effectiveness of our approach, dubbed pins (pruning with principled importance and self-regularization), we conducted extensive experiments with various pre-trained language models on a wide variety of tasks, including natural language understanding on glue @xcite ), question answering on squad @xcite , named entity recognition on conll 2003 @xcite , and data-to-text generation on webnlg @xcite , dart @xcite , and e2e @xcite . experimental results show that pins provides more accurate models at different sparsity levels. detailed analysis shed further light on some intriguing properties of models pruned by pins. by exploiting the resulting high sparsity, we show that the storage/inference can be reduced/accelerated by 8.9x and 2.7x using csr format and a sparsityaware inference runtime @xcite ) on consumer-level cpus foot_0 . in summary, our contributions are: in this section, we review the necessary background on transformer-based pre-trained language models and popular importance criteria for iterative pruning. ## methodology instead of heuristically defining the importance criterion as in prior pruning methods, we take a step back and rethink the design of the importance criterion for model pruning from an optimization perspective. from our analysis, we draw an equivalence between finding the optimal stepwise pruning decision and solving an equality-constrained 0-1 integer linear programming problem. we further show that the optimal solution to this problem leads to a new importance criterion for model pruning. moreover, we propose a simple yet effective self-regularization scheme to facilitate the generalization ability of the sparse model. we elucidate our analysis in section 3.1 and describe our self-regularization scheme in section 3.2. ## experiments in this section, we compare pins with state-ofthe-art pruning algorithms and perform detailed analysis to understand the effectiveness of pins. ## conclusion we present pins, a new iterative pruning method that hinges on a principled weight importance criterion to deliver the optimal stepwise pruning decision. integrated with a self-regularization scheme tailored to pruning-during-adaptation, pins allows for provably better generalization ability. empirical experiments and analyses confirm the effectiveness of our method and shed further light on the different sparsity patterns produced by pins and other existing methods.