A Unified Paths Perspective for Pruning at Initialization

1 INTRODUCTION . A wealth of recent work has been dedicated to characterizing the training dynamics and generalization bounds of neural networks under a linearized approximation of the network depending on its parameters at initialization ( Jacot et al. , 2018 ; Arora et al. , 2019 ; Lee et al. , 2019a ; Woodworth et al. , 2020 ) . This approach makes use of the Neural Tangent Kernel , and under infinite width assumptions , the training dynamics of gradient descent over the network become analytically tractable . In this paper , we make use of the Neural Tangent Kernel theory with the goal of approximating the effects of various initialization pruning methods on the resulting training dynamics and performance of the network . Focusing on networks with homogeneous activation functions ( ReLU , Leaky-ReLU , Linear ) , we introduce a novel decomposition of the Neural Tangent Kernel which separates the effects of network architecture from effects due to the data on the training dynamics of the network . We find the data-independent factor of the Neural Tangent Kernel to have a particularly nice structure as a symmetric matrix representing the covariance of path values in the network which we term the Path Kernel . We subsequently show that the Path Kernel offers a data-independent approximation of the network ’ s convergence dynamics during training . To validate the empirical benefits of this theoretical approach , we turn to the problem of pruning at initialization . While the problem of optimally pruning deep networks is nearly as old as deep networks themselves ( Reed , 1993 ) , interest in this problem has experienced a revival in recent years . This revival is likely a product of a number of underlying factors , but much of the recent interest could be ascribed to the Lottery Ticket Hypothesis ( Frankle & Carbin , 2018 ) which states that sparse , trainable networks–which achieve task performance that matches or exceeds those of its dense counterparts–can be found at initialization . The Lottery Ticket Hypothesis implies that the over-parameterization of neural networks is incidental in finding a trainable solution , the topology of which often exists at initialization . However , finding these lottery ticket networks currently requires some amount of iterative re-training of the network at increasing levels of sparsity which is inefficient and difficult to analyze theoretically . The resurgence of interest in optimal pruning has spurred the development of a number of recent approaches for pruning deep neural networks at initialization ( Lee et al. , 2019b ; Liu & Zenke , 2020 ; Wang et al. , 2020 ; Tanaka et al. , 2020 ) in supervised , semi-supervised , and unsupervised settings , borrowing theoretical motivation from linearized training dynamics ( Jacot et al. , 2018 ) , mean-field isometry ( Saxe et al. , 2013 ) , and saliency ( Dhamdhere et al. , 2019 ) . While each of these methods have their own theoretical motivations , little work has been dedicated to formally describing the effect of these pruning methods on the expected performance of the pruned network . Also , the diversity in theoretical motivations that give rise to these pruning methods makes it difficult to observe their similarities . In this paper , we observe that a number of initialization pruning approaches are implicitly dependent on the path covariance structure captured by the Path Kernel which , in turn , affects the network ’ s training dynamics . We show that we can approximate these training dynamics in general , and our approximation results for a number of initialization pruning approaches suggests that it is possible to estimate , prior to training , the efficacy of a particular initialization pruning approach on a given architecture by investigating the eigenstructure of its Path Kernel . Motivated by our theoretical results and the unification of a number of initialization pruning methods in this Path Kernel framework , we investigate the close relationship between the SynFlow ( Tanaka et al. , 2020 ) pruning approach and our path decomposition . This leads to our suggestion of two new initialization pruning approaches which we predict to perform well under various assumptions on the stability of the Path Kernel and the input distribution of the data . We then validate these predictions empirically by comparing the performance of these pruning approaches across a number of network architectures . The insights on initialization pruning provided by the Path Kernel decomposition are only one of a number of potential application domains which could benefit from this path-centric framework . Importantly , the coviariance structure over paths encoded by the Path Kernel is general and may be computed at any point in time , not just at initialization . We anticipate that this representation will provide insight into other application areas like model interpretation , model comparison , or transfer learning across domains . The sections of the paper proceed as follows . We start with a brief introduction to the Neural Tangent Kernel in Section 2 before introducing the Path Kernel decomposition in Section 3 and its relationship to approximations of network convergence properties . In Section 4 , we reformulate in this path framework three popular initialization pruning approaches and introduce two additional initialization pruning approaches inspired by this path decomposition . We validate these convergence approximations and the behavior of these pruning approaches in Section 5 and conclude with a discussion of the results and opportunities for future work . 2 THE NEURAL TANGENT KERNEL . Recent work by Jacot et al . ( 2018 ) has shown that the exact dynamics of infinite-width network outputs through gradient descent training corresponds to kernel gradient descent in function space with respect the Neural Tangent Kernel . More formally , for a neural network f parameterized by θ and loss function ` : RK × RK → R , let L = ∑ ( x∈X , y∈Y ) ` ( ft ( x , θ ) , y ) denote the empirical loss function . Here , X is the training set , Y is the associated set of class labels . For multiple inputs , denote f ( X , θ ) ∈ RNK the outputs of the network where K is the output dimension and N is the number of training examples . In continuous-time gradient descent , the evolution of parameters and outputs can be expressed as θ̇t = −η∇θf ( X , θt ) ᵀ∇f ( X , θt ) L ( 1 ) ḟ ( X , θt ) = ∇θf ( X , θt ) θ̇t = −ηΘt ( X , X ) ∇f ( X , θ ) L ( 2 ) where the matrix Θt ( X , X ) ∈ RNK×NK is the Neural Tangent Kernel at time step t , defined as the covariance structure of the Jacobian of the parameters over all training samples : Θt ( X , X ) = ∇θf ( X , θt ) ∇θf ( X , θt ) ᵀ . ( 3 ) For infinitely wide networks , the NTK exactly captures the output space dynamics through training , and Θt ( X , X ) remains constant throughout . Lee et al . ( 2019a ) have shown that neural networks of any depth tend to follow the linearized training dynamics as predicted by the NTK . Moreover , we can approximate the outputs of the network linearly through a one-step Taylor expansion given input x as : f− ( x , θt ) = f ( x , θ0 ) +∇θf ( x , θ0 ) ωt ( 4 ) where ωt = θt − θ0 is the change in parameters from their initial values . The first term of Equation 4 is constant while the second term captures the dynamics of the initial outputs during training . Substituting f− for f in Equations 1 and 2 , the dynamics of the linearized gradient flow become ω̇t = −η∇θf ( X , θ0 ) ᵀ∇f− ( X , θt ) L ( 5 ) ḟ− ( x , θt ) = −ηΘ ( x , X ) ∇f− ( X , θt ) L ( 6 ) Under MSE loss , the above ODEs have closed-form solutions as ωt = −∇θf ( X , θ0 ) ᵀΘ−10 ( I − e−ηΘ0t ) ( f ( X , θ0 ) −Y ) ( 7 ) f− ( X , θt ) = ( I − e−ηΘ0t ) Y + e−ηΘ0tf ( X , θ0 ) . ( 8 ) In other words , through the tangent kernel and initial outputs of the network , we can compute the training convergence of a linearized neural network before running any gradient descent steps . We will show in Section 3 that , through the Path Kernel , we can reliably approximate the linearized convergence rate of the network in the absence of data and without computing the full NTK . 3 THE PATH KERNEL . We will now provide a reformulation of the output behavior of networks in terms of activated paths and their values . This reformulation provides access to a unique decomposition of the NTK that separates the data-dependent output dynamics of the network from those dependent on the architecture and initialization . This reformulation of network behavior in terms of active paths is motivated by Meng et al . ( 2019 ) wherein the authors show that gradient descent on networks with homogenous activations may be computed completely in path space . The decomposition provided in this section makes explicit the relationships between the pruning-at-initialization based approaches described in Section 4 and allows for estimation of the convergence behavior of networks at initialization and how pruning affects this behavior . Let θ ∈ Rm denote the network parameters in vector form and let x ∈ Rd denote an input . Assume the network has K output nodes . We define a path from input to output as a binary vector p such that pj = 1 when θj is an edge along path p , otherwise pj = 0 . Denote by P the enumerable set of all paths and Ps→k the subset of paths that go from input node s to output node k. Let P = |P| be the number of paths . Given the enumeration in P , we will abuse notation slightly and refer to paths p by their index p in P . The value of a path p can be calculated as a product of parameters with binary exponents vp ( θ ) = ∏m j=1 θ pj j such that vp is the product along the θ weights that are in path p. The activation status of a path p is ap ( x , θ ) = ∏ { j | pj=1 } I ( opj ( x , θ ) > 0 ) . Here opj is the output of the hidden node which path pj passes through immediately after parameter j . We can then define the output of the network at node k as fk ( x , θ ) = d∑ s=1 ∑ p∈Ps→k vp ( θ ) ap ( x , θ ) xs . ( 9 ) The derivative of ap ( x , θ ) with respect to θ will goes to zero in expectation . Using the chain rule , we can break the derivative of the output of the network with respect to θ into two parts , the partial derivative of the output with respect to the path values and the partial derivative of the path values with respect to the parameters : ∇θf ( x , θ ) = ∂f ( x , θ ) ∂θ = ∂f ( x , θ ) ∂v ( θ ) ∂v ( θ ) ∂θ = Jfv ( x ) J v θ . ( 10 ) Note the inner product structure formed over paths as ∂f ( x , θ ) ∂v ( θ ) = J f v ( x ) ∈ RK×P and ∂v ( θ ) ∂θ = Jvθ ∈ RP×m . The change in output with respect to path values now only depends on the activation status of each path leading to the output and the input to that path . Each entry of this matrix has value ( Jfv ( x ) ) k , p = ( ∂f ( x , θ ) ∂v ( θ ) ) k , p = d∑ s=1 I ( p ∈ Ps→k ) ap ( x , θ ) xs . Similarly , the change in path values with respect to parameters is a function of only the parameters and their relational structure through paths : ( Jvθ ) p , j = { vp θj pj = 1 0 otherwise . Given this reparameterization of the output of the network in terms of activated paths and their values , we see the NTK similarly decomposes nicely along this structure : Θ ( x , x ) = ∇θf ( x , θ ) ∇θf ( x , θ ) ᵀ = Jfv ( x ) J v θ ( J v θ ) ᵀJfv ( x ) ᵀ = Jfv ( x ) ΠθJ f v ( x ) ᵀ where Πθ = Jvθ J v θ ᵀ is the Path Kernel . The Path Kernel Πθ is positive semidefinite with maximal values on the diagonal Πθ ( p , p ) = ∑m j=1 ( vp ( θ ) θj ) 2 and off-diagonal elements Πθ ( p , p′ ) =∑m j=1 ( vp ( θ ) θj ) ( vp′ ( θ ) θj ) . We can view Πθ as a covariance matrix on the weighted paths defined by the network architecture and parameter initialization . Note that Jfv ( x ) entirely captures the dependence of f on the input by choosing which paths are active and re-weighting by the input , while Πθ is completely determined by the architecture and initialization . Therefore , we can expand our one-sample NTK to the entire training set through the appropriate expansion of dimensions such that Jfv ( X ) ∈ RNK×P . In the following section , we will show that the Path Kernel decomposition of the NTK allows us to approximate , at initialization , the convergence behavior of the network during training . Additionally , we find we can compute the trace of the path kernel efficiently through an implicit computation over the parameter gradients of a particular loss function . This trace computation serves as an approximation to full eigenstructure of the NTK .

Summary: The paper studies the problem of neural network pruning at initialization through the lens of neural tangent kernels (NTK). As a result, the paper delivers a unified perspective on SNIP, GRASP, and SynFlow. Based on the framework, the paper provides a method to approximate the convergence dynamics of pruned models. The paper also motivates a SynFlow-variant from the theoretical framework.