Robustness to Pruning Predicts Generalization in Deep Neural Networks

1 INTRODUCTION . The gap between learning-theoretic generalization bounds for highly overparameterized neural networks and their empirical generalization performance remains a fundamental mystery to the field ( Zhang et al. , 2016 ; Jiang et al. , 2020 ; Allen-Zhu et al. , 2019 ) . While these models are already being successfully used in many applications , improving our understanding of how neural networks perform on unseen data is crucial for safety-critical use cases . By understanding which factors drive generalization in neural networks we may further be able to develop more efficient and performant network architectures and training methods . Numerous theoretically and empirically motivated attempts have been made to identify generalization measures , that is , properties of the trained model , training procedure and training data that distinguish models that generalize well from those that do not ( Jiang et al. , 2020 ) . A number of generalization measures have attempted to quantify Occam ’ s razor , i.e . the principle that simpler models generalize better than complex ones ( Neyshabur et al. , 2015 ; Bartlett et al. , 2017 ) . This has proven to be non-trivial , as many measures , particularly norm-based measures , grow with the size of the model and thus incorrectly predict that larger networks generalize worse than smaller networks . Other approaches have tried to establish a connection between model compression and generalization ( Arora et al. , 2018 ; Zhou et al. , 2019 ) . While both of these approaches are theoretically elegant and yield tighter bounds than bounds that are based on the size of uncompressed networks , they nonetheless grow with the size of the original network . Recent empirical studies ( Jiang et al. , 2020 ; 2019 ) , on the other hand , identify three classes of generalization measures that do seem predictive of generalization : measures that estimate the flatness of local minima , the speed of the optimization , and the margin of training samples to decision boundaries . While these measures are correlated with generalization , their failure to fully explain the test performance of the model demonstrate a need for other notions of model simplicity to explain generalization in neural networks . In this paper , we leverage the empirical observation that large fractions of trained neural networks ’ parameters can be pruned – that is , set to 0 – without hurting the models ’ performance ( Gale et al. , 2019 ; Zhu & Gupta , 2018 ; Han et al. , 2015 ) . Based on this insight , we introduce a new measure of a model ’ s simplicity which we call prunability : the smallest fraction of weights we can keep while pruning a network without hurting its training loss . In a range of empirical studies , we demonstrate that a model ’ s prunability is indeed highly informative of its generalization . In particular , we find that the larger the fraction of parameters that can be pruned without hurting a model ’ s training loss , the better the model will generalize . Overall , we show that the smaller the fraction of parameters a model actually “ uses ” - the simpler it is - the better the network ’ s generalization performance . In summary , our main contributions are thus the following : 1 . We introduce a new generalization measure called prunability that captures a model ’ s simplicity ( Sec . 4 ) and show that across a large set of models this measure is highly informative of a network ’ s generalization performance ( Sec . 5.4.1 ) . 2 . We show that even in a particularly challenging setting in which we observe a test loss double descent ( Nakkiran et al. , 2020 ; He et al. , 2016 ) , prunability is informative of models ’ test performance and is competitive with an existing strong generalization measures ( Sec . 5.4.2 ) . 3 . Lastly , we investigate whether the success of prunability can be explained by its relationship to flat local minima . We find that while prunability makes similar predictions to some existing measures that estimate the flatness of minima , it differs from them in important ways , in particular exhibiting a stronger correlation with generalization performance than these flatness-based measures ( Sec . 5.4.1 and 5.4.3 ) . 2 RELATED WORK . Zhang et al . ( 2016 ) demonstrate that neural networks can perfectly memorize randomly labeled data while still attaining good test set performance on correctly labeled data , a phenomenon that complexity measures such as the VC-Dimension and Rademacher complexity fail to explain . More recently , increasing the parameter count of models has been shown to improve their generalization performance , yielding ‘ double descent ’ curves in which the test loss first decreases , then increases and then decreases again as the parameter count is increased ( He et al. , 2016 ; Nakkiran et al. , 2020 ; Belkin et al. , 2019 ) . In contrast , many existing generalization measures grow monotonically with the model ’ s size and thus fail to capture double descent , making it a particularly interesting setting to study new generalization measures ( Maddox et al. , 2020 ) . While generalization measures have been studied for a long time , Jiang et al . ( 2020 ) recently brought new empirical rigor to the field . They perform a large-scale empirical study by generating a large set of trained neural networks with a wide range of generalization gaps . Proposing a range of new evaluation criteria , they test how informative of generalization previously proposed generalization measures actually are . In this paper , we evaluate our new generalization measure in the same framework and use the strongest measures as baselines for comparison against our proposed measure , the neural network ’ s prunability . A common theme across many generalization measures is that they try to formalize Occam ’ s Razor , the idea that simpler models generalize better than more complex models . Rissanen ( 1986 ) formalizes this principle as a model ’ s Minimum Description length which was later applied to small neural networks by Hinton & Van Camp ( 1993 ) . Similarly , Neyshabur et al . ( 2015 ; 2017 ) and numerous other approaches suggest that networks with smaller parameter norms generalize better . We are aware of two existing generalization measures that are based on compression or pruning . Arora et al . ( 2018 ) derive a PAC-Bayesian generalization bound based on the size of a network following a particular compression method based on noise stability . Zhou et al . ( 2019 ) , on the other hand , formulate a generalization bound that is non-vacuous even for large network architectures based on the description length of a compressed network which in turn also grows with the size of the original network . While these approaches are theoretically elegant , all such measures grow with the number of parameters in a given network and thus do not capture the phenomenon that more highly overparameterized networks tend to generalize better . Three classes of generalization measures seem to be particularly predictive of generalization . First , measures that estimate the flatness of local minima , either based on random perturbations of the weights ( Neyshabur et al. , 2017 ; Dziugaite & Roy , 2017 ; Jiang et al. , 2020 ; McAllester , 1999 ) or directly based on the curvature of the loss surface ( Maddox et al. , 2020 ) . In a purely empirical study , Morcos et al . ( 2018 ) find that a model ’ s robustness to clamping the activations of individual units to fixed values is predictive of its generalization performance . Second , margin-based measures that use the distance of training samples to decision boundaries ( Jiang et al. , 2019 ) . Lastly , Jiang et al . ( 2020 ) ; Ru et al . ( 2020 ) have proposed measures that estimate the speed of optimization . A large fraction of a network ’ s parameters can be removed without impairing performance ( Gale et al. , 2019 ; Blalock et al. , 2020 ) ; in fact , such pruning often improves generalization performance ( Thodberg , 1991 ) . There are many approaches to selecting parameters for removal ; we focus on randomly selected parameters ( random pruning ) and removing the smallest-magnitude parameters ( magnitude pruning ) , though other approaches exist in the literature ( Han et al. , 2015 ; Renda et al. , 2020 ; Lee et al. , 2019 ; Wang et al. , 2020 ; Blalock et al. , 2020 ) . Random pruning is applied by the regularization method dropout ( Srivastava et al. , 2014 ) , which is used to improve network generalization performance . Magnitude pruning can be used in an iterative pruning and retraining procedure ( Frankle & Carbin , 2019 ) to identify sparse subnetworks that , for some sparsity levels , generalize better than their parent network . Both approaches are therefore natural candidates for a notion of network simplicity , which we will leverage in our proposed generalization measure . While research on dropout and the Lottery Ticket Hypothesis suggest that pruning during training and generalization are connected , we study the distinct question of whether pruning after training and generalization are connected . 3 NOTATION . Following the notation of Jiang et al . ( 2020 ) , we denote by D the distribution over inputs X and their labels y and let κ denote the number of classes . By S we denote a given data set containing m i.i.d . tuples { ( X1 , y1 ) , . . . , ( Xm , ym ) } drawn from D where Xi ∈ X is the input data and yi ∈ { 1 , . . . , κ } the class labels . We denote a feedforward neural network by fw : X → Z , its weight parameters by w and the number of its parameters by ω . We denote the weight tensor of the ith layer of the network by Wi and we then have that w = vec ( W1 , . . . , Wd ) , where d is the depth of the network and vec denotes the vectorization operator . Let L be the 1-0 classification error under the data distribution D : L ( fw ) : = E ( X , y ) ∼D [ fw ( X ) 6= y ] , and let L̂ be the empirical 1-0 error on S : L̂ ( fw ) : = 1m ∑m i=1 [ fw ( Xi ) 6= yi ] . We define the generalization gap as g ( fw ) : = L ( fw ) − L̂ ( fw ) . When training a network , we denote each hyperparameter by θi ∈ Θi for i = 1 , . . . n where n is the total number of hyperparameters . We denote a particular hyperparameter configuration by θ : = ( θ1 , θ2 , . . . , θn ) ∈ Θ , where Θ : = Θ1 × Θ2 × · · · ×Θn . 4 PRUNABILITY . A generalization measure should satisfy three straightforward desiderata : 1 ) It should correlate with the generalization gap across a broad range of models . 2 ) It should not suffer trivially from Goodhart ’ s law , that is , incorporating it into the optimization procedure of a model should lead to improved generalization . 3 ) It should ideally have a plausible theoretical justification to allow for further study of its properties . A measure based on magnitude- and random pruning presents a promising direction for satisfying all three desiderata . First , it is well-established that we can remove large fractions of a neural network ’ s parameters without negatively impacting its train performance , while also improving its test performance ( Blalock et al. , 2020 ; Thodberg , 1991 ) . Second , training models using dropout , i.e . random pruning , is widely known to improve generalization ( Srivastava et al. , 2014 ) , as does the iterative magnitude pruning procedure of Frankle & Carbin ( 2019 ) ( up to certain sparsity levels ) . Finally , pruning has a theoretical grounding in the minimum description length principle and in PAC-Bayesian bounds . The connection between dropout and generalization has been studied in the PAC-Bayesian framework ( McAllester , 1999 ) , which allows us to quantify the worst-case generalization performance of random predictors . PAC-Bayesian bounds define model complexity as the KL divergence be- tween some fixed prior distribution P , and the probability distribution over functions induced by the random predictor ( often called the posterior Q , although this may not correspond to a Bayesian posterior ) to obtain generalization bounds of the following form . Theorem 1 ( McAllester ( 1999 ) ) . For any δ > 0 , data distributionD , prior P , with probability 1−δ over the training set , for any posterior Q the following bound holds : Ev∼Q [ L ( fv ) ] ≤ Ev∼Q [ L̂ ( fv ) ] + √ KL ( Q||P ) + log ( mδ ) 2 ( m− 1 ) McAllester ( 2013 ) shows that in the case where both the prior and posterior are distributions over networks to which dropout with a given dropout probability α is applied , the KL term becomes : KL ( Q||P ) = 1−α2 ||w|| 2 . We can directly use this formalization to derive a generalization bound for networks that are randomly pruned to the largest α , s.t . Ev∼Q [ L̂ ( fv ) ] < L̂ ( fw ) × ( 1 + β ) , i.e. , such that the training loss does not increase more than a relative tolerance β . We have to adjust the bound to account for the search for α . We make use of the fact that we search over a fixed number c of possible values for α and use union bound which will change the log term in the bound to log ( cmδ ) . We thus obtain the following bound : Ev∼Q [ L ( fv ) ] ≤ Ev∼Q [ L̂ ( fv ) ] + √ 1−α 2 ||w|| 2 2 + log ( m δ ) + 5 2 ( m− 1 ) . ( 1 ) While this bound will be vacuous for many large neural networks , it does provide us with some intuition for why we should expect a model ’ s generalization to be linked to its robustness to pruning . In particular , this bound predicts the relative generalization performance of models from which different fractions of parameters can be removed without adversely affecting their training loss . It is this prediction that we evaluate with the experiments described below . We provide the details on the derivation in Appendix C. Given a trained model and a pruning algorithm , we define the model ’ s prunability as the smallest fraction of the model ’ s parameters that we can keep while pruning without increasing the model ’ s training loss by more than a given relative tolerance β . We choose the fraction of parameters pruned rather than the number of parameters remaining based on the empirical observation that naive parameter counting is a poor predictor of generalization performance . Definition 1 ( Prunability ) . Let L̂ denote the cross-entropy loss . For a given model fw , a training data set Strain , a relative bound on the change in train loss β and a pruning method φ , we thus define prunability as µprunability ( fw ) = arg min α φ ( fw , α ) s.t . L̂ ( φ ( fw , α ) , Strain ) ≤ ( 1 + β ) × L̂ ( fw , Strain ) , ( 2 ) where φ ( fw , α ) sets the fraction 1− α of the weights of fw to zero . In our experiments we evaluate robustness to pruning both with regards to magnitude pruning , in which the parameters with the smallest absolute magnitude are set to 0 , and random pruning , where random parameters are set to 0 . Magnitude pruning can remove a greater number of model weights without affecting performance than random pruning ( Blalock et al. , 2020 ) , and can therefore be considered a more effective pruning strategy . We hypothesize that using a more effective pruning method might yield a measure of prunability that is more predictive of generalization . Per default , we refer to a model ’ s prunability with regards to magnitude pruning in the following . Note that while iterative pruning is usually used when the aim is to achieve the best compression rate , we are only concerned with the generalization of the model at the end of training and thus apply one-shot pruning . In our experiments we use a tolerance β of 0.1 and note that the particular choice of β did not seem to have a large impact on the predictiveness of prunability . See Appendix A.6.1 for pseudo-code of how prunability is computed for a given model .

Summary: In the present work, the authors tackle the highly debated (and sometimes confusing) problem of finding a good simplicity/complexity measure able to predict generalization performance of deep networks. A novel measure called 'prunability' is introduced and compared with some of the many alternatives in the literature. This property measures how networks are able to retain low training loss when a fraction of the weights is set to zero, and is clearly related to common training practices (e.g. dropout) that seems to yield better generalization performance in practice. The experimental settings and the evaluation methods for this new metric are inspired by recent extensive studies on deep networks performance. The authors are able to show that prunability is in fact associated with good generalization and seems able to capture some non-trivial phenomena (double-descent), but they also find it to be inferior to pre-existing (margin based) measures. Moreover, the close relationship to perturbation robustness and flatness measures is investigated, but the results are not fully conclusive.