A Signal Propagation Perspective for Pruning Neural Networks at Initialization

1 INTRODUCTION . Deep learning has made great strides in machine learning and been applied to various fields from computer vision and natural language processing , to health care and playing games ( LeCun et al. , 2015 ) . Despite the immense success , however , it remains challenging to deal with the excessive computational and memory requirements of large neural network models . To this end , lightweight models are often preferred , and network pruning , a technique to reduce parameters in a network , has been widely employed to compress deep neural networks ( Han et al. , 2016 ) . Nonetheless , designing pruning algorithms has been often purely based on ad-hoc intuition lacking rigorous underpinning , partly because pruning was typically carried out after training the model as a post-processing step or interwoven with the training procedure , without adequate tools to analyze . Recently , Lee et al . ( 2019 ) have shown that pruning can be done on randomly initialized neural networks in a single-shot prior to training ( i.e. , pruning at initialization ) . They empirically showed that as long as the initial random weights are drawn from appropriately scaled Gaussians ( e.g. , Glorot & Bengio ( 2010 ) ) , their pruning criterion called connection sensitivity can be used to prune deep neural networks , often to an extreme level of sparsity while maintaining good accuracy once trained . However , it remains unclear as to why pruning at initialization is effective , how it should be understood theoretically and whether it can be extended further . In this work , we first look into the effect of initialization on pruning , and find that initial weights have critical impact on connection sensitivity , and therefore , pruning results . Deeper investigation shows that connection sensitivity is determined by an interplay between gradients and weights . Therefore when the initial weights are not chosen appropriately , the propagation of input signals into layers of these random weights can result in saturating error signals ( i.e. , gradients ) under backpropagation , and hence unreliable connection sensitivity , potentially leading to a catastrophic pruning failure . This result leads us to develop a signal propagation perspective for pruning at initialization , and to provide a formal characterization of how a network needs to be initialized for reliable connection sensitivity measurements and in turn effective pruning . Precisely , we show that a sufficient condition to ensure faithful1 connection sensitivity is layerwise dynamical isometry , which is defined as all singular values of the layerwise Jacobians being concentrated around 1 . Our signal propagation perspective is inspired by the recent literature on dynamical isometry and mean field theory ( Saxe et al. , 2014 ; Poole et al. , 2016 ; Schoenholz et al. , 2017 ; Pennington et al. , 2017 ) , in which the general signal propagation in neural networks is studied . We extend this result to understanding and improving pruning at initialization . Moreover , we study signal propagation in the pruned sparse networks and its effect on trainability . We find that pruning neural networks can indeed break dynamical isometry , and hence , hinders signal propagation and degrades the training performance of the resulting sparse network . In order to address this issue , we propose a simple , yet effective data-free method to recover the layerwise orthogonality given the sparse topology , which in turn improves the training performance of the compressed network significantly . Our analysis further reveals that in addition to signal propagation , the choice of pruning method and sparsity level can influence trainability in sparse neural networks . Perfect layerwise dynamical isometry can not always be ensured in the modern networks that have components such as ReLU nonlinearities ( Pennington et al. , 2017 ) and/or batch normalization ( Yang et al. , 2019 ) . Even in such cases , however , our experiments on various modern architectures ( including convolutional and residual neural networks ) indicate that connection sensitivity computed based on layerwise dynamical isometry is robust and consistently outperforms pruning based on other initialization schemes . This indicates that the signal propagation perspective is not only important to theoretically understand pruning at initialization , but also it improves the results of pruning for a range of networks of practical interest . Furthermore , this signal propagation perspective for pruning poses another important question : how informative is the error signal computed on randomly initialized networks , or can we prune neural networks even without supervision ? To understand this , we compute connection sensitivity scores with different unsupervised surrogate losses and evaluate the pruning results . Interestingly , our results indicate that we can in fact prune networks in an unsupervised manner to extreme sparsity levels without compromising accuracy , and it often compares competitively to pruning with supervision . Moreover , we test if pruning at initialization can be extended to obtain architectures that yield better performance than standard pre-designed architectures with the same number of parameters . In fact , this process , which we call neural architecture sculpting , compares favorably against hand-designed architectures , taking network pruning one step further towards neural architecture search . 2 PRELIMINARIES . Pruning at initialization . The principle behind conventional approaches for network pruning is to find unnecessary parameters , such that by eliminating them the complexity of the model is reduced while minimizing the impact on what is learned ( Reed , 1993 ) . Naturally , a typical pruning algorithm starts after convergence to a minimum or training to some degree . This pretraining requirement has been left unattended until Lee et al . ( 2019 ) recently showed that pruning can be performed on untrained networks at initiailzation prior to training . They proposed a method called SNIP which relies on a new saliency criterion , namely connection sensitivity , defined as follows : sj ( w ; D ) = |gj ( w ; D ) |∑m k=1 |gk ( w ; D ) | , where gj ( w ; D ) = ∂L ( c w ; D ) ∂cj ∣∣∣∣ c=1 . ( 1 ) Here , sj is the saliency of the parameter j , w ∈ Rm is the network parameters , c ∈ { 0 , 1 } m is the auxiliary indicator variables representing the connectivity of network parameters , m is the total number of parameters in the network , and D is a given dataset . Also , gj is the derivative of the loss L with respect to cj , which turns out to be an infinitesimal approximation of the change in the 1 The term faithful is used to describe signals propagating in a network isometrically with minimal amplification or attenuation , and borrowed from Saxe et al . ( 2014 ) , the first work to introduce dynamical isometry . loss with respect to removing the parameter j . Designed to be computed at initialization , pruning is performed by keeping top-κ ( where κ denotes a desired sparsity level ) salient parameters based on the above sensitivity scores . Dynamical isometry and mean field theory . The success of training deep neural networks is due in large part to the initial weights ( Hinton & Salakhutdinov , 2006 ; Glorot & Bengio , 2010 ; Pascanu et al. , 2013 ) . In essence , the principle behind these random weight initializations is to have the mean squared singular value of a network ’ s input-output Jacobian close to 1 , so that on average , an error vector will preserve its norm under backpropagation ; however , this is not sufficient to prevent amplification or attenuation of an error vector on worst case . A stronger condition that having as many singular values as possible near 1 is called dynamical isometry ( Saxe et al. , 2014 ) . Under this condition , error signals backpropagate isometrically through the network , approximately preserving its norm and all angles between error vectors . Alongside dynamical isometry , mean field theory is used to develop a theoretical understanding of signal propagation in neural networks with random parameters ( Poole et al. , 2016 ) . Precisely , the mean field approximation states that preactivations of wide , untrained neural networks can be captured as a Gaussian distribution . Recent works revealed a maximum depth through which signals can propagate at initialization , and verified that networks are trainable when signals can travel all the way through them ( Schoenholz et al. , 2017 ; Yang & Schoenholz , 2017 ; Xiao et al. , 2018 ) . 3 SIGNAL PROPAGATION PERSPECTIVE TO PRUNING RANDOM NETWORKS . Problem setup . Consider a fully-connected , feed-forward neural network with weight matrices Wl ∈ RN×N , biases bl ∈ RN , pre-activations hl ∈ RN , and post-activations xl ∈ RN , for l ∈ { 1 . . .K } up to K layers . Now , the feed-forward dynamics of a network can be written as , xl = φ ( hl ) , hl = Wlxl−1 + bl , ( 2 ) where φ : R→ R is an elementwise nonlinearity , and the input is denoted by x0 . Given the network configuration , the parameters are initialized by sampling from a probability distribution , typically a zero mean Gaussian with scaled variance ( LeCun et al. , 1998 ; Glorot & Bengio , 2010 ) . 3.1 EFFECT OF INITIALIZATION ON PRUNING . It is observed in Lee et al . ( 2019 ) that pruning results tend to improve when initial weights are drawn from a scaled Gaussian , or so-called variance scaling initialization ( LeCun et al. , 1998 ; Glorot & Bengio , 2010 ; He et al. , 2015 ) . As we wish to better understand the role of these random initial weights in pruning , we will examine the effect of varying initialization on the pruning results . In essence , variance scaling schemes introduce normalization factors to adjust the variance σ of the weight sampling distribution , which can be summarized as σ → αψlσ , where ψl is a layerwise scalar that depends on an architecture specification such as the number of output neurons in the previous layer ( e.g. , fan-in ) , and α is a global scalar throughout the network . Notice in case of a network with layers of the same width , the variance can be controlled by a single scalar γ = αψ as ψl = ψ for all layers l. In particular , we take both linear and tanh multilayer perceptron networks ( MLP ) of layers K = 7 and width N = 100 on MNIST with σ = 1 as the default , similar to Saxe et al . ( 2014 ) . We initialize these networks with different γ , compute the connection sensitivity , prune it , and then visualize layerwise the resulting sparsity patterns c as well as the corresponding connection sensitivity used for pruning in Figure 1 . It is seen in the sparsity patterns that for the tanh network , unlike the linear case , more parameters tend to be pruned in the later layers than the earlier layers . As a result , this limits the learning capability of the subnetwork critically when a high sparsity level is requested ; e.g. , for κ̄ = 90 % , only a few parameters in later layers are retained after pruning . This is explained by the connection sensitivity plot . The sensitivity of parameters in the nonlinear network tends to decrease towards the later layers , and therefore , choosing the top-κ parameters globally based on the sensitivity scores results in a subnetwork in which retained parameters are distributed highly non-uniformly and sparsely towards the end of the network . This result implies that the initial weights have a crucial effect on the connection sensitivity , and from there , the pruning results .

Summary: This paper analyzes how signals propagate through randomly initialized neural networks that have undergone a kind of pruning/sparsification. The pruning method utilizes a metric called 'connection sensitivity', which has been used in prior work and which measures the infinitessimal impact of turning off specific parameters. The distribution of singular values in the layer-to-layer Jacobian matrices for pruned networks becomes increasingly pathological as the depth increases. This observation motivates the concept of 'layerwise dynamical isometry' (LDI), a slight generalization of the concept of 'dynamical isometry' that has been studied in prior work. Several methods for approximately obtaining differing amounts of LDI are investigated in a series of in-depth experiments that show a strong correlation between increased signal propagation and improved trainability of sparse networks.