Progressive Weight Pruning Of Deep Neural Networks Using ADMM

Deep neural networks ( DNNs ) , although achieving human-level performance in many domains , have very large model size that hinders their broader applications on edge computing devices . Extensive research work has been conducted on DNN model compression or pruning . However , most of the previous work has taken heuristic approaches . This work proposes a progressive weight pruning approach based on ADMM ( Alternating Direction Method of Multipliers ) , a powerful technique to deal with non-convex optimization problems with potentially combinatorial constraints . Motivated by dynamic programming , the proposed method reaches extremely high pruning rate by using partial prunings with moderate pruning rates . Therefore , it resolves the accuracy degradation and long convergence time problems when pursuing extremely high pruning ratios . It achieves up to 34× pruning rate for ImageNet data set and 167× pruning rate for MNIST data set , significantly higher than those reached by existing work in the literature . Under the same number of epochs , the proposed method also achieves better convergence and higher compression rates . The codes and pruned DNN models are avilable in the link : bit.ly/2zxdlss . 1 INTRODUCTION . Deep neural networks ( DNNs ) have achieved human-level performance in many application domains such as image classification ( Krizhevsky et al. , 2012 ) , object recognition ( LeCun et al. , 1998 ; He et al. , 2016 ) , natural language processing ( Hinton et al. , 2012 ; Dahl et al. , 2012 ) , etc . At the same time , the networks are growing deeper and bigger for higher classification/recognition performance ( i.e. , accuracy ) ( Simonyan & Zisserman , 2015 ) . However , the very large DNN model size increases the computation time of the inference phase . To make matters worse , the large model size hinders DNN ’ deployments on edge computing , which provides the ubiquitous application scenarios of DNNs besides cloud computing applications . As a result , extensive research efforts have been devoted to DNN model compression , in which DNN weight pruning is a representative technique . Han et al . ( 2015 ) is the first work to present the DNN weight pruning method , which prunes the weights with small magnitudes and retrains the network model , heuristically and iteratively . After that , more sophisticated heuristics have been proposed for DNN weight pruning , e.g. , incorporating both weight pruning and growing ( Guo et al. , 2016 ) , L1 regularization ( Wen et al. , 2016 ) , and genetic algorithms ( Dai et al. , 2017 ) . Other improvement directions of weight pruning include trading-off between accuracy and compression rate , e.g. , energy-aware pruning ( Yang et al. , 2017 ) , incorporating regularity , e.g. , channel pruning ( He et al. , 2017 ) , and structured sparsity learning ( Wen et al. , 2016 ) . While the weight pruning technique explores the redundancy in the number of weights of a network model , there are other sources of redundancy in a DNN model . For example , the weight quantization ( Leng et al. , 2017 ; Park et al. , 2017 ; Zhou et al. , 2017 ; Lin et al. , 2016 ; Wu et al. , 2016 ; Rastegari et al. , 2016 ; Hubara et al. , 2016 ; Courbariaux et al. , 2015 ) and clustering ( Zhu et al. , 2017 ; Han et al. , 2016 ) techniques explore the redundancy in the number of bits for weight representation . The activation pruning technique ( Jung et al. , 2018 ; Sharify et al. , 2018 ) leverages the redundancy in the intermediate results . While our work focuses on weight pruning as the major DNN model compression technique , it is orthogonal to the other model compression techniques and might be integrated under a single ADMM-based framework for achieving more compact network models . The majority of prior work on DNN weight pruning take heuristic approaches to reduce the number of weights as much as possible , while preserving the expressive power of the DNN model . Then one may ask , how can we push for the utmost sparsity of the DNN model without hurting accuracy ? and what is the maximum compression rate we can achieve by weight pruning ? Towards this end , Zhang et al . ( 2018b ) took a tentative step by proposing an optimization-based approach that leverages ADMM ( Alternating Direction Method of Multipliers ) , a powerful technique to deal with nonconvex optimization problems with potentially combinatorial constraints . This direct ADMM-based weight pruning technique can be perceived as a smart DNN regularization where the regularization target is dynamically changed in each ADMM iteration . As a result it achieves higher compression ( pruning ) rate than heuristic methods . Inspired by Zhang et al . ( 2018b ) , in this paper we propose a progressive weight pruning approach that incorporates both an ADMM-based algorithm and masked retraining , and takes a progressive means targeting at extremely high compression ( pruning ) rates with negligible accuracy loss . The contributions of this work are summarized as follows : • We make a key observation that when pursuing extremely high compression rates ( say 150× for LeNet-5 or 30× for AlexNet ) , the direct ADMM-based weight pruning approach ( Zhang et al. , 2018b ) can not produce exactly sparse models upon convergence , in that many weights to be pruned are close to zero but not exactly zero . Certain accuracy degradation will result from this phenomenon if we simply set these weights to zero . • We propose and implement the progressive weight pruning paradigm that reaches an extremely high compression rate through multiple partial prunings with progressive pruning rates . This progressive approach , motivated by dynamic programming , helps to mitigate the long convergence time by direct ADMM pruning . • Extensive experiments are performed by comparing with many state-of-the-art weight pruning approaches and the highest compression rates in the literature are achieved by our progressive weight pruning framework , while the loss of accuracy is kept negligible . Our method achieves up to 34× pruning rate for the ImageNet data set and 167× pruning rate for the MNIST data set , with virtually no accuracy loss . Under the same number of epochs , the proposed method achieves notably better convergence and higher compression rates than prior iterative pruning and direct ADMM pruning methods . We provide codes ( both Caffe and TensorFlow versions ) and pruned DNN models ( both for the ImageNet and MNIST data sets ) in the link : bit.ly/2zxdlss . 2 THE PROGRESSIVE WEIGHT PRUNING FRAMEWORK OF DNNS . This section introduces the proposed progressive weight pruning framework using ADMM . Section 2.1 describes the overall framework . Section 2.2 discusses the ADMM-based algorithm for DNN weight pruning ( Zhang et al. , 2018b ) , which we will improve and incorporate into the progressive weight pruning framework . Section 2.3 proposes a direct improvement of masked retraining to restore accuracy . Section 2.4 provides the motivations and details of the proposed progressive weight pruning framework . 2.1 THE OVERALL FRAMEWORK . Masked ADMM regularization Masked retraining Thresholding mask updating Pruned modelPre-trained model Figure 1 : The overall progressive weight pruning framework including masked ADMM-based algorithm , thresholding mask updating , and masked retraining steps . The overall framework of progressive weight pruning is shown in Figure 1 . It applies the ADMMbased pruning algorithm on a pre-trained ( uncompressed ) network model . Then it defines thresholding masks , with which the weights smaller than thresholds are forced to be zero . To restore accuracy , the masked retraining step is applied , that only updates nonzero weights specified by the thresholding masks . The ADMM-based algorithm , thresholding mask updating , and masked retaining steps are performed for several rounds , and each round is considered as a partial pruning , progressively pushing for the utmost of the DNN model pruning . Note that in our progressive weight pruning framework , we change the ADMM-based algorithm into a “ masked ” version that reuses the partially pruned model by masking the gradients of the pruned weights , thereby preventing them from recovering to nonzero weights and thus accelerating convergence . 2.2 ADMM-BASED PRUNING ALGORITHM . Our ADMM-based pruning algorithm takes a pre-trained network as the input and outputs a pruned network model satisfying some sparsity constraints . Consider anN -layer DNN , where the collection of weights in the i-th ( convolutional or fully-connected ) layer is denoted by Wi and the collection of biases in the i-th layer is denoted by bi . The loss function associated with the DNN is denoted by f ( { Wi } Ni=1 , { bi } Ni=1 ) . The DNN weight pruning problem can be formulated as : minimize { Wi } , { bi } f ( { Wi } , { bi } ) , subject to Wi ∈ Si , i = 1 , . . . , N , ( 1 ) where Si = { Wi | card ( Wi ) ≤ li } , i = 1 , . . . , N and li is the desired number of weights in the i-th layer of the DNN . It is clear that S1 , . . . , SN are nonconvex sets , and it is in general difficult to solve optimization problems with nonconvex constraints . The problem can be equivalently rewritten in a format without constraints , namely minimize { Wi } , { bi } f ( { Wi } , { bi } ) + N∑ i=1 gi ( Wi ) , ( 2 ) where gi ( · ) is the indicator function of Si , i.e. , gi ( Wi ) = { 0 if card ( Wi ) ≤ li , +∞ otherwise . ( 3 ) The ADMM technique ( Boyd et al. , 2011 ) can be applied to solve the weight pruning problem by formulating it as : minimize { Wi } , { bi } f ( { Wi } , { bi } ) + N∑ i=1 gi ( Zi ) , subject to Wi = Zi , i = 1 , . . . , N. Through the augmented Lagrangian , the ADMM technique decomposes the weight pruning problem into two subproblems , and solving them iteratively until convergence . The first subproblem is : minimize { Wi } , { bi } f ( { Wi } , { bi } ) + N∑ i=1 ρi 2 ‖Wi − Zki + Uki ‖2F . ( 4 ) This subproblem is equivalent to the original DNN training plus an L2 regularization term , and can be effectively solved using stochastic gradient descent with the same complexity as the original DNN training . Note that we can not prove global optimality of the solution to subproblem ( 4 ) , just as we can not prove optimality of the solution to the original DNN training problem . On the other hand , the second subproblem is : minimize { Zi } N∑ i=1 gi ( Zi ) + N∑ i=1 ρi 2 ‖Wk+1i − Zi + U k i ‖2F . Since gi ( · ) is the indicator function of the set Si , the globally optimal solution to this subproblem can be explicitly derived as in Boyd et al . ( 2011 ) : Zk+1i = ΠSi ( W k+1 i + U k i ) , ( 5 ) where ΠSi ( · ) denotes the Euclidean projection onto the set Si . Note that Si is a nonconvex set , and computing the projection onto a nonconvex set is a difficult problem in general . However , the special structure of Si = { Wi | card ( Wi ) ≤ li } allows us to express this Euclidean projection analytically . Namely , the optimal solution ( 5 ) is to keep the li largest elements of Wk+1i + U k i and set the rest to zero ( Boyd et al. , 2011 ) . Here we introduce set Pi for weights that were pruned to be zero . In every layer , Pi is a subset of Si . By introducing set Pi , we introduce progressive pruning . We will show more detail in Algorithm 1 . Finally , we update the dual variable Ui as Uk+1i = U k i + W k+1 i − Z k+1 i . ( 6 ) This concludes one iteration of the ADMM . In the context of deep learning , the ADMM-based algorithm for DNN weight pruning can be understood as a smart DNN regularization technique ( see Eqn . ( 4 ) ) , in which the regularization target ( in the L2 regularization term ) is dynamically updated in each ADMM iteration . This is one reason that the ADMM-based algorithm for weight pruning achieves higher performance than heuristic methods and other regularization techniques ( Wen et al. , 2016 ) , and the Projected Gradient Descent technique ( Zhang et al. , 2018a ) .

Summary: This paper focus on weight pruning for neural network compression. The proposed method is based on ADMM optimization method for neural network loss with constraint on the l_0 norm of weights, proposed in Zhang et al. 2018b. Two improvements, masked retraining and progressive pruning, are introduced. Masked retraining set the weights to zero at early stages and stop updating those weights. Progressing pruning keeps a buffer of partial pruning results and select the best performed model for further pruning. The proposed method achieves 30x compression rate for AlexNet and VGG for ImageNet.