Title: IN SEARCH OF THEORETICALLY GROUNDED PRUNING

Abstract: Deep learning relies on resource-heavy linear algebra operations which can be prohibitively expensive when deploying to constrained embedded and mobile devices, or even when training large-scale networks. One way to reduce a neural network's resource requirements is to sparsify its weight matrices -a process often referred to as pruning. It is typically achieved by removing least important weights as measured by some salience criterion, with pruning by magnitude being the most popular option. This, however, often makes close to random judgments. In this paper we aim to closely investigate the concept of model weight importance, with a particular focus on the magnitude criterion and its most suitable substitute. To this end we identify a suitable Statistical framework and derive deep model parameter asymptotic theory to use with it. Thus, we derive a statisticallygrounded pruning criterion which we compare with the magnitude pruning both qualitatively and quantitatively. We find this criterion to better capture parameter salience, by accounting for its estimation uncertainty. This results in improved performance and easier post-pruned re-training.

1. INTRODUCTION:

Deep learning is powered by resource intensive computations, to the extent that only models that have been carefully modified to regulate the system resources they consume can be deployed to mobile and wearable technology (Sze et al., 2017) . Currently, the management of its energy, latency, and memory demands is one of the most pressing practical challenges. Many argue that fully realizing the potential of deep learning requires a shift towards a structured consideration of these resource requirements and their inclusion in the architecture design and model selection process (LeCun et al., 1990; Han et al., 2015; Yang et al., 2017; Chen et al., 2018) .

One way to manage the trained model's resource requirements is to manipulate its size by altering its parameter matrices. Weight matrices can be shrunk by using approximations, such as SVD or Tucker decompositions (Bhattacharya & Lane, 2016; Kim et al., 2015) , or sparsification (Dong et al., 2017; Castellano et al., 1997) . Alternatively, floating-point precision of values can be reduced, sometimes all the way to binary (Gupta et al., 2015; Rastegari et al., 2016) . Finally, specific elements of the matrices can be dropped from the model. This requires the practitioner to specify some importance (salience) criterion which determines the order in which parameters are supposed to be pruned.

State-of-the-art model compression frameworks often make use of pruning with the magnitude salience criterion (Li et al., 2016; Yang et al., 2017; He et al., 2018; Yang et al., 2018) , which states that a weight's importance is its magnitude (Han et al., 2015) . Evidence suggests this criterion may have behavior very close to pruning weights at random (Molchanov et al., 2016 ) -a result we further observe in our experiments in this paper. Furthermore, the relative performance of pruning criteria varies considerably and without a discernible pattern across domains and data sets (Augasta & Kathirvalavakumar, 2013) . Such results collectively call for a closer investigation of the concept of parameter importance, the salience criteria, and their role in the pruning process.

In this paper we set out to understand and improve the pruning process. Our contributions are:

1. We identified suitable Statistical framework for parameter salience determination. 2. We derived deep model's parameter distribution compatible with the framework. 3. We experimentally verified the superior accuracy and easier post-pruning re-training recovery delivered by the Statistical framework.

We emphasize that in this work, we explicitly do not consider methods that induce sparsity during training or otherwise trade optimization target for extra sparsity.

2. LIMITS OF MAGNITUDE PRUNING FROM A STATISTICAL VIEW:

The Magnitude criterion defines parameter salience as their magnitude and intuitively justifies this as the size of their contribution to the network's prediction. While simple to implement, this method makes several implicit assumptions that we enumerate below. Crucially, it assumes all weights to be estimated with the same uncertainty, which is highly unlikely to hold.

Norm-regularization. First implicit requirement of magnitude pruning is for the norm-based regularization to be a part of the model's loss function. Unless the network is norm-regularized, weights that don't get updated remain at their initialized state and thus have random magnitude (salience) values. Thus, it fails to determine weight importance unless aided by the environment.

Normalization. The second implicit assumption is that the data and intermediate activations have been normalized, e.g. through the use of Batch Normalization. Take, for instance, two weights, one estimated at 0.1 the other at 0.01. Magnitude pruning would judge the former as more important because it is larger, arguing that this implies bigger contribution. However, this may not be the case if the former's activation, on average, has 100x lower magnitude.

Uncertainty. Finally, when weights of a neural network obtained during training are seen as mean estimates from some distribution, magnitude pruning does not take into account the uncertainty, the variance, with which they were estimated. In the example above, if the larger weight also has larger associated standard deviation, it is arguably less important to the overall prediction.

3. THEORETICAL FOUNDATION FOR WALD PRUNING:

The Wald statistic is a Statistical tool used to test hypotheses about model parameters in general multivariate models. It requires the model's parameter's asymptotic distribution to be approximately normal and estimable. We derive Deep Learning's distribution in this form by recognizing it to be a member of the M-Estimator family and utilizing the associated asymptotic theory1 .

3.1. THE ASYMPTOTIC DISTRIBUTION OF PARAMETERS:

Let Θ be the parameter space for a general deep model with an objective (loss) function Q n and parameters θ, where n is the number of data points and Q n = 1 n m(y, x, θ)-an average of some per-example loss function m(y, x, θ), where (y, x) is an example from the dataset D. Now, let us assume the following: ID Identification Condition: Q has a unique minimum at θ 0 in the parameter space Θ . CMP Compactness Condition: Θ is compact (closed and bounded).

DOM Domination Condition: Esup θ∈Θ m (y, x, θ) < ∞.

WKC Weak Continuity Condition: m is weakly continuous. IID IID Data: {y i , x i } n i=1 is independently identically distributed. I Interior condition: θ 0 is in the interior of (not at the edge of) Θ. D Differentiability condition: m is twice continuously differentiable for any (y, x). Then, asymptotically, as the sample size n grows to infinity, the model parameters are distributed as:

SD Score Distribution: 1 √ n n t=1 s (y t , x t , θ 0 ) → d N (0, Σ), Σ

θ ∼ N θ 0 , (E [H(y t , x t , θ 0 )]) -1 Σ (E [H(y t , x t , θ 0 )]) -1

(1)

The following three assumptions are fundamental to the application of the theory to deep models 4 :

Score Distribution assumption (SD). First, the assumption that is driving the functional form of the parameter distribution above is the estimable normal Score Distribution (SD). Since the object of interest is an average, we can derive its distribution by an application of the Central Limit Theorem (CLT). The specific choice of CLT is domain specific:

1. In applications for which we can assume independence of examples and stable, finite first and second moments of the distribution of gradients, the basic Lindeberg-Lvy CLT can be applied (Billingsley, 1995) . 2. When the stability of the second moments can't be assumed, the Lyapunov CLT can be used instead (Hayashi, 2000) . 3. For time series applications, we can call on the Gordin's CLT, which relaxes the requirement of example independence to an existence of a limit to the time-series' memory (Hayashi, 2000) . That is it assumes past observations to become irrelevant as time passes.

Identification and Hessian Non-singularity conditions (ID & HN).:

If the model's minima are isolated, the ID can be trivially satisfied if we restrict the parameter space Θ to contain a single minimum. Then, the weight distribution is derived for the given minimum (the one model converged to). Furthermore, the HN is satisfied because the isolated minimum requires the Hessian to be positive definite and hence non-singular. Statistical theory, however, suggests that the minima are likely to be connected, violating the isolation requirement. This was observed in practice by Sagun et al. (2017) and Draxler et al. (2018) . In such case, Gill & King (2004) suggest working with an approximation of a weight distribution, as opposed to the exact formulation.

3.2. DERIVING THE WALD STATISTIC-BASED PRUNING CRITERION:

The Wald test is a general framework for testing the significance of explanatory features (Hayashi, 2000) . In this context, a lack of significance means that features can be set to 0 with little loss of model performance. The test assumes that the difference between the optimal and the estimated parameters is approximately Normally distributed. We apply the test by considering γ-many weights to be prunable. The null hypothesis H 0 , which states that some weights can be set to 0, and Wald's salience measure W would take the following form:

H 0 : a (θ 0 ) = 0 W = n * a θ T ∂a (θ) ∂θ Υ ∂a (θ) ∂θ T -1 a θ ∼ χ 2 γ (2)

In this equation, Υ is the parameters' covariance matrix and a here is defined as a(θ 0 ) = p T θ 0 , where p is a vector that contains γ-many 1s and is otherwise 0. If the null hypothesis is true, elements of the parameter vector θ 0 picked out by p have to be 0.

The salience criterion based on the distribution in Equation 1 is then:

EQUATION

3 So that for any consistent estimator θ we have n t=1 H(yt, xt, θ)

→ d E [H(y, x, θ0)]

. 4 The rest are met close to trivially. Therefore we isolate them in the appendix.

Avar θ = E [H(y t , x t , θ 0 )] -1 Σ E [H(y t , x t , θ 0 )] -1

The θ are the weight estimates, Σ is the estimated score (gradient) covariance matrix.

In the pruning's single weight importance setup γ = 1. The function a θ will one-hot select the parameter at position k, where k is the investigated parameter's index. This means that the inner portion of the sandwich form in Equation 3 picks out the kk-th (diagonal) element of the Avar θ square matrix. Finally, the W k , the salience of the k-th parameter is therefore defined as:

W k = n * θk Avar θ kk -1 θk ∼ χ 2 1 (4)

This criterion, which we call Wald criterion in a nod to its roots in the statistical theory, bares close resemblance to other pruning criteria previously developed and adopted in the deep learning literature:

1. Similarly to magnitude pruning, the magnitude of a weight estimate is captured by θk . 2. Similarly to Optimal Brain Damage (LeCun et al., 1990) , the salience measure uses Hessians but it is not specific to the mean square error loss function. 3. Similarly to N2PS (Tong & Tanaka, 2015) , it uses gradient-derived information but only after convergence, when values are stable.

In addition to being derived using statistical theory, Wald pruning can also be justified with intuitions that were used to justify the criteria above, as it incorporates all three of them. Therefore, beside providing a theory-based alternative, this novel use of the Wald statistic in deep model pruning reconciles the differences between previous pruning approaches by appropriately combining them into a single holistic method.

3.3. ASYMPTOTIC DISTRIBUTION PARAMETER ESTIMATION:

Finally, in order to compute the Wald-based salience measure, we need to be able to estimate the asymptotic variance, Avar θ , as given in Equation 3. Below we provide the two estimators needed to compute it -the estimator of the Hessians (H) and of the score covariance matrix ( Σ).

First, we look at the covariance estimator. Since the necessary conditions of the SD assumption imply the existence of the first and the second score moments, we have, also, justified the existence and estimability of the score variance.

EQUATION

Where = 1 holds because the model is assumed to have converged. Then, since θ → p θ 0 , we can write the sample score covariance estimator as:

EQUATION

)

Where all is defined as above.

Secondly, we focus on the Hessian estimator. Since θ → p θ 0 , the Local Dominance condition further implies that:

EQUATION

where all is defined as above.

Therefore we can rewrite the estimated deep model parameter distribution as:

θ ∼ N θ0, E H(yt, xt, θ) -1 1 n n t=1 s yt, xt, θ s yt, xt, θ T E H(yt, xt, θ) -1 (8)

where all is defined as above.

This is a fully estimable empirical distribution. It can be used to test weight exclusion hypotheses, as in the pruning setup, but it can also be used to test weight equality restrictions, as would required by quantization applications. Its usefulness, therefore, reaches beyond the present setup.

4. EXPERIMENT RESULTS:

In this section, we present the experimental results of pruning in several setups using the magnitude and the Wald statistics. We aim to quantify the method's relative performance right after pruning and during as well as after model re-training.

4.1. METHOD:

We adapt Han et al. (2015) 's setup. We use a three-step pruning procedure to prune p% (i.e., pruning threshold) of weights: (a) a network is trained until convergence; (b) salience score is computed for all parameters and the least salient p% are set to 0 and frozen; (c) the network is retrained to recover some of the lost performance. We consider two modes of pruning-one-shot, where all p% of parameters are pruned at once, and iterative, where parameters are pruned in smaller groups and steps (b) and (c) are repeated until the desired sparsity is reached. In our experiments we consider p = 60%, 80%, 90%, 92%, 94%, 96%, 98%, 99%.

To approximate Hessian computation for our models, we note that optimizing a model with a Softmax output and the cross-entropy loss function is equivalent to Maximum Likelihood problem5 . In this case, the approximation H(x) = J(x) T J(x) holds exactly (Hayashi (2000, p. 510) ; Newey & Mcfadden (1994 , p. 2147) ).

Some networks contain Dropout (Srivastava et al., 2014) and Batch Normalization (Ioffe & Szegedy, 2015) layers which have to be adapted for pruned networks. For Dropout, which can be seen as randomized pruning with probability t, we update the Dropout probability to t √ p (Han et al., 2015) .

For Batch Normalization, we chose to reset the learned mean and variance parameters, letting the network to re-learn them during retraining.

We use the MNIST dataset (LeCun et al., 1998) for our experiments, which is a popular 10-class handwritten digit classification problem, and use TensorFlow (Abadi et al., 2016) to implement the models and perform training.

The purpose of the experiments that follow is to aid the understanding of the pruning decisions and their implications for the post-pruning recovery process. For that reason we select the data set MINST, as it is straightforwardly interpretable and well suited for demonstrating intuitions. While other, larger, data sets exist, critically, MINST allows us to reason, with high degree of confidence, about specific input layer parameters. That would not be practical with more complex data sets. Using this data set, We train architectures LeNet-300-100 and LeNet5, as they are tailored to this data set. Again, while larger architectures exist, their use would be impractical as their depth would obscure, rather than illuminate the core intuitions of this paper.

4.2. EXPERIMENTS:

In what follows we provide three core experiments that examine Walt-based Pruning with magnitude and random pruning criteria.

Salience Map Comparison. In this first experiment, we compare the input layer salience values (i.e. the salience map) for a two-layer MLP under three different criteria. We deliberately design this experiment under MNIST so that it is straightforward to reason qualitatively about what input layer regions should be more important than others.

We make use of the LeNet-300-100 (LeCun et al., 1998) , which is a two-layer MLP model, with 300 and 100 hidden units in each layer respectively, followed by the output layer with 10 neurons for each digit class.

Random The Figure 1 shows the salience maps of the input (784 by 300) fully-connected layer. A reasonable expectation is that the background of the MNIST images ought to be disregarded. It is the same in all examples, thus weights connected to it are at their initialized values.

λ L2 = 0.0001 Magnitude Stat λ L2 = 0.5

The Wald pruning exhibits the reasonable behavior across all regularization setups. The magnitude does not. It requires an extreme regularization of around λ = 0.5 to pick up the same pattern. This, however, comes at a significant cost. The best performing regularization constant λ = 0.0001 delivered 94.3% accuracy, while the λ = 0.5 scored only 67.9%. We conclude that only very severe norm-based regularization induces the magnitude to behave in a manner significantly differing from random. This means, that magnitude criterion delivers non-random decisions only in degenerate setups. Wald criterion, on the other hand, correctly predicts background irrelevance irrespective of how much help it receives from the network regularizer.

Compression Frontier -Fully Connected Network. Next, we investigate what the accuracy implications of the criteria are. The Figure 2 shows how the model's accuracy changes compared to the full model (fraction 0.0) as more parameters are pruned (L 2 regularization at λ = 0.0001). We show results for the aforementioned iterative and one-shot pruning schedules.

As was also observed by Han et al. (2015) , the one-shot pruning is strictly inferior in performance. This is to be expected as the criteria were derived for individual / small groups of parameters. The one-shot setup seriously violates this setup. In either setup, the Wald criterion pushes the frontier by roughly 0.8% of compression, or equivalently by circa 2.5% of accuracy retention. This is reflected in savings of around 2130 parameters.

The margin of outperformance might feel underwhelming at first, but it has to be recognized that the fully connected layers of LeNet-300-100 have a significant degree of flexibility, whereby weights within a layer are relatively easily substitutable6 . Consequently, a lot of inefficient pruning is compensated for by more aggressive7 re-training. Furthermore, the state of the art frameworks compete for every tenth of a percent improvement in either metric (Yang et al., 2017; 2018) . Therefore, 0.8% pruning shift and 2.5% accuracy improvement are very significant when viewed in the appropriate scale.

Compression Frontier -Convolutional Network. The natural extension is to test the same pruning-accuracy setup on a convolutional network. We choose to work with LeNet5 (LeCun et al., 1998) for its similarity in size to the LeNet-300-100. It is a model which consists of 2 convolutional and pooling layers followed by 3 fully-connected layers.

Figure 3 , as did Figure 2 , shows the accuracy change resulting from varying levels of compression.

One-shot pruning, again, outperforms the iterative alternative. This time, however, the importance of the individual weight setup is revealed to be much stronger for the Wald criterion than to the magnitude. This is to be expected if we recognize that the magnitude pruning is likely to have close to random behavior.

Relative to the fully connected setup, we observe larger accuracy drops at the far end of pruning, in line with Han et al. (2015) .

Finally, the frontier shift occurs later, indicating that convolution's weight magnitudes are probably more informative about their importance. Further experiments are needed to make that conclusion.

Recovery post-pruning. Overall, the above results suggest convergence of pruned model performance during the re-training process. We test this implication in this final experiment. Figure 4 shows how training loss evolves in the re-training process at an intermediate 90% pruning ratio.

Indeed, the prediction was confirmed for the iterative schedule. The Wald criterion started with a head start, which it gradually, but not completely, lost to flexibility. The one-shot schedule hurt the Wald criterion much more. This reinforces our confidence in the above justifications.

4.3. DISCUSSION:

Wald pruning is superior. The Wald criterion successfully internalizes the weights' uncertainty by scaling them by their variances. Consequently, it delivers superior performance with margin that is comparable in scale to the state of the art relevance levels (Yang et al., 2017; 2018) . Training loss Stat (iterative) @ 0.90 Magnitude (iterative) @ 0.90 Stat (one-shot) @ 0.90 Magnitude (one-shot) @ 0.90 Models recover considerably in re-training. The model can recover to original levels of accuracy even when a large fraction of parameters was pruned close to randomly thanks to the easy substitutability of model parameters.

Iterative pruning is preferable. One-shot pruning severely violates the assumptions under which the criteria were computed. Both performances are impacted, with Wald criterion being much more sensitive. Unlike for magnitude pruning, one might straightforwardly derive the Wald criterion for groups of parameters. This would, however require computing such criterion for all possible pruning group compositions, which would be too expensive and defeat the purpose.

5. CONCLUSION:

In this paper we analyzed the use of weight salience pruning with a focus on the magnitude criterion. We argued it fails to account for parameter uncertainty in its salience measure. Therefore, we derived a Statistics-founded, Wald, alternative, which internalizes the consideration. The novel criterion delivers superior accuracy and requires less intensive re-training post-pruning. The over-performance is on scale relevant to the current state of the art.

covariance matrix. The proof follows exactly the same lines as that for the logit with log likelihood objective function in Hayashi (2000, p. 510) . This means that for as long as the previous layer does not contain perfectly correlated nodes, the loss function hessian will be bounded and thus estimable. Such feature correlation would imply an obvious redundancy that can be diagnosed and remedied by dropping all but one of the correlated features. A re-training would be needed afterwards.

:

A detailed argument for deep learning's membership in the family and extended justification of the assumptions are presented in the Appendix. Details behind the distribution result are available inHayashi (2000).

2 This is the distribution of the average of loss gradients with respect to the parameter vector evaluated at the true parameter vector and the individual samples.

For all deep models with Cross-entropy loss and Softmax output layer, there exists a probability model specification, such that its Maximum Likelihood objective function matches one-to-one the deep network's loss function.

Activations are not fixed, re-training can re-learn those that were lost if sufficient number of parameters in the reminder of the network were retained.

Larger learning rate.

We change the sign in order to stay in the maximization paradigm.

Figures and Tables:

Figure 1: Salience map for the first layer (768 inputs, 300 outputs) of the LeNet-300-100 model for 3 different criteria with L 2 regularization at λ = 0.0001 and λ = 0.5. Note that less salient areas correspond to edges of the input image, which are empty in MNIST. Unless regularized heavily, Magnitude criterion is indistinguishable from Random.

Figure 2: Accuracy versus fraction of parameters pruned of the LeNet-300-100 model for both (Stat) and magnitude pruning, in both one-shot and iterative modes. Wald pushes out the possibility frontier by circa 0.8%, that is roughly 2130 parameters.

Figure 3: Accuracy versus fraction of parameters pruned of the LeNet5 model for bothWald (Stat) and magnitude pruning, in both one-shot and iterative modes. Wald was derived for individual/small groups of weights, using it en-mass is detrimental.

Figure 4: Train loss evolution during re-training of the LeNet5 model after pruning 90% of the weights. When used properly, Wald gains performance advantage at the beginning of the re-training process which it retains throughout.

Peer review (iclr2019):
 Comments * all conditions are not discussed, for instance, isn*t it quite strong to impose that the parameter space .Theta is compact? * LDH condition: what is .mathcal{N}, it isn*t defined * The statement on top of page 3: "Then, asymptotically, as the sample size n grows to infinity, the model parameters are distributed as (...)" is awkward: is it standard? is it a new result? please consider stating it as a Proposition/Theorem or something, and explain where it comes from. * Actually, I am skeptical about the above statement itself: it is stated that "the model parameters are distributed as (...)". Does that mean that the parameters are random, in a Bayesian fashion? please state it then. Or the alluded to distribution is that of parameters* estimators? then the statement should be amended, and it should be explained what estimators are considered. * footnote 3: shouldn*t the considered sum be divided by n? * the paper is full of flawed equations where, I think, parameters are mixed up with estimators. * For instance, Equations (2) and (3) should express the same quantity, , however the former displays the parameters’ covariance matrix, while the latter uses an estimator of it, .hat{Avar}. * Same in page 4 with the estimator of the Hessians (H) and of the score covariance matrix (.hat Σ). * Same for Equation (6): .hat .theta is suddenly plugged-in in place of .theta_0, all of this conserving the equality sign. * Same for Equation (7) Typos, minor comments and un-precise statements * the footnotes superscripts are misleading for footnotes 2 and 3; they could be better placed. * Lindeberg-Lvy should be Lindeberg-Lévy * after equations (6,7,8) you write "Where all is defined as above": this is fairly un-precise; as you use LaTeX, you should conveniently use some .label and .ref. * end of section 4.1: MINST dataset, a new dataset? * top of section 4.2: there are even typos in the proposed name for methodology! : "Walt-based Pruning" * Many references in the bibliography are incomplete, so that it is impossible to see where they come from (at least seven of them)