An Operator Theoretic View On Pruning Deep Neural Networks

1 INTRODUCTION . A surprising , but well-replicated , result in the study of deep neural network ( DNN ) optimization is that it is often possible to significantly “ prune ” the number of parameters after training with little effect on performance ( Janowsky , 1989 ; Mozer & Smolensky , 1989a ; b ; LeCun et al. , 1989 ; Karnin , 1990 ; Han et al. , 2015 ; Blalock et al. , 2020 ) . Because DNNs have become increasingly large and require considerable storage , memory bandwidth , and computational resources ( Han et al. , 2015 ) , finding a reduced form is desirable and has become an area of major research . In addition to the practical considerations , DNN pruning has become of theoretical interest , as the search for subnetworks that can be pruned with minimal effect on performance led to new insight on the success of large DNNs , via the development of the Lottery Ticket Hypothesis ( LTH ) ( Frankle & Carbin , 2019 ; Frankle et al. , 2020 ) . A considerable body of literature has been devoted to expanding , interrogating , and critiquing this theory , furthering the understanding of sparse DNNs . While many different methods exist for choosing which parameters to prune , including ones that take into account the gradient and the Hessian of the loss function ( LeCun et al. , 1989 ; Hassibi & Stork , 1993 ; Yu et al. , 2018 ; Lee et al. , 2019 ; 2020 ) , the simple approach of pruning based on a global threshold of parameters ’ magnitudes has been found to be robust and competitive with more complicated , state-of-the-art methods ( Blalock et al. , 2020 ) . In addition , early work on the LTH found that magnitude pruning could be successfully performed early during training ( although , iteratively pruning based on magnitude was found to yield considerably better results ) ( Frankle & Carbin , 2019 ) . Subsequent research has related the point in training time when iterative magnitude pruning can be successfully applied to a phenomenon known as “ linear mode connectivity ” ( Frankle et al. , 2020 ) . However , why magnitude pruning does not work before the emergence of linear mode connectivity , is not well understood . Therefore , there is a need to develop new , theoretically driven tools for studying magnitude pruning , especially during training pre-convergence ( Frankle et al. , 2021 ) . DNN training can be viewed as a discrete dynamical system , evolving the associated parameters along a trajectory . Despite the fact that this perspective is familiar ( indeed it is often taught in introductions to gradient descent ) , the complex dependence of the trajectory on the choice of optimization algorithm , architecture , activation functions , and data has kept dynamical systems theory from providing much insight into the behavior of DNNs . However , the recent development of Koopman operator theory ( Koopman , 1931 ; Koopman & Neumann , 1932 ; Mezić , 2005 ; Budišić et al. , 2012 ) , a rigorous , data-driven framework for studying nonlinear dynamical systems , has been successfully applied to DNNs and machine learning , more generally ( Dietrich et al. , 2020 ; Dogra & Redman , 2020 ; Tano et al. , 2020 ; Manojlović et al. , 2020 ; Mohr et al. , 2021 ; Naiman & Azencot , 2021 ) . This recent success motivated us to ask whether we could leverage Koopman operator theory to provide insight into the pruning of DNNs . We first made use of Koopman methods to define a new class of pruning algorithms . One method in this class , which we refer to as Koopman magnitude pruning , is shown to be equivalent to magnitude pruning at convergence . Early during training , we again find that the two methods are equivalent . The dynamical systems perspective of by Koopman pruning allows us to derive insight into the parameter dynamics present when the equivalence holds and when it does not . We find that the breaking of the equivalence occurs at a similar time that non-trivial structure emerges within the DNN . Finally , we extend our work beyond magnitude based pruning , showing that gradient pruning is equivalent to another form of Koopman pruning . Thus , the Koopman framework unifies magnitude and gradient pruning algorithms . Our contributions are the following : • A new class of theoretically motivated pruning methods ; • New insight on how training dynamics may impact the success ( or lack thereof ) of magnitude pruning pre-convergence ; • A unifying framework for magnitude and gradient based pruning . 2 KOOPMAN OPERATOR THEORY BASED PRUNING . Koopman operator theory is a dynamical systems theory that lifts the nonlinear dynamics present in a finite dimensional state space , to an infinite dimensional function space , where the dynamics are linear ( Koopman , 1931 ; Koopman & Neumann , 1932 ; Mezić , 2005 ) . While gaining an infinite number of dimensions would appear to make the problem intractable , numerical methods have been developed in the past two decades to find good finite dimensional approximations of the relevant spectral objects ( Mezić & Banaszuk , 2004 ; Mezić , 2005 ; Rowley et al. , 2009 ; Schmid , 2010 ; Tu et al. , 2014 ; Williams et al. , 2015 ; Mezić , 2020 ) . Given the powerful and interpretable tools that exist for studying linear systems , Koopman operator theory has found tremendous success in a number of applied settings . Below we discuss the basics of Koopman mode decomposition , a key element in Koopman operator theory , and use it to motivate a new class of pruning algorithms . We broadly refer to these methods as Koopman based pruning . For more details on Koopman operator theory , we refer the interested reader to Budišić et al . ( 2012 ) . 2.1 KOOPMAN MODE DECOMPOSITION . The central object of interest in Koopman operator theory is the Koopman operator , U , an infinite dimensional linear operator that describes the time evolution of observables ( i.e . functions of the underlying state space variables ) that live in the function space F . That is , after t > 0 amount of time , which can be continuous or discrete , the value of the observable g ∈ F , which can be a scalar or a vector valued function , is given by Utg ( p ) = g [ Tt ( p ) ] . ( 1 ) Here T is the nonlinear dynamical map evolving the system and p is the initial condition or location in state space . In the case of training a DNN , T is determined by the optimization algorithm and its associated hyperparameters ( e.g . learning rate ) , architecture , activation functions , and order in which the training data is presented . If the optimizer is ( a variant of ) stochastic gradient descent , then each epoch will have its own T , since the dynamic map depends on the ordering of the training data ( see Sec . 2.3 ) . The DNN parameter values that are set at initialization , θ ( 0 ) , is the associated initial condition [ i.e . p = θ ( 0 ) ] . The observable , g , could be chosen to be the identity function or some other function that might be relevant to a given DNN , such as its loss or gradient . However , note that Eq . 1 must hold for all g ∈ F . For the remainder of the paper , it will be assumed that the state space being considered is of finite dimension , and that F is the space of square-integrable functions . The action of the Koopman operator on the observable g can be decomposed as Ug ( p ) = R∑ i=0 λiφi ( p ) vi , ( 2 ) where the φi are eigenfunctions of U , with λi ∈ C as their eigenvalues and vi as their Koopman modes ( Mezić , 2005 ) . The decomposition of Eq . 2 is called the Koopman mode decomposition . The Koopman mode decomposition is powerful because , for a discrete dynamical system , the value of g at t ∈ Z+ time steps in the future is given by g [ Tt ( p ) ] = Utg ( p ) = R∑ i=0 λtiφi ( p ) vi . ( 3 ) From Eq . 3 , we see that the dynamics of the system in the directions vi , scaled by φi ( p ) , are given by the magnitude of the corresponding λi . Assuming that |λi| ≤ 1 for all i , finding the long time behavior of g amounts to considering only the φi ( p ) vi whose λi ≈ 1 . Given that the Koopman operator is linear , and it is known that even simple DNNs can have parameters that exhibit nonlinear dynamics ( Saxe et al. , 2014 ) , it may seem that applying Koopman methods to the training of DNNs is an over simplification . However , the Hartman-Grobman Theorem , a foundational result in dynamical systems theory , guarantees local neighborhoods around hyperbolic fixed points where the linearization captures the dynamical properties of the full , nonlinear system ( i.e . the two are conjugate ) ( Wiggins , 2003 ) . Additional work by Lan & Mezić ( 2013 ) extended the neighborhood where this conjugacy holds to the entirety of the basin . Therefore , as long as the DNN is in a basin of attraction of some fixed point , we can assume that the Koopman mode decomposition will provide an accurate representation to the true , nonlinear dynamics . There exist a number of different methods for computing the Koopman mode decomposition . While they differ in their assumptions on the dynamics of the system , they are all data-driven and require multiple snapshots of some g. In the case of DNN training , we can take g to be the identity function , and collect τ + 1 column vectors of the values the DNN parameters , θ , each taken at a different training step . That is , to compute the Koopman mode decomposition , the data matrix D = [ θ ( 0 ) , θ ( 1 ) , ... , θ ( τ ) ] ∈ RN× ( τ+1 ) must be constructed . For more details on the method used in this paper to compute the Koopman modes , and its computational demands , see Appendix A . 2.2 KOOPMAN BASED PRUNING . Any combination of DNN optimizer , architecture , activation function , and ordering of training data has an associated Koopman operator , U , which governs the flow of DNN parameters during training . This Koopman operator , in turn , has various directions along which the flow exponentially grows and shrinks ( i.e . the φi [ θ ( 0 ) ] vi ) . Previous work has shown that , over a non-trivial window of training , approximations of the Koopman operator can well capture the weight and bias dynamics of DNNs ( Dogra & Redman , 2020 ; Tano et al. , 2020 ) , suggesting that the φi [ θ ( 0 ) ] vi may be useful for pruning . If we order the modes by the real part of their eigenvalues , i.e . re ( λ1 ) > re ( λ2 ) > ... > re ( λR ) , and assume that λ1 = 1 ( a reasonable assumption in the case of stable dynamical systems , which we find corroborated in our experiments ) , then φ1 [ θ ( 0 ) ] v1 corresponds to the predicted fixed point of training . That is , the first mode identifies the parameter values that the DNN would take after training reached convergence . All additional modes would identify parameters that changed at increasingly larger exponential decays . If φ1 [ θ ( 0 ) ] v1 is accurate in its approximation of the fixed point , then pruning parameters based on their magnitude in φ1 [ θ ( 0 ) ] v1 could prove to be useful . We refer to this as Koopman magnitude pruning ( KMP ) to denote its similarity to global magnitude pruning ( GMP ) , a popular and robust approach to pruning ( Blalock et al. , 2020 ) . We make this connection explicit in Section 3 . Alternatively , the decaying modes , i.e . φ2 [ θ ( 0 ) ] v2 , φ3 [ θ ( 0 ) ] v3 ... , might be expected to convey useful information about which parameters are most coherently dynamic . We refer to this strategy as Koopman gradient pruning ( KGP ) to denote its similarity to gradient pruning . We make this connection explicit in Section 5 . Algorithm 1 General form for class of Koopman based pruning methods . 1 . Construct the data matrix D = [ θ ( 0 ) , θ ( 1 ) , ... , θ ( τ ) ] from τ + 1 snapshots of the parameter values θ during training . 2 . Compute the Koopman mode decomposition from D to obtain R Koopman modes , ( λi , φi , vi ) . 3 . Using a scoring function , s , that is a function of ( λi , φi , vi ) , create a pruning mask , m , that compresses the network by an amount c. In the case of KMP , s ( θ , λi , φi , vi ) = |φ1 [ θ ( 0 ) ] v1| , assuming that λ1 = 1 . 4 . Create the pruned DNN , f ( x ; θ m ) . These algorithms , KMP and KGP , are two of many possible methods for pruning based on the Koopman mode decomposition . Our methods can be extended to layer magnitude pruning , by splitting φ1 [ θ ( 0 ) ] v1 into subvectors and pruning each layer based on thresholds computed from each subvector . Algorithm 1 provides a general form for approaches in the class of Koopman based pruning . Note that we follow standard practice in the pruning literature and denote a DNN by f ( x ; θ ) , where x is its input and θ is its N parameters . The goal of pruning is to find a mask , m ∈ { 0 , 1 } N , such that the accuracy of f ( x ; θ m ) is comparable to the original DNN . Here denotes the usual element-wise multiplication . To determine which parameters to prune ( i.e . which mi = 0 ) , a scoring function , s , is used . In the case of GMP , s ( θi ) = |θi| . The DNN is compressed an amount c by keeping only the 100c % of parameters with the largest scores . If c = 4 , then only the parameters with scoring function values that are in the top 25 % are kept .

Summary: The authors studied network pruning from the perspective of dynamical system theory. They show that a new type of pruning method, named Koopman pruning, unifies magnitude pruning and gradient-based pruning to a degree. It also clarifies aspects of magnitude-based pruning before convergence in training.