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: In this paper, the authors studied and formalized the effect of initialization to connection-sensitivity-based pruning. The authors first pointed out that a previously studied pruning criterion -- connection sensitivity (CS) -- is a normalized magnitude of gradients. Based on signal propagation theory, to achieve a 'faithful' (with minimal amplification) CS, the gradients must be also faithful. Then by using relation of Jacobians and gradient, the authors proved that orthogonally initial weights guarantees faithful on linear networks and certain distribution property on nonlinear network can achieve layerwise dynamic isometry, which is to ensure faithful signal propagation. Based on these findings, the authors proposed an initialization setup for improving pruning performance, with the goal to ensure dynamic isometry by orthogonal initialization and approximation.