Title: Aiming towards the minimizers: fast convergence of SGD for overparametrized problems
Venue: NeurIPS 2023
Accepted: yes
Keywords: Polyak-Lojasiewicz condition

SGD

interpolation

fast convergence
Abstract: Modern machine learning paradigms, such as deep learning, occur in or close to the interpolation regime, wherein the number of model parameters is much larger than the number of data samples. 
In this work, we propose a regularity condition within the interpolation regime which endows the stochastic gradient method with the same worst-case iteration complexity as the deterministic gradient method, while using only a single sampled gradient (or a minibatch) in each iteration. In contrast, all existing guarantees require the stochastic gradient method to take small steps, thereby resulting in a much slower linear rate of convergence. Finally, we demonstrate that our condition holds when training sufficiently wide feedforward neural networks with a linear output layer.

Aiming towards the minimizers: fast convergence of SGD for overparameterized problems

Abstract

Modern machine learning paradigms, such as deep learning, occur in or close to the interpolation regime, wherein the number of model parameters is much larger than the number of data samples. In this work, we propose a regularity condition within the interpolation regime which endows the stochastic gradient method with the same worst-case iteration complexity as the deterministic gradient method, while using only a single sampled gradient (or a minibatch) in each iteration. In contrast, all existing guarantees require the stochastic gradient method to take small steps, thereby resulting in a much slower linear rate of convergence. Finally, we demonstrate that our condition holds when training sufficiently wide feedforward neural networks with a linear output layer.

1 Introduction

Recent advances in machine learning and artificial intelligence have relied on fitting highly overparameterized models, notably deep neural networks, to observed data; e.g. [39, 37, 16, 23]. In such settings, the number of parameters of the model is much greater than the number of data samples, thereby resulting in models that achieve near-zero training error. Although classical learning paradigms caution against overfitting, recent work suggests ubiquity of the 'double descent' phenomenon [3], wherein significant overparameterization actually improves generalization. The stochastic gradient method is the workhorse algorithm for fitting overparametrized models to observed data and understanding its performance is an active area of research. The goal of this paper is to obtain improved convergence guarantees for SGD in the interpolation regime that better align with its performance in practice.

Classical optimization literature emphasizes conditions akin to strong convexity as the phenomena underlying rapid convergence of numerical methods. In contrast, interpolation problems are almost never convex, even locally, around their solutions [24]. Furthermore, common optimization problems have complex symmetries, resulting in nonconvex sets of minimizers. Case in point, standard formulations for low-rank matrix recovery [4] are invariant under orthogonal transformations while ReLU neural networks are invariant under rebalancing of adjacent weight matrices [9]. The PolyakŁojasiewicz (PŁ) inequality, introduced independently in [26] and [34], serves as an alternative to strong convexity that holds often in applications and underlies rapid convergence of numerical algorithms. Namely, it has been known since [34] that gradient descent convergences under the PŁ condition at a linear rate O (exp( -t/κ )) , where κ is the condition number of the function. 1 In contrast, convergence guarantees for the stochastic gradient method under PŁ-the predominant algorithm in practice-are much less satisfactory. Indeed, all known results require SGD to take

le-1

batch_size=64

batch_size=32

batch_size=120

le-1

batch_size=64

batch_size=250

batch_size=128

5e-2

batch_size=500

batch_size=256

batch_size=1000

batch_size=512

batch_size=1000

training loss

training loss

batch_size=2000

2e-2

1e-2

1e-2

1e-3

0

500

1000

1500

2000

2500

3000

3500

4000

0

20000

40000

60000

80000

100000

number of iterations number of iterations

shorter steps than gradient descent to converge at all [2, 12], with the disparity between the two depending on the condition number κ . The use of the small stepsize directly translates into a slow rate of convergence. This requirement is in direct contrast to practice, where large step-sizes are routinely used. As a concrete illustration of the disparity between theory and practice, Figure 1 depicts the convergence behavior of SGD for training a neural network on the MNIST data set. As is evident from the Figure, even for small batch sizes, the linear rate of convergence of SGD is comparable to that of GD with an identical stepsize η = 0 . 1 . Moreover, experimentally, we have verified that the interval of stepsizes leading to convergence for SGD is comparable to that of GD; indeed, the two are off only by a factor of 20 . Using large stepsizes also has important consequences for generalization. Namely, recent works [19, 5] suggest that large stepsizes bias the iterates towards solutions that generalize better to unseen data. Therefore understanding the dynamics of SGD with large stepsizes is an important research direction. The contribution of our work is as follows.

In this work, we highlight regularity conditions that endow SGD with a fast linear rate of convergence exp( -t/κ ) both in expectation and with high probability, even when the conditions hold only locally. Moreover, we argue that the conditions we develop are reasonable because they provably hold on any compact region when training sufficiently wide feedforward neural networks with a linear output layer.

1.1 Outline of main results.

We will focus on the problem of minimizing a loss function L under the following two assumptions. First, we assume that L grows quadratically away from its set of global minimizers S :

where B r ( w 0 ) is a ball of radius r around the initial point w 0 . This condition is standard in the optimization literature and is implied for example by the PŁ-inequality holding on the ball B r ( w 0 ) ; see Section A.2. Secondly, and most importantly, we assume there there exist constants θ, ρ > 0 satisfying the aiming condition: 2

Here, proj S ( w ) denotes a nearest point in S to w and dist( w,S ) denotes the distance from w to S . The aiming condition ensures that the negative gradient -∇L ( w ) points towards S in the sense that -∇L ( w ) correlated nontrivially with the direction proj S ( w ) -w . At first sight, the aiming condition appears similar to quasar-convexity, introduced in [14] and further studied in [15, 22, 20]. Namely a function L is quasar-convex relative to a fixed point ¯ w ∈ S if the estimate (Aiming) holds with proj S ( w ) replaced by ¯ w . Although the distinction between aiming and quasar-convexity may

appear mild, it is significant. As a concrete example, consider the function L ( x, y ) = 1 2 ( y -ax 2 ) 2 for any a > 0 . It is straightforward to see that L satisfies the aiming condition on some neighborhood of the origin. However, for any neighborhood U of the origin, the function L is not quasar-convex on U relative to any point ( x, ax 2 ) ∈ U ; see Section C. More generally, we show that (Aiming) holds automatically for any C 3 smooth function L satisfying (QG) locally around the solution set. Indeed, we may shrink the neighborhood to ensure that θ is arbitrarily close to 2 . Secondly, we show that (Aiming) holds for sufficiently wide feedforward neural networks with a linear output layer.

Our first main result can be summarized as follows. Roughly speaking, as long as the SGD iterates remain in B r ( w 0 ) , they converge to S at a fast linear rate O (exp( -tαθ 2 /β )) with high probability.

Theorem 1.1 (Informal) . Consider minimizing the function L ( w ) = E ℓ ( w,z ) , where the losses ℓ ( · , z ) are nonnegative and have β -Lipschitz gradients. Suppose that the minimal value of L is zero and both regularity conditions (QG) and (Aiming) hold. Then as long as the SGD iterates w t remain in B r ( w 0 ) , they converge to S at a linear rate O (exp( -tαθ 2 /β )) with high probability.

The proof of the theorem is short and elementary. The downside is that the conclusion of the theorem is conditional on the iterates remaining in B r ( w 0 ) . Ideally, one would like to estimate this probability as a function of the problem parameters. With this in mind, we show that for a special class of nonlinear least squares problems, including those arising when fitting wide neural networks, this probability may be estimated explicitly. The end result is the following unconditional theorem.

Theorem 1.2 (Informal) . Consider the loss L ( w ) = 1 n ∑ n i =1 ( f ( w,x i ) -y i ) 2 , where f ( w, · ) is a fully connected neural network with l hidden layers and a linear output layer. Let λ ∞ be the minimal eigenvalue of the Neural Tangent Kernel of an infinitely wide neural network. Then with high probability both conditions (QG) and (Aiming) hold on a ball of radius r around the initial point w 0 ∼ N (0 , I ) with θ = 1 and α = λ ∞ / 2 , as long as the network width m satisfies m = ˜ Ω( nr 6 l +2 /λ 2 ∞ ) . If in addition r = Ω(1 /δ √ λ ∞ ) , with probability 1 -δ , SGD with stepsize η = Θ(1) converges to a zero-loss solution at the fast rate O (exp( -tλ ∞ )) . This parameter regime is identical as for gradient descent to converge in [24], with the only exception of the inflation of r by 1 /δ .

A key part of the argument is to estimate the probability that the iterates remain in a ball B r ( w 0 ) . A naive approach is to bound the length of the iterate trajectory in expectation, but this would then require the radius r to expand by an additional factor of 1 /λ ∞ , which in turn would increase m multiplicatively by λ -6 l -2 ∞ . We avoid this exponential blowup by a careful stopping time argument and the transition to linearity phenomenon that has been shown to hold for sufficiently wide neural networks [24]. While Theorem 1.2 is stated with a constant failure probability, there are standard ways to remove the dependence. One option is to simply set δ 1 = Ω(1) and rerun SGD logarithmically many times from the the same initialization w 0 and return the final iterate with smallest function value. Section 4 outlines a more nuanced strategy based on a small ball assumption, which entirely avoids computation of the function values of the full objective.

1.2 Comparison to existing work.

We next discuss how our results fit within the existing literature, summarized in Table 1. Setting the stage, consider the problem of minimizing a smooth function L ( w ) = E ℓ ( w,z ) and suppose for simplicity that its minimal value is zero. We say that L satisfies the Polyak-Łojasiewicz (PŁ) inequality if there exists α > 0 satisfying

for all w ∈ R d . In words, the gradient ∇L ( w ) dominates the function value L ( w ) , up to a power. Geometrically, such functions have the distinctive property that the gradients of the rescaled function √ L ( w ) are uniformly bounded away from zero outside the solution set. See Figures 2a and 2b for an illustration. Using the PŁ inequality, we may associate to L two condition numbers, corresponding to the full objective and its samples, respectively. Namely, we define ¯ κ ≜ ¯ β/α and κ ≜ β/α , where ¯ β is a Lipschitz constant of the full gradient ∇L and β is a Lipschitz constant of the sampled gradients ∇ ℓ ( · , z ) for all z . Clearly, the inequality ¯ κ ≤ κ holds and we will primarily be interested in settings where the two are comparable.

The primary reason why the PŁ condition is useful for optimization is that it ensures linear convergence of gradient-type algorithms. Namely, it has been known since Polyak's seminal work

(a)

L

(

w

)

satisfying the PŁ condition

(b) gradients of

√

L

(

w

)

are far from zero

For stochastic optimization problems under the PŁ condition, the story is more subtle, since the rates achieved depend on moment bounds on the gradient estimator, such as:

for A,B,C ≥ 0 . In the setting where C > 0 -the classical regime- stochastic gradient methods converge sublinearly at best, due to well-known lower complexity bounds in stochastic optimization [31]. On the other hand, in the setting where C = 0 -interpolation problems-stochastic gradient methods converge linearly when equipped with an appropriate stepsize, as shown in [2, Theorem 1], [22, Corollary 2], [38], and [13, Theorem 4.6]. Although linear convergence is assured, the rate of linear converge under the PŁ condition and interpolation is an order of magnitude worse than in the deterministic setting. Namely, the three papers[2, Theorem 1], [22, Corollary 2] and [13, Theorem 4.6] obtain linear rates on the order of exp( -t/ ¯ κκ ) . On the other hand, in the case A = C = 0 , which is called the strong growth property, the paper [38, Theorem 4] yields the seemingly better rate exp( -t/B ¯ κ ) . The issue, however, is that B can be extremely large. As an illustrative example, consider the loss functions ℓ ( w,z ) = 1 2 dist 2 ( w,Q z ) where Q z are smooth manifolds. A quick computation shows that equality ∥∇ ℓ ( w,z ) ∥ 2 = 2 ℓ ( w,z ) holds. Therefore, locally around the intersection of ∩ z Q z , the estimate (1.1) with A = C = 0 is exactly equivalent to the PŁ-condition with B = 1 /α . As a further illustration, Figure 3 shows the possible large value of the constant B along the SGD iterates for training a neural network on MNIST. Another related paper is [36]: assuming so-called small gradient confusion and requiring a stronger version of PL condition (i.e., each individual loss l i is µ -PL), the paper [36] showed a slow rate of convergence exp( -t/n 2 κ ) , where n is the dataset size.

40k

fficient

30k

coef

20k

G

S

10k

0

0

1

2

3

4

5

6

7

8

9

10

epochs (x100)

The purpose of this work is to understand whether we can improve stepsize selection and the convergence rate of SGD for nonconvex problems under the (local) PŁ condition and interpolation. Unfortunately, the PŁ condition alone appears too weak to yield improved rates. Instead, we take inspiration from recent work on accelerated deterministic nonconvex optimization, where the recently introduced quasar-convexity condition has led to improved rates [15]. We note that quasar-convexity is a very restrictive assumption in the interpolation regime because it requires the solution set to be an affine subspace; see the discussion in Appendix E. Recent work has also shown that quasar convexity can lead to accelerated sublinear rates of convergence for certain stochastic optimization problems [20, Theorem 4.4] (and the concurrent work [11, Corollary 3.3]), but to the best of our knowledge, there are no works that analyze improved linear rates of convergence. Thus, in this work, we fill the gap in the literature, by providing a rate that matches that of deterministic gradient descent and allows for a large stepsize. Moreover, in contrast to most available results, we only assume that regularity conditions hold on a ball-the common setting in applications. The local nature of the assumptions requires us to bound the probability of the iterates escaping.

2 Main results

Throughout the paper, we will consider the stochastic optimization problem

where P is a probability distribution that is accessible only through sampling and ℓ ( · , z ) is a differentiable function on R d . We let S denote the set of minimizers of L . We impose that L satisfies the following assumptions on a set W . The two main examples are when W is a ball B r ( w 0 ) and when W is a tube around the solution set:

Assumption 1 (Running assumptions) . Suppose that there exist constants α, β, θ ≥ 0 and a set W ⊂ R d satisfying the following.

(Interpolation) The losses ℓ ( w,z ) are nonnegative, the minimal value of L is zero, and the set of minimizers S ≜ argmin L is nonempty.

(Smoothness) For almost every z ∼ P , the loss ℓ ( · , z ) is differentiable and the gradient ∇ ℓ ( · , z ) is β -Lipschitz continuous on W .

(Quadratic growth) The estimate holds:

(Aiming) For all w ∈ W there exists a point ¯ w ∈ proj( w,S ) such that

We define the condition number κ ≜ β/α .

As explained in the introduction, the first three conditions (1)-(3) are classical in the literature. In particular, both quadratic growth (3) on a ball W = B r ( w 0 ) and existence of solutions in W follow from a local PŁ-inequality. In order to emphasize the local nature of the condition, following [24] we say that L is α -PŁ ∗ on W if the inequality (PŁ) holds for all w ∈ W . We recall the proof of the following lemma in Section A.2.

Lemma 2.1 (PŁ ∗ condition implies quadratic growth) . Suppose that L is differentiable and is α -PŁ ∗ on a ball B 2 r ( w 0 ) . Then as long as L ( w 0 ) < 1 2 αr 2 , the intersection S ∩ B r ( w 0 ) is nonempty and

The aiming condition (4) is very closely related to quasar-convexity , which requires (2.2) to hold for all w ∈ R d and a distinguished point ¯ w ∈ S that is independent of w . This distinction may seem mild, but is in fact important because aiming holds for a much wider class of problems. As a concrete example, consider the function L ( x, y ) = 1 2 ( y -ax 2 ) 2 for any a > 0 . It is straightforward to see that L satisfies the aiming condition on some neighborhood of the origin. However, for any neighborhood U of the origin, the function L is not quasar-convex on U relative to any point ( x, ax 2 ) ∈ U ; see Section C. We now show that (4) is valid locally for any C 3 -smooth function satisfying quadratic growth, and we may take θ arbitrarily close to 2 by shrinking r . Later, we will also show that problems of learning wide neural networks also satisfy the aiming condition.

Theorem 2.2 (Local aiming) . Suppose that L is C 2 -smooth and ∇ 2 L is L -Lipschitz continuous on the tube S r . Suppose moreover that L satisfies the quadratic growth condition (2.1) and r < 6 α 5 L . Then the aiming condition (2.2) holds with parameter θ = 2 -5 Lr 3 α . An analogous statement holds if S r is replaced by a ball B r ( ¯ w 0 ) for some ¯ w 0 ∈ S .

Algorithm 1 SGD ( w 0 , η, T )

Initialize: Initial w 0 ∈ R d , learning rate η > 0 , iteration counter T ∈ N .

For t = 1 , . . . , T -1 do:

Return: w T .

2.1 SGD under regularity on a tube S r

Convergence analysis for SGD (Algorithm 1) is short and elementary in the case W = S r and therefore this is where we begin. We note, however, that the setting W = B r ( w 0 ) is much more realistic, as we will see, but also more challenging.

The converge analysis of SGD proceeds by a familiar one-step contraction argument.

Lemma 2.3 (One-step contraction on a tube) . Suppose that Assumption 1 holds on a tube W = S 2 r and fix a point w ∈ S r . Define the updated point w + = w -η ∇ f ( w,z ) where z ∼ P . Then for any stepsize η < θ β , the estimate holds:

Using the one step guarantee of Lemma 2.3, we can show that SGD iterates converge linearly to S if Assumption 1 holds on a tube W = S 2 r . The only complication is to argue that the iterates are unlikely to leave the tube if we start in a slightly smaller tube S r ′ for some r ′ < r . We do so with a simple stopping time argument.

Theorem 2.4 (Convergence on a tube) . Suppose that Assumption 1 holds relative to a tube W = S 2 r for some constant r > 0 . Fix a stepsize η > 0 satisfying η < θ β . Fix a constant δ 1 > 0 and a point

w 0 ∈ S √ δ 1 r . Then with probability at least 1 -δ 1 , the SGD iterates { w t } t ≥ 0 remain in W . Moreover, with probability at least 1 -δ 1 -δ 2 , the estimate dist 2 ( w t , S ) ≤ ε · dist 2 ( w 0 , S ) holds after

Thus as long SGD is initialized at a point w 0 ∈ S √ δ 1 r , with probability at least 1 -δ 1 -δ 2 , the iterates remain in S r and converge at linear rate O ( 1 δ 2 exp( -tθ 2 α/β )) . Note that the dependence on δ 2 is logarithmic, while the dependence on δ 1 appears linearly in the initialization requirement w 0 ∈ S √ δ 1 r . One simple way to remove the dependence on δ 1 is to simply rerun the algorithm from the same initial point logarithmically many times and return the point with the smallest function value. An alternative strategy that bypasses evaluating function values will be discussed in Section 4.

2.2 SGD under regularity on a ball B r ( w 0 )

Next, we describe convergence guarantees for SGD when Assumption 1 holds on a ball W = B r ( w 0 ) . The key complication is the following. While w t are in the ball, the distance dist 2 ( w t , S ) shrinks in expectation. However, the iterates may in principle quickly escape the ball B r ( w 0 ) , after which point we lose control on their progress. Thus we must lower bound the probability that the iterates w t remain in the ball. To this end, we will require the following additional assumption.

Assumption 2 (Uniform aiming) . The estimate

The intuition underlying this assumption is as follows. We would like to replace ¯ w in the aiming condition (2.2) by an arbitrary point v ∈ B r ( w 0 ) ∩ S , thereby having a condition of the form ⟨∇L ( w ) , w -v ⟩ ≥ θ · L ( w ) . The difficulty is that this condition may not be true for the main problem we are interested in- training wide neural networks. Instead, it suffices to lower bound the inner product by θ L ( w ) -ρ · dist( w,S ) where ρ is a small constant. This weak condition provably holds for wide neural networks, as we will see in the next section. The following is our main result.

Theorem 2.5 (Convergence on a ball) . Suppose that Assumptions 1 and 2 hold on a ball W = B 3 r ( w 0 ) . Fix constants δ 1 ∈ (0 , 1 3 ) and δ 2 ∈ (0 , 1) , and assume dist 2 ( w 0 , S ) ≤ δ 2 1 r 2 . Fix a stepsize η < θ β and suppose ρ ≤ ( θ -βη ) αr . Then with probability at least 1 -5 δ 1 , all the SGD iterates { w t } t ≥ 0 remain in B r ( w 0 ) . Moreover, with probability at least 1 -5 δ 1 -δ 2 , the estimate dist 2 ( w t , S ) ≤ ε · dist 2 ( w 0 , S ) holds after

Thus as long as ρ is sufficiently small and the initial distance satisfies dist 2 ( w 0 , S ) ≤ δ 2 1 r 2 , with probability at least 1 -5 δ 1 -δ 2 , the iterates remain in B r ( w 0 ) and converge at a fast linear rate O ( 1 δ 2 exp( -tθ 2 α/β )) . While the dependence on δ 2 is logarithmic, the constant δ 1 linearly impacts the initialization region. Section 4 discusses a way to remove this dependence. As explained in Lemma 2.1, both quadratic growth and the initialization quality holds if L is α -PŁ ∗ on the ball B r ( w 0 ) , and r is sufficiently big relative to 1 /α .

3 Consequences for nonlinear least squares and wide neural networks

We next discuss the consequences of the results in the previous sections to nonlinear least squares and training of wide neural networks. To this end, we begin by verifying the aiming (2.2) and uniform aiming (2.4) conditions for nonlinear least squares. The key assumption we will make is that the nonlinear map's Jacobian ∇ F has a small Lipschitz constant in operator norm.

Theorem 3.1. Consider a function L ( w ) = 1 2 ∥ F ( w ) ∥ 2 , where F : R d → R n is C 1 -smooth. Suppose that there is a point w 0 satisfying dist( w 0 , S ) ≤ r and such that on the ball B 2 r ( w 0 ) , the gradient ∇L is β -Lipschitz, the Jacobian ∇ F is L -Lipschitz in the operator norm, and the quadratic growth condition (2.1) holds. Then as long as L ≤ 2 α r √ β , the aiming (2.2) and uniform aiming (2.4) conditions

We next instantiate Theorem 3.1 and Theorem 2.5 for a nonlinear least squares problem arising from fitting a wide neural network. Setting the stage, an l -layer (feedforward) neural network f ( w ; x ) , with parameters w , input x , and linear output layer is defined as follows:

Here, m i is the width (i.e., number of neurons) of i -th layer, α ( i ) ∈ R m i denotes the vector of i -th hidden layer neurons, w := { W (1) , W (2) , . . . , W ( l ) , W ( l +1) } denotes the collection of the parameters (or weights) W ( i ) ∈ R m i × m i -1 of each layer, and σ is the activation function, e.g., sigmoid , tanh , linear activation. We also denote the width of the neural network as m := min i ∈ [ l ] m i , i.e., the minimal width of the hidden layers. The neural network is usually randomly initialized, i.e., each individual parameter is initialized i.i.d. following N (0 , 1) . Henceforth, we assume that the activation functions σ are twice differentiable, L σ -Lipschitz, and β σ -smooth. In what follows, the order notation Ω( · ) and O ( · ) will suppress multiplicative factors of polynomials (up to degree l ) of the constants C , L σ and β σ .

Given a dataset D = { ( x i , y i ) } n i =1 , we fit the neural network by solving the least squares problem

We assume that all the the data inputs x i are bounded, i.e., ∥ x i ∥ ≤ C for some constant C .

Our immediate goal is to verify the assumptions of Theorem 3.1, which are quadratic growth and (uniform) aiming. We begin with the former. Quadratic growth is a consequence of the PŁ-condition. Namely, define the Neural Tangent Kernel K ( w 0 ) = ∇ F ( w 0 ) ∇ F ( w 0 ) ⊤ at the random initial point w 0 ∼ N (0 , I ) and let λ 0 be the minimal eigenvalue of K ( w 0 ) . The value λ 0 has been shown to be positive with high probability in [10, 8]. Specifically, it was shown that, under a mild non-degeneracy condition on the data set, the smallest eigenvalue λ ∞ of NTK of an infinitely wide neural network is positive (see Theorem 3.1 of [10]). Moreover, if the network width satisfies m = Ω( n 2 · 2 O ( l ) λ 2 ∞ log nl ϵ ) , then with probability at least 1 -ϵ the estimate λ 0 > λ ∞ 2 holds [8, Remark E.7]. Of course, this is worst case bound and for our purposes we will only need to ensure that λ 0 is positive. It will also be important to know that ∥ F ( w 0 ) ∥ 2 = O (1) , which indeed occurs with high probability as shown in [18]. To simplify notation, let us lump these two probabilities together and define

Next, we require the following theorem, which shows two fundamental properties on B r ( w 0 ) when the width m is sufficiently large: (1) the function w ↦→ f ( w,x ) is nearly linear and (2) the function L satisfies the PŁ condition with parameter λ 0 / 2 .

Theorem 3.2 (Transition to linearity [25] and the PŁ condition [24]) . Given any radius r > 0 , with probability 1 -p -2 exp( -ml 2 ) -(1 /m ) Θ(ln m ) of initialization w 0 ∼ N (0 , I ) , it holds:

In the same event, as long as the width of the network satisfies m = ˜ Ω ( nr 6 l +2 λ 2 0 ) , the function L is PŁ ∗ on B r ( w 0 ) with parameter λ 0 / 2 .

Note that (3.1) directly implies that the Lipschitz constant of ∇ F is bounded by ˜ O ( r 3 l √ m ) on B r ( w 0 ) , and can therefore be made arbitrarily small. Quadratic growth is now a direct consequence of the PŁ ∗ condition while (uniform) aiming follows from an application of Theorem 3.1.

Theorem 3.3 (Aiming and quadratic growth condition for wide neural network) . With probability at least 1 -p -2 exp( -ml 2 ) -(1 /m ) Θ(ln m ) with respect to the initialization w 0 ∼ N (0 , I ) , as long as

the following are true:

the quadratic growth condition (2.1) holds on B r ( w 0 ) with parameter λ 0 / 2 and the intersection B r ( w 0 ) ∩ S is nonempty,

aiming (2.2) and uniform aiming (2.4) conditions hold in B r ( w 0 ) with θ = 1 and ρ = ˜ O ( r 3 l +2 √ m ) ,

the gradient of each function ℓ i ( w ) ≜ ( f ( w,x i ) -y i ) 2 is β -Lipschitz on B r ( w 0 ) with β = O (1) .

Please see Appendix D for a numerical verification of an estimate of the aiming condition. It remains to deduce convergence guarantees for SGD by applying Theorem 2.5.

Corollary 3.4 (Convergence of SGD for wide neural network) . Fix constants δ 1 ∈ (0 , 1 3 ) , δ 2 ∈ (0 , 1) , ε > 0 and t ∈ N . There is a stepsize η = Θ(1) such that the following is true. With probability at least 1 -p -δ 1 -δ 2 -2 exp( -ml 2 ) -(1 /m ) Θ(ln m ) , as long as

all the SGD iterates { w t } t ≥ 0 remain in B r ( w 0 ) and the estimate dist 2 ( w t , S ) ≤ ε · dist 2 ( w 0 , S ) holds after t ≥ 1 λ 0 log ( 1 εδ 2 ) iterations.

Thus, the width requirements for SGD to converge at a fast linear rate are nearly identical to those for gradient descent [24], with the exception being that the requirement r = Ω ( 1 √ λ 0 ) is strengthened to r = Ω ( 1 δ 1 √ λ 0 ) . That is, the radius r needs to shrink by the probability of failure.

4 Boosting to high probability

A possible unsatisfying feature of Theorems 2.4 and 2.5 and Corollary 3.4 is that the size of the initialization region shrinks with the probability of failure δ 1 . A natural question is whether this requirement may be dropped. Indeed, we will now see how to boost the probability of success to be independent of δ 1 . A first reasonable idea is to simply rerun SGD a few times from an initialization region corresponding to δ 1 = 1 / 2 . Then by Hoeffding's inequality, after very trials, at least a third of them will be successful. The difficulty is to determine which trial was indeed successful. The fact that the solution set is not a singleton rules out strategies based on the geometric median of means [31, p. 243], [29]. Instead, we may try to estimate the function value at each of the returned points. In a classical setting of stochastic optimization, this is a very bad idea because it amounts to mean estimation, which in turn requires O (1 /ε 2 ) samples. The saving grace in the interpolation regime is that ℓ ( w, · ) is a nonnegative function of the samples . While estimating the mean of nonnegative random variables still requires O (1 /ε 2 ) samples, detecting that a nonnegative random variable is large requires very few samples! This basic idea is often called the small ball principle and is the basis for establishing generalization bounds with heavy tailed data [28]. With this in mind, we will require the following mild condition, stipulating that the empirical average 1 m ∑ m i =1 ℓ ( w,z i ) to be lower bounded by L ( w ) with high probability over the iid samples z i .

Assumption 3 (Detecting large values) . Suppose that there exist constants c 1 > 0 and c 2 > 0 such that for any w ∈ R d , integer m ∈ N , and iid samples z 1 , . . . , z m ∼ P , the estimate holds:

Importantly, this condition does not have anything to do with light tails. A standard sufficient condition for Assumption 3 is a small ball property.

Assumption 4 (Small ball) . There exist constants τ > 0 and p ∈ (0 , 1) satisfying

The small ball property simply asserts that ℓ ( w, · ) should not put too much mess on small values relative to its mean L ( w ) . Bernstein's inequality directly shows that Assumption 4 implies Assumption 3. We summarize this observation in the following theorem.

Lemma 4.1. Assumption 4 implies Assumption 3 with c 1 = pτ 2 and c 2 = p 4 .

A valid bound for the small ball probabilities is furnished by the Paley-Zygmund inequality [33]:

Thus if the ratio E ℓ ( w,z ) 2 L ( w ) is bounded by some D > 0 , then the small ball condition holds with p = (1 -τ ) 2 /D where τ ∈ [0 , 1] is arbitrary.

The following lemma shows that under Assumption 3, we may turn any estimation procedure for finding a minimizer of L that succeeds with constant probability into one that succeeds with high probability. The procedure simply draws a small batch of samples z 1 , . . . , z m ∼ P and rejects those trial points w i for which the empirical average 1 m ∑ m j =1 ℓ ( w i , z j ) is too high.

Lemma 4.2 (Rejection sampling) . Let w 1 , . . . , w k be independent random variables satisfying P ( L ( w i ) ≤ ϵ ) ≥ 1 / 2 . For each i = 1 , . . . , k draw m samples z 1 , . . . , z m i.i.d. ∼ P . For any λ > 1 , define admissible indices I = { i ∈ [ k ] : 1 m ∑ m j =1 ℓ ( w i , z j ) ≤ λϵ } . Then with probability 1 -exp( -k 16 ) -k exp( -c 2 m ) -λ -k/ 4 , the set I is nonempty and L ( w i ) ≤ λϵ c 1 for any i ∈ I .

We may now simply combine SGD with rejection sampling to obtain high probability guarantees. Looking at Lemma 4.2, some thought shows that the overhead for high probability guarantees is dominated by c -1 2 . As we saw from the Paley-Zygmond inequality, we always have c -1 2 ≲ D where D upper bounds the ratios E ℓ ( w,z ) 2 L ( w ) . It remains an interesting open question to investigate the scaling of small ball probabilities for overparametrized neural networks.

5 Conclusion

Existing results ensuring convergence of SGD under interpolation and the PŁ condition require the method to use a small stepsize, and therefore converge slowly. In this work we isolated conditions that enable SGD to take a large stepsize and therefore have similar iteration complexity as gradient descent. Consequently, our results align theory better with practice, where large stepsizes are routinely used. Moreover, we argued that these conditions are reasonable because they provably hold when training sufficiently wide feedforward neural networks with a linear output layer.

References

Hedy Attouch, Jérôme Bolte, and Benar Fux Svaiter. 'Convergence of descent methods for semi-algebraic and tame problems: proximal algorithms, forward-backward splitting, and regularized Gauss-Seidel methods'. In: Mathematical Programming 137.1-2 (2013), pp. 91129.

Raef Bassily, Mikhail Belkin, and Siyuan Ma. 'On exponential convergence of sgd in nonconvex over-parametrized learning'. In: arXiv preprint arXiv:1811.02564 (2018).

Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. 'Reconciling modern machinelearning practice and the classical bias-variance trade-off'. In: Proceedings of the National Academy of Sciences 116.32 (2019), pp. 15849-15854.

Samuel Burer and Renato DC Monteiro. 'A nonlinear programming algorithm for solving semidefinite programs via low-rank factorization'. In: Mathematical Programming 95.2 (2003), pp. 329-357.

Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. 'Gradient descent on neural networks typically occurs at the edge of stability'. In: arXiv preprint arXiv:2103.00065 (2021).

Dmitriy Drusvyatskiy, Alexander D Ioffe, and Adrian S Lewis. 'Curves of descent'. In: SIAM Journal on Control and Optimization 53.1 (2015), pp. 114-138.

Dmitriy Drusvyatskiy and Adrian S. Lewis. 'Error Bounds, Quadratic Growth, and Linear Convergence of Proximal Methods'. In: Math. of Oper. Res. 43.3 (2018), pp. 919-948. DOI: 10.1287/moor.2017.0889 . eprint: https://doi.org/10.1287/moor.2017.0889 . URL: https://doi.org/10.1287/moor.2017.0889 .

Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. 'Gradient Descent Finds Global Minima of Deep Neural Networks'. In: International Conference on Machine Learning . 2019, pp. 1675-1685.

Simon S Du, Wei Hu, and Jason D Lee. 'Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced'. In: Advances in neural information processing systems 31 (2018).

Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. 'Gradient Descent Provably Optimizes Over-parameterized Neural Networks'. In: International Conference on Learning Representations . 2018.

Qiang Fu, Dongchu Xu, and Ashia Wilson. 'Accelerated Stochastic Optimization Methods under Quasar-convexity'. In: arXiv preprint arXiv:2305.04736 (2023).

Guillaume Garrigos and Robert M Gower. 'Handbook of convergence theorems for (stochastic) gradient methods'. In: arXiv preprint arXiv:2301.11235 (2023).

Robert Gower, Othmane Sebbouh, and Nicolas Loizou. 'Sgd for structured nonconvex functions: Learning rates, minibatching and interpolation'. In: International Conference on Artificial Intelligence and Statistics . PMLR. 2021, pp. 1315-1323.

Moritz Hardt, Tengyu Ma, and Benjamin Recht. 'Gradient descent learns linear dynamical systems'. In: arXiv preprint arXiv:1609.05191 (2016).

Oliver Hinder, Aaron Sidford, and Nimit Sohoni. 'Near-Optimal Methods for Minimizing Star-Convex Functions and Beyond'. In: Proceedings of Thirty Third Conference on Learning Theory . Ed. by Jacob Abernethy and Shivani Agarwal. Vol. 125. Proceedings of Machine Learning Research. PMLR, Sept. 2020, pp. 1894-1938. URL: https://proceedings.mlr. press/v125/hinder20a.html .

Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. 'Gpipe: Efficient training of giant neural networks using pipeline parallelism'. In: Advances in neural information processing systems 32 (2019).

Aleksandr Davidovich Ioffe. 'Metric regularity and subdifferential calculus'. In: Russian Mathematical Surveys 55.3 (2000), p. 501.

Arthur Jacot, Franck Gabriel, and Clément Hongler. 'Neural tangent kernel: Convergence and generalization in neural networks'. In: Advances in neural information processing systems 31 (2018).

Stanisław Jastrz˛ ebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. 'Three factors influencing minima in sgd'. In: arXiv preprint arXiv:1711.04623 (2017).

Jikai Jin. 'On the convergence of first order methods for quasar-convex optimization'. In: arXiv preprint arXiv:2010.04937 (2020).

Hamed Karimi, Julie Nutini, and Mark Schmidt. 'Linear convergence of gradient and proximalgradient methods under the polyak-łojasiewicz condition'. In: Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16 . Springer. 2016, pp. 795-811.

Ahmed Khaled and Peter Richtárik. 'Better theory for SGD in the nonconvex world'. In: arXiv preprint arXiv:2002.03329 (2020).

Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. 'Big transfer (bit): General visual representation learning'. In: Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part V 16 . Springer. 2020, pp. 491-507.

Chaoyue Liu, Libin Zhu, and Mikhail Belkin. 'Loss landscapes and optimization in overparameterized non-linear systems and neural networks'. In: Applied and Computational Harmonic Analysis 59 (2022), pp. 85-116.

Chaoyue Liu, Libin Zhu, and Misha Belkin. 'On the linearity of large non-linear models: when and why the tangent kernel is constant'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 15954-15964.

Stanislaw Lojasiewicz. 'A topological property of real analytic subsets'. In: Coll. du CNRS, Les équations aux dérivées partielles 117.87-89 (1963), p. 2.

Z.-Q. Luo and P. Tseng. 'Error bounds and convergence analysis of feasible descent methods: a general approach'. In: Annals of Operations Research 46.1 (Mar. 1993), pp. 157-178. ISSN: 1572-9338. DOI: 10.1007/BF02096261 . URL: https://doi.org/10.1007/BF02096261 .

Shahar Mendelson. 'Learning without concentration'. In: Journal of the ACM (JACM) 62.3 (2015), pp. 1-25.

STANISLAV MINSKER. 'Geometric median and robust estimation in Banach spaces'. In: Bernoulli 21.4 (2015), pp. 2308-2335. ISSN: 13507265, 15739759. URL: http://www.jstor. org/stable/43590532 (visited on 05/14/2023).

Ion Necoara, Yu Nesterov, and Francois Glineur. 'Linear convergence of first order methods for non-strongly convex optimization'. In: Mathematical Programming 175 (2019), pp. 69107.

Arkadij Semenoviˇ c Nemirovskij and David Borisovich Yudin. 'Problem complexity and method efficiency in optimization'. In: (1983).

Samet Oymak and Mahdi Soltanolkotabi. 'Overparameterized nonlinear learning: Gradient descent takes the shortest path?' In: International Conference on Machine Learning . PMLR. 2019, pp. 4951-4960.

Raymond EAC Paley and Antoni Zygmund. 'A note on analytic functions in the unit circle'. In: Mathematical Proceedings of the Cambridge Philosophical Society . Vol. 28. 3. Cambridge University Press. 1932, pp. 266-272.

B. T. Poljak. 'Gradient methods for minimizing functionals'. In: Ž. Vyˇ cisl. Mat i Mat. Fiz. 3 (1963), pp. 643-653. ISSN: 0044-4669.

Quentin Rebjock and Nicolas Boumal. 'Fast convergence to non-isolated minima: four equivalent conditions for C 2 functions'. In: arXiv preprint arXiv:2303.00096 (2023).

Karthik Abinav Sankararaman, Soham De, Zheng Xu, W Ronny Huang, and Tom Goldstein. 'The impact of neural network overparameterization on gradient confusion and stochastic gradient descent'. In: International conference on machine learning . PMLR. 2020, pp. 84698479.

Mingxing Tan and Quoc Le. 'Efficientnet: Rethinking model scaling for convolutional neural networks'. In: International conference on machine learning . PMLR. 2019, pp. 6105-6114.

Sharan Vaswani, Francis Bach, and Mark Schmidt. 'Fast and faster convergence of sgd for overparameterized models and an accelerated perceptron'. In: The 22nd international conference on artificial intelligence and statistics . PMLR. 2019, pp. 1195-1204.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. 'Understanding deep learning (still) requires rethinking generalization'. In: Communications of the ACM 64.3 (2021), pp. 107-115.

A Missing proofs

A.1 Proof of Theorem 2.2

The proof relies on the following elementary lemma.

Lemma A.1 (Aiming for smooth functions) . Let L : R → R be a differentiable function. Fix two points w, ¯ w ∈ R d satisfying ∇L ( ¯ w ) = 0 and L ( ¯ w ) = 0 . Suppose that the Hessian ∇ 2 L exists and is L -Lipschitz continuous on the segment [ w, ¯ w ] . Then we have

d

Proof. Define the function g ( t ) = L ( w + t ( ¯ w -w )) . The theorem is evidently equivalent to

In order to establish this estimate, we first note that g ′′ is Lipschitz continuous with constant ˆ L := L ∥ w -¯ w ∥ 3 , as follows from a quick computation. Taylor's theorem with remainder applied to g and g ′ , respectively, then gives

where | E 1 | ≤ ˆ L 6 and | E 2 | ≤ ˆ L 2 . Combining the two estimates yields

thereby completing the proof.

Turning to the proof of Theorem 2.2, an application of Lemma A.1 guarantees

for all w ∈ S r and any ¯ w ∈ proj S ( w ) . If the quadratic growth condition (3) is satisfied on the tube S r , then we have the upper bound 5 Lr 6 ∥ w -¯ w ∥ 2 ≤ 5 Lr 3 α L ( w ) . Theorem 2.2 follows. It remains to prove Lemma A.1.

A.2 Proof of Lemma 2.1

In this section, we verify the classical result that the PŁ ∗ condition on a ball implies quadratic growth. We begin with the following lemma estimating the distance of a single point to sublevel set of a function; this result is a special instance of the descent principle [6, Lemma 2.5], whose roots can be traced back to [17, Basic Lemma, Chapter 1].

)

Lemma A.2 (Descent principle) . Fix a differentiable function L : R d → [0 , ∞ ) and a ball B r ( w 0 and define S = { w : L ( w ) = 0 } . Suppose that L satisfies the PŁ ∗ condition on B r ( w 0 ) with parameter α . Then as long as L ( w 0 ) < 1 2 r 2 α , the intersection S ∩ B r ( w 0 ) is nonempty and the estimate holds:

Proof. Define the function f ( w ) = √ L ( w ) and observe that for any w ∈ B r ( w 0 ) with L ( w ) > 0 we have ∥∇ f ( w ) ∥ 2 = ∥∇L ( w ) ∥ 2 4 L ( w ) ≥ α 2 . Therefore an application of the descent principle [6, Lemma 2.5] implies that the set S ∩ B r ( w 0 ) is nonempty and the estimate dist( w 0 , S ) ≤ 1 √ α/ 2 · f ( w 0 ) holds. Squaring both sides completes the proof.

We may now complete the proof of Lemma 2.1 by extending from a single point w 0 to a neighborhood of w 0 as follows. First, Lemma A.2 ensures that B r ( w 0 ) intersects S at some point ¯ w 0 and the inequality L ( w 0 ) ≥ α 2 dist 2 ( w 0 , S ) holds. Fix a point w ∈ B r ( w 0 ) . Then clearly L satisfies the PŁ ∗ condition on B r ( w ) with parameter α . Let us now consider two cases: L ( w ) < 1 2 r 2 α and L ( w ) ≥ 1 2 r 2 α . In the former case, Lemma A.2 implies the claimed estimate L ( w ) ≥ α 2 · dist 2 ( w,S ) . In the remaining case L ( w ) ≥ 1 2 r 2 α , we compute

Rearranging completes the proof of Lemma 2.1 .

A.3 Proof of Lemma 2.3

Fix a point w ∈ S r and let ¯ w ∈ proj S ( w ) be a point satisfying the aiming condition (2.2). Observe that Lipschitz continuity of ∇ ℓ ( · , z ) and interpolation ensures

Therefore for every τ ∈ [0 , 1 /β ] , the point w τ := w -τ ∇ ℓ ( w,z ) satisfies

Therefore the gradient ∇ ℓ ( · , z ) is β -Lipschitz on the entire line segment { w τ : 0 ≤ τ ≤ 1 /β } . The descent lemma therefore guarantees ∥∇ ℓ ( w,z ) ∥ 2 ≤ 2 βℓ ( w,z ) . Therefore upon taking expectations we obtain the second moment bound: E ∥∇ ℓ ( w,z ) ∥ 2 ≤ 2 β L ( w ) . Next, we compute

where (A.1) follows from (2.2) while (A.2) follows from (2.1). The proof is complete.

A.4 Proof of Theorem 2.4

Define the stopping time τ = inf { t ≥ 1 : w t / ∈ S r } and set U t = dist 2 ( w t , S ) . Note that we may equivalently write τ = inf { t ≥ 1 : U t > r 2 } . Now, multiplying (2.3) through by 1 τ>t yields the estimate

An application of Theorem B.2 with q = (1 -αη ( θ -βη )) completes the proof.

A.5 Proof of Theorem 2.5

We begin with the following simple lemma that bounds the second moment of the gradient estimator.

Lemma A.3. For any point w ∈ B r ( w 0 ) , we have E ∥ ℓ ( w,z ) ∥ 2 ≤ 2 β · L ( w ) .

Proof. Let ¯ w 0 ∈ proj S ( w 0 ) be arbitrary. Observe that Lipschitz continuity of ∇ ℓ ( · , z ) on B 3 r ( w 0 ) and interpolation ensure

Therefore for every τ ∈ [0 , 1 /β ] , the point w τ := w -τ ∇ ℓ ( w,z ) satisfies

Therefore the gradient ∇ ℓ ( · , z ) is β -Lipschitz on the entire line segment { w τ : 0 ≤ τ ≤ 1 /β } . The descent lemma therefore guarantees ∥∇ ℓ ( w,z ) ∥ 2 ≤ 2 βℓ ( w,z ) . Taking the expectation of both sides completes the proof.

Next we prove the following lemma that simultaneously estimates (1) one step progress of the iterates towards S and (2) how far the iterates move away from the the center w 0 .

Lemma A.4. Fix a point w ∈ B r/ 3 ( w 0 ) and choose ¯ w 0 ∈ proj S ( w 0 ) ∩ B r/ 3 ( w 0 ) . Assume η ≤ θ β and define w + = w -η ∇ ℓ ( w,z ) where z ∼ P . Then the following estimates hold:

Proof. Fix any point v ∈ S and observe that

Taking the expectation with respect to z and using Lemma A.3, we deduce

We will use this estimate multiple times for different vectors v .

Choose any point ¯ w ∈ proj S ( w ) satisfying the aiming condition (2.2). Define U := dist 2 ( w,S ) and U + := dist 2 ( w + , S ) . Then setting v := ¯ w in (A.6), we deduce

Using the aiming condition (2.2) we therefore deduce

where the last inequality follows from quadratic growth. This establishes (A.4).

Next, set v := ¯ w 0 in (A.6). Defining V = ∥ w -¯ w 0 ∥ 2 , V + = ∥ w + -¯ w 0 ∥ 2 , and using Assumption 2 we therefore deduce

where the last estimate follows from the inequality βη ≤ θ . This establishes (A.5).

For each t ≥ 0 , let 1 t denote the indicator of the event E t := { w 0 , . . . , w t ∈ B r/ 3 ( w 0 ) } . Define the random variables U t = dist 2 ( w t , S ) and V t = ∥ w t -¯ w 0 ∥ 2 . We may now multiply (A.4) by 1 E t . Noting that 1 E t +1 ≤ 1 E t we may iterate the bound yielding

Similarly, multiplying (A.4) by 1 E t and using (A.7), we deduce

Iterating the recursion gives

We now lower bound the probability of escaping from the ball. Note that within event E t , we have

Therefore, as desired.

A.6 Proof of Theorem 3.1

We first prove the following lemma, which does not require quadratic growth and relies on the Lipschitz continuity of the Jacobian ∇L .

Lemma A.5. Consider a function L ( w ) = 1 2 ∥ F ( w ) ∥ 2 where F : R d → R n is C 1 -smooth and the Jacobian ∇ F is L -Lipschitz on the ball B r ( w 0 ) . Then the following estimates

hold for all w,u,v ∈ B r ( w 0 ) satisfying F ( u ) = F ( v ) = 0 .

Proof. Fix any w ∈ B r ( w 0 ) and u, v ∈ B r ( w 0 ) satisfying F ( u ) = F ( v ) = 0 . We first prove (A.11). To this end, we compute

where the last estimate follows from the Cauchy-Schwarz inequality. Next, the fundamental theorem of calculus and Lipschitz continuity of ∇ F yields

where (A.9) follows from Markov's inequality and (A.10) follows from (A.8).

Next, we estimate the probability that U t remains small within the event E t . To that end, let define the constant C t := (1 -η ( θ -βη ) α ) t U 0 /δ 2 . Then Markov's inequality yields

Finally, we unconditionally bound the probability that U t remains small:

where ∥ E ∥ ≤ L ∥ v -u ∥ ( ∥ u -w ∥ + ∥ v -u ∥ ) ≤ 8 r 2 L . Thus we have proved (A.11). In order to see (A.12), we compute

where the last inequality follows from Cauchy-Schwarz and Lipschtiz continuity of the Jacobian ∇ F . Thus (A.12) holds.

Lemma A.5 quickly yields the following corollary.

Corollary A.6. Consider a function L ( w ) = 1 2 ∥ F ( w ) ∥ 2 where F : R d → R n is C 1 -smooth and the Jacobian ∇ F is L -Lipschitz on the ball B r ( w 0 ) . Suppose moreover that

where S = { w : F ( w ) = 0 } . Then the estimates

hold for all w ∈ B r ( w 0 ) , v ∈ B r ( w 0 ) ∩ S , and ¯ w ∈ B r ( w 0 ) ∩ proj S ( w ) .

Proof. We will apply Lemma A.5. Setting u := ¯ w in Lemma A.12 and using quadratic growth yields the estimate:

which establishes (A.13). Adding this estimate to (A.12) with u = ¯ w yields (A.14).

Theorem 3.1 follows immediately from the corollary. Indeed, note that let ¯ w ∈ proj S ( w ) satisfies ∥ w -¯ w ∥ ≤ 2 r . Therefore we may apply Corollary A.6 and use the estimate L ( w ) ≤ β 2 ∥ w -¯ w ∥ 2 .

A.7 Proof of Theorem 3.3

Theorem 3.2 implies that with probability at least 1 -p -2 exp( -ml 2 ) -(1 /m ) Θ(ln m ) , the Jacobian of ∇ F is Lipschitz with constant L = ˜ O ( r 3 l √ m ) on B r ( w 0 ) and the PŁ-condition holds on B 2 r ( w 0 ) with parameter λ 0 / 2 . Using the assumption r = Ω( 1 √ λ 0 ) , we may apply Lemma 2.1 to deduce quadratic growth. Next, we will apply Theorem 3.1 in order to deduce (uniform) aiming with parameters θ = 1 . To this end, it suffices to ensure 2 rL √ β/λ 0 ≤ 1 , which follows from the prerequisite assumption m = Ω( r 6 l +2 α 2 ) .

To see the last claim, for any w ∈ B r ( w 0 ) we compute the Hessian

Therefore,

Setting ¯ w 0 ∈ proj S ( w 0 ) ∩ B r ( w 0 ) , observe that

where the last inequality follows from transition to linearity (3.1). Thus f ( · , x i ) is Lipschitz continuous on the ball B r ( w 0 ) . Therefore, we may also bound

Returning to (A.15), we thus have ∥∇ 2 ℓ i ( w ) ∥ op = ˜ O ( r 6 L +2 m ) . Taking into account that m = ˜ Ω ( nr 6 l +2 λ 2 0 ) we deduce r 6 L +2 m = O (1) thereby completing the proof.

A.8 Proof of Corollary 3.4

The goal is to apply Theorem 2.5. To this end, we begin by applying Theorem 3.3. We may then be sure that with high probability Assumptions 1 and 2 hold on a ball B r ( w 0 ) with α = λ 0 / 2 , θ = 1 , and ρ = ˜ O ( r 3 l +2 √ m ) . We may then choose η = 1 2 β = Θ(1) . In order to make sure that ρ ≤ ( θ -βη ) αr , it suffices to be in the regime m = Ω( r 6 l +2 λ 2 0 ) . Finally it remains to ensure that dist 2 ( w 0 , S ) ≤ δ 2 1 r 2 . To do so, using quadratic growth, we have dist 2 ( w 0 , S ) ≤ 4 λ 0 L ( w 0 ) = O ( 1 λ 0 ) . Thus it suffices to let r = Ω( 1 δ 1 √ λ 0 ) . An application of Theorem 2.5 completes the proof.

A.9 Proof of Lemma 4.2

Define the set J = { i ∈ [ k ] : f ( w i ) ≤ ϵ } . Then Hoeffding inequality ensures that the inequality |J | ≥ k 4 holds with probability at least 1 -exp( -k/ 16) . Conditioned on this event, consider any index i ∈ [ k ] satisfying L ( w i ) > λϵ c 1 . Then from (3), we know that with probability at least 1 -exp( -c 2 m ) , we have

and therefore i / ∈ I . Let us now estimate the probability that I is nonempty conditioned on |J | ≥ k 4 . To this end, by Markov's inequality for any i ∈ J , we have

Consequently, the probability that the set I is empty is at most λ -k/ 4 .

B Auxiliary results on stopping times

Theorem B.1 (Stopping time argument) . Let { U t } t ≥ 0 be a sequence of nonnegative random variables and define the stopping time τ = inf { t ≥ 0 : U t > u } for some constant u > 0 . Suppose that

Then the estimate P [ τ ≤ t ] ≤ E U 0 + ∑ t -1 i =0 ζ i u holds for all t ≥ 1 .

Proof. Observe that by Markov's inequality, we have

Let us therefore bound the expectation of the stopped random variable V t := U t ∧ τ . Letting E t [ · ] denote the conditional expectation E [ · | U 1: t ] , we successively compute

Taking the expectation with respect to U 1: t , applying the tower rule, and iterating the recursion we deduce E [ V t ] ≤ E [ U 0 ] + ∑ t -1 i =0 ζ i , thereby completing the proof.

Theorem B.2 (Stopping time with contractions) . Let { U t } t ≥ 0 be a sequence of nonnegative random variables and define the stopping time τ = inf { t ≥ 0 : U t > u } for some constant u > 0 . Suppose that there exists q ∈ (0 , 1) such that

Then as long as E U 0 ≤ δ 1 u , the event { τ = ∞} occurs with probability at least 1 -δ 1 . Moreover, with probability at least 1 -δ 1 -δ 2 , the estimate U t ≤ εU 0 holds after t ≥ 1 1 -q log ( 1 δ 2 ε ) iterations.

Proof. Define the stopping time τ = inf { t ≥ 1 : U t > u } . An application of Lemma B.1 therefore implies P [ τ < t ] ≤ E U 0 u ≤ δ 1 . Taking the limit as t →∞ , we deduce that the event { τ = ∞} occurs with probability at least 1 -δ 1 . Next taking the expectation with respect to U 1: t in (B.2) and applying the tower rule gives

Taking into account that 1 τ>t ≥ 1 τ>t +1 we may iterate the recursion thereby yielding

Now, setting ε ′ := εU 0 , Markov's inequality yields

Finally observe

Therefore we deduce

as claimed.

C A bad example

Lemma C.1. Consider the objective L ( x, y ) = 1 2 ( y -ax 2 ) 2 for any a > 0 . Then L is C ∞ near (0 , 0) and the function L satisfies the PŁ inequality (PŁ) with constant α = 1 . Therefore L satisfies the aiming condition on some neighborhood of the origin. However, for any neighborhood U of the origin, the function L is not quasar-convex on U relative to any point ( x, ax 2 ) ∈ U .

Proof. To see the validity of the PŁ-condition, observe that

It follows immediately from Lemma 2.2 that L satisfies the aiming condition (2.2) on some neighborhood of the origin. Next we verify the failure of quasar convexity. Consider an arbitrary point ( x, ax 2 ) on the parabola. For any ( u, v ) ∈ R 2 , we compute

In particular, setting z γ = (0 , γ ) we obtain

̸

Note the right hand side is negative for any 0 < γ < ax 2 . Thus, letting γ > 0 tend to zero, we deduce that L is not quasar-convex relative to ( x, ax 2 ) ∈ U with x = 0 . Let us consider now the setting x = 0 . Then we compute

The right side is negative if au 2 < v < 2 au 2 and therefore ∇L is not quasar-convex relative to (0 , 0) on any neighborhood of (0 , 0) .

MNIST

0

500

1000

1500

2000

2500

3000

3500

4000

iterations

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

2.00

Aiming coefficient

0

500

1000

1500

2000

2500

3000

3500

4000

iterations

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Aiming coefficient

CIFAR-10

D Numerical verification of aiming condition on neural networks

In principle, the aiming condition is difficult to verify, as it involves proj S ( w ) - the nearest point of w to the solution set S , which we have no way of computing. Instead, we replace proj S ( w ) by the last iterate ¯ w of an SGD run, and compute an estimate of the aiming condition coefficient by the quotients ˆ θ t ≜ ⟨∇ L ( w t ) , w t -¯ w ⟩ /L ( w t ) along the iterate path. Note that this quotient is not the same quotient as would appear in the definition of quasar-convexity because ¯ w is a random point that depends on the iterate path taken by SGD.

Settings: We conduct the experiments on two datasets, MNIST and CIFAR-10. For MNIST, we trained a 3 -hidden layer fully-connected neural network (width = 1000 ). First, we train the network until convergence (loss function value smaller than 10 -4 ), and record the parameters of this trained network as the optimal solution ¯ w . For CIFAR-10, we train a two-layer CNN for 4000 iterations, and record the final parameter setting as ¯ w . Then, for each experiment, we took a second run, and at each iteration w t we compute the estimate of the aiming condition coefficient using the equation mentioned above.

Results: Figure 4 show the plots of against the iteration number. In both cases, we observe that the estimate stays positive, which suggests that aiming condition holds.

E Aiming Condition vs. quasar-convexity

Following [15], a nonnegative function L is called quasar-convex relative to a point ¯ w with L ( ¯ w ) = 0 if the inequality ⟨∇L ( w ) , w -¯ w ⟩ ≥ θ · L ( w ) holds for all w near ¯ w . Aiming, in contrast, stipulates the analogous inequality but with ¯ w crucially replaced by proj S ( w ) , where S is the set of zero loss solutions. We argue that in the interpolation setting, quasar-convexity is a very restrictive condition because it requires the set of zero-loss solutions to be an affine subspace, which is usually not the case (see for example [24]). To see this, note that quasar-convexity implies that S is star-convex relative to ¯ w (Observation 3 in [15]). That is, there exists ϵ > 0 such that for any w ∈ B ϵ ( ¯ w ) ∩ S , the line segment joining w and ¯ w is fully contained in S . This is a very strong conclusion. For example, a smooth-manifold S is star convex near ¯ w if and only if it coincides with a linear subspace around ¯ w . The PL-condition in turn implies by Theorem 2.16 in [35] that S is a smooth manifold. Therefore, if loss functions for training wide neural networks were quasar-convex, then the set of interpolating solutions would form a linear space on any compact set, which is certainly not true. In contrast, as we have proved, any C 3 -smooth function satisfying PL automatically satisfies the aiming inequality near a point ¯ w ∈ S . In summary, quasar-convexity does not hold when training wide neural networks (since otherwise it would imply that interpolating solutions form a linear subspace), while aiming provably holds-one of our main results.

Peer reviews:
comments: SUMMARY:
This work proposes a new condition called the aiming condition, which looks similar to quasar-convexity but provides fundamentally different convergence guarantees for SGD. Under the aiming condition, along with several other regularity conditions, SGD can achieve the same sample complexity as GD. It is then shown that wide neural networks enjoys this property with high probability.

STRENGTHS:
- The aiming condition is impressive as a condition for SGD to achieve the same sample complexity as GD.
- The presentation of the results is clear. This work embodies a lot of results, which could potentially make the paper hard to follow. However, the introduction part provides a clear roadmap.

WEAKNESSES:
- I am expecting more discussions on the comparison of the aiming condition against existing conditions (e.g. quasar-convexity).This involves two aspects: 1. I would like to see some more examples where the aiming condition is satisfies but quasar-convexity does not hold; 2. It would be interesting to provide some intuition about why quasar-convexity cannot provide a similar result.
- This work does not seem to be enough strong technically.
- This work has some minor problems presentations. For example, the same sentence appears twice in Line 62 and Line 160; Furthermore, in Line 251, the radius needs to be scaled up rather than shrunk.

QUESTIONS:
I am mainly curious about the comparison of quasar-convexity and the aiming condition (see "Weaknesses" above).

LIMITATIONS:
This work, with a theorical nature, does not have potential negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work shows a regularization condition for SGD in the interpolation regime which allows it to have same fast linear convergence rate as deterministic gradient descent. Hence, the theory presented in this paper supports the practical observation that with the same (large) learning rate, mini-batch SGD has almost the same convergence rate of SGD. This goes against the traditional approaches for SGD convergence under PL condition that require SGD to have smaller step-size than GD, hence making SGD's convergence slower in theory. The authors here present the conditions in which SGD have similar iteration complexity as GD.\n\nSTRENGTHS:\n1) The main strength of the work lies in it's novelty. The motivation of the problem seems clear from the introduction. \n2) This is an important problem since it is crucial to have a large learning rate for SGD as it helps improving generalization. Hence, the current theories that analyze SGD under PL-condition and requires it to have a smaller l.r, may suffer from having good generalization. The theory presented by the authors allows the same large learning rate to be used as GD. Hence, the variance of SGD coupled with the alowable large learning rate can boost generalization.\n\nWEAKNESSES:\n1) Looking at the theorems, I still feel the assumptions made are too strong too hold in a non-convex landscape especially Aiming. For theorem-1.2, which consdiers a non-convex loss landscape on $w$, it is unclear how aiming can hold on a ball of radius $r$ and it's relation to the curvature at that point (given by minimal eigenvalue of NTK). The authors make sure that the iterates remain inside a ball $B_{r}(w_{0})$ near the minima but the radius depends on the eigenvalue of NTK at initializaiton! This makes a very strong assumption that the initial point is set to be very close to the true minima $w_{0}$. More so, if initialization is made far away from the minima, there is a high probability that aiming does not hold. \n\n2) From the theorems, it is unclear that how the stochasticity from the mini-batch gradients effect the convergence or what is it's effect in ensuring that the iterates don't escpae $B_{r}(w_{0})$. For theorem-2.3, I assume that it will take longer iterations if the variance in mini-batch gradient is high. It is not very clear from the inequality as how the relation is. \n\nA minor correction:\n\n3) On line-42, it might be better to cite a different source such as [1] to refer to the generalization effect of large learning rate, as the authors of \"Edge of stability\" themselves mention that they don't claim the generalization effect of large learning rate but rather focus on the stability effect. See https://youtu.be/6xeh6gfESuc?t=3016\n\n[1] Li, Yuanzhi, Colin Wei, and Tengyu Ma. \"Towards explaining the regularization effect of initial large learning rate in training neural networks.\" Advances in Neural Information Processing Systems 32 (2019).\n\nQUESTIONS:\nSee points 1 and 2.\n\nLIMITATIONS:\nThe main limitation of this work is a strong assumption on the locality of the SGD analysis. Some further clarification would be helpful.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents a set of conditions under which the convergence rate of SGD with large step size is similar to that of gradient descent (the deterministic setting). This is in contrast to prior work, where the convergence of over-parameterized SGD under PL condition require small step size, and converge slowly as a result of this small step size.\n\nSTRENGTHS:\n1. The paper is well-motivated and clearly organized. This work provides a theoretical result that improve previous analyses of the linear convergence rate of overparameterized SGD.\n2. The paper thoroughly compares its result to prior works mentioned. \n3. The paper attempts verifies assumptions that it makes.\n\nWEAKNESSES:\n1. The paper states that its goal is to improve stepsize selection and convergence rate of SGD for nonconvex problems under certain conditions (see lines 133-134). However, the paper does not directly do this.\n2. This work claims to be the first to present an analysis of the convergence rate of over-parameterized SGD with large step size which is similar the convergence rate attained by gradient descent,. However, prior work touches upon this as well (see “Observations in simplified settings” in The Impact of Neural Network Overparameterization on Gradient Confusion and Stochastic Gradient Descent by Sankararaman, De, Xu, Huang, and Goldstein). Perhaps it is worth connecting both works to further demonstrate this paper’s contribution. (Please correct me if I am mistaken.)\n3. While the contribution of this paper is predominantly theoretical, I believe that this work could benefit from some empirical evaluation as well. Current experiments are fairly limited (MNIST is the only dataset used and is easy to learn. Additionally, model architecture is also quite limited. Using other datasets e.g. CIFAR10/100 or running the experiments for networks with varying widths might be useful).\n\nQUESTIONS:\nThere is no figure 2\n\nLIMITATIONS:\nThe authors have described the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper studies the convergence of SGD with large step size. It is shown that under some regularity conditions, SGD enjoys a fast linear convergence rate, both in expectation and with high probability. These results can be applied to show fast convergence of SGD for wide enough feed-forward networks.

STRENGTHS:
The paper is well organized and easy to follow, and especially the authors have compared in detail the assumptions/results with those in the existing papers.

It is indeed an important question to study the convergence of SGD with large step sizes, and the results in this paper seem novel and interesting.

WEAKNESSES:
1. The main application of the results is for wide enough neural networks in the NTK regime, which seems restrictive.
2. It would be helpful if the authors can highlight the technical novelties.
3. It is discussed around line 40-42 that large batch sizes could be beneficial for generalization. However, the current results don't seem to have such implications. In particular, I don't think the regime studied in the current paper is a case of the "edge of stability" regime. 
4. Is it possible to verify the proposed conditions for common neural network models, at least empirically? Otherwise it's not clear if the developed theory reflects what happens in practice.

QUESTIONS:
Please see the question above.

LIMITATIONS:
The authors have adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09