Title: Approximate Heavy Tails in Offline (Multi-Pass) Stochastic Gradient Descent
Venue: NeurIPS 2023
Accepted: yes
Keywords: SGD

heavy-tails

wasserstein convergence
Abstract: A recent line of empirical studies has demonstrated that SGD might exhibit a heavy-tailed behavior in practical settings, and the heaviness of the tails might correlate with the overall performance. In this paper, we investigate the emergence of such heavy tails. Previous works on this problem only considered, up to our knowledge, online (also called single-pass) SGD, in which the emergence of heavy tails in theoretical findings is contingent upon access to an infinite amount of data. Hence, the underlying mechanism generating the reported heavy-tailed behavior in practical settings, where the amount of training data is finite, is still not well-understood. Our contribution aims to fill this gap. In particular, we show that the stationary distribution of offline (also called multi-pass) SGD exhibits ‘approximate’ power-law tails and the approximation error is controlled by how fast the empirical distribution of the training data converges to the true underlying data distribution in the Wasserstein metric. Our main takeaway is that, as the number of data points increases, offline SGD will behave increasingly ‘power-law-like’. To achieve this result, we first prove nonasymptotic Wasserstein convergence bounds for offline SGD to online SGD as the number of data points increases, which can be interesting on their own. Finally, we illustrate our theory on various experiments conducted on synthetic data and neural networks.

Approximate Heavy Tails in Offline (Multi-Pass) Stochastic Gradient Descent

Abstract

A recent line of empirical studies has demonstrated that SGD might exhibit a heavy-tailed behavior in practical settings, and the heaviness of the tails might correlate with the overall performance. In this paper, we investigate the emergence of such heavy tails. Previous works on this problem only considered, up to our knowledge, online (also called single-pass) SGD, in which the emergence of heavy tails in theoretical findings is contingent upon access to an infinite amount of data. Hence, the underlying mechanism generating the reported heavy-tailed behavior in practical settings, where the amount of training data is finite, is still not wellunderstood. Our contribution aims to fill this gap. In particular, we show that the stationary distribution of offline (also called multi-pass) SGD exhibits 'approximate' power-law tails and the approximation error is controlled by how fast the empirical distribution of the training data converges to the true underlying data distribution in the Wasserstein metric. Our main takeaway is that, as the number of data points increases, offline SGD will behave increasingly 'power-law-like'. To achieve this result, we first prove nonasymptotic Wasserstein convergence bounds for offline SGDto online SGD as the number of data points increases, which can be interesting on their own. Finally, we illustrate our theory on various experiments conducted on synthetic data and neural networks.

1 Introduction

Many machine learning problems can be cast as the following population risk minimization problem:

where x ∈ R d denotes the model parameters, µ z is the data distribution over the measurable space ( Z , Z ) and f : R d × Z → R is a loss function. The main difficulty in addressing this problem is that µ z is typically unknown. Suppose we have access to an infinite sequence of independent and

identically distributed (i.i.d.) data samples D := { Z 1 , Z 2 , . . . } from µ z . In that case, we can resort to the online stochastic gradient descent (SGD) algorithm, which is based on the following recursion:

where k denotes the iterations, Ω k := { b ( k -1) + 1 , b ( k -1) + 2 , . . . , bk } is the batch of data-points at iteration k , η denotes the step size, b denotes the batch size such that | Ω k | = b , and Z i with i ∈ Ω k , denotes the i -th data sample at the k -th iteration. This algorithm is also called 'single-pass' SGD, as it sees each data point Z i in the infinite data sequence D only once.

While drawing i.i.d. data samples at each iteration is possible in certain applications, in the majority of practical settings, we only have access to a finite number of data points, preventing online SGD use. More precisely, we have access to a dataset of n i.i.d. points D n := { Z 1 , . . . , Z n } 1 , and given these points the goal is then to minimize the empirical risk ˆ F ( n ) , given as follows:

To attack this problem, one of the most popular approaches is the offline version of (2), which is based on the following recursion:

where Ω ( n ) k ⊂ { 1 , . . . , n } denotes the indices of the (uniformly) randomly chosen data points at iteration k with | Ω ( n ) k | = b ≤ n and · ( n ) emphasizes the dependence on the sample size n . Analogous to the single-pass regime, this approach is also called the multi-pass SGD, as it requires observing the same data points multiple times.

Despite its ubiquitous use in modern machine learning applications, the theoretical properties of offline SGD have not yet been well-established. Among a plethora of analyses addressing this question, one promising approach has been based on the observation that parameters learned by SGD can exhibit heavy tails . In particular, [¸ SGN + 19] have empirically demonstrated a heavy-tailed behavior for the sequence of stochastic gradient noise:

Since then, further studies have extended this observation to other sequences appearing in machine learning algorithms [ZFM + 20, MM19, ZLMU22]. These results were recently extended by [BSE + 21], who showed that the parameter sequence ( X ( n ) k ) k ≥ 1 itself can also exhibit heavy tails for large η and small b . These empirical investigations all hinted a connection between the observed heavy tails and the generalization performance of SGD: heavier the tails might indicate a better generalization performance.

Motivated by these empirical findings, several subsequent papers theoretically investigated how heavytailed behavior can emerge in stochastic optimization. In this context, [GSZ21] have shown that when online SGD is used with a quadratic loss (i.e., f ( x, z = ( a, y )) = 2 -1 ( a ⊤ x -y ) 2 ), the distribution of the iterates ( X k ) k ≥ 0 can converge to a heavy-tailed distribution, even with exponentially light-tailed data (e.g., Gaussian). More precisely, for a fixed step-size η > 0 , they showed that, under appropriate assumptions, there exist constants c ∈ R + and α > 0 , such that the following identity holds:

where X ∞ ∼ π and π denotes the stationary distribution of online SGD (2). This result illustrates that the distribution of the online SGD iterates follows a power-law decay with the tail index α : P ( ∥ X ∞ ∥ > t ) ≈ t -α for large t . Furthermore, α depends monotonically on the algorithm

0

1400

0

1400

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0

4600

Linear regression parameter norms ||

w

||

Frequency

n: 20

n: 250

n: 1000

+2*

0

2

e

4

0

2

e

4

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

0

2

e

4

Neural network parameter norms

||

w

||

Frequency

n = 1000

n = 2000

n = 5000

+2*

hyperparameters η , b , and certain properties of the data distribution µ z . This result is also consistent with [DMN + 21, Example 1], which shows a simple instance from linear stochastic approximation such that for any fixed step-size η , there exists p c > 0 such that for any p ≥ p c , lim k → + ∞ E [ ∥ X k ∥ p ] = + ∞ . In a concurrent study, [HM21] showed that this property is, in fact, not specific for quadratic problems and can hold for more general loss functions and different choices of stochastic optimizers, with the constraint that i.i.d. data samples are available at every iteration 2 .

Another line of research has investigated the theoretical links between heavy tails and the generalization performance of SGD. Under different theoretical settings and based on different assumptions, [SSDE20, BSE + 21, LW¸ S22, RBG + 22, HSKM22, RZG¸ S23] proved generalization bounds (i.e., bounds on | ˆ F ( n ) ( x ) -F ( x ) | ) illustrating that the heavy tails can be indeed beneficial for better performance. However, all these results rely on exact heavy tails, which, to our current knowledge, can only occur in the online SGD regime where there is access to an infinite sequence of data points. Hence the current theory (both on the emergence of heavy tails and their links to generalization performance) still falls short in terms of explaining the empirical heavy-tailed behavior observed in offline SGD as reported in [¸ SGN + 19, MM19, ZFM + 20, BSE + 21].

Main problematic and contributions. Although empirically observed, it is currently unknown how heavy-tailed behavior arises in offline SGD since the aforementioned theory [GSZ21, HM21] requires infinite data. Our main goal in this paper is hence to develop a theoretical framework to catch how and in what form heavy tails may arise in offline SGD, where the dataset is finite.

We build our theory based on two main observations. The first observation is that, since we have finitely many data points in offline SGD, the moments (of any order) of the iterates X ( n ) k can be bounded (see, e.g., [CGZ19]), hence in this case we cannot expect exact power-law tails in the form of (5) in general. Our second observation, on the other hand, is that since the finite dataset D n can be seen as the first n elements of the infinite sequence D , as n increases the offline SGD recursion (4) should converge to the online SGD recursion (2) in some sense. Hence, when online SGD shows a heavy-tailed behavior (i.e., the case where n → + ∞ ), as we increase n , we can expect that the tail behavior of offline SGD should be more and more 'power-law-like'.

To further illustrate this observation, as a preliminary exploration, we run offline SGD in a 100-dimensional linear regression problem, as well as a classification problem on the MNIST dataset, using a fully-connected, 3-layer neural network. 3

In Figure 1, we plot the histograms of the parameter norms after 1000 (and 5000

10

20

50

100

500

n

0.5

1.0

1.5

2.0

(

n

)

= 0.01

10

20

50

100

500

n

0.5

1.0

1.5

2.0

= 0.025

10

20

50

100

500

n

0.5

1.0

1.5

2.0

= 0.04

Online

Figure 2: Estimated tail indices

resp.) iterations. We observe the following consistent patterns: (i) the means and standard deviations of the parameters' estimated stabilizing distributions increase with n , and (ii) , the quantity and magnitude of the parameters far from the bulk of the distribution increase with n . In other words, as we increase the number of samples n in offline SGD, the behavior of the iterates becomes heavier-tailed.

As another example, we utilize the tail-estimator from [GSZ21] 4 to compare the estimated tail indices α for offline and online SGD. The results are plotted in Figure 2 and correspond to 10 random initializations with different step-sizes (for further details, see Tail estimation). We observe that, as n increases, the estimated tail indices for offline SGD (shown in green) get closer to the true tail index corresponding to online SGD (the horizontal blue line).

Our main contribution in this paper is to make these observations rigorous and explicit. In particular, we extend the tail estimation results for online SGD [GSZ21, HM21] to offline SGD, and show that the stationary distribution of offline SGD iterates will exhibit 'approximate' heavy tails. Informally, for both quadratic and a class of strongly convex losses, we show that, with high probability, there exist constants c 1 , c 2 , such that for large enough t we have the following tail bound:

where α is the tail index of online SGD as given in (5), X ( n ) ∞ ∼ π ( n ) denotes a sample from the stationary distribution π ( n ) of offline SGD (4), µ ( n ) z = n -1 ∑ n i =1 δ Z i denotes the empirical measure of the data points in the data sample D n = { Z 1 , . . . , Z n } , and W 1 denotes the Wasserstein-1 distance (to be defined formally in the next section).

Our result indicates that the tail behavior of offline SGD is mainly driven by two terms: (i) a persistent term ( c/t α ) that determines the power-law decay, and (ii) a vanishing term W 1 ( µ z , µ ( n ) z ) /t , as we take n →∞ . In other words, compared to (5), we can see that, with high-probability, X ( n ) ∞ exhibits an approximate power-law decay behavior, with a discrepancy controlled by W 1 ( µ z , µ ( n ) z ) .

Fortunately, the term W 1 ( µ z , µ ( n ) z ) has been well-studied in the literature. By combining our results with Wasserstein convergence theorems for empirical measures [FG15], we further present nonasymptotic tail bounds for offline SGD. Our main takeaway is as follows: as n increases, offline SGD will exhibit approximate power-laws, where the approximation error vanishes with a rate given by how fast the empirical distribution of the data points µ ( n ) z converges to the true data distribution µ z as n →∞ .

To prove our tail bounds, as an intermediate step, we prove nonasymptotic Wasserstein convergence results for offline SGD in the form: W 1 ( π, π ( n ) ) ≲ W 1 ( µ z , µ ( n ) z ) , with high probability, recalling that π and π ( n ) respectively denote online and offline SGD stationary distributions. These results, we believe, hold independent interest for the community. Finally, we support our theory with various experiments conducted on both synthetic quadratic optimization problems and real neural network settings.

2 Preliminaries and Technical Background

2.1 Notation and preliminary definitions

We let I d denote the d × d identity matrix. We denote [ n ] = { 1 , . . . , n } , P n,b = { I ⊂ [ n ] , | I | = b } . For a vector x ∈ R d , ∥ x ∥ denotes the Euclidean norm of x , and for matrix A , the norm ∥ A ∥ = sup ∥ x ∥ =1 ∥ A x ∥ denote its spectral norm. We denote with σ min ( A ) , σ max ( A ) the smallest and largest singular values of A , respectively. P ( R d ) is the set of all probability measures on R d , and L ( X ) denotes the probability law of a random variable X . Wedenote the set of all couplings of two measures µ and ν with Γ( µ, ν ) . Finally, for two functions f ( n ) and g ( n ) defined on R + , we denote f = O ( g ) , if there exist constants c ∈ R + , n 0 ∈ N , such that f ( n ) ≤ cg ( n ) , ∀ n > n 0 . The Dirac measure concentrated at x is defined as follows: for a measurable set E , δ x ( E ) = 1 if x ∈ E , δ x ( E ) = 0 , otherwise.

For p ∈ [1 , ∞ ) , the Wassersteinp distance between two distributions µ and ν on R d is defined as:

We correspondingly define the Wassersteinp distance between two distributions µ and ν on R d × d as

.

2.2 Heavy tails in online SGD

We first formally define heavy-tailed distributions with a power-law decay, used throughout our work. Definition 1 (Heavy-tailed distributions with a power-law decay) . A random vector X is said to follow a distribution with a power-law decay if lim t →∞ t α P ( ∥ X ∥ ≥ t ) = c 0 , for constants c 0 > 0 and α > 0 , the latter known as the tail index.

The tail index α determines the thickness of the tails: larger values of α imply a lighter tail.

Recently, [GSZ21] showed that online SGD might produce heavy tails even with exponentially lighttailed data. More precisely, they focused on a quadratic optimization problem, where z ≡ ( a, q ) ∈ R d +1 and f ( x, z ) = 2 -1 ( a ⊤ x -q ) 2 . With this choice, the online SGD recursion reads as follows:

where ( Z k ) k ≥ 1 ≡ ( a k , q k ) k ≥ 1 , A k := b -1 ∑ i ∈ Ω k a i a ⊤ i , and b k := b -1 ∑ i ∈ Ω k a i q i . In this setting, they proved the following theorem.

Theorem 1 ([GSZ21]) . Assume that ( a k , q k ) k ≥ 1 are i.i.d. random variables such that a 1 ∼ N(0 , σ 2 I d ) , with σ 2 > 0 and q i has a continuous density with respect to the Lebesgue measure on R with all its moments being finite. In addition suppose that E [log ∥ I d -η A 1 ∥ ] < 0 . Then, there exists a unique α > 0 such that E [ ∥ I d -η A 1 ∥ α ] = 1 and the iterates ( X k ) k ≥ 0 in (7) converge in distribution to a unique stationary distribution π such that if X ∞ ∼ π and α / ∈ N :

The authors of [GSZ21] also showed that the tail index α is monotonic with respect to the algorithm hyperparameters η and b : a larger η or smaller b indicates heavier tails, i.e., smaller α 5 .

In a concurrent study, [HM21] considered a more general class of loss functions and choices of stochastic optimization algorithms. They proved a general result, where we translate it for online SGD applied on strongly convex losses in the following theorem.

Theorem 2 ([HM21]) . Assume that for every z ∈ Z , f ( · , z ) is twice differentiable and is strongly convex. Consider the recursion in (2) with a sequence of i.i.d. random variables ( Z k ) k ≥ 1 and let R and r be two non-negative functions defined as: for any z ∈ Z ,

R ( z ) := sup x ∈ R d ∥ ∥ I d -η ∇ 2 f ( x, z ) ∥ ∥ , and r ( z ) := lim inf ∥ x ∥→∞ σ min ( I d -η ∇ 2 f ( x, z ) ) . Further assume that E [ R ( Z 1 ) + ∥∇ f ( x ⋆ , Z 1 ) ∥ ] < + ∞ , for some x ⋆ ∈ R d , E [log R ( Z 1 )] < 0 , and P ( r ( Z 1 ) > 1) > 0 . Then, the iterates ( X k ) k ≥ 1 in (2) admit a stationary distribution π . Moreover, there exist α, β > 0 such that E [ r ( Z ) α ] = 1 , E [ R ( Z ) β ] = 1 , and for any ε > 0 :

The result illustrates that power laws can arise in general convex stochastic optimization algorithms. We note that in general, Theorem 2 does not require each f ( · , z ) to be strongly convex, in fact, it can accommodate non-convex functions as well. However, we are not aware of any popular non-convex machine learning problem that can satisfy all the assumptions of Theorem 2. Nevertheless, for better illustration, in Section F, we show that an ℓ 2 -regularized logistic regression problem with random regularization coefficients falls into the scope of Theorem 2.

The two aforementioned theorems are limited to the online setting, and it is unclear whether they apply to a finite number of data points n . In the subsequent section, we aim to extend these theorems to the offline setting.

3 Main results

3.1 Quadratic objectives

We first focus on the quadratic optimization setting considered in Theorem 1 as its proof is simpler and more instructive. Similarly as before, let ( a i , q i ) i ≥ 1 be i.i.d. random variables in R d +1 , such that z i ≡ ( a i , q i ) . Following the notation from (4), we can correspondingly define the offline SGD recursion as follows:

where A ( n ) k := b -1 ∑ i ∈ Ω ( n ) k a i a ⊤ i , b ( n ) k := b -1 ∑ i ∈ Ω ( n ) k a i q i , and Ω ( n ) k is as defined in (4). Now, ( A ( n ) k , b ( n ) k ) k ≥ 1 are i.i.d. random variables with respect to the empirical measure:

where we enumerate all possible choices of minibatch indices P n,b = { I ⊂ { 1 , . . . , n } : | I | = b } as P n,b = { S ( n ) 1 , S ( n ) 2 , . . . , S ( n ) ( n b ) } and ( A ( n ) i , b ( n ) i } ( n b ) i =1 are i.i.d. random variables defined as follows:

We denote by µ A , b the common distribution of ( A ( n ) i , b ( n ) i ) ( n b ) i =1 . Note that it does not depend on n but only b .

With these definitions at hand, we define the two marginal measures of µ A , b and µ ( n ) A , b with µ A , µ b , and µ ( n ) A , µ ( n ) b , respectively. Before proceeding to the theorem, we have to define Wasserstein ergodicity. A discrete-time Markov chain ( Y k ) k ≥ 0 is said to be (Wasserstein-1) geometrically ergodic if it admits a stationary distribution π Y and there exist constants c ≥ 0 and ρ ∈ (0 , 1) , such that

The chain is simply called (Wasserstein-1) ergodic if lim k →∞ W 1 ( L ( Y k ) , π Y ) = 0 .

We can now state our first main contribution, an extension of Theorem 1 to the offline setting:

Theorem 3. Let n ≥ 1 and ϵ n ∈ [0 , 1] be the probability that ( X ( n ) k ) k ≥ 0 is not ergodic (in the Wasserstein sense) 6 with a stationary distribution π ( n ) having finite q -th moment with q > 1 . Assume that the conditions of Theorem 1 hold with α > 1 . Then, for any ϵ > 0 , there exist constants ˜ c 1 , ˜ c 2 , and t 0 > 0 such that for all ζ ∈ (0 , 1] , t > t 0 , with probability larger than 1 -ϵ n -ζ , the following inequalities hold:

where c is given in (8) .

This result shows that whenever the online SGD recursion admits heavy tails, the offline SGD recursion will exhibit 'approximate' heavy tails. More precisely, as can be seen in (13), the tails will have a global power-law behavior due to the term t -α , where α is the tail index determined by online SGD. On the other hand, the power-law behavior is only approximate, as we have an additional term, which vanishes at a rate O ( n -1 / 2 ) . Hence, as n increases, the power-law behavior will be more prominent, which might bring an explanation for why heavy tails are still observed even when n is finite.

To establish Theorem 3, as an intermediate step, we show that π ( n ) converges to π in the Wasserstein-1 metric. The result is given in the following theorem and can be interesting on its own.

Theorem 4. Under the setting of Theorem 3, for p, q > 0 with 1 /p +1 /q = 1 , p > d/ 2 , q < α , the following holds with probability larger than 1 -ϵ n :

where c 0 = ( E [ ∥ X ( n ) 0 ∥ q ]) 1 /q +1 with X ( n ) 0 ∼ π ( n ) , and δ A = E [ ∥ I d -η A 1 ∥ ] .

This result shows that the stationary distribution of offline SGD will converge to the one of online SGD as n → ∞ , and the rate of this convergence is determined by how fast the empirical distribution of the dataset converges to the true data distribution, as measured by W p ( µ A , b , µ ( n ) A , b ) . Once (14) is established, to obtain Theorem 3, we use the relationship that W 1 ( π, π ( n ) ) ≥ W 1 ( L ( ∥ X ( n ) ∞ ∥ ) , L ( ∥ X ∞ ∥ )) = ∫ ∞ 0 | P ( ∥ X ( n ) ∞ ∥ > t ) -P ( ∥ X ∞ ∥ > t ) | d t (see Lemmas 1 and 3 in the Appendix). Finally, by exploiting the assumptions made on µ A , b in Theorem 1, we can invoke [FG15] (see Lemma 2 in Section A.1) and show that W p ( µ A , b , µ ( n ) A , b ) behaves as O ( n -1 / 2 ) .

3.2 Strongly convex objectives

We now focus on more complex loss functions, as in Theorem 2. For simplicity, we consider b = 1 . The general case, while similar, would complicate the notation and, in our opinion, make our results more difficult to digest. We start by rewriting offline SGD with respect to an empirical measure µ ( n ) z :

where ( ˆ Z ( n ) k ) k ≥ 1 are i.i.d. random variables associated with the empirical measure defined for any A ∈ Z as:

We first make the following assumption.

Assumption 1. (a) There exists L > 0 such that ∥∇ f ( x, z ) -∇ f ( x, z ′ ) ∥ ≤ L ( ∥ x ∥ +1)( ∥ z -z ′ ∥ ) , for any z, z ′ ∈ Z and x ∈ R d .

(b) The data distribution µ z has finite q -th moments with some q ≥ 1 .

Assumption 1-(a) is a Lipschitz-like condition that is useful for decoupling x and z , and is commonly used in the analysis of optimization and stochastic approximations algorithms; see, e.g., [BWMP12]. We can now state our second main contribution, an extension of Theorem 2 to the offline setting.

Theorem 5. Let ( X k ) k ≥ 0 and ( X ( n ) k ) k ≥ 0 be defined as in (2) and (4) with b = 1 and n ≥ 1 . Assume Assumption 1 and the conditions of Theorem 2 with β > 1 hold. Further assume that R ( Z 1 ) is nondeterministic (see Theorem 2), ( X k ) k ≥ 0 is geometrically ergodic and ϵ n ∈ [0 , 1] is the probability that ( X ( n ) k ) k ≥ 0 is not ergodic (in the Wasserstein sense). Then, for every ε > 0 , there exist constants c, c α , c β , and t 0 ≥ 0 , such that for all t > t 0 , the following inequalities hold with probability larger than 1 -ϵ n :

This theorem shows that the approximate power-laws will also be apparent for more general loss functions: for large enough n , the tails of offline SGD will be upper- and lower-bounded by two power-laws with exponents α and β . Note that the ergodicity of ( X k ) k ≥ 0 can be established by using similar tools provided in [DDB20]. On the other hand, the assumption on R ( Z 1 ) being nondeterministic and E [ R β ( Z 1 )] = 1 with β > 1 jointly indicate that δ R = E [ R ( Z 1 )] < 1 , which is a form of contraction on average, hence the strong convexity condition is needed.

Similarly to Theorem 3, in order to prove this result, we first obtain a nonasymptotic Wasserstein convergence bound between π ( n ) and π , stated as follows:

SampleQuantiles

2

GaussianQuantiles

-2

Z

3

n

2

50

100

Sample

-1

0

250

500

750

2

1000

=3

E-

2

-1

0

1

2

GaussianQuantiles

Theorem 6. Under the setting of Theorem 5, the following holds with probability greater than 1 -ϵ n :

The convergence rate implied by Theorem 6 is implicitly related to the finite moments of the datagenerating distribution µ z . Using the results of [FG15] ( Lemma 2 in the Appendix), we can determine the convergence rate r , where with high probability it will hold that W 1 ( π, π ( n ) ) = O ( n -r ) , with r ∈ (0 , 1 / 2] . Let the data have finite moments up to order q . Then, for example, if q > 4 d , we get r = 1 / 2 , i.e., the fastest rate. As another example: if d > 2 and q = 3 , we obtain a rate r = 1 /d . In other words, the moments of the data-generating distribution determine whether we can achieve a dimension-free rate or fall victim to the curse of dimensionality. Thus, if the data do not admit higher-order moments, our results show that we need a large number of data points n to be able to observe power-law tails in offline SGD.

4 Experiments

In this section, we support our theoretical findings on both synthetic and real-world problems. Our primary objective is to validate the tail behavior of offline SGD iterates by adapting previous online SGD analyses [GSZ21, SSG19], in which the authors investigated how heavy-tailed behavior relates to the step size η and batch size b , or their ratio η/b 7 .

4.1 Linear regression

Experimental setup. In the linear regression examples, we examine the case with a quadratic loss and Gaussian data. In this simple scenario, it was observed that online SGDcould exhibit heavytailed behavior [GSZ21]. The model can be sum-

10

20

50

100

200

500

n

0.0

0.2

0.4

|

(

n

)

|

Batch size: 1

10

20

50

100

200

500

n

Batch size: 5

10

20

50

100

200

500

n

Batch size: 20

0.001

0.003

0.005

0.007

0.009

1.15

1.368

1.587

1.806

2.024

marized as follows: with initial parameters ¯ X 0 ∼ N ( 0 , σ 2 x I d ) , features a i ∼ N ( 0 , σ 2 I d ) , and targets y i | a i , ¯ X 0 ∼ N ( a ⊤ i ¯ X 0 , σ 2 y ) , for i = 1 , . . . , n , and σ, σ x , σ y > 0 . In our experiments, we fix σ = 1 , σ x = σ y = 3 , use either d = 1 or d = 100 , and simulate the statistical model to obtain { a i , y i } n i =1 .

In order to examine the offline version of the model (i.e., for offline SGD), we utilize a finite dataset with n points rather than observing new samples at each iteration. We then analyze the same experimental setting with an increasing number of samples n and illustrate how similar heavy-tailed behavior emerges in offline SGD.

0.00

0.02

0.04

/

b

1.0

1.5

2.0

Layer: 1

0.00

0.02

0.04

/

b

Layer: 2

0.00

0.02

0.04

/

b

Layer: 3

0.00

0.02

0.04

/

b

Layer: 1

0.00

0.02

0.04

/

b

Layer: 2

0.00

0.02

0.04

/

b

Layer: 3

Data %:

0.25

0.50

0.75

1.00

Data %:

0.25

0.50

0.75

1.00

MNIST data

CIFAR10 data

Preliminary tail inspection. We begin by analyzing how the tail behavior differs between online and offline SGD (with varying n in the latter). Specifically, in Figure 3, we depict a 1-dimensional linear regression task by analyzing the QQ-plots of the estimated stabilizing distributions after 1000 iterations. We observe that online SGD with a sufficiently large η exhibits heavy, non-Gaussian tails, underscoring the prevalence of heavy tails even in such simple scenarios. Furthermore, we also observe that offline SGD exhibits increasingly heavier tails as the sample size n increases, as our theoretical results suggest.

Tail estimation. We now set d = 100 and run the corresponding offline and online SGD recursions. We then use a tail-index estimator [MMO15], which assumes that the recursions both converge to an exact heavy-tailed distribution. While this is true for online SGD due to Theorem 1, offline SGD will only possess approximate heavy-tails, and the power-law behavior might not be apparent for small n . Hence, for small n , we expect the estimated tail index for offline SGD will be inaccurate and get more and more accurate as we increase n .

We illustrate this in Figure 2, in which we plot the range of estimated offline SGD tail indices (marked in green) corresponding to 10 random initializations, while varying n and η . We can see that, across all learning rates, the variance of the tail estimation decreases as n increases and that the estimated values get closer to the estimated tail index for online SGD (marked as the horizontal blue line).

Further analyses. We now run online and offline SGD recursions, varying η from 0 . 001 to 0 . 01 , b from 1 to 20, and n from 1 to 500 . Each hyperparameter configuration is run 1600 times with distinct initializations, yielding 1600 estimated samples from the stationary distributions, used to estimate the tail indexes. In order to estimate the tail-indexes α (online SGD) and α ( n ) (offline SGD) of the respective stationary distributions, we follow the procedure as explained in [GSZ21]. Finally, in Figure 4, we plot the absolute difference of the estimated indexes, | ˆ α ( n ) -ˆ α | .

We find that a larger number of data samples leads to a smaller discrepancy between the online and offline approximations across all batch sizes b and step sizes η . This trend is consistently observed. Moreover, we observe that larger values of η lead to a smaller discrepancy on average, across all n . The conclusion here is that, as n increases, the power-law tails in offline SGD become more apparent and the estimator can identify the true tail index corresponding to online SGD even for moderately large n , which confirms our initial expectations.

4.2 Neural networks (NN)

To test the applicability of our theory in more practical scenarios, we conduct a second experiment using fully connected (FC) NNs (3 layers, 128 neurons, ReLU activations), as well as larger architectures, such as LeNet (60k parameters, 3 convolutional layers, 2 FC layers) 8 [LBBH98], and AlexNet (62.3M parameters, 5 convolutional layers, 3 FC layers)[KSH17]. The models are trained for 10 , 000 iterations using cross-entropy loss on the MNIST and CIFAR-10 datasets. We vary the learning rate from 10 -4 to 10 -1 , and the batch size b from 1 to 10, with offline SGD utilizing 25% , 50% , and 75% of the training data.

We proceed similarly to the linear regression experiment and again replicate the method presented in [GSZ21]: we estimate the tail index per layer, and plot the corresponding results, with different colors representing different data proportions ( 1 . 00 indicates the full data set). The results for the fully connected network are presented in Figure 5 (on CIFAR-10 & MNIST), LeNet (on MNIST)

0.00

0.02

0.04

/

b

1.0

1.5

2.0

2.5

Layer: 1

0.00

0.02

0.04

/

b

Layer: 2

0.00

0.02

0.04

/

b

Layer: 3

0.00

0.02

0.04

/

b

Layer: 4

0.00

0.02

0.04

/

b

Layer: 5

Data %:

0.25

0.50

0.75

1.00

Data %:

0.25

0.50

0.75

1.00

in Figure 7, and AlexNet (on CIFAR-10) in Figure 6. The LeNet CIFAR-10 and AlexNet MNIST results can be found in the Appendix D.

Our observations show that the estimated ˆ α ( n ) has a strong correlation with ˆ α : using a reasonably large proportion of the data enables us to estimate the tail index that is measured over the whole dataset. Moreover, although the dependence of ˆ α ( n ) on η/b varies between layers, the

1.5

2.0

2.5

Layer: 1

Layer: 2

Layer: 3

Layer: 4

0.000

0.025

/

b

1.5

2.0

2.5

Layer: 5

0.000

0.025

/

b

Layer: 6

0.000

0.025

/

b

Layer: 7

0.000

0.025

/

b

Layer: 8

Data %:

0.25

0.50

0.75

1.00

Data %:

0.25

0.50

0.75

1.00

measured ˆ α ( n ) 's correlate well with the ratio η/b across all datasets and NN architectures 9 . While our theory does not directly cover this setup, our results show that similar behavior is also observed in more complicated scenarios.

5 Conclusion

We established a relationship between the data-generating distributions and stationary distributions of offline and online SGD. This enabled us to develop the first theoretical result illuminating the heavy-tailed behavior in offline SGD. We extended previous results encompassing both quadratic losses, as well as more sophisticated strongly convex loss functions. Through an experimental study, we validated our theoretical findings in a variety of settings.

Limitations. There are two main limitations to our work. Firstly, our analysis focuses predominantly on the 'upper' tails of the distributions. Studying the bulk of the distribution of the iterates may reveal new findings. For example, research in this direction could connect our findings to learning theory by leveraging existing generalization bounds (e.g., see Corollary 2 in [HSKM22] for a link between 'lower' tails and generalization error). Secondly, it would be of great interest to extend our results to the non-convex deep learning settings. Finally, since this is a theoretical paper studying online and offline SGD, our work contains no direct potential negative societal impacts.

Acknowledgments. We thank Benjamin Dupuis for the valuable feedback. AD would like to thank the Isaac Newton Institute for Mathematical Sciences for support and hospitality during the programme The mathematical and statistical foundation of future data-driven engineering when work on this paper was undertaken. Umut ¸ Sim¸ sekli'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) and the European Research Council Starting Grant DYNASTY - 101039676.

References

[BDM + 16] Dariusz Buraczewski, Ewa Damek, Thomas Mikosch, et al. Stochastic models with power-law tails. The equation X= AX+ B. Cham: Springer , 2016.

[Bog06] VI Bogachev. Measure theory (second volume), 2006.

[BSE + 21] Melih Barsbey, Milad Sefidgaran, Murat A Erdogdu, Gael Richard, and Umut Simsekli. Heavy tails in sgd and compressibility of overparametrized neural networks. Advances in Neural Information Processing Systems , 34:29364-29378, 2021.

[BWMP12] A. Benveniste, S.S. Wilson, M. Metivier, and P. Priouret. Adaptive Algorithms and Stochastic Approximations . Stochastic Modelling and Applied Probability. Springer Berlin Heidelberg, 2012.

[CGZ19] Bugra Can, Mert Gurbuzbalaban, and Lingjiong Zhu. Accelerated linear convergence of stochastic momentum methods in wasserstein distances. In International Conference on Machine Learning , pages 891-901. PMLR, 2019.

[DDB20] Aymeric Dieuleveut, Alain Durmus, and Francis Bach. Bridging the gap between constant step size stochastic gradient descent and markov chains, 2020.

[DMN + 21] Alain Durmus, Eric Moulines, Alexey Naumov, Sergey Samsonov, Kevin Scaman, and Hoi-To Wai. Tight high probability bounds for linear stochastic approximation with fixed stepsize. Advances in Neural Information Processing Systems , 34:30063-30074, 2021.

[FG15] Nicolas Fournier and Arnaud Guillin. On the rate of convergence in wasserstein distance of the empirical measure. Probability Theory and Related Fields , 162(3):707-738, 2015.

[GSZ21] Mert Gurbuzbalaban, Umut Simsekli, and Lingjiong Zhu. The heavy-tail phenomenon in sgd. In International Conference on Machine Learning , pages 3964-3975. PMLR, 2021.

[HM21] Liam Hodgkinson and Michael Mahoney. Multiplicative noise and heavy tails in stochastic optimization. In International Conference on Machine Learning , pages 4262-4274. PMLR, 2021.

[HSKM22] Liam Hodgkinson, Umut Simsekli, Rajiv Khanna, and Michael Mahoney. Generalization bounds using lower tail exponents in stochastic optimizers. In International Conference on Machine Learning , pages 8774-8795. PMLR, 2022.

[KSH17] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM , 60(6):84-90, 2017.

[LBBH98] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998.

[LW¸ S22] Soon Hoe Lim, Yijun Wan, and Umut ¸ Sim¸ sekli. Chaotic regularization and heavy-tailed limits for deterministic gradient descent. arXiv preprint arXiv:2205.11361 , 2022.

[MM19] Charles H Martin and Michael W Mahoney. Traditional and heavy-tailed self regularization in neural network models. arXiv preprint arXiv:1901.08276 , 2019.

[MMO15] Mohammad Mohammadi, Adel Mohammadpour, and Hiroaki Ogata. On estimating the tail index and the spectral measure of multivariate α -stable distributions. Metrika , 78(5):549-561, 2015.

[PZ19] Victor M Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual review of statistics and its application , 6:405-431, 2019.

[RBG + 22] Anant Raj, Melih Barsbey, Mert Gürbüzbalaban, Lingjiong Zhu, and Umut ¸ Sim¸ sekli. Algorithmic stability of heavy-tailed stochastic gradient descent on least squares. arXiv preprint arXiv:2206.01274 , 2022.

[RZG¸ S23] Anant Raj, Lingjiong Zhu, Mert Gürbüzbalaban, and Umut ¸ Sim¸ sekli. Algorithmic stability of heavy-tailed sgd with general loss functions. arXiv preprint arXiv:2301.11885 , 2023.

[SCW20] James Stewart, Daniel K Clegg, and Saleem Watson. Calculus: early transcendentals . Cengage Learning, 2020.

[ ¸ SGN + 19] Umut ¸ Sim¸ sekli, Mert Gürbüzbalaban, Thanh Huy Nguyen, Gaël Richard, and Levent Sagun. On the heavy-tailed theory of stochastic gradient descent for deep neural networks. arXiv preprint arXiv:1912.00018 , 2019.

[SSDE20] Umut Simsekli, Ozan Sener, George Deligiannidis, and Murat A Erdogdu. Hausdorff dimension, heavy tails, and generalization in neural networks. Advances in Neural Information Processing Systems , 33:5138-5151, 2020.

[SSG19] Umut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning , pages 5827-5837. PMLR, 2019.

[Vil09] Cédric Villani. The wasserstein distances. In Optimal transport , pages 93-111. Springer, 2009.

[ZFM + 20] Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems , 33:21285-21296, 2020.

[ZLMU22] Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. Strength of minibatch noise in sgd, 2022.

Approximate Heavy Tails in Offline (Multi-Pass) Stochastic Gradient Descent

APPENDIX

The organization of the appendix is as follows:

In Section A, we provide the background material required for the proof methodology.

In Section B, we present technical results used to conclude the results of our work.

In Section C, we give the proofs of Theorems 3-6.

In Section D, we provide additional experimental results.

In Section E, we justify the tail estimator choice.

A Additional Technical Background

A.1 Existing results

Lemma 1 (Explicit solution of Wasserstein distance [PZ19]) . Let µ 1 , µ 2 ∈ P ( R ) be two probability measures on R , and denote their cumulative distribution functions by F 1 ( x ) and F 2 ( x ) respectively. Then, the Wassersteinp distance between µ 1 and µ 2 has an explicit formula:

where F -1 1 and F -1 2 denote the quantile functions. In the case when p = 1 , by applying the change of variables, one can obtain the following:

Now, for q > 0 , α > 0 , γ > 0 and µ ∈ P R d , let:

Furthermore, consider an i.i.d. sequence ( X k ) k ≥ 1 of µ -distributed random variables and, for n ≥ 1 , define the empirical measure by:

We can now proceed to state the following result.

Lemma 2 ([FG15]) . Let µ ∈ P ( R d ) and let p > 0 . Assume one of the three following conditions: (1) α > p, γ > 0 , E α,γ ( µ ) < ∞ ,

(3) q > 2 p, M q ( µ ) < ∞ .

Then for all n ≥ 1 , all x ∈ (0 , ∞ ) , where

and

The positive constants C and c depend only on p, d and either on α, γ, E α,γ ( µ ) (under (1)) or on α, γ, E α,γ ( µ ) , ε (under (2)) or on q, M q ( µ ) , ε (under (3)).

B Technical Lemmas

Lemma 3. Let X and Y be two random vectors in R d . Then, we have

Proof. By the definition of the Wasserstein distance, we have that where γ ∗ is the coupling that attains the infimum (the proof of the existence of such coupling for any Wassersteinp distance, where p ≥ 1 , can be found in, e.g., [Bog06, Theorem 8.10.45]). Then, by the reverse triangle inequality, we have

where T : R d × R d ↦→ R + × R + is the map ( x, y ) ↦→ ( ∥ x ∥ , ∥ y ∥ ) and T # γ ∗ is the pushforward measure such that for a measurable set B ⊂ R + × R + , we have T # γ ∗ ( B ) = γ ∗ ( T -1 ( B )) . As T # γ ∗ is a coupling between ∥ X ∥ and ∥ Y ∥ , we have that

This concludes the proof.

C Proofs

This section contains the proofs of our theoretical findings.

C.1 Proof of Theorem 4

Proof. First, note that the unique stationary distributions of ( X k ) k ≥ 0 and ( X ( n ) k ) k ≥ 0 are denoted respectively by π and π ( n ) . Denote by E n the event on which ( X ( n ) k ) k ≥ 0 has a stationary distribution π ( n ) such that P ( E n ) ≥ 1 -ϵ n (where ϵ n ∈ [0 , 1] is the probability that ( X ( n ) k ) k ≥ 0 is not ergodic, in the Wasserstein sense). Given E n , let ( A k , A ( n ) k ) k ≥ 1 and ( b k , b ( n ) k ) k ≥ 1 be the sequences of i.i.d. optimal couplings for µ A and µ ( n ) A , µ b and µ ( n ) b respectively. Therefore, by construction, for any k ∈ N , E [ ∥ A k -A ( n ) k ∥ ] = W 1 ( µ A , µ ( n ) A ) , and E [ ∥ b k -b ( n ) k ∥ ] = W 1 ( µ b , µ ( n ) b ) . Based on these two sequences, we consider the processes ( X k ) k ≥ 0 , ( X ( n ) k ) k ≥ 0 , ( Y k ) k ≥ 0 defined by the recursions:

Note that ( X k ) k ≥ 0 corresponds to the online SGD recursion (7), and ( X ( n ) k ) k ≥ 0 to the offline SGD recursion (9). In addition, since these two Markov chains are started at stationarity, we have:

With these definitions at hand and this observation, using triangle inequality, we can obtain:

Now, by [GSZ21, Theorem 8], we have that Y k is geometrically ergodic with respect to its stationary distribution π , i.e., there exist constants c ρ > 0 , ρ ∈ (0 , 1) such that the following inequality holds:

Therefore, we have:

as W 1 ( L ( X k ) , L ( Y k )) = W 1 ( π, L ( Y k )) ≤ c ρ e -ρk W 1 ( π, L ( Y 0 )) . Rearranging the above terms, we get W 1 ( π, π ( n ) )(1 -c ρ e -ρk ) ≤ W 1 ( L ( Y k ) , L ( X ( n ) k )) implying:

To bound the right-hand side of (25), we consider the following difference by using the recursion definitions:

Writing out the recursion, we can obtain:

where for any sequence ( a i ) i ≥ 0 , we let ∏ k i = j a i = 1 when j > k . Now, taking the norm of both sides, using the triangle inequality, and taking expectations given E n , we obtain:

We shall now analyze each of the three summands separately. For the first term in (28), as we start with Y 0 = X ( n ) 0 , it equals zero. For the second term in (28), we have:

where (a) follows from the linearity of expectation, (b) follows as A 1 , . . . , A i -1 are i.i.d. and independent from A i , . . . , A k +1 , and (c) is obtained by applying the Hölder's inequality, where p and q are such that 1 /p +1 /q = 1 , with p, q ∈ (1 , + ∞ ) , q < α and p > d/ 2 (used later in the proof), where α is the tail index.

As the chain ( X ( n ) k ) k ≥ 0 starts from its stationary distribution π ( n ) , we have that ( E [ ∥ X ( n ) i ∥ q ]) 1 /q = ( E [ ∥ X ( n ) 0 ∥ q ]) 1 /q , ∀ i . Using the fact that ( A i -A ( n ) i ) k +1 i =1 are i.i.d. random variables (the randomness arises from the empirical measure µ ( n ) A ), E [ ∥ I d -η A 1 ∥ α ] = 1 and α > 1 , we have that:

Therefore, we can obtain the following bound:

where the last equality follows from (29) and the power series sum formula. Using the same arguments, we can bound the third term in (28):

Combining the above, we can obtain that:

where the penultimate inequality follows from Jensen's inequality: it implies that W p ( µ, ν ) ≤ W q ( µ, ν ) , for p ≤ q (see, e.g., [Vil09]). The last inequality follows as W p ( µ A , µ ( n ) A ) ≤ W p ( µ A , b , µ ( n ) A , b ) , and W p ( µ b , µ ( n ) b ) ≤ W p ( µ A , b , µ ( n ) A , b ) . Together with (25), we have:

By taking the limit as k →∞ , we can conclude the proof of Theorem 4.

C.2 Proof of Theorem 3

Proof. Using results from Theorem 4, we now proceed to bound the Wassersteinp distance between the probability laws of A 1 and A ( n ) 1 . Reminiding ourselves that each element in the i.i.d. sequence ( a k ) k ≥ 1 follows a N(0 , σ 2 I d ) , we denote this measure with µ a . Correspondingly, we denote its empirical measure with µ ( n ) a = n -1 ∑ n j =1 δ a j . Given E n , we let ( a k , a ( n ) k ) k ≥ 1 be the sequence of i.i.d. optimal couplings for µ a and µ ( n ) a , so that by construction, for any k ∈ N , ( E [ ∥ a k -a ( n ) k ∥ 2 p ]) 1 / 2 p = W 2 p ( µ a , µ ( n ) a ) . Now:

where (a) follows from Minkowski's inequality and the fact that ( a k ) k ≥ 1 are i.i.d. random variables, as well as ( a ( n ) k ) k ≥ 1 . Inequality (b) follows due to Minkowski's inequality, and (c) follows using the generalized Hölder's inequality with 1 / 2 p +1 / 2 p = 1 /p .

Now, we can proceed to utilize Lemma 2 and (30). In order to apply Lemma 2 to the difference of measures between A 1 and A ( n ) 1 (i.e., between a 1 and a ( n ) 1 ), as E [ ∥ X ( n ) 0 ∥ q ] 1 /q < ∞ for all q < α , we require q < α and 1 /p +1 /q = 1 (due to Hölder's inequality).

First, we select p large and q small enough for both q < α and p > d/ 2 . Now, in order to satisfy the requirements of Lemma 2, we set ϵ ∗ := √ 1 nC 1 log 4 c 1 ζ . This choice allows us to obtain that, with probability greater than 1 -ζ/ 2 , ( E [ ∥ a 1 -a ( n ) 1 ∥ 2 p ]) 1 / 2 p = W 2 p ( µ a , µ ( n ) a ) < ϵ ∗ . Therefore, we have that, with probability greater than 1 -ζ/ 2 :

where the positive constants C 1 and c 1 depend only on p , d , µ a 1 . Therefore, with probability greater than 1 -ζ/ 2 :

In order to bound ( E [ ∥ b k -b ( n ) k ∥ p ]) 1 /p , we can apply Lemma 2 again to obtain constants C 4 and c 5 such that, with probability greater than 1 -ζ/ 2 :

with C 2 = √ 1 C 1 (( E [ ∥ a 1 ∥ 2 p ]) 1 / 2 p +( E [ ∥ a ( n ) 1 ∥ 2 p ]) 1 / 2 p ) and c 3 = 4 c 1 .

Using (32), (33) with Theorem 4 (see (25) and (30)), and denoting C 6 = C 2 ( E [ ∥ X ( n ) 0 ∥ q ]) 1 /q + C 4 and c 7 = max { c 3 , c 5 } we can obtain, with probability greater than 1 -ζ :

Therefore, we obtain that, with probability greater than 1 -ζ , W 1 ( π, π ( n ) ) = O ( n -1 / 2 ) .

Due to [GSZ21] we have lim t →∞ t α P ( ∥ X ∞ ∥ > t ) ∈ (0 , ∞ ) . Now, let ϵ > 0 . Then, there exists t 0 s.t. for all t ≥ t 0 , | t α P ( ∥ X ∞ ∥ > t ) -c | ≤ ϵ , for some c ∈ (0 , ∞ ) . From (34), we know that with probability greater than 1 -ζ , there exist constants ˜ c 1 , ˜ c 2 such that W 1 ( π, π ( n ) ) ≤ ˜ c 1 √ log ˜ c 2 ζ n -1 / 2 .

Now, using Lemma 3, we can obtain that, with probability greater than 1 -ϵ n -ζ :

where we have used that lim t →∞ P ( ∥ X ∞ ∥ > t ) = 0 and lim t →∞ P ( ∥ X ( n ) ∞ ∥ > t ) = 0 , in order to select t ′ and t ′′ large enough for the last inequality to hold. Now, we have that

Therefore, by choosing t ′′ = 2 t ′ , we can obtain:

Similarly, for any ϵ > 0 and t > t 0 , we have t α P ( ∥ X ∞ ∥ > t ) ≤ ϵ + c , and using similar arguments we obtain:

Finally, choosing as before t ′′ = 2 t ′ , we can obtain:

Substituting for ¯ t = 2 t ′ , we can obtain:

Combining (35) and (37), we obtain that with probability greater than 1 -ζ -ϵ n , for any ϵ > 0 and t > t 0 :

This concludes the proof of Theorem 3.

Proof of results in Section 3.2

Let us now consider a setting where f ( · , z ) is not necessarily a quadratic function. As before, we assume the data comes from an unknown data distribution µ z on R d +1 . Furthermore, with ( ˆ Z ( n ) k ) k ≥ 1 , we denote the i.i.d. random variables associated with the empirical measure µ ( n ) z = n -1 ∑ n j =1 δ ˆ Z ( n ) j , as defined in (15). We consider the case b = 1 for simplicity.

C.3 Proof of Theorem 6

Proof. First, we utilize the tail index limits from [HM21] in the strongly convex setting from Theorem 2. We use α and β in order to construct the proof and provide the same argument as in Section 3.1.

Denote by E n the event on which ( X ( n ) k ) k ≥ 0 has a stationary distribution π ( n ) such that P ( E n ) ≥ 1 -ϵ n . Given E n , let ( Z k , ˆ Z ( n ) k ) k ≥ 1 be the sequences of i.i.d. optimal couplings for µ z and µ ( n ) z . Therefore, by construction, for any k ∈ N , E [ ∥ Z k -ˆ Z ( n ) k ∥ ] = W 1 ( µ z , µ ( n ) z ) . Based on this sequence, we consider the processes ( X k ) k ≥ 0 , ( X ( n ) k ) k ≥ 0 , ( Y k ) k ≥ 0 defined by the recursions:

with the unique stationary distributions of ( X k ) k ≥ 0 and ( X ( n ) k ) k ≥ 0 being denoted by π and π ( n ) respectively. Note again that ( X k ) k ≥ 0 corresponds to the online SGD recursion (2), and ( X ( n ) k ) k ≥ 0 to the offline SGD recursion (4), both with b = 1 . Under our assumptions, by the definition of geometric ergodicity, there exist constants c ρ > 0 and ρ ∈ (0 , 1) such that the following inequality holds:

Analogously to (25), we can conclude:

The proof then utilizes Taylor's expansion with the remainder part in integral form (see, e.g., [SCW20]) in the following way:

Now, taking the norm of both sides and using the triangle inequality together with Assumption 1, we can obtain the following:

Now, where again R ( z ) = sup x ∈ R d ∥ ∥ I d -η ∇ 2 f ( x, z ) ∥ ∥ . This allows us to conclude the following: ∥ Y k +1 -X ( n ) k +1 ∥ ≤ R ( Z k +1 ) ∥ Y k -X ( n ) k ∥ + ηL ( ∥ X ( n ) k ∥ +1) ∥ Z k +1 -ˆ Z ( n ) k +1 ∥ .

Denoting B ( n ) k := ∥ X ( n ) k ∥ +1 , we have:

where again, for any sequence ( a i ) i ≥ 0 , ∏ k i = j a i = 1 , for j > k . Furthermore, as β > 1 and E [ R ( Z 1 ) β ] = 1 , we have:

Now, we can take the expectation given E n and use that R ( Z 1 ) , ..., R ( Z k +1 ) , and B ( n ) 0 , . . . , B ( n ) k are i.i.d. Furthermore, using (43), (44), and proceeding as in the proof of Theorem 4, we can obtain:

Finally, denoting E [ ∥ X ( n ) 0 ∥ ] + 1 = c 0 , we obtain that, with probability larger than 1 -ϵ n :

W

1

(

π, π

(

n

)

ρk

1

η

(

n

)

-Taking the limit as k →∞ , we can conclude the proof of Theorem 6.

C.4 Proof of Theorem 5

Proof. Due to [HM21] we know that there exist α, β such that β < α and:

Using Lemma 3, we can obtain:

where we have used that lim t →∞ P ( ∥ X ∞ ∥ > t ) = 0 and lim t →∞ P ( ∥ X ( n ) ∞ ∥ > t ) = 0 , in order to select t ′ and t ′′ large enough for the last inequality to hold. Now, using Theorem 2, we can obtain:

)

c

0

L

(1

-

c

ρ

e

)

1

δ

W

1

(

µ

z

, µ

z

)

.

-

-

R

for some constant c α . Therefore, by choosing t ′′ = 2 t ′ , we can obtain:

Using similar arguments, we can obtain:

for a constant c β . Choosing t ′′ = 2 t ′ and substituting 1 2 ¯ t = t ′ , we obtain:

Using Theorem 6 and (46)-(47), we can conclude the proof of Theorem 5.

0.00

0.02

0.04

/

b

1.0

1.5

2.0

2.5

Layer: 1

0.00

0.02

0.04

/

b

Layer: 2

0.00

0.02

0.04

/

b

Layer: 3

0.00

0.02

0.04

/

b

Layer: 4

0.00

0.02

0.04

/

b

Layer: 5

Data %:

0.25

0.50

0.75

1.00

Data %:

0.25

0.50

0.75

1.00

1

2

Layer: 1

Layer: 2

Layer: 3

Layer: 4

0.000

0.025

/

b

1

2

Layer: 5

0.000

0.025

/

b

Layer: 6

0.000

0.025

/

b

Layer: 7

0.000

0.025

/

b

Layer: 8

Data %:

0.25

0.50

0.75

1.00

Data %:

0.25

0.50

0.75

1.00

D Further experimental results

In this section, we present the remaining experimental findings.

D.1 Neural Network Experiments

First, we showcase the remaining plots from Section 4. As outlined in 4.2, we conducted model training for 10 , 000 iterations, utilizing the cross-entropy loss and employing the MNIST and CIFAR10 datasets. The learning rates ranged from 10 -4 to 10 -1 , while the batch sizes varied between 1 and 10. We employed offline SGD with a subset of the data amounting to 25% , 50% , and 75% . In Figure 8, we exhibit the estimated tail indices for the LeNet architecture with the CIFAR-10 dataset, while Figure 9 shows the estimated tail indices for the AlexNet architecture implemented with the MNIST dataset.

The inclusion of these plots serves the purpose of completeness, as the overall conclusions remain unchanged; a strong correlation between ˆ α ( n ) and ˆ α is exhibited, as well as a notable correlation between ˆ α ( n ) 's and the ratio η/b . These conclusions hold true across all datasets and neural network architectures, providing further substantiation for our theoretical propositions.

D.2 Further tail examination

From [GSZ21], we have that P ( ∥ X ∞ ∥ > t ) ≈ t -α , so the log-log tail histogram of the stationary distribution should follow a linear line with slope -α in the tails (for large t ). We first examine the histograms of the estimated stationary distributions for both the linear regression setting (as in Sec. 4.1) and the NN setting (as in Sec. 4.2), analyzing their behavior as the number of samples n increases. These are depicted in Figures 10-11. Then, we re-run the experiments and examine the histograms on a log-log scale to see whether the linear slope -α becomes more apparent. For both cases, as n increases, we observe that this heavy-tailed phenomenon becomes increasingly apparent, and the tails follow a clearer linear trend. The linear behavior in the tails can be observed in Figures 12-13.

In the entirety of our experiments, encompassing diverse learning rates and layers within both the linear regression and NN settings, a consistent observation emerges: an increase in the number of samples employed in offline SGD leads to a convergence of behavior towards that of online SGD. As previously mentioned in Section 1, while we do not anticipate an exact power-law tailed behavior, it is noteworthy that the tails exhibit progressively more characteristics resembling a power-law

distribution. Specifically, the log-log plots demonstrate a linear trend in the tails. This empirical observation aligns with our theoretical findings, again reinforcing the consistency between the two.

0.2

100

0.2

100

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

0.2

100

n = 20

n = 250

n = 1000

10

2

10

2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

10

2

First layer

n:

0.5

0.75

1.0

10

1

10

3

10

1

10

3

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

10

1

10

3

Second layer

n:

0.5

0.75

1.0

10

0

10

1

10

2

10

0

10

1

10

2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

10

0

10

1

10

2

Third layer

n:

0.5

0.75

1.0

10

2

10

1

10 0

10 0

10 1

10 2

10

2

10

1

10 0

10

2

10

1

10 0

n = 20

n = 250

n = 1000

10

2

First layer

10

2

10

2

Second layer

Third layer

% = 0.5

% = 0.75

% = 1.0

10

2

10

2

10

2

% = 0.5

% = 0.75

% = 1.0

10

2

10

2

10

2

% = 0.5

% = 0.75

% = 1.0

10

2

10

3

10

2

10

1

10

0

10

2

10

2

10

1

10

0

10

1

% = 0.5

% = 0.75

10

2

10

1

10

0

% = 1.0

=0.015

=0.03

=0.04

=0.045

E Estimator choice justification

Analytical computation of the true tail exponent, even in the linear regression setting, is, to our knowledge, unfortunately not possible. Therefore, an estimator choice is required. In this section, we provide our rationale for selecting the tail estimator from [MMO15] by highlighting its strengths.

2

1.5

Estimate

1

0.5

0.5

1

1.5

2

True α

ε

|

∈ 1 ≤ i ≤ r . For t ≥ 0 , w tε -α → Y m i ( t ) , as ε → 0 , in the sense of finite-dimensional distributions, where Y = ( Y y ( t )) t ≥ 0 is a continuous-time Markov chain on a state space { m 1 , m 2 , . . . , m r } with the infinitesimal generator Q = ( q ij ) r i,j =1 with [SSG19]), which, in our opinion, contains convincing estimation results (e.g., small error bars regardless of the magnitude of the true tail index).

(a) Figure 3: (a) An objective with two local minima (b) Illustration of the tail-index estimator ˆ α Theorem 2 ( [42]) . Let w 0 S i , for some The main arguments are as follows: (i) the estimator's theoretical framework is established through its convergence in distribution to the true tail index (via a Central Limit Theorem result, detailed in Theorem 2.3 [MMO15]), which is further complemented by its demonstrated asymptotic consistency (as per Corollary 2.4 [MMO15]). (ii) The estimator has already been used in various articles and its qualities have been thoroughly examined. For example, this can be observed in Figure 14 (taken from

1

q ij = α ∣ ∣ ∣ ∣ | s j -q ii = -∑ F Strongly Convex Problem Example

∣

∣

j

/negationslash

1

=

i

1

-

q

m

ij

.

(13)

(14)

This process admits a density π satisfying Q T π = 0 . A consequence of this theorem is that equilibrium probabilities p i are typically larger for 'wide valleys". To see this consider the special case illustrated in Figure 3(a) with r = 2 local minima Consider the following one-dimensional logistic regression where the regularization parameter λ ∼ Exp ( µ ) , i.e., it follows an Exponential distribution with mean 1 /µ - its probability density function is as follows:

m

1

< s

1

= 0

< m

2

s

1

m

separated by a local maximum at

m

|

1

|

+

m

2

|

= 0

1

|

1

-

m

.

,

i

α

∣

∣

For this example,

+

|

m

2

|

We see that π 2 > π 1 , that is in the equilibrium the process spends more time on the wider walley. In particular, the ratio π 2 π 1 = ( m 2 | m 1 | ) α grows with an exponent α when the ratio m 2 | m 1 | of the width of the valleys grows. Consequently, if the gradient noise is indeed α -stable distributed, As before, let ( a i , q i ) i ≥ 1 be i.i.d. random variables in R 2 , such that z i ≡ ( a i , q i ) . In this scenario, where a ∼ N (0 , σ 2 ) and y ∈ { 0 , 1 } , the loss function equals:

these results directly provide theoretical evidence for the wide-minima behavior of SGD.

noise

We first consider a fully-connected network (FCN) on the MNIST and

The order of the total number of parameters in

p

CIFAR10 datasets. For this model, we vary the depth (i.e. the number of layers) in the set { 2 , 3 , . . . , 10 } , the width (i.e. the number of neurons per layer) in the set { 2 , 4 , 8 , . . . , 1024 } , and the minibatch size ranging from 1 to full batch. We then consider a convolutional neural network (CNN) architecture (AlexNet) on the CIFAR10 and CIFAR100 datasets. We scale the number of filters in each convolutional layer in range { 2 , 4 , . . . , 512 } . We randomly split the MNIST dataset into train and test parts of sizes 60 K and 10 K, and CIFAR10 and CIFAR100 datasets into train and test parts of sizes 50 K and 10 K, respectively. Now, the second derivative of the loss with respect to the parameter x equals ∇ 2 ℓ ( x, z ) = a 2 e ax (e ax +1) 2 + λ . Note that |∇ 2 ℓ ( x, z ) | ≤ a 2 / 4 + λ, ∀ x ∈ X . In the SGD case, from [HM21], we have r ( z ) = lim inf ∥ x ∥→∞ σ min ( I d -γ ∇ 2 ℓ ( x, z ) ) and R ( z ) = sup x ∥ ∥ I d -γ ∇ 2 ℓ ( x, z ) ∥ ∥ . In other words, as we are in the one-dimensional setting, we have:

a variety of

scenarios.

range from several thousands to tens of millions.

Now, we require P ( r ( Z 1 ) > 1) > 0 and E [ R ( Z 1 )] < 1 . As λ ∼ Exp ( µ ) , we have that P ( r ( Z 1 ) > 1) > 0 . For the latter condition, we can calculate the required expectation:

m

2

>

m

|

1

|

, and

-

i

|

α

|

s

j

Furthermore, we have:

We can see that, for example, both expressions are less than 1 in modulus for µ = 0 . 1 , σ 2 = 1 , and γ = 0 . 1 . It is important to highlight that the aforementioned calculations necessitate the condition 1 2 σ 2 -µ 4 > 0 , or equivalently σ 2 µ < 2 . This condition can be interpreted as a stability criterion, indicating that the mean of the penalization term λ must increase as the variance of the data grows. In other words, a larger variance necessitates a higher mean value for λ .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the approximate heavy-tailed behavior of stochastic gradient descent (SGD) in local minima in practical settings and its correlation with overall performance. It shows that the stationary distribution of offline SGD exhibits approximate power-law distribution, and the approximation error is controlled by how fast the empirical distribution of the training data converges to the true underlying data distribution in the Wasserstein metric. The main contribution is to fill the gap in understanding the underlying mechanism generating the reported heavy-tailed behavior in practical settings, where the amount of training data is finite. The paper also proves nonasymptotic Wasserstein convergence bounds for offline SGD to online SGD as the number of data points increases. The theory is verified with experiments\n\nSTRENGTHS:\n- **Originality**: The paper investigates the heavy-tailed behavior of SGD in practical settings where the amount of data is finite, which has not been well understood before. It proves a nonasymptotic Wasserstein convergence bound for offline SGD which reduces to online SGD as the number of data points increases. The difference from a true heavy-tailed distribution with a finite number of data is given by the Wasserstein distance.\n- **Quality**: The paper provides rigorous theoretical analysis and proof to support its claims. It also conducts experiments to verify its theory.\n- **Clarity**: The paper is well-written and easy to follow, with clear explanations of technical terms and concepts. The authors also provide many clear schematic illustrations to help readers understand their concepts and results.\n- **Significance**: The theory of this paper is of both theoretical interest and direct practical significance. The results are easy to understand and bounds on the distance from a true heavy-tailed distribution are given explicitly with a clear form.\n\nWEAKNESSES:\n- Possible evaluation method for the values of estimated tail indices is not provided. In Figure 5, the estimated tail indices depend on the step size $\\eta$, the batch size $b$, or their ratio $\\eta/b$, and of course the number of data $n$. It will be strong if an evaluation method for the tail indices, even approximately, can be provided.\n\nQUESTIONS:\n- In Line 267, the authors state that the tail indices depend on the ratio of the step size $\\eta$ to the batch size $b$. However, a mere dependence on the ratio $\\eta/b$ has been found broken in online SGD with a large $\\eta$ as discussed in Sec.6.6 of Ref.[1]. A modification that is linear in the batch size $b$ is found for a simple 1-dimensional example. Could the authors elaborate on this result?\n\n- Still in Sec.6.6 of Ref.[1], the tail index is estimated with an explicit expression in terms of $\\eta$ and $b$ for a 1D example. Is it possible to compare the estimated tail indices predicted by your theorem with the explicit expression given in Ref.[1]?\n\n[1] **Strength of Minibatch Noise in SGD**, Liu Ziyin, Kangqiao Liu, Takashi Mori, Masahito Ueda, *ICLR 2022*\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis manuscript considers the problem of multi-pass stochastic gradient descent with a finite batch-size and strongly convex objective. The key result is to show that when the stationary distribution of the parameters is heavy-tailed in the infinite data (one-pass) limit, then the stationary distribution of the finite batch limit is approximately heavy-tailed, in the sense of having a heavy-tailed component plus a correction depending on the 1-Wasserstein distance between empirical and population data distribution (hence decaying at worst as $\\sim n^{-1/2}$). The authors also present numerical experiments that suggest the relevance of this theoretical result to more practical settings.\n\nSTRENGTHS:\nUnderstanding SGD in its different flavours is an important problem in theoretical machine learning, and results on multi-pass SGD are scarce compared to the one-pass case. Therefore, the contribution is timely and significant. The technical part is sound. Moreover, the paper is well-written and the thread is relatively easy to follow: the motivation and goals are clearly stated and the main result addresses them. Finally, the numerical simulations supporting a broader scope of applicability are nice.\n\nWEAKNESSES:\nThere are two immediate shortcomings of the work. First, the setting is rather restricted (strongly convex goals). Second, while the authors cite related literature connecting heavy-tails with generalization, this is not explicitly explored in the context of the paper. Overall, one question that remains in the end of the reading is: the stationary distribution of the weights are approximately heavy-tailed, but so what?\n\nQUESTIONS:\n- **[Q1]**: As a reader who is not very familiar with this line of work ([HM21, HSKM22, MM], etc.) I wonder how to conciliate these results to the classical literature [Fabian 1968, Kushner 1981, Pflug 1986] showing asymptotic normality of the stationary distribution under certain conditions. For instance, considering online SGD in the simple least-squares setting. Since the gradient is proportional to the residual $r_{i} = y_{i}-\\langle a_{i}, x\\rangle$:\n$$\n  x^{k+1}  = x^{k} + \\gamma r_{k} a_{k}\n$$\n\nif the residuals are Gaussian close to the minimum (e.g. $y_{i}=\\angle a_{i}, x_{\\star}\\rangle + z_{i}$ for some $x_{\\star}$ and $z_{i}\\sim\\mathcal{N}(0,\\sigma^2)$), won't the stationary distribution be Gaussian? What I am missing?\n\n- **[Q2]**: Theorems 2 and 3 have an explicit dependence on the sample size $n$, but how do they depend on the batch size b?\n\n**Minor comments**\n\n- The following sentences are strong and misleading:\n> *Previous works on this problem only considered, up to our knowledge, online (also called single-pass) SGD, which requires an infinite amount of data.*\n\n> *However, all these results rely on exact heavy tails, which, to our current knowledge, can only occur in the online SGD regime where there is access to an infinite sequence of data points.*\n\nWhile it is true that in one-pass SGD the maximum number of steps is limited by the availability of data, this doesn't necessarily means that it requires an infinite amount of data. Note that the number of steps required generically depend on the task under consideration, and in some practical scenarios where data is abundant only a few epochs are required, see e.g. Table 2.2. of [[Kaplan et al. 2020]](https://arxiv.org/pdf/2005.14165.pdf). The authors even acknowledge this fact in L29-30.\n\n- Please add the details of the plots in the caption in Fig. 2 & 5: the reported tail indices correspond to what simulation?\n\n- L128: You mean $[n] = \\{1,\\cdots, n\\}$?\n\n**References**\n\n[[Fabian 1968]](https://projecteuclid.org/journals/annals-of-mathematical-statistics/volume-39/issue-4/On-Asymptotic-Normality-in-Stochastic-Approximation/10.1214/aoms/1177698258.full) Vaclav Fabian. *On Asymptotic Normality in Stochastic Approximation*. Ann. Math. Statist. 39 (4) 1327 - 1332, August, 1968. https://doi.org/10.1214/aoms/1177698258\n\n[[Kushner 1981]](https://epubs.siam.org/doi/10.1137/0319007) Harold J. Kushner and Hai Huang. *Asymptotic Properties of Stochastic Approximations with Constant Coefficients*. SIAM Journal on Control and Optimization, 19(1), 1981\n\n[[Pflug 1986]](https://epubs.siam.org/doi/10.1137/0324039) Georg Ch. Pflug. *Stochastic minimization with constant step-size: Asymptotic laws*. SIAM Journal on Control and Optimization, 24(4):655–666, 1986\n\nLIMITATIONS:\nI have discussed some of the limitations in the \"Weaknesses\" part.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors study the heavy-tail distribution for the parameters in offline stochastic gradient descent algorithm (SGD). Theoretical results are provided for a quadratic loss and a strongly convex problem, while numerical results cover more realistic cases such as fully connected NN or CNN. In the theoretical part of the paper, the authors link the tail of the distribution of the offline SGD to the online one. To be more specific, the authors show that the qualitative difference between the online and the offline tail is bounded by the Wasserstein distance between the generating and empirical distribution of the data. The author further bound the Wasserstein distance between the distribution of the online and the offline parameters by the same distance. The theoretical result is qualitatively supported by the numerical experiment.\n\nSTRENGTHS:\n* Originality and significance: the authors extend the analysis of heavy tail in SGD from the online settings to offline settings. Personally, I find it really interesting of having a quantitative non-asymptotic bound of the difference between the two cases. It could make the existing and future online SGD analysis more relevant.\n* Quality: In my personal opinion, the authors perform convincing theoretical and numerical analysis.\n* Clarity: The paper is well written. The authors present their results in a self-contained way, with multiple examples and intuitions helping illustrating the idea behind their main results. And I appreciate that the code is attached to the paper.\n\nWEAKNESSES:\nIt appears to me that the fig.1 and fig.3 are not as convincing as the rest of the paper. The axes are unlabeled in these figures. The fact that the distribution has a heavy tail is not very obvious in these figures. For example, to my eye, the n=250 histogram is more power-law comparing to the n=1000 histogram.\n\nThe weaknesses and questions are well addressed by the authors in their rebuttal. The authors provided detailed explanations and additional figures.\n\nQUESTIONS:\nI believe that it would be nice if the authors could include results, expressed in numbers, which characterizes the difference between the experimental tail and a real power-law tail.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

comments: SUMMARY:
This paper considers offline SGD and proves an approximate power-law tail behavior of the stochastic gradient for strongly convex objectives, confirming the heavy-tail heuristic encountered in practices. Explicit tail estimations are obtained, and as a intermediate result, nonasymptotic Wasserstein convergence results for offline SGD are proved.

STRENGTHS:
1. The paper provides a good analysis of the behavior of offline SGD, and contributes to our understanding of heavy-tail phenomena in SGD training.
2. I have gone over the proofs and believe it's basically correct.
3. The paper is well-written and comfortable to read.

WEAKNESSES:
I do not see apparent weaknesses in this paper.

QUESTIONS:
I have no additional questions.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07