Heavy Tails in SGD and Compressibility of Overparametrized Neural Networks
Year: 2021

Abstract:
Neural network compression techniques have become increasingly popular as they can drastically reduce the storage and computation requirements for very large networks. Recent empirical studies have illustrated that even simple pruning strategies can be surprisingly effective, and several theoretical studies have shown that compressible networks (in specific senses) should achieve a low generalization error. Yet, a theoretical characterization of the underlying causes that make the networks amenable to such simple compression schemes is still missing. In this study, focusing our attention on stochastic gradient descent (SGD), our main contribution is to link compressibility to two recently established properties of SGD: (i) as the network size goes to infinity, the system can converge to a mean-field limit, where the network weights behave independently [DBDFŞ20], (ii) for a large stepsize/batch-size ratio, the SGD iterates can converge to a heavy-tailed stationary distribution [HM20, GŞZ21]. Assuming that both of these phenomena occur simultaneously, we prove that the networks are guaranteed to be ‘`p-compressible’, and the compression errors of different pruning techniques (magnitude, singular value, or node pruning) become arbitrarily small as the network size increases. We further prove generalization bounds adapted to our theoretical framework, which are consistent with the observation that the generalization error will be lower for more compressible networks. Our theory and numerical study on various neural networks show that large step-size/batch-size ratios introduce heavy tails, which, in combination with overparametrization, result in compressibility.

heading: 1 Introduction
text: 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].

heading: 2 Preliminaries and Technical Background
text: 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 cannot 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].

heading: 3 Compressibility and the Heavy-Tailed Mean-Field Regime
text: 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.

heading: 3.1 The heavy-tailed mean-field regime
text: 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.

heading: 3.2 Analysis of compression algorithms
text: In this section, we will analyze the compression errors of three different compression schemes under the HML condition. All three methods are based on pruning, which we formally define as follows. Let x be a vector of length d, and consider its sorted version in descending order with respect to the magnitude of its entries, i.e., |xi1 | ≥ |xi2 | ≥ · · · ≥ |xid |, where {i1, . . . , id} = {1, . . . , d}. For any k ≤ d, the k-best term approximation of x, denoted as x(k) = (x(k)1 , . . . , x (k) d ), is defined as follows:
for l ∈ J1, dkeK, x(k)il := xil and for l /∈ J1, dkeK, x (k) il
:= 0. Informally, we keep the k-largest entries of x with the largest magnitudes, and ‘prune’ the remaining ones. Current results pertain to one-shot pruning with no fine-tuning; other settings are left for future work (cf. [EKT20]).
In this section, we consider that we have access to a sample from the stationary distribution of the SGD, i.e. w{∞}, and for conciseness we will simply denote it by w. We then consider a compressed network ŵ (that can be obtained by different compression schemes) and measure the performance of the compression scheme by its ‘relative `p-compression error’ (cf. [AUM11, GCD12]), defined as: ‖ŵ −w‖p/‖w‖p. Importantly for the following results, in the supplement, we further prove that a small compression error also implies a small perturbation on the network output.
Magnitude pruning. Magnitude pruning has been one of the most common and efficient algorithms among all the network pruning strategies [HPTD15, BOFG20, KLG+21]. In this section, we consider the global and layer-wise magnitude pruning strategies under the HML condition.
More precisely, given a network weight vector w ∈ Rd and a remaining parameter ratio κ ∈ (0, 1), the global pruning strategy compresses w by using w(κd), i.e., it prunes the smallest (in magnitude) (1− κ)d entries of w. Also, note that 1/κ corresponds to the frequently used metric compression rate [BOFG20]. On the other hand, the layer-wise pruning strategy applies the same approach to each layer separately, i.e., given layer-wise remaining parameter ratios κl ∈ (0, 1), we compress each layer weight wl ∈ Rdl by using w(κldl)l . The following result shows that the compression error of magnitude pruning can be made arbitrarily small as the network size grows.
Theorem 1. Assume that the recursion (2) satisfies the HML condition.
(i) Global magnitude pruning: if the weights of all layers have identical asymptotic distributions µ?l ≡ µ? with tail index α?l = α, for all l ∈ J1, LK, then for every > 0, ε > 0, κ ∈ (0, 1), and p ≥ α, there exists d0 ∈ N, such that ‖w(κd) −w‖p ≤ ε‖w‖p holds with probability at least 1− , for d ≥ d0. (ii) Layer-wise magnitude pruning: for every > 0, εl > 0 and κl ∈ (0, 1), where l ∈ J1, LK, and p ≥ maxl αl, there exists dl,0 ∈ N, such that ‖w(κldl)l − wl‖p ≤ εl‖wl‖p holds with probability at least 1− , for dl ≥ dl,0.
This result shows that any pruning ratio of the weights is achievable as long as the network size is sufficiently large and the network weights are close enough to an i.i.d. heavy-tailed distribution. Empirical studies report that global magnitude pruning often works better than layer-wise magnitude pruning [BOFG20], except when it leads to over-aggressive pruning of particular layers [WZG20].
The success of this strategy under the HML condition is due to a result from compressed sensing theory, concurrently proven in [GCD12, AUM11], which informally states that for a large vector of i.i.d. heavy-tailed random variables, the norm of the vector is mainly determined by a small fraction of its entries. We also illustrate this visually in the supplementary document.
An important question here is that, to achieve a fixed relative error εl, how would the smallest κl differ with varying tail-indices αl, i.e., whether “heavier tails imply more compression”. In our experiments, we illustrate this behavior positively: heavier-tailed weights are indeed more compressible. We partly justify this behavior for a certain range of p in the supplement; however, a more comprehensive theory is needed, which we leave as future work.
Singular value pruning. In recent studies, it has been illustrated that the magnitudes of the eigenvalues of the sample covariance matrices (for different layers) can decay quickly, hence pruning the singular values of the weight matrices, i.e., only keeping the largest singular values and corresponding singular vectors, is a sensible pruning method. Exploiting the low-rank nature of fully connected and convolutional layer weights in network compression has been investigated theoretically and empirically [AGNZ18, YLWT17]. Here we will present a simple scheme to demonstrate our results.
More precisely, for the weight matrix at layer l, Wl, consider its singular value decomposition, Wl = UΣV >, and then, with a slight abuse of notation, define W[κlhl−1]l := UΣ (κlhl−1)V>, where Σ(κlhl−1) is the diagonal matrix whose diagonal entries contain the dκlhl−1e largest singular values (i.e., prune the diagonal of Σ). Accordingly, denote w[κlhl−1]l := vec(W [κlhl−1] l ).
The next theorem shows that under the HML condition with an additional requirement that the limiting distributions are SαS , the eigenvalues of the (properly normalized) sample covariance matrices will be indeed compressible and the pruning strategy achieves negligible errors as the network size grows.
Theorem 2. Assume that the recursion (2) satisfies the HML condition, L ≥ 3, and for all l ∈ J2, L− 1K, µ?l ≡ SαlS(σl) with some σl > 0. Then, for every > 0, εl > 0, and κl ∈ (0, 1), there exists {hl,0,}Ll=1, such that the following inequalities hold for every hl ≥ hl,0 and p ≥ maxl αl/2:
‖λ(κlhl−1)l − λl‖p ≤ ε2l ‖λl‖p, and ‖w [κlhl−1] l −wl‖ ≤ εl‖wl‖, (5)
with probability at least 1− , where λl ∈ Rhl−1 denotes the vector of eigenvalues corresponding to the sample covariance matrix h−2/αll W > l Wl.
We shall note that the proof of Theorem 2 in fact only requires the limiting distributions to be ‘regularly varying’ [TTR+20] and symmetric around zero, which covers a broad range of heavy-tailed distributions beyond the α-stable family [BDM+16]. The sole reason why we require the SαS condition here is to avoid introducing further technical notation.
In Theorem 2, p needs to be greater than or equal to α/2, in contrast to the condition p ≥ α in Theorem 1. The reason is that Theorem 2 is based on pruning the eigenvalues of the normalized covariance matrix and moreover in general if a matrix A ∈ Rn×m has elements [A]i,j that are independent and identically distributed from the symmetric α-stable distribution, then by [TTR+20, Theorem 2.7], as m→∞, the eigenvalues of m−2/αATA weakly converge to independent random variables, that are identically distributed from a positive stable distribution with tail index α/2.
Node pruning. The last pruning strategy that we consider is a structured pruning strategy, that corresponds to the removal of the whole columns of a fully connected layer weight matrix. Even though below we consider pruning based on the norms of the weight layer columns, the same arguments apply for pruning rows; see supplement for further discussion.
The idea in column pruning is that, for a given layer l, we first sort the columns of the weight matrix Wl ∈ Rhl×hl−1 with respect to their `p-norms for a given p ≥ maxl αl. Then, we remove the entire columns that have the smallest `p-norms. More precisely, let Wl(i) ∈ Rhl be the i-th column of Wl, for i ∈ J1, hl−1K, and suppose that ‖Wl(i1)‖p ≥ ‖Wl(i2)‖p ≥ · · · ≥ ‖Wl(ihl−1)‖p, where {i1, . . . , ihl−1} = {1, . . . , hl−1}. Then, we define the k-best column approximation of Wl, denoted as W{k},pl ∈ Rhl×hl−1 , as follows: for j ∈ J1, dkeK, W {k},p l (ij) := Wl(ij) and for j /∈ J1, dkeK, W {k},p l (ij) := 0. Denote also w {k},p l := vec ( W {k},p l ) .
Theorem 3. Assume that the recursion (2) satisfies the HML condition. Then, for every > 0, εl > 0, κl ∈ (0, 1), where l ∈ J2, LK, and p ≥ maxl αl, there exists hl−1,0 ∈ N, such that ‖w{κlhl−1},pl −wl‖p ≤ εl‖wl‖p holds with probability at least 1− , for every hl−1 ≥ hl−1,0.
This theorem indicates that we can remove entire columns in each layer, without considerably affecting the network weights, as long as the network is large enough. In other words, effectively the widths of the layers can be reduced. Structured pruning schemes are commonly used in CNNs where filters, channels, or kernels can be pruned by norm-based or other criteria [LKD+17, HZS17, HPTD15].

heading: 3.3 A note on the limiting independence structure in the HML condition
text: We conclude this section by discussing the particular independence condition in the limit, which appears in Condition 1. We shall underline that the element-wise independence is not a necessity and under weaker conditions, we can still obtain Theorems 1-3 with identical or almost identical proofs. For instance, the proof of Theorem 3 remains the same when the columns of Wl are i.i.d. vectors with dependent components; hence, the element-wise independence is indeed not needed, but is used for the clarity of presentation. More generally, in all three theorems, the main requirement is to ensure a weak dependence between the components of the weight vector. More precisely, for Theorems 1, 2, 3, we respectively need (1) the entries of Wl or (2) its singular values, or (3) the `p-norms of its columns to be stationary and ergodic with a heavy-tailed distribution. Under this condition, the same proof strategies will still work by invoking [SD15, Theorem 1], instead of [GCD12, Proposition 1].

heading: 4 Generalization Bounds
text: So far, we have shown that the heavy-tailed behavior in the weights of a network together with their independence result in compressibility. In this section, we further show that these phenomena bring forth a better generalization performance bound in the network. More precisely, we establish a generalization bound such that if a network is more compressible, then the corresponding compressed network has a smaller generalization error bound. Throughout the section, we will focus on layer-wise magnitude pruning; yet, our results can be easily extended to the other pruning strategies. Note that similar results have already been proven in [AGNZ18, SAM+20, SAN20, HJTW21]; yet, they cannot be directly used in our specific setting, hence, we prove results that are customized to our setup. More precisely, although these works, similar to our result, are based on the assumption of the “compressibility” of the network and moreover the generalization gap of the population risk with respect to the empirical margin risk is considered in [AGNZ18, HJTW21], they cannot be applied directly to the `p-compressibility based strategies discussed in the previous section. We should emphasize that our results, as well as the previous mentioned works, are merely upper bounds on the generalization performance of the network.
Our generalization bounds are derived by applying the previously developed techniques as in [AGNZ18, NBS18]. In particular, we follow the approach of [AGNZ18] by further adapting the technique for the magnitude pruning strategy and allowing the compressed network weights to take unbounded continuous values. As in [AGNZ18], we consider the 0-1 loss function with margin γ ≥ 0, `γ : Y × Y 7→ {0, 1}, for the multiclass classifier fw, given as follows:
`γ(y, fw(x)) =
{ 1, if fw(x)[y]−max
j 6=y fw(x)[j] ≤ γ,
0, otherwise. (6)
Still denoting w = w{∞}, the population and empirical risks associated with `γ are denoted as Rγ(w) and R̂γ(w), respectively. By having the dataset S ∼ µ⊗nZ , we assume the weights w are sampled according to the stationary distribution Pw|S of SGD. Denote the joint distribution of (S,w) by PS,w := µ⊗nZ Pw|S .
In the following theoretical results, we will assume that we have access to a random compressible neural network that is amenable to the layer-wise magnitude pruning strategy with high probability. This assumption is essentially the outcome of Theorem 1 under the HML condition together with an additional uniformity requirement on dl,0 over S1.
H 1. For ε ≥ 0, > 0, and {κl}Ll=1 : κl ∈ (0, 1), there exists {dl,0}Ll=1 : dl,0 ∈ N independent of S, such that for dl ≥ dl,0, l ∈ J1, LK, the relation ‖ŵl − wl‖ ≤ ε‖wl‖ holds for all l ∈ J1, LK simultaneously with probability at least 1− , where ŵl := w(κldl)l and the probability is w.r.t. PS,w.
In the following result, we relate the population risk of the compressed network to the empirical margin risk of the original (uncompressed) network. For notational convenience, for δ, τ > 0, let
R(δ) := inf{R : P(‖w‖ ≥ R) ≤ δ}, and L(τ, δ) := √ 2BL ( 2R(δ)/ √ L )L−1 /τ, (7)
where the probability is with respect to the joint distribution PS,w.
Theorem 4. Assume H1 holds. Then for n : n/ log(n) ≥ 10L, {dl}Ll=1 : dl ≥ dl,0, and any δ, τ > 0, with probability at least 1− 2e−κd/2 − δ − ,
R0(ŵ) ≤ R̂γ(δ,τ)(w) + ( 12L(τ, δ)R(δ) + √ d )√ (κ+ κ) log(n)/n, (8)
where R(δ) and L(τ, δ) are defined in (7), κ := 1d ∑L l=1dκldle,2
κ := (2hb(κ)− κ log(κ))/ log(n), and γ(δ, τ) := τ + √ 2B
τ
( R(δ)/ √ L )L( (1 + ε)L − 1 ) .
1Note that the uniformity assumption is mild and can be avoided by combining Theorem 1 with Egoroff’s theorem with additional effort, as in [ŞSDE20].
2The binary entropy hb(κ) (in nats) is defined as −κ log(κ)− (1− κ) log(1− κ) for κ ∈ [0, 1], with the convention hb(0) = hb(1) = 0. Note that 0 ≤ hb(κ) ≤ log(2).
This result establishes a generalization bound such that, for a fixed relative compression error ε, if a network is more compressible, i.e. if κ is smaller, then its corresponding compressed network has a smaller generalization error bound, as the bound is scaled by the factor ≈√κ. In our proof, we prove an intermediate result, inspired by [NTS15] and stated in the supplement, which bounds the perturbation of the network output after compression. This guarantees that the risks of the original and pruned networks can be made arbitrarily close, as long as the relative compression error is small enough.
To have a better understanding of the constants in Theorem 4, in the next result, we consider a special case where the weights follow a stable distribution, and we make the above bound more explicit.
Corollary 1. Assume that for l ∈ J1, LK and i ∈ J1, dlK, the conditional distribution of wl,i i.i.d.∼ SαlS(σl) with αl ∈ (1, 2). Further, assume that the scale parameters satisfy the following property:
σ2 := ∑L
l=1 (dl/d)(σl/σαl)
2 = [ (4−1/α √ L/3)σ0d −(1/2+1/α) ]2 , (9)
where σαl := (2Γ(−αl) cos((2− αl)π/2))1/αl , α:= minl αl, and σ0 is a constant, and also {αl}Ll=1 and σ are independent from S. Then, for every ε > 0 and κl ∈ (0, 1), l ∈ J1, LK, there exists dl,0 ∈ N, such that for dl ≥ dl,0, n : n/ log(n) ≥ 10L, and every τ > 0, with probability at least 1−3d−α/(2L),
R0(ŵ) ≤ R̂γ(w) + ( aσL0 /τ + 1 )√ (κ+ κ)d log(n)/n, (10)
where {ŵl}l = {w(κldl)l }l, γ := τ + bεσL0 √ d/τ , a := 6 √ 2B2LL3/2, and bε :=√
2B ( (1 + ε)L − 1 ) .
To simplify our presentation, we have set the ‘scale’ of the distribution as a decreasing function of dimension, which intuitively states that the typical magnitude of each entry of the network will get smaller as the network grows. We observe that for a fixed ε and d, the bound improves as pruning ratio, 1− κ, increases. This result is of interest in particular since it is observed experimentally (and in part, theoretically) that heavier-tailed weights are more compressible, and hence due to this result have better generalization bounds. This provides an alternative perspective to the recent bounds that aim at linking heavy tails to generalization through a geometric approach [ŞSDE20].
Finally, in the supplement, we further show that the uncompressed network also inherits this good generalization performance bound, which is consistent with the results of [HJTW21, KLG+21]: if a network is more “compressible”, not only the generalization performance for the compressed network but also for the original network improves. The generalization bound adapted to `p- compressibility based strategies, discussed in the previous section, is highlighted in Section S2 of the supplement.

heading: 5 Experiments
text: In this section, we present experiments conducted with neural networks to investigate our theory. We use three different model architectures: a fully connected network with 4 hidden layers (FCN4), a fully connected network with 6 hidden layers (FCN6), and a convolutional neural network with 8 convolutional layers (CNN). Hidden layer widths were 2048 for both FCN models. All networks include ReLU activation functions and none include batch normalization, dropout, residual blocks, or any explicit regularization term in the loss function. Each model is trained on MNIST [LCB10] and CIFAR10 [Kri09] datasets under various hyperparameter settings, using the default splits for training and evaluation. The total number of parameters were approximately 14M for FCN4-MNIST, 19M for FCN4-CIFAR10, 23M for FCN6-MNIST, 27M for FCN6-CIFAR10, and 9M for both CNN models. All models were trained with SGD until convergence with constant learning rates and no momentum. The convergence criteria comprised 100% training accuracy and a training negative log-likelihood
less than 5× 10−5. The training hyperparameter settings include two batch-sizes (b = 50, 100) and various learning rates (η) to generate a large range of η/b values. See the supplement for more details.
By invoking [GŞZ21, Corollary 11] which shows the ergodic average of heavy-tailed SGD iterates converges to a multivariate stable distribution, after convergence, we kept running SGD for 1000 additional iterations to obtain the average of the parameters to be used in this estimation. The tail index estimations were made by the estimator proposed in [MMO15, Corollary 2.4], which has been used by other recent studies that estimate tail index in neural network parameters [ŞSDE20, GŞZ21, ZFM+20]. The ergodic averaging does not change the tail index of the parameters, and was employed to facilitate tail index estimation by enabling the utilization of the aforementioned estimator. The rest of the experiments have been conducted without ergodic averaging. We also observe that the results with/without ergodic averaging are virtually identical in both tail index estimation tasks and pruning experiments. In all pruning methods, the parameters were centered before pruning is conducted with the median value to conform with the tail index estimation.
Training hyperparameters and layer statistics. We first verify that models trained with higher learning rate to batch-size ratio (η/b) lead to heavier-tailed parameters, replicating the results presented in [GŞZ21]. For each trained model, we compute the mean of separately estimated tail indices for all layers, so that each model is represented by a single value. This averaging is a heuristic without a clear theoretical meaning, and has been also used by recent works in the literature [GŞZ21, MM19]. Results
presented in Figure 1 demonstrates that higher η/b leads to heavier tails (lower α̂).
Also of interest are the distribution of the resulting parameters from training. Figure 2 (left) demonstrates a representative example of the parameters from an CNN layer trained on MNIST, where two overlaid histograms representing empirical distributions of a random partition of the parameters are almost identical. On the right in the same figure, this distribution is compared to a simulated symmetric α-stable distribution that has the same tail index (≈ 1.95). The figure demonstrates that the two distributions have similar qualitative properties.
Tail index and prunability. In this section we examine whether networks with heavier-tailed layers, trained with higher η/b ratios, are more prunable and whether they generalize better. As a baseline test, we first examine whether neural network layers which are heavier-tailed can be pruned more given a fixed maximum relative compression error. Figure 3 demonstrates for an error of 0.1 that this is indeed the case. We next test our hypothesis that posit models with heavy-tailed parameters to be more prunable. Both the results pertaining to layer-wise magnitude pruning and singular value pruning, demonstrated in Figures 4 and 5 show that this is indeed the case. Here, relative test accuracy
stands for test accuracy after pruning / unpruned test accuracy. The results show that models with heavier-tailed parameters (shown with darker colors) are starkly more robust against pruning. Similar results with global magnitude pruning can be seen in the supplement.
For structured pruning, we prune 3×3 kernel parameters in CNN models. The results (Figure 6) show a similar, hypothesis conforming pattern. Results for FCNs, presented in the supplement, were underwhelming; perhaps unsurprisingly as structured pruning is not as commonly used in FCNs [O’N20]. More successful attempts could be due to alternative scoring methods [SAM+20]; our approach might require wider layers to conform with theoretical conditions.
Tail index and generalization. Following our theoretical results, we examine whether the heavier-tailed networks lead to better generalization performance. Consistent with our hypothesis, Figure 7 shows that models with the highest tail index have consistently the worst test performances. The same conclusion applies to generalization performance as training accuracy is 100% for all models. See supplementary material for additional experiment results and discussion of the findings.

heading: 6 Conclusion
text: We investigated the conditions under and the mechanism through which various pruning strategies can be expected to work, and confirmed the relationship between pruning and generalization from our theoretical approach. Future directions for research include formally identifying settings in which HML condition holds, eliciting the relationship between α̂ and κmore clearly, and further examination of structured pruning in FCNs. Extending our work to other pruning strategies and network structures is also another important future direction. For example, it has recently been observed that gradients also exhibit heavy-tailed behavior [ŞGN+19, ZKV+20, ZFM+20]; we suspect that our theory might be applicable in the case of gradient-based pruning as well. The extension of our analyses to other network structures such as recurrent layers or ResNet blocks [HZRS15] would also be valuable.

heading: Acknowledgements
text: The authors are grateful to A. Taylan Cemgil, Valentin De Bortoli, and Mohsen Rezapour for fruitful discussions on various aspects of the paper. MAE is partially funded by CIFAR AI Chairs program, and CIFAR AI Catalyst grant, NSERC Grant [2019-06167]. UŞ’s research is supported by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute).

Review summaries:
This paper gives a theoretical characterization of the underlying cause that makes the neural networks amenable to common compression schemes in the over-parameterized regime. It is done by connecting the dynamics of the learning algorithms such as SGD to the heavy tails of the weights of a network, which leads to better compressibility. The main contribution is to provide sufficient conditions for the compressibility from the limiting distribution of SGD, and show how the compressibility of a network can lead to better generalization.  

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.

The main contribution of this paper is to link the two properties of compressibility and SGD. (i) as the network size goes to infinity, the system can converge to a mean-field limit, where the network weights behave independently, (ii) for a large step-size/batch-size ratio, the SGD iterates can converge to a heavy-tailed stationary distribution. In the case where these two phenomena occur simultaneously, we prove that the networks are guaranteed to be ‘lp-compressible’. Further experiments on fully connected and convolutional neural networks and show that the results are in strong accordance with the theory.

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.