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: The paper looks at the compressibility and generalization ability of multi-layer neural networks under the heavy tailed mean-field limit. The mean-field condition is invoked to simplify the analysis of the distribution of weight parameters as the distributions of individual weights are independent, while the heavy-tailed condition is needed to ensure an abundance of small weights (weights at the tail) whose removal would result in bounded relative error. Compression error is defined in weight space as the relative distance between the compressed weight vector (weight vector where small weights are set to zero) and the original weight vector.  The paper shows that arbitrarily low compression errors can be achieved under a range of pruning strategies (global pruning, layer-wise pruning, singular value pruning, and neuron pruning) if the network is sufficiently large. Another condition is for the p-norm used to calculate the error to have p larger than the heavy tail coefficients.