An Operator Theoretic View on
Pruning Deep Neural Networks

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 the LTH, 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 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; Mezić, 2005; Budišić 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; Manojlović 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 can be equivalent. The dynamical systems perspective afforded by Koopman pruning allows us to gain 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.

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; Mezić, 2005). While gaining an infinite number of dimensions would appear to make the problem intractable, numerical methods have been developed over the past two decades to find good finite dimensional approximations of the relevant spectral objects (Mezić & Banaszuk, 2004; Mezić, 2005; Rowley et al., 2009; Schmid, 2010; Tu et al., 2014; Williams et al., 2015; Mezić, 2020). Given the powerful and interpretable tools that exist for studying linear systems, Koopman operator theory has found success in a number of applied settings. Below we discuss the basics of Koopman mode decomposition, a key element of 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 Budišić et al. (2012).

We start by examining how global and Koopman magnitude pruning compare in the long-training time limit. We consider a DNN, initialized with parameters $\theta(0)$, optimized via SGD on training data $x$, with the order of the $i^{\text{th}}$ epoch’s data denoted by $x_{i}$. As the number of epochs goes to $\infty$, the DNN will converge to $\theta^{*}$, a local minimum in the loss landscape. This requires that $||\theta(t)-\theta^{*}||_{2}\rightarrow 0$, as $t\rightarrow\infty$. Therefore, for $\epsilon\in\mathbb{R}^{+}$, there exists $t_{\epsilon}\in\mathbb{N}$, such that $||\theta(t)-\theta^{*}||_{2}<\epsilon$, for all $t>t_{\epsilon}$. Here, $||\cdot||_{2}$ is the $l^{2}$ norm. Note that this additionally implies that at $t_{\epsilon}$ the DNN is in the basin of attraction of $\theta^{*}$.

Previous work on approximating the Koopman operator has found that if a system is sufficiently near a fixed point (and if the trajectory is sufficiently “nice”), then the computed Koopman mode decomposition should well approximate the true dynamics. That is, for $\epsilon\in\mathbb{R}^{+}$, there exists $t_{\epsilon}^{\text{Koop}}\in\mathbb{N}$, such that $||\tilde{\theta}^{*}-\theta^{*}||_{2}<\epsilon$, for all $t>t_{\epsilon}^{\text{Koop}}$. Here, $\tilde{\theta}^{*}=\phi_{1}[\theta(t_{\epsilon}^{\text{Koop}})]\textbf{v}_{1}$ is the computed Koopman mode corresponding to $\lambda=1$.

Assuming that a DNN has trained for at least $\max(t_{\epsilon},t_{\epsilon}^{\text{Koop}})$ training steps, then

If $\epsilon$ is sufficiently small, then $m(|\tilde{\theta}^{*}|,c)\approx m(|\theta(t)|,c)$, and the two methods are thus equivalent.

To test whether this result holds in practice, we used ShrinkBench, a recent open source toolkit for standardizing pruning results and avoiding bad practices that have been identified in the pruning literature (Blalock et al., 2020). The results of pruning a custom ShrinkBench convolutional DNN called “MnistNet” (Blalock et al., 2020), pre-trained on MNIST ¹¹1https://github.com/JJGO/shrinkbench-models/tree/master/mnist, and ResNet-20, pre-trained on CIFAR-10 ²²2https://github.com/JJGO/shrinkbench-models/tree/master/cifar10, are presented in Fig. 1. Each DNN was trained from a point near convergence for a single epoch, so that the parameter trajectories could be captured. All hyperparameters of the DNNs were the same as the off-the-shelf implementation of ShrinkBench, except that we allowed for pruning of the classifier layer. Training experiments were repeated independently three times, each with a different random seed. Our code has been made publicly available ³³3https://github.com/william-redman/Koopman_pruning.

We found that, for both MnistNet and ResNet-20, KMP and GMP produced nearly identical results, both immediately after pruning and after one epoch of refinement (Fig. 1). These results support the claim developed above that KMP and GMP are equivalent in the long training time limit.

Motivated by the equivalence of Koopman and global magnitude pruning in the long training time limit, we next examined how the two methods compared before training converged. In this regime, we might expect KMP to outperform GMP, since computing the Koopman mode decomposition requires multiple snapshots of the training data. This provides KMP with considerably more dynamic information than GMP, which is based on a single snapshot.

As previous work has found that ResNet-20 trained on CIFAR-10 does not achieve linear mode connectivity until at least 5 epochs of training (Frankle et al., 2020), we examined pruning performance at epoch 1 and epoch 20 to capture the effect that linear mode connectivity had. We find that, despite the additional information, KMP performs similarly to GMP at both epochs (Fig. 2, right column).

There are two possible explanations for this similarity in performance. The first is that KMP and GMP prune the same parameters, making them equivalent methods (even before convergence). The second is that KMP and GMP prune distinct parameters, but those that remain result in a subnetwork with similar accuracy. To test these two hypotheses, we compare the overlap between the parameters not pruned by GMP and KMP at different compressions and different epochs (see Appendix B for more details on this metric). The right column of Fig. 3 shows that the two methods largely prune the same parameters. From epochs 1-10, the mean overlap is roughly $>95\%$, even for a compression of 64. At epoch 20, there is a decrease in mean overlap and an increase in standard deviation. This suggests that the two methods are equivalent at the earliest epochs of training, with the equivalence breaking at some subsequent epoch.

Because KMP makes use of dynamical systems theory, the equivalence of KMP and GMP during (at least) the first 10 epochs has a strong implication on the trajectory in parameter space that is induced by training: at the end of each epoch, the largest parameters of the DNN must be the same ones that are expected to be the largest at convergence, if the order of training data were to remain the same. In this sense, the expected identity of the largest parameters are “static”. The overlap between masks found using KMP from increasingly distant epochs shows a decrease with time (Appendix B, Fig. 5), suggesting that SGD slowly re-orders the identity of the largest parameters epoch-to-epoch, making these initial largest parameters inaccurate.

By epoch 20, KMP and GMP stop being equivalent and, on average, start pruning a larger number of different parameters. This implies that, even if the training data was ordered the exact same way for all the remaining epochs, the identity of (some of) the largest parameters are expected to evolve. To contrast this with the “static” dynamics found during epochs 1–10, we call this regime “active”. We speculated that this change from “static” to “active” might be caused by emerging structure in the parameter evolution. To test this, we pruned ResNet-20 at epochs 1 and 20 using layer shuffle pruning (LSP). This method preserves the amount of compression in each layer, but randomly chooses which parameters to remove. Previous work has found that many state-of-the-art methods designed to prune at initialization do not do better than LSP (Frankle et al., 2021). We find that after only a single epoch of training, both GMP and KMP do no better, and in fact, do worse, than LSP for nearly all compressions tested. However, by epoch 20, both GMP and KMP were better than LSP for compression values of less than 16. These results support the idea that non-trivial structure is emerging in the DNN at a similar time that KMP and GMP stop being equivalent. When this structure is not present (e.g. at epoch 1), the two methods prune the same parameters.

We also examined the ability of KMP and GMP to prune MnistNet, trained on MNIST, after epochs 1 and 20. While we found similar results to ResNet-20 (Figs. 2 and 3, left column), there are subtle differences in the behavior of the overlap between pruning masks. For MnistNet, the overlap is smallest at epochs 1 and 5, suggesting that the two methods are not equivalent at the earliest part of training (Fig. 3, left). That this is different to what we found for ResNet-20 may be due to the fact that the two models are at different phases of training at equivalent epochs. Previous work has shown that, shortly after initialization, linear mode connectivity emerges in LeNet, trained on MNIST (Frankle et al., 2020). This does not happen for ResNet-20, trained on CIFAR-10, until at least epoch 5. Additionally, we find that, for MnistNet, KMP and GMP outperform LSP for compression above 8 as early as epoch 1. These MnistNet results further support the fact that, as structure beyond what LSP can find emerges, the dynamics associated with the identity of the largest parameters goes from being “static” to “active”.

Finding what is driving this “static” to “active” transition, and whether it is related to the emergence of linear mode connectivity and/or other noted transitions during training (Gur-Ari et al., 2018; Achille et al., 2019; Kalimeris et al., 2019), will be the focus of future work.

Thus far, we have focused only on magnitude based pruning. In this section, we extend our work to another method that has found success and is of general interest: gradient based pruning (Yu et al., 2018; Lee et al., 2019; 2020; Blalock et al., 2020). In practice, this way of pruning is accomplished by scoring based on the gradient of the loss function, $\nabla_{\theta}L$, as propagated through the network. For instance, global gradient pruning (GGP) scores using the magnitude of the gradient times the parameter values, $s(\theta)=\big{|}\nabla_{\theta}L\cdot\theta(t)\big{|}$ (Blalock et al., 2020). Given that the gradient is applied to update the parameters of the DNN, $\nabla_{\theta}L$ is directly related to the parameters’ dynamics.

From the Koopman mode decomposition, we have that the parameters at training step $t$, $\theta(t)$, are given by

where the $\tilde{\theta}_{i}$ are the scaled Koopman modes (i.e. $\tilde{\theta}_{i}=\phi_{i}(p)\textbf{v}_{i}$). As before, we use $\tilde{\theta}^{*}$ to denote the scaled Koopman mode that has $\lambda=1$. Importantly, Eq. 5 tells us that the future state of the parameters is given by the sum of $\tilde{\theta}_{i}$ times their decaying eigenvalues. While there are $R-1$ modes in this sum, those with large norms (i.e. $|\tilde{\theta}_{i}|\gg 0$) are expected to make the most substantial contribution. For simplicity, we consider only the largest modes that also have real, positive eigenvalues. These modes have easily interpretable dynamics (i.e. exponential decay) and are directly related to $\nabla_{\theta}L$. Pruning using the magnitude of these modes is what we call Koopman gradient pruning (KGP).

The Koopman spectrum of a ResNet-20 trained on CIFAR-10 for 20 epochs is shown in Fig. 4. In addition to the presence of an eigenvalue close to $1$, corresponding to $\tilde{\theta}^{*}$, there is also an eigenvalue that satisfies the above criteria. We refer to its corresponding mode as $\tilde{\theta}_{2}$.

As expected from the discussion above, pruning ResNet-20 after 20 epochs of training using KGP is very similar to pruning using the scoring function $s(\theta)=|\nabla_{\theta}L|$, a method we refer to as just gradient pruning (JGP) (Fig. 4, right). We note that we only expect the two methods to be “nearly” equivalent (and not truly equivalent, as in the case of KMP and GMP in the long training time limit) because, KGP finds exponential decaying modes from dynamic data, whereas JGP is the average gradient computed over a batch of input data (Blalock et al., 2020).

These results highlight the fact that, while gradient and magnitude based pruning methods have seemingly little to do with each other, they are in fact, related in the Koopman operator framework. Whereas magnitude pruning is, at certain phases of training, equivalent to pruning using $|\tilde{\theta}^{*}|$, gradient pruning is equivalent to pruning using the exponentially decaying modes.

Inspired by the widespread success Koopman operator theory has found in a number of applied fields, and its recent application to problems in machine learning, we extended Koopman tools to study deep neural network pruning (Janowsky, 1989; Mozer & Smolensky, 1989a; b; LeCun et al., 1989; Karnin, 1990; Han et al., 2015; Blalock et al., 2020). By making use of the Koopman mode decomposition, we defined a new class of pruning algorithms (Algorithm 1) that are theoretically motivated and data-driven in implementation. We found that both global magnitude and gradient pruning (Blalock et al., 2020) have equivalent (or nearly equivalent) methods at certain regimes of training in the Koopman framework (Figs. 1, 4). In this way, Koopman pruning provides a unified perspective on these “disparate” methods.

Applying KMP to DNNs in the early part of training revealed that it could be equivalent to GMP (Fig. 3). From a dynamical systems perspective, this equivalence implies that the largest parameters, at the end of each epoch, are the same ones that are expected to remain the largest, if training continued using the same ordering of the data. We refer to this behavior as “static”. Given that we found the identity of the largest parameters evolved with training (Appendix B, Fig. 5), pruning too early may confine the sparsified DNN to a subspace where it stays in the static regime and therefore is unable find a good local minimum. This is additionally supported by the fact that layer shuffle pruning (LSP) outperforms GMP during this part of training, suggesting that the structure present in the GMP mask is more “harmful” than a randomly structured mask (Fig. 2).

However, at a certain subsequent point in training, the masks produced by GMP and KMP stop overlapping as completely (Fig. 3). This suggests that there is some transition in the dynamics, with the largest parameters becoming “active”. As this transition seems to occur at a similar time that GMP starts to outperform LSP, and at a similar time that linear mode connectivity has been reported to emerge (Frankle et al., 2020), we speculate that it is related to the sparsified DNN still being able to evolve and find a good local minima in the subspace it is restricted to.

By making use of ShrinkBench, our results can be easily replicated and avoid some of the bad practices that have been recently identified as being pervasive in the DNN pruning literature (Blalock et al., 2020). Our custom implementations are integrated into the ShrinkBench framework, allowing for additional investigation into Koopman pruning.

While dynamical systems theory is a natural language with which to frame DNN optimization, the complex dependencies on optimizer, architecture, activation function, and training data have historically kept efforts in this direction to a minimum. This has led to a reliance on heuristic methods, such as iterative magnitude pruning, whose basis of success is still not clear (Elesedy et al., 2021). Other groups have attempted to more principally examine DNN behavior by studying spectral objects, such as the spectrum of the Hessian matrix (Gur-Ari et al., 2018) and the spectrum of the weight matrices (Martin et al., 2021). This work has been influential and led to new insight, however the connection between these metrics and the (future) dynamics of DNN training is not direct. An exception to this is the recent work on using inertial manifold theory to study the Lottery Ticket Hypothesis, which is similar in spirit to our Koopman based pruning methods (Zhang et al., 2021).

Given the data-driven tools that have been developed for Koopman operator theory (Mezić & Banaszuk, 2004; Mezić, 2005; Rowley et al., 2009; Tu et al., 2014; Williams et al., 2015; Mezić, 2020) and the rich theoretical literature that connects Koopman objects to the state space geometry (Mezić, 2005; Arbabi & Mezić, 2017; Mauroy et al., 2013; Mezić, 2020), recent work has focused on bring Koopman operator theory to the study of DNNs (Dietrich et al., 2020; Dogra & Redman, 2020; Tano et al., 2020; Manojlović et al., 2020; Mohr et al., 2021; Naiman & Azencot, 2021). This work has shown that Koopman methods can:

• •

  Approximate the parameter dynamics during training and lead to speed-ups in wall clock time when substituted for standard optimizers (Dogra & Redman, 2020; Tano et al., 2020);

• •

  Predict the number of layers necessary for a hierarchical support vector machine to achieve a given level of performance (Manojlović et al., 2020; Mohr et al., 2021);

• •

  Shed light on the behavior of sequence neural models (Naiman & Azencot, 2021).

Our current work further illustrates the power that Koopman operator theory can bring. In particular, we found that the Koopman mode decomposition encodes dynamical information about training that could be leveraged to prune DNNs, enables a framework that unifies disparate pruning algorithms, and provides insight on how training dynamics impact the success of magnitude pruning. To the best of our knowledge, these could not have been done using any other linearization method (e.g. principal component analysis). We believe that continued exchange between the Koopman and machine learning communities will lead to further advances.

Forethought to ensure reproducibility was taken in both the design and description of our methods and results. We made use of ShrinkBench, a recent open-source framework for developing and testing new pruning methods (Blalock et al., 2020), to implement all our experiments. This not only allowed us to report comparisons between our methods and existing methods across a number of different compressions, but will also enable easy sharing of our algorithms, as they are integrated into the ShrinkBench framework. Additionally, we have provided general pseudo-code for implementing Koopman based pruning (Algorithm 1) and a description of the method we used to compute the Koopman mode decomposition (Appendix A), to make it straightforward for others to understand and implement.