Balancing training time vs. performance with Bayesian Early Pruning

1 INTRODUCTION . Deep neural networks ( DNNs ) are known to be overparameterized ( Allen-Zhu et al. , 2019 ) as they usually have more learnable parameters than needed for a given learning task . So , a trained DNN contains many ineffectual parameters that can be safely pruned or zeroed out with little/no effect on its predictive accuracy . Pruning ( LeCun et al. , 1989 ) is an approach to alleviating overparameterization of a DNN by identifying and removing its ineffectual parameters while preserving its predictive accuracy on the validation/test dataset . Pruning is typically applied to the DNN after training to speed up testing-time evaluation . For standard image classification tasks with MNIST , CIFAR-10 , and ImageNet datasets , it can reduce the number of learnable parameters by up to 50 % or more while maintaining test accuracy ( Han et al. , 2015 ; Li et al. , 2017 ; Molchanov et al. , 2017 ) . In particular , the overparameterization of a DNN also leads to considerable training time being wasted on those DNN elements ( e.g. , connection weights , neurons , or convolutional filters ) which are eventually ineffectual after training and can thus be safely pruned . Our work in this paper considers early pruning of such DNN elements by identifying and removing them throughout the training process instead of after training.1 As a result , this can significantly reduce the time incurred by the training process without compromising the final test accuracy ( upon convergence ) much . Recent work ( Section 5 ) in foresight pruning ( Lee et al. , 2019 ; Wang et al. , 2020 ) show that pruning heuristics applied at initialization work well to prune connection weights without significantly degrading performance . In contrast to these work , we prune throughout the training procedure , which improves performance after convergence of DNNs , albeit with somewhat longer training times . In this work , we pose early pruning as a constrained optimization problem ( Section 3.1 ) . A key challenge in the optimization is accurately modeling the future efficacy of DNN elements . We achieve this through the use of multi-output Gaussian process which models the belief of future efficacy conditioned upon efficacy measurements collected during training ( Section 3.2 ) . Although the posed optimization problem is NP-hard , we derive an efficient Bayesian early pruning ( BEP ) approximation algorithm , which appropriately balances the inherent training time vs. performance tradeoff in pruning prior to convergence ( Section 3.3 ) . Our algorithm relies on a measure of network element efficacy , termed saliency ( LeCun et al. , 1989 ) . The development of saliency functions is an active area of research with no clear optimal choice . To accomodate this , our algorithm is agnostic , and therefore 1In contrast , foresight pruning ( Wang et al. , 2020 ) removes DNN elements prior to the training process . flexible , to changes in saliency function . We use BEP to prune neurons and convolutional filters to achieve practical speedup during training ( Section 4 ) .2 Our approach also compares favorably to state-of-the-art works such as SNIP ( Lee et al. , 2019 ) , GraSP ( Wang et al. , 2020 ) , and momentum based dynamic sparse reparameterization ( Dettmers & Zettlemoyer , 2019 ) . 2 PRUNING . Consider a dataset of D training examples X = { x1 , . . . , xD } , Y = { y1 , . . . , yD } and a neural network Nvt parameterized by a vector of M pruneable network elements ( e.g . weight parameters , neurons , or convolutional filters ) vt , [ vat ] a=1 , ... , M , where vt represent the network elements after t iterations of stochastic gradient descent ( SGD ) for t = 1 , . . . , T . Let L ( X , Y ; Nvt ) be the loss function for the neural network Nvt . Pruning aims at refining the network elements vt given some sparsity budget B and preserving the accuracy of the neural network after convergence ( i.e. , NvT ) , which can be stated as a constrained optimization problem ( Molchanov et al. , 2017 ) : minm∈ { 0,1 } M |L ( X , Y ; Nm vT ) − L ( X , Y ; NvT ) | s.t . ||m||0 ≤ B ( 1 ) where is the Hadamard product and m is a pruning mask . Note that we abuse the Hadamard product for notation simplicity : for a = 1 , .. , M , ma × vaT corresponds to pruning vaT if ma = 0 , and keeping vaT otherwise . Pruning a network element refers to zeroing the network element or the weight parameters which compute the network element . Any weight parameters which reference the output of the pruned network element are also zeroed since the element outputs a constant 0 . The above optimization problem is difficult due to the NP-hardness of combinatorial optimization . This leads to the approach of using saliency function s which measures efficacy of network elements at minimizing the loss function . A network element with small saliency can be pruned since it ’ s not salient in minimizing the loss function . Consequently , pruning can be done by maximizing the saliency of the network elements given the sparsity budget B : maxm∈ { 0,1 } M ∑M a=1m as ( a ; X , Y , NvT , L ) s.t . ||m||0 ≤ B ( 2 ) where s ( a ; X , Y , NvT , L ) measures the saliency of vaT at minimizing L after convergence through T iterations of SGD . The above optimization problem can be efficienctly solved by selecting the B most salient network elements in vT . The construction of the saliency function has been discussed in many existing works : Some approaches derived the saliency function from first-order ( LeCun et al. , 1989 ; Molchanov et al. , 2017 ) and second-order ( Hassibi & Stork , 1992 ; Wang et al. , 2020 ) Taylor series approximations of L. Other common saliency functions include L1 ( Li et al. , 2017 ) or L2 ( Wen et al. , 2016 ) norm of the network element weights , as well as mean activation ( Polyak & Wolf , 2015 ) . In this work , we use a first-order Taylor series approximation saliency function defined for neurons and convolutional filters3 ( Molchanov et al. , 2017 ) , however our approach remains flexible to arbitrary choice of saliency function on a plug-n-play basis . 3 BAYESIAN EARLY PRUNING . 3.1 PROBLEM STATEMENT . As has been mentioned before , existing pruning works based on the saliency function are typically done after the training convergence ( i.e. , ( 2 ) ) to speed up the testing-time evaluation , which waste considerable time on training these network elements which will eventually be pruned . To resolve this issue , We extend the pruning problem definition ( 2 ) along the temporal dimension , allowing network elements to be pruned during the training process consisting of T iterations of SGD . 2Popular deep learning libaries do not accelerate sparse matrix operations over dense matrix operations . Thus , pruning network connections can not be easily capitalized upon with performance improvements . It is also unclear whether moderately sparse matrix operations ( i.e. , operations on matrices generated by connection pruning ) can be significantly accelerated on massively parallel architectures such as GPUs ( see Yang et al . ( 2018 ) Fig . 7 ) . See Section 5 in Buluç & Gilbert ( 2008 ) for challenges in parallel sparse matrix multiplication . 3Implementation details of this saliency function can be found in Appendix A.1 . Let sat , s ( a ; X , Y , Nvt , L ) be a random variable which denotes the saliency of network element vat after t iterations of SGD , st , [ sat ] a=1 , ... , M for t = 1 , . . . , T , and sτ1 : τ2 , [ st ] t=τ1 , ... , τ2 be a vector of saliency of all the network elements between iterations τ1 and τ2 . Our early pruning algorithm is designed with the goal of maximizing the saliency of the unpruned elements after iteration T , yet allowing for pruning at each iteration t given some computational budget Bt , c for t = 1 , . . . , T : ρT ( mT−1 , BT , c , Bs ) , maxmT mT · sT ( 3a ) s.t . ||mT ||0 ≤ Bs ( 3b ) mT ≤̇mT−1 ( 3c ) BT , c ≥ 0 ( 3d ) ρt ( mt−1 , Bt , c , Bs ) , maxmtEp ( st+1|s̃1 : t ) [ ρt+1 ( mt , Bt , c−||mt||0 , Bs ) ] ( 4a ) s.t . mt ≤̇mt−1 ( 4b ) where Bs is the trained network sparsity budget , s̃1 : t is a vector of observed values for s1 : t , m0 is an M -dimensional 1 ’ s vector , andmt ≤̇mt−1 represents an element-wise comparison between mt and mt−1 : mat ≤ mat−1 for a = 1 , . . . , M . At each iteration t , the saliency st is observed and mt ∈ { 0 , 1 } M in ρt represents a pruning decision performed to maximize the expectectation of ρt+1 conditioned upon saliency measurements s1 : t collected up to and including iteration t. This recursive structure terminates with base case ρT where the saliency of the unpruned elements is maximized after T iterations of training . In the above early pruning formulation4 , constraints ( 3c ) and ( 4b ) ensure pruning is performed in a practical manner whereby once a network element is pruned , it can no longer be recovered in a later training iteration . We define a trained network sparsity budget , Bs ( 3b ) , which may differ significantly from initial network size ||m0||0 ( e.g. , in the case where the network is trained on GPUs , but deployed on resource constrained edge or mobile devices ) . We also constrain a total computational effort budget Bt , c which is reduced per training iteration t by the number of unpruned network elements ||mt||0 . We constrain BT , c ≥ 0 ( 3d ) to ensure training completion within the specified computational budget . Here we assume that a more sparse pruning maskmt corresponds to lower computational effort during training iteration t due to updating fewer network elements . Finally , ( 3a ) maximizes the saliency with a pruning maskmT constrained by a sparsity budget Bs ( 3b ) . Our early pruning formulation balances the saliency of network elements after convergence against the total computational effort to train such network ( i.e. , mT · sT vs. ∑T t=1||mt||0 ) . This appropriately captures the balancing act of training-time early pruning whereby the computational effort is saved by early pruning network elements while preserving the saliency of the remaining network elements after convergence . 3.2 MODELING THE SALIENCY WITH MULTI-OUTPUT GAUSSIAN PROCESS . To solve the above early pruning problem , we need to model the belief p ( s1 : T ) of the saliency for computing the predictive belief p ( st+1 : T |s̃1 : t ) of the future saliency in ( 4a ) . At the first glance , one may consider to decompose the belief : p ( s1 : T ) , ∏M a=1 p ( s a 1 : T ) and model the saliency s a 1 : T , [ sat ] t=1 , ... , T of each network element independently . Such independent models , however , ignore the co-adaptation and co-evolution of the network elements which have been shown to be a common occurrence in DNN ( Hinton et al. , 2012 ; Srivastava et al. , 2014 ; Wang et al. , 2020 ) . Also , modeling the correlations between the saliency of different network elements explicitly is non-trivial since considerable feature engineering is needed for representing diverse network elements such as neurons , connections , or convolutional filters . To resolve such issues , we use multi-output Gaussian process ( MOGP ) to jointly model the belief p ( s1 : T ) of all saliency measurements . To be specific , we assume that the saliency sat of the a-th network element at iteration t is a linear mixture5 of Q independent latent functions { uq ( t ) } Qq=1 : sat , ∑Q q=1 γ a q uq ( t ) . As shown in ( Álvarez & Lawrence , 2011 ) , if each uq ( t ) is an independent GP with prior zero mean and covariance kq ( t , t′ ) , then the resulting distribution over p ( s1 : T ) is a multivariate Gaussian distribution with prior zero mean and covariance determined by the mixing 4In contrast to PruneTrain ( Lym et al. , 2019 ) , our problem definition balances training time vs. performance under an additional constraint on the trained network size ( 3b ) . We discuss this further in Section 5 . 5Among the various types of MOGPs ( see Álvarez & Lawrence ( 2011 ) for a detailed review . ) , we choose this linear model such that the correlations between sat and sa ′ t′ can be computed analytically . weights : cov [ sat , s a′ t′ ] = ∑Q q=1 γ a q γ a′ q kq ( t , t ′ ) . This explicit covariance between sat and s a′ t′ helps to exploit the co-evolution and co-adaptation of network elements within the neural networks . To capture the horizontal asymptote trend of sa1 , . . . , s a T as visualized in Appendix A.2 , we turn to a kernel used for modeling decaying exponential curves known as the “ exponential kernel ” ( Swersky et al. , 2014 ) and set kq ( t , t′ ) , βq αq ( t+t′+βq ) αq where αq and βq are hyperparameters of MOGP and can be learned via maximum likelihood estimation ( Álvarez & Lawrence , 2011 ) . Then , given a vector of observed saliency s̃1 : t , the MOGP regression model can provide a Gaussian predictive distribution for any future saliency st′ . Thus , the predictive mean µat′|1 : t , E [ s a t′ | s̃1 : t ] of the saliency sat′ and the predictive ( co ) variance σaa ′ t′|1 : t , cov [ s a t′ , s a′ t′ | s̃1 : t ] between the saliency sat′ and sa ′ t′ can be computed analytically , as detailed in Appendix A.3 .

Summary: This paper introduces a method for pruning during the training process in order to filter out unimportant/redundant components of the network continuously to speed up training and perform gradual pruning over the training process. The proposed approach is novel in the sense that the vast amount of prior work on pruning has focused on either (i) pruning on network initialization (e.g., SNIP, etc.) or (ii) pruning after the network has been fully trained (e.g., Magnitude Pruning, among many others). The introduced method uses the Taylor-series based saliency criterion (of Molchanov et al., 2017) and uses a multi-output Gaussian process to predict future saliencies and to determine whether a parameter can be safely removed early on during training.