Picking Winning Tickets Before Training by Preserving Gradient Flow

1 INTRODUCTION . Deep neural networks exhibit good optimization and generalization performance in the overparameterized regime ( Zhang et al. , 2016 ; Neyshabur et al. , 2019 ; Arora et al. , 2019 ; Zhang et al. , 2019b ) , but both training and inference for large networks are computationally expensive . Network pruning ( LeCun et al. , 1990 ; Hassibi et al. , 1993 ; Han et al. , 2015b ; Dong et al. , 2017 ; Zeng & Urtasun , 2019 ; Wang et al. , 2019 ) has been shown to reduce the test-time resource requirements with minimal performance degradation . However , as the pruning is typically done to a trained network , these methods don ’ t save resources at training time . Moreover , it has been argued that it is hard to train sparse architectures from scratch while maintaining comparable performance to their dense counterparts ( Han et al. , 2015a ; Li et al. , 2016 ) . Therefore , we ask : can we prune a network prior to training , so that we can improve computational efficiency at training time ? Recently , Frankle & Carbin ( 2019 ) shed light on this problem by proposing the Lottery Ticket Hypothesis ( LTH ) , namely that there exist sparse , trainable sub-networks ( called “ winning tickets ” ) within the larger network . They identify the winning tickets by taking a pre-trained network and removing connections with weights smaller than a pre-specificed threshold . They then reset the remaining weights to their initial values , and retrain the sub-network from scratch . Hence , they showed that the pre-trained weights are not necessary , only the pruned architecture and the corresponding initial weight values . Nevertheless , like traditional pruning methods , the LTH approach still requires training the full-sized network in order to identify the sparse sub-networks . Can we identify sparse , trainable sub-networks at initialization ? This would allow us to exploit sparse computation with specified hardware for saving computation cost . ( For instance , Dey et al . ( 2019 ) demonstrated 5x efficiency gains for training networks with pre-specified sparsity . ) At first glance , a randomly initialized network seems to provide little information that we can use to judge the importance of individual connections , since the choice would seem to depend on complicated training dynamics . However , recent work suggests this goal is attainable . Lee et al . ( 2018 ) proposed the first algorithm for pruning at initialization time : Single-shot Network Pruning ( SNIP ) , which uses a connection sensitivity criterion to prune weights with both small magnitude and small gradients . Their empirical results are promising in the sense that they can find sparse , trainable sub-networks at initialization . However , connection sensitivity is sub-optimal as a criterion because the gradient of each weight might change dramatically after pruning due to complicated interactions between weights . Since SNIP only considers the gradient for one weight in isolation , it could remove connections that are important to the flow of information through the network . Practically , we find that this blocking of information flow manifests as a reduction in the norm of the gradient . Therefore , we aim to prune connections in a way that accounts for their role in the network ’ s gradient flow . Specifically , we take the gradient norm after pruning as our criterion , and prune those weights whose removal will result in least decrease in the gradient norm after pruning . Because we rely on preserving the gradient flow to prune the network , we name our method Gradient Signal Preservation ( GraSP ) . Our approach is easy to implement and conceptually simple . Moreover , the recently introduced Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) provides tools for studying the learning dynamics in the output space . Building on the analysis of Arora et al . ( 2019 ) , we show that our pruning criterion tends to keep those weights which will be beneficial for optimization . We evaluate GraSP on CIFAR-10 , CIFAR-100 ( Krizhevsky , 2009 ) , Tiny-ImageNet and ImageNet ( Deng et al. , 2009 ) with modern neural networks , such as VGGNet ( Simonyan & Zisserman , 2014 ) and ResNet ( He et al. , 2016 ) . GraSP significantly outperforms SNIP in the extreme sparsity regime . 2 RELATED WORK AND BACKGROUND . In this section , we review the literature on neural network pruning including pruning after training , pruning during training , dynamic sparse training and pruning before training . We then discuss propagation of signals in deep neural networks and recent works in dynamical isometry and meanfield theory . Lastly , we also review the Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) , which builds up the foundation for justifying our method in Section 4 . 2.1 NETWORK PRUNING . After training . Most pruning algorithms ( LeCun et al. , 1990 ; Hassibi et al. , 1993 ; Dong et al. , 2017 ; Han et al. , 2015b ; Li et al. , 2016 ; Molchanov et al. , 2016 ) operate on a pre-trained network . The main idea is to identify those weights which are most redundant , and whose removal will therefore least degrade the performance . Magnitude based pruning algorithms ( Han et al. , 2015b ; a ) remove those weights which are smaller than a threshold , which may incorrectly measure the importance of each weight . In contrast , Hessian-based pruning algorithms ( LeCun et al. , 1990 ; Hassibi et al. , 1993 ) compute the importance of each weight by measuring how its removal will affect the loss . More recently , Wang et al . ( 2019 ) proposed a network reparameterization based on the Kroneckerfactored Eigenbasis for further boosting the performance of Hessian-based methods . However , all the aforementioned methods require pre-training , and therefore aren ’ t applicable at initialization . During training . There are also some works which attempt to incorporate pruning into the training procedure itself . Srinivas & Babu ( 2016 ) proposed generalized dropout , allowing for tuning the individual dropout rates during training , which can result in a sparse network after training . Louizos et al . ( 2018 ) proposed a method for dealing with discontinuity in training L0 norm regularized networks in order to obtain sparse networks . Both methods require roughly the same computational cost as training the full network . Dynamic Sparse Training Another branch of pruning algorithms is Dynamic Sparse Training methods , which will dynamically change the weight sparsity during training . Representative works , such as Bellec et al . ( 2018 ) ; Mocanu et al . ( 2018 ) ; Mostafa & Wang ( 2019 ) ; Dettmers & Zettlemoyer ( 2019 ) , follow a prune-redistribute-regrowth cycle for pruning . Among them , Dettmers & Zettlemoyer ( 2019 ) , proposed the sparse momentum algorithm , which dynamically determines the sparse mask based on the mean momentum magnitude during training . However , their method requires maintaining the momentum of all the weights during training , and thus does not save memory . These techniques generally achieve higher accuracy compared with fixed sparse connectivity , but change the standard training procedure and therefore do not enjoy the potential hardware acceleration . Before training . Pruning at initialization is more challenging because we need to account for the effect on the training dynamics when removing each weight . There have been several attempts to conduct pruning before training . Frankle & Carbin ( 2019 ) ; Frankle et al . ( 2019 ) proposed and validated the Lottery Ticket Hypothesis ( LTH ) , namely that the network structure found by traditional pruning algorithms and the corresponding initialization are together sufficient for training the subnetwork from scratch . Lee et al . ( 2018 ) proposed the SNIP algorithm , which was the first attempt to directly identify trainable and sparse sub-networks at initialization time . Their method was based on connection sensitivity , which aims to preserve the loss after pruning , and achieved impressive results . Concurrently to our work , Lee et al . ( 2019b ) studied the pruning problem from a signal propagation perspective , and proposed to use an orthogonal initialization to ensure faithful signal propagation . Though their work shares the same spirit as our GraSP algorithm , they focused on the weight initialization scheme , while we focus on the pruning criterion . 2.2 SIGNAL PROPAGATION AT INITIALIZATION . Our pruning criteria shares the same spirit as recent works in dynamical isometry and mean-field theory ( Saxe et al. , 2013 ; Xiao et al. , 2018 ; Yang & Schoenholz , 2017 ; Poole et al. , 2016 ) where they derived initialization scheme theoretically by developing a mean field theory for signal propagation and by characterizing singular values of the input-output Jacobian matrix . Particularly , Xiao et al . ( 2018 ) successfully train 10,000-layers vanilla ConvNets with specific initialization scheme . Essentially , this line of work shows that the trainability of a neural network at initialization is crucial for final performance and convergence . While they focus on address the trainability issue of very deep networks , we aim to solve the issue of sparse neural networks . Besides , they measure the trainability by examining the input-output Jacobian while we do that by checking the gradient norm . Though different , gradient norm is closely related to the input-output Jacobian . 2.3 NEURAL TANGENT KERNEL AND CONVERGENCE ANALYSIS . Jacot et al . ( 2018 ) analyzed the dynamics of neural net training by directly analyzing the evolution of the network ’ s predictions in output space . Let L denote the cost function , X the set of all training samples , Z = f ( X ; θ ) ∈ Rnk×1 the outputs of the neural network , and k and n the output space dimension and the number of training examples . For a step of gradient descent , the change to the network ’ s predictions can be approximated with a first-order Taylor approximation : f ( X ; θt+1 ) = f ( X ; θt ) − ηΘt ( X , X ) ∇ZL , ( 1 ) where the matrix Θt ( X , X ) is the Neural Tangent Kernel ( NTK ) at time step t : Θt ( X , X ) = ∇θf ( X ; θt ) ∇θf ( X ; θt ) > ∈ Rnk×nk , ( 2 ) where ∇θf ( X ; θ ) denotes the network Jacobian over the whole training set . Jacot et al . ( 2018 ) showed that for infinitely wide networks , with proper initialization , the NTK exactly captures the output space dynamics throughout training . In particular , Θt ( X , X ) remains constant throughout training . Arora et al . ( 2019 ) used the NTK to analyze optimization and generalization phenomena , showing that under the assumptions of constant NTK and squared error loss , the training dynamics can be analyzed in closed form : ‖Y − f ( X ; θt ) ‖2 = √√√√ n∑ i=1 ( 1− ηλi ) 2t ( u > i Y ) 2 ± ( 3 ) where Y ∈ Rnk×1 is all the targets , Θ = UΛU > = ∑n i=1 λiuiu > i is the eigendecomposition , and is a bounded error term . Although the constant NTK assumption holds only in the infinite width limit , Lee et al . ( 2019a ) found close empirical agreement between the NTK dynamics and the true dynamics for wide but practical networks , such as wide ResNet architectures ( Zagoruyko & Komodakis , 2016 ) .

Summary: This paper proposes a novel one-shot-pruning algorithm which improves the training of sparse networks by maximizing the norm of the gradient at initialization. The utility of training sparse neural networks and shortcomings of dense-to-sparse algorithms like Pruning, LotteryTicket are nicely motivated at introduction. The pruning criterion is motivated by the first order approximation of the change in the gradient norm when a single connection is removed, though the results show that removing many connections together with GraSP increases the total gradient norm therefore allowing the loss to decrease faster. Experiments suggest employing such pruning algorithm improves final performance over two baselines: random and SNIP.