Layer-wise Model Pruning based on Mutual Information (2021)

Abstract: Inspired by mutual information (MI) based feature selection in SVMs and logistic regression, in this paper, we propose MI-based layer-wise pruning: for each layer of a multi-layer neural network, neurons with higher values of MI with respect to preserved neurons in the upper layer are preserved. Starting from the top softmax layer, layer-wise pruning proceeds in a top-down fashion until reaching the bottom word embedding layer. The proposed pruning strategy offers merits over weight-based pruning techniques: (1) it avoids irregular memory access since representations and matrices can be squeezed into their smaller but dense counterparts, leading to greater speedup; (2) in a manner of top-down pruning, the proposed method operates from a more global perspective based on training signals in the top layer, and prunes each layer by propagating the effect of global signals through layers, leading to better performances at the same sparsity level. Extensive experiments show that at the same sparsity level, the proposed strategy offers both greater speedup and higher performances than weight-based pruning methods (e.g., magnitude pruning, movement pruning).

## 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. 1 as a result, training and using existing huge models not only requires rich hardware resources, but also consumes high environmental costs @xcite . 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 @xcite @xcite . 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 @xcite , first-order derivative @xcite and second-order derivative information @xcite , and it has been successfully applied to a large variety of model architectures @xcite @xcite and downstream tasks @xcite . 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 @xcite in svms and logistic regression, we propose mi based layerwise 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 weightbased 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 @xcite , movement pruning @xcite and l 0 pruning @xcite , the proposed method yields greater speedup along with better performances for the same sparsity levels on generative nlp tasks of wmt'14 @xmath0, and discriminative nlp tasks of squad v1.1 @xcite , mnli @xcite and sst-5 @xcite . 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-theart performance of 43.9 bleu score for @xmath1. ## 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 @xcite to maintain a sourcetarget vocabulary of 37,000. we use transformers @xcite 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 ## 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. ## 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.