Layer-wise Model Pruning based on Mutual Information

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3079–3090 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 3079 Layer-wise Model Pruning based on Mutual Information Chun Fan♠⋆▲, Jiwei Li♦♣, Xiang Ao▼, Fei Wu♦, Yuxian Meng♣ and Xiaofei Sun♣ ♣ Shannon.AI, ♦Zhejiang University ▲National Biomedical Imaging Center, Peking University ♠Computer Center of Peking University, ⋆Peng Cheng Laboratory ▼Key Lab of Intelligent Information Processing of Chinese Academy of Sciences {yuxian_meng, jiwei_li, xiaofei_sun}@shannonai.com fanchun@pku.edu.cn, aoxiang@ict.ac.cn, wufei@zju.edu.cn Abstract Inspired by mutual information (MI) based fea- ture selection in SVMs and logistic regression, in this paper, we propose MI-based layer-wise pruning: for each layer of a multi-layer neu- ral network, neurons with higher values of MI with respect to preserved neurons in the up- per 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 prun- ing techniques: (1) it avoids irregular mem- ory 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 perspec- tive based on training signals in the top layer, and prunes each layer by propagating the ef- fect 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 of- fers both greater speedup and higher perfor- mances than weight-based pruning methods (e.g., magnitude pruning, movement pruning). 1 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 exist- ing 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 com- putation intensity and memory footprint of large models at the lowest cost of accuracy on end tasks 1For 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. (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 (Mc- Carley, 2019; Gordon et al., 2020). While weight-based methods have been success- fully applied to a wide range of neural models for model pruning, they come with the following short- comings: …

Content truncated to 500 words