MLPrune: Multi-Layer Pruning for Automated Neural Network Compression

1 INTRODUCTION . Deep neural networks have proven very successful in many artificial intelligence tasks such as computer vision , natural language processing , and robotic control . In exchange for such success , modern architectures are usually composed of many stacked layers parameterized with a large number of learnable weights . This contrasts classical networks which have hundreds or even thousands times fewer parameters . As a consequence , modern architectures require considerable memory storage and intensive computation . This is problematic in applications that need to run on small embedded systems or that require low-latency to make safety-critical decisions . Fortunately , we can compress these large networks with little to no loss of accuracy , by exploiting the fact that many redundancies exist within their parameters . A popular approach is to quantize the parameters using lower precision , thus encoding the network with fewer bits ( Courbariaux et al. , 2016 ; Gong et al. , 2014 ; Rastegari et al. , 2016 ; Wu et al. , 2016 ; Zhu et al. , 2016 ) . On the other hand , pruning techniques ( Dong et al. , 2017 ; Guo et al. , 2016 ; Han et al. , 2015a ; b ; Li et al. , 2016 ; Wen et al. , 2016 ) aim at removing redundant connections . It has been widely exploited due to its simplicity and efficacy . In this paper , we follow this line of work and propose a new pruning approach . Pruning algorithms differ on the pruning criteria employed . Popular pruning methods typically rely on heuristics , such as weight magnitudes . However , a small magnitude does not necessarily mean unimportance ( Hassibi & Stork , 1993 ; LeCun et al. , 1990 ) ; if the input neuron has a large expected value , a small weight could still have a large effect on its output neuron . As a consequence , magnitudebased pruning might delete important parameters , or preserve unimportant ones . Furthermore , magnitude-based pruning requires manually setting the compression ratio of each layer . These ratios are not easy to tune , as different layers ( and architectures ) typically have different sensitivities for compression . Directly comparing the parameters ’ magnitudes from all layers and pruning the smallest one would not solve the problem , since magnitudes are not calibrated across layers . Because of this , existing approaches usually involve time-consuming trial and error processes that require domain expertise . In this paper we employ the Hessian matrix as a principled pruning criterion , as it characterizes the local curvature of the training loss . The Hessian was exploited in the past , but was either assumed to be computable ( Hassibi & Stork , 1993 ) , which is not possible in practice for modern neural networks with millions of parameters , or simplistic approximations ( e.g. , diagonal ( LeCun et al. , 1990 ) ) were utilized resulting in poor performance . Here we first note that the Fisher Information matrix is close to the Hessian under certain conditions . We then use a Kronecker-factored Approximate Curvature ( K-FAC ) ( Martens & Grosse , 2015 ) method to efficiently estimate the Fisher matrix . In our experiments , the overhead of estimating the Hessian approximation is very small compared with the pre-train and re-train time . Therefore , our method is efficient to use , while exhibiting better compression performance compared with previous second-order pruning . Importantly , this criterion is also calibrated across layers , allowing us to directly compare parameters from different layers and set one global compression ratio for the whole network , rather than one ratio per layer . This makes our method significantly easier to use when compared to magnitude-based pruning . Our method achieves state-of-the-art results on several benchmark datasets and model architectures . Furthermore , it can be easily applied to new customized network architectures , without costly tuning many compression hyper-parameters . 2 BACKGROUND . In this section , we first introduce the notation used in the paper . Since our work is inspired by the Optimal Brain Surgeon ( Hassibi & Stork , 1993 ) , we also briefly review the idea and its limitations . 2.1 NOTATION . Suppose we have a neural network with L fully-connected layers1 . Let p ( y|x , Θ ) be the output distribution and L be the loss function . y denotes the target outputs , x the inputs , and Θ the network ’ s weights . Each layer l is associated with its input al−1 ∈ Rdl−1 , a weight matrix Wl ∈ Rdl−1×dl , a binary mask matrix Γl ∈ Rdl−1×dl and its output sl ∈ Rdl . The elements γ ( i , j ) l ∈ { 0 , 1 } in the binary matrix Γl indicate whether the corresponding weight is retained or pruned away . We use subscript to denote the layer index , and superscript to denote the element index in a matrix . We also distinguish a matrix ( vector ) from a scalar by capital boldface2 . The weights in a layer l can be computed as θ ( i , j ) l = w ( i , j ) l γ ( i , j ) l , and thus the forward-pass can be written as , sl = ( Wl Γl ) al−1 , al = h ( sl ) , ( 1 ) where denotes element-wise product , and h is the Relu activation in our experiments . We also use W to denote the concatenation of [ vec { W1 } , vec { W2 } , · · · , vec { WL } ] , which is a column vector with the same size as the total number of parameters in this neural network , similarly for Γ and Θ . Index q denotes the index number of a parameter in W ( or Γ , Θ ) that is associated with w ( i , j ) l ( or γ ( i , j ) l , θ ( i , j ) l ) . Parameters , weights and connections are used interchangeably . 2.2 A REVIEW ON OPTIMAL BRAIN SURGEON ( OBS ) . Given a neural network with parameters Θ , the local surface of the training loss can be characterized by its Taylor expansion : δL = ∂L ∂Θ δΘ + 1 2 δΘTHδΘ +O ( ||δΘ||3 ) , ( 2 ) where δΘ is a small change of the parameter values ( e.g. , setting some parameters to zero ) and H is the Hessian matrix defined as ∂2L/∂Θ2 . Note that at convergence the first-order term vanishes and the higher-order term can be neglected . OBS aims to find an index q , such that when changing θq from its original value to zero , the change in training loss δL is minimized : min q ( min δΘ ( 1 2 δΘTHδΘ ) ) , s.t . eTq δΘ + δΘ = 0 , ( 3 ) where eTq is the unit vector in the weight space corresponding to θq . The inner minimization problem can be solved by the method of Lagrangian multipliers . Hassibi & Stork ( 1993 ) applied OBS on a small network with 18,000 weights and showed good results . However , since the exact solution involves the inverse of the Hessian matrix , original OBS is intractable for modern neural networks that contain millions of parameters . 1Convolutional layers can be transformed into fully-connected layers if we vectorize filters within each channel and expand the input feature map with repeated patches 2with minor abuse of notation for al−1 and sl 3 MULTI-LAYER PRUNING . Our method follows a similar pipeline when compared to conventional pruning methods . We first pre-train a neural network until converged , and then apply Multi-Layer Pruning ( including surgeon ) upon it . This will give us a compressed network with slightly worse performance . Finally , the network is re-trained to recover the model performance . In the following paragraph , we first introduce the problem formulation and its ideal solution in Section 3.1 . We then show how to make approximations in Section 3.2 . 3.1 MULTI-LAYER PRUNING AND SURGEON . Our objective is to prune as much as possible while maintaining good model performance , e.g . high accuracy in classification . This is similar to the objective of Minimum Description Length ( MDL , ( Rissanen , 1978 ) ) min Θ D ( data|Θ ) +D ( Θ ) . ( 4 ) In our case , the first term in Eq . ( 4 ) is equivalent to the training loss , and the second term is the description length denoting how many bits we need to encode the model . Here , we use the number of parameters as an indicator of such description length . We also introduce a relative importance λ , which characterizes the trade-off between sparsity and model performance , and controls the final size of a compressed model . Our objective can then be written as Ψ = min W , Γ L ( y|x , W Γ ) + λ L∑ l=1 ∑ ( i , j ) γ ( i , j ) l , ( 5 ) where Γ is the binary mask matrix defined in Section 2.1 . This objective can not be directly optimized with SGD as γ ( i , j ) l is contrained to be binary . Instead , we consider an easier case where we first pre-train the model to the local minimum of L , then update one γ ( i , j ) l from 1 to 0 if such change will decrease Ψ , or equivalently , if such change will not increase L more than λ . Since the parameter for forward-pass is θ ( i , j ) l = γ ( i , j ) l w ( i , j ) l , updating γ ( i , j ) l is equivalent to prune a parameter θ ( i , j ) l . Similar to Hassibi & Stork ( 1993 ) , we also consider a surgeon operation after pruning . This will change other preserved parameters accordingly , so as to minimize the loss increment caused by the pruning . More specifically , after pruning one parameter θ ( i , j ) l and applying the surgeon , the increment of L will be , min δΘ ( 1 2 δΘTHδΘ ) , s.t . δθ ( i , j ) l + w ( i , j ) l = 0 . ( 6 ) δθ ( i , j ) l is the corresponding scalar value in δΘ , and H is the Hessian matrix of L w.r.t ( W Γ ) . Solving such a minimization problem with Lagrangian multipliers results in the following update ∆Lq = 1 2 ( w ( i , j ) l ) 2 [ H−1 ] ( q , q ) , δΘ∗ = [ − w ( i , j ) l [ H−1 ] ( q , q ) H−1 ] ( . , q ) . ( 7 ) Here , δΘ∗ is the surgeon applied on the remaining Θ ( and W accordingly ) . ∆Lq is the increment of loss after pruning and the surgeon operation . Again , superscript ( q , q ) denotes the element at the qth row and qth column , ( . , q ) denotes the qth column in that matrix , and q is the index of the parameter in H that associates with w ( i , j ) l . The estimation of H −1 will be introduced in section 3.2 . Theoretically , we can calculate ∆Lq for each parameter and prune the one with the smallest ∆L , but this will become impractically time-consuming . Therefore , at each pruning step , we calculate ∆Lq for each remaining parameter , and prune all parameters with ∆Lq < λ . This is fine as long as such assumption holds , that H does not change significantly when we prune those unimportant parameters one by one . In our experiments , we simply set λ to be the pth percentile of all ∆Lq from all layers . Therefore , our method can jointly determine the importance of parameters from all layers and prune multiple layers simultaneously . After pruning and doing surgeon , we apply SGD to those retained parameters to re-train the network , W←W − α∂L ( W Γ ) ∂ ( W Γ ) Γ . ( 8 )

Summary: The paper proposes a multi-layer pruning method called MLPrune for neural networks, which can automatically decide appropriate compression ratios for all the layers. It firstly pre-trains a network. Then it utilizes K-FAC to approximate the Fisher matrix, which in turn approximates the exact Hessian matrix of training loss w.r.t model weights. The approximated Hessian matrix is then used to estimate the increment of loss after pruning a connection. The connections from all layers with the smallest loss increments are pruned and the network is re-trained to the final model.