Heavy Tails in SGD and Compressibility of Overparametrized Neural Networks

1 Introduction . With the increasing model sizes in deep learning and with its increasing use in low-resource environments , network compression is becoming ever more important . Among many network compression techniques , network pruning has been arguably the most commonly used method [ O ’ N20 ] , and it is rising in popularity and success [ BOFG20 ] . Though various pruning methods are successfully used in practice and their theoretical implications in terms of generalization are increasingly apparent [ AGNZ18 ] , a thorough understanding of why and when neural networks are compressible is lacking . * Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A common conclusion in pruning research is that overparametrized networks can be greatly compressed by pruning with little to no cost at generalization , including with simple schemes such as magnitude pruning [ BOFG20 , O ’ N20 ] . For example , research on iterative magnitude pruning [ FC19 ] demonstrated the possibility of compressing trained deep learning models by iteratively eliciting a much sparser substructure . While it is known that the choice of training hyperparameters such as learning rate affects the performance of such pruning strategies [ FDRC20 , HJRY20 , RFC20 ] , usually such observations are low in granularity and almost never theoretically motivated . Overall , the field lacks a framework to understand why or when a pruning method should be useful [ O ’ N20 ] . Another strand of research that highlights the importance of understanding network compressibility includes various studies [ AGNZ18 , SAM+20 , SAN20 , HJTW21 , KLG+21 ] that presented generalization bounds and/or empirical evidence that imply that the more compressible a network is , the more likely it is to generalize well . The aforementioned bounds are particularly interesting since classical generalization bounds increase with the dimension and hence become irrelevant in high dimensional deep learning settings , and fall short of explaining the generalization behavior of overparametrized neural networks . These results again illustrate the importance of understanding the conditions that give rise to compressibility given their implications regarding generalization . In this paper , we develop a theoretical framework to address ( i ) why and when modern neural networks can be amenable to very simple pruning strategies and ( ii ) how this relates to generalization . Our theoretical results are based on two recent disparate discoveries regarding deep neural networks trained with the stochastic gradient descent ( SGD ) algorithm . The first one is the emergence of heavytailed stationary distributions , which appear when the networks are trained with large learning rates and/or small batch-sizes [ HM20 , GŞZ21 ] . The second one is the propagation of chaos phenomenon , which indicates that , as the network size goes to infinity , the network weights behave independently [ MMN18 , SS20 , DBDFŞ20 ] . We show that , assuming both of the aforementioned phenomena occur simultaneously , fully connected neural networks will be provably compressible in a precise sense , and the compression errors of ( i ) unstructured global or layer-wise magnitude pruning , ( ii ) pruning based on singular values of the weight matrices , ( iii ) and node pruning can be made arbitrarily small for any compression ratio as the dimension increases . Our formulation of network compressibility in terms of ‘ ` p-compressibility ’ enables us to access results from compressed sensing literature [ GCD12 , AUM11 ] to be used in neural network analysis . Moreover , we prove generalization bounds adapted to our framework that agree with existing compression-based generalization bounds [ AGNZ18 , SAM+20 , SAN20 ] and confirm that compressibility implies better generalization . We conduct experiments on fully connected and convolutional networks and show that the results are in strong accordance with our theory . Our study reveals an interesting phenomenon : depending on the algorithm hyperparameters , such as learning rate and batch-size , the resulting neural networks might possess different compressibility properties . Under the decoupling effect of propagation of chaos that emerges with overparametrization [ DBDFŞ20 ] , the networks become compressible in a way that they are amenable to simple pruning strategies as the tails get heavier , which is shown to depend on the step-size/batch-size ratio [ GŞZ21 ] . Finally , when compressible , the networks become likely to generalize better . In this sense , our results also provide an alternative perspective to the recent theoretical studies that suggest that heavy tails can be beneficial in terms of generalization [ ŞSDE20 , ZFM+20 ] . 2 Preliminaries and Technical Background . Notation . Matrices and vectors are denoted by upper- and lower-case bold letters , respectively , e.g . X and x . A sequence of n scalars x1 , . . . , xn is shown as { xi } ni=1 . Similar notations are used for sequences of matrices { Xi } li=1 and vectors { xi } l i=1 , whose entries are indexed with convention Xi= [ Xi ] m , n and xi= ( xi,1 , xi,2 , . . . , xi , n ) , respectively . The set of integers { m , . . . , n } is denoted by Jm , nK . We denote the ` p ( semi- ) norm of a vector x ∈ Rd as ‖x‖p= ( ∑d i=1 |xi|p ) 1/p for all p ∈ ( 0 , ∞ ) and ‖x‖ implies ‖x‖2 . For a matrix A ∈ Rn×m , ‖A‖ denotes its Frobenius norm . Fully connected neural networks . In the entirety of the paper , we consider a multi-class classification setting . We denote the space of data points by Z = X × Y , where X ⊂ RdX is the space of features and Y = { 1 , 2 , . . . , dY } is the space of the labels . Similar to prior art ( e.g. , [ NBS18 ] ) , we focus our attention on the bounded feature domain X = XB : = { x ∈ RdX ; ‖x‖ ≤ B } . We denote a fully connected neural network with L layers by a collection of weight matrices { Wl } Ll=1 , such that Wl ∈ Rhl×hl−1 , where hl denotes the number of hidden units at layer l with h0 = dX and hL = dY . Accordingly , the prediction function fw ( x ) : XB 7→ RdY , with elements fw ( x ) : = ( fw ( x ) [ 1 ] , . . . , fw ( x ) [ dY ] ) , corresponding to the neural network is defined as follows : fw ( x ) = WLφ ( WL−1φ ( · · ·φ ( W1x ) ) ) , ( 1 ) where φ : R → R is the rectified linear unit ( ReLU ) activation function , i.e. , φ ( x ) = max ( 0 , x ) , and it is applied element-wise when its input is a vector . For notational convenience , let us define dl : = hl × hl−1 and d : = ∑L l=1 dl . Furthermore , let wl denote the vectorized weight matrix of layer l , i.e. , wl : = vec ( Wl ) ∈ Rdl , where vec denotes vectorization . Finally , let w be the concatenation of all the vectorized weight matrices , i.e. , w : = [ w1 , . . . , wL ] ∈ Rd . We assume that L ≥ 2 . Risk minimization and SGD . In order to assess the quality of a neural network represented by its weights w , we consider a loss function ` : Y × Y 7→ R+ , such that ` ( y , fw ( x ) ) measures the loss incurred by predicting the label of x as arg maxj fw ( x ) [ j ] , when the true label is y . By following a standard statistical learning theoretical setup , we consider an unknown data distribution µZ over Z , and a training dataset with n elements , i.e. , S = { z1 , . . . , zn } , where each zi = : ( xi , yi ) i.i.d.∼ µZ . We then denote the population and empirical risks as R ( w ) : = E ( x , y ) ∼µZ [ ` ( y , fw ( x ) ) ] and R̂ ( w ) : = 1 n ∑n i=1 ` ( yi , fw ( xi ) ) . Since µZ is unknown , we can not directly attempt to minimizeR in practice . One popular approach to address this problem is the empirical risk minimization strategy , where the goal is to solve the following optimization problem : minw∈Rd R̂ ( w ) . To tackle this problem , SGD has been one of the most popular optimization algorithms , which is based on the following simple recursion : w { k+1 } = w { k } − η∇R̃k+1 ( w { k } ) , where ∇R̃k ( w ) : = ( 1/b ) ∑ i∈Ωk ∇ ` ( yi , fw ( xi ) ) . ( 2 ) Here , w { k } denotes the weight vector at iteration k ∈ N+ , η > 0 is the step-size ( or learning rate ) , ∇R̃k ( · ) is the stochastic gradient , and Ωk is the mini-batch with size b , drawn with or without replacement from J1 , nK . Heavy-tailed distributions and the α-stable family . In this study , we will mainly deal with heavy-tailed random variables . While there exist different definitions of heavy tails in the literature , here , we call a random variable heavy-tailed if its distribution function has a power-law decay , i.e. , P ( X ≥ x ) ∼ cx−α as x→∞ , for some c > 0 and α ∈ ( 0 , 2 ) . Here , the tail index α determines the tail thickness of the distribution , i.e. , as α get smaller , the distribution becomes heavier-tailed . An important subclass of heavy-tailed distributions is the family of stable distributions . A random variable X has symmetric α-stable distribution , denoted by SαS ( σ ) , if its characteristic function is equal to E [ exp ( iwX ) ] = exp ( −|σw|α ) , where α ∈ ( 0 , 2 ] is again the tail index and σ ∈ ( 0 , ∞ ) is the scale parameter . An important property of SαS is that whenever α < 2 , E|X|p is finite if and only if p < α . This implies that the distribution has infinite variance as soon as α < 2 . In addition to their wide use in applied fields [ Nol20 ] , recently , SαS distributions have also been considered in deep learning theory [ SSG19 , PFF20 , ZFM+20 , ŞSDE20 ] and optimization [ WGZ+21 ] . 3 Compressibility and the Heavy-Tailed Mean-Field Regime . In this section , we will present our first set of theoretical results . We first identify a sufficient condition , then we prove that , under this condition , the compression errors of different pruning techniques become arbitrarily small as the network size increases . 3.1 The heavy-tailed mean-field regime . Due to the peculiar generalization behavior of neural networks trained with SGD , recent years have witnessed an extensive investigation of the theoretical properties of SGD in deep learning [ Lon17 , DR17 , KL18 , ZZL19 , AZLL19 , ŞSDE20 , Neu21 , CDE+21 , BLGŞ21 ] . We will now mention two recently established theoretical properties of SGD , and our main contribution is studying the compressibility properties of the network under the assumption that these two seemingly unrelated ‘ phenomena ’ occur simultaneously . The heavy tail phenomenon . Several recent studies [ MM19 , SSG19 , ŞGN+19 , ŞZTG20 , ZFM+20 , ZKV+20 , CWZ+21 ] have empirically illustrated that neural networks can exhibit heavy-tailed behavior when optimized with SGD . Theoretically investigating the origins of this heavy-tailed behavior , [ HM20 ] and [ GŞZ21 ] later proved several theoretical results on online SGD , with rather surprising implications : due to the ‘ multiplicative ’ structure of the gradient noise ( i.e. , ∇R̃k ( w ) − ∇R̂ ( w ) ) , as k → ∞ , the distribution of the SGD iterates w { k } can converge to a heavy-tailed distribution with infinite second-order moment , and perhaps more surprisingly , this behavior can even emerge in simple linear regression with Gaussian data . The authors of [ GŞZ21 ] further showed that , in the linear regression setting , the tail index α is monotonic with respect to the step-size η and batch-size b : larger η and/or smaller b result in smaller α ( i.e. , heavier tails ) . Propagation of chaos and mean-field limits . Another interesting property of SGD appears when the size of the network goes to infinity . Recently , it has been shown that , under an appropriate scaling of step-sizes , the empirical distribution of the network weights converges to a fixed meanfield distribution , and the SGD dynamics can be represented as a gradient flow in the space of probability measures [ MMN18 , CB18 , SS20 , DBDFŞ20 ] . Moreover , [ DBDFŞ20 ] showed that a propagation of chaos phenomenon occurs in this setting , which indicates that when the network weights are initialized independently ( i.e. , a-priori chaotic behavior ) , the weights stay independent as the algorithm evolves over time ( i.e. , the chaos propagates ) [ Szn91 ] . Our focus in this paper is the setting where both of these phenomena occur simultaneously : we assume that as the size of the network goes to infinity , the SGD iterates become independent and their distribution converges to a heavy-tailed distribution . To formalize this setting , let us introduce the required notations . For an integer ml < dl , let wl , ( 1 : ml ) denote the first ml coordinates of wl , and for ml ≥ dl set wl , ( 1 : ml ) = wl . Furthermore , parametrize the dimension of each layer dl with a parameter ρ ∈ R , i.e. , dl = dl ( ρ ) , such that for all l = 1 , . . . , L , we have : dl ( ρ ) ∈ N+ , and limρ→∞ dl ( ρ ) =∞ . ( 3 ) This construction enables us to take the dimensions of each layer to infinity simultaneously [ NP21 ] . Condition 1 ( Heavy-tailed mean-field limit condition ) . The SGD recursion ( 2 ) satisfies the heavytailed mean-field limit ( HML ) condition , if the dimensions obey ( 3 ) and if there exist heavy-tailed probability measures { µ ? l } Ll=1 on R with tail indices { αl } Ll=1 , such that µ ? l ( { 0 } ) = 0 for all l and for any m1 , . . . , mL ∈ N+ , the joint distribution of ( w { k } 1 , ( 1 : m1 ) , . . . , w { k } L , ( 1 : mL ) ) weakly converges to ( µ ? 1 ) ⊗m1 ⊗ · · · ⊗ ( µ ? L ) ⊗mL , ( 4 ) as ρ , k →∞ , where µ⊗ ν denotes the product measure and µ⊗n denotes the n-fold product measure . Informally , the HML condition states that in the infinite size and infinite iteration limit , the entries of the weight matrices will become independent , and the distribution of the elements within the same layer will be identical and heavy-tailed . We acknowledge that the above condition may not be always satisfied , and theoretical evidence for this condition is provided in rather seemingly simpler settings such as linear regression [ GŞZ21 ] and/or two-layer neural networks [ DBDFŞ20 ] . However , we will empirically illustrate the behavior described by the HML condition in Section 5 in certain scenarios . Note that , in ( 4 ) , the particular form of independence in the limit is not crucial and we will discuss weaker alternatives in Section 3.3 . We further note that [ DBDFŞ20 ] proved that a similar form of ( 4 ) indeed holds for SGD applied on single hidden-layered neural networks , where the limiting distributions possess second-order moments ( i.e. , not heavy-tailed ) and the independence is column-wise . Recently , [ PFF20 ] investigated the infinite width limits of fully connected networks initialized from a SαS distribution and proved heavy-tailed limiting distributions . On the other hand , heavy-tailed propagation of chaos results have been proven in theoretical probability [ JMW08 , LMW20 ] ; however , their connection to SGD has not been yet established . We believe that ( 4 ) can be shown to hold under appropriate conditions , which we leave as future work .

Summary: This paper analyze the compressibility of over-parameterized neural networks using the heavy-tail property of SGD. The authors show that when the network weights are approximately sampled from a heavy-tailed distribution, the weights are more compressible. The authors also provide a generalization guarantee showing that more compressible networks have smaller generalization gap.