Aiming towards the minimizers: fast convergence of SGD for overparametrized problems
Authors: Chaoyue Liu
Dmitriy Drusvyatskiy
Misha Belkin
Damek Davis
Yian Ma
Venue: NeurIPS 2023 poster (2023)
Keywords: Polyak-Lojasiewicz condition
SGD
interpolation
fast convergence
TL;DR: We develop mathematial conditions under which SGD have as fast convergence rate as full batch GD in terms of number of iterations, and show that these condtions are satisfied by wide neural networks

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
Chaoyue Liu", Dmitriy Drusvyatskiy", Yian Ma", Damek Davis" "** and Mikhail Belkin"
Halicioglu Data Science Institute, University of California San Diego Mathematics Department, University of Washington
School Operations Research and Information Engineering, Cornell University
Abstract
Modern machine learning paradigms, such as deep learning, occur in close the interpolation regime, wherein the number model parameters much larger than the number of data samples. In this work, we propose regularity condition within the interpolation regime which endows the stochastic gradient method with the worst-case iteration complexity as the deterministic gradient method, while using only a single sampled gradient (or minibatch) in each iteration. 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 overparam-
eterized models, notably deep neural networks, observed data; e.g. [39 37 the number of parameters the model is much greater than the number of data samples, thereby re- sulting 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 overparameetrizzatio actually improves generalization. The stochastic gradient method is the workhorse algorithm for fitting overparametrized models to observed data and understanding its performance an active area of research. The goal of this paper is obtain improved convergence guarantees for SGD in the interpolation regime that better align with its performance in practice.
16 23 In such settings,
Classical optimization literature emphasizes conditions akin strong convexity the phenomena underlying rapid convergence of numerical methods. In contrast, interpolation problems are almost convex, 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] invariant under orthogonal transformations while ReLU neural networks are invariant under rebalancing of adjacent weight matrices (9. The Polyak- Lojasiewiez (PL) 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 Pt. condition at línear rate (fexp(-t/<)) where x. is the condition number of the function! In contrast, convergence guarantees for the stochastic gradient method under the-the predominant algorithm in practice- are much less satisfactory. Indeed, all known results require SGD to take
In particular, for smooth function with L-pischhit gradient satisfying a PL inequality with constant a, the condition number is x = 3/a.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

16.1
(00--21-000 batch (bstch
betth une une-1000
batch 2io-22000
202
borning
Se
502 1000
3300
2000
2300
3000
2000
6000
20020
90000
60000
#000C
1020000
Figure 1: Convergence plot of SGD when training fully connected neural network with 3 hidden layers and 1000 neurons in each MNIST (left) and a ResNet-28 on CIFAR- (right). MNIST has 60k images and the convergence behavior stabilizes with small batchsize of m = 64; CIFAR-10 has 60k images and the convergence behavior stabilizes with batchsize of m 512, although early on in training (first 20k iterations) convergence behavior looks identical for all batchsizes.
shorter steps than gradient descent to converge at all (2 12]. with the disparity between the two depending on the condition number K. The use of the small stepsize directly translates into slow rate of convergence. This requirement is in direct contrast to practice, where large step-sizes are routinely used. As concrete illustration of the disparity between theory and practice. Figure 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 7j = 0.1. Moreover, experimentally, we have verified that the interval of stepsizes leading convergence for SGD is comparable to that of GD; indeed, the off only by a factor of 20. Using large stepsizes also has important consequences for generalization. Namely, recent works [T9 suggest that large stepsizes bias the iterates towards solutions that generalize hetter to unseen data. Therefore understanding the dynamics of SGD with large stepsizes an important research direction. The contribution of our work as follows.
In this work, we reghlighttegularity conditions that endow SGD with fast linear rate of convergence cxp(-t/x) 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 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:
[(w) 2 a (w,
Vw E B, (wo);
(QG)
2
where (wo) is ball of radius around the initial point wo- This condition is standard in the optimization literature and implied for example by the Pt.-inequality holding on the ball (wo): see Section A.2] Secondly, and most importantly, we assume there there exist constants 0, p 0
satisfying the aiming condition
(VL(w), w projs(w)) 2 0. C(w)
Yw B,
(Aiming)
Here, projg (w) denotes a nearest point in S to w and dist(w, S) denotes the distance from to S. The aiming condition ensures that the negative gradient - (L((w) points towards the sense that -V[(w) correlated nontrivially with the direction projs(w) condition appears similar to quasar-convexitg introduced in (14) and further studied in [15 22 20). Namely a function L is quasar-conee relative to fixed point E S if the estimate AAiming holds with projs (w)) replaced by w. Although the distinction between aiming and quasar-convexity may
At first sight, the aiming
'If projs (w) is not singleton, (w) in the expression should be replaced with any element w from projs(w).

---

appear mild, it is significant. As concrete example, consider the function [(x, y) - (v )2²) for any a > 0. is straightforward 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 (xt, ax²) U: see Section C More generally, we show that (Aiming) holds automatically for any C smooth function L satisfying (QG) locally around the solution set. indeed, we may shrink the neighborhood to ensure that I 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, (wo), they converge to S at a fast linear rate ((exp(-
(oxxpp(tt2///0 ttttttb2//p/ /B)) with high probability. Theorem 1.1 (Informal). Consider minimizing the function [(w) = El(w, 2), where the losses CG 2) are nonnegative and have B-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 We remain in B, B,(wa) (wa), they converge to S at linear rate O(xx(-laaf /3)) with high probability.
The proof of the theorem is short and elementary. The downside is that the conclusion of the theorem conditional on the iterates remaining (wo)). Ideally, one would like to estimate this probability as a function of the problem parameters. With this in mind, we show that for 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. 1 xf) yi)², where fww,-) is a fully connected neural network with I hidden layers and a linear output layer. Let doo 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 Wy [0,, I) with 0 = and a = X2 as long as the network width m satisfies m = 61r (X²x). in addition r = 91(1 (1/1////)))) with probability 1 - 5, SGD with stepsize 1j- = ((1) converges to a zero-loss solution at the fast rate ((xp(-t)>=))) This parameter regime identical as for gradient descent to converge in (24), with the only exception of the inflation of r by 1/8.
A key part of the argument is to estimate the probability that the iterates remain in a ball (wo). 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/Ase. which in turn would increase m multiplicatively by 62- 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[[.2) [.2 is stated with a constant failure probability, there are standard ways to remove the dependence. One option is to simply set 61 221 and rerun SGD logarithmically many times from the the same initialization We and return the final iterate with smallest function value. Section 4 outlines more nuanced strategy based on 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 I Setting the stage, consider the problem of minimizing a smooth function [(w) = Efww, 2) and suppose for simplicity that its minimal value is zero. We say that L satisfies the Polyak-Lojasiewice (PL)
inequality if there exists a > 0 satisfying
|| 2 2a L(w),
(PL.)
for all w E Rd. words, the gradient V((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 /C(w) are uniformly bounded away from zero outside the solution set. See Figures 2a and 2b for an illustration. Using the Pt. inequality, we may associate to L two condition numbers, corresponding to the full objective and its samples, respectively. Namely, we define R é 3/a and k - 3/a, where / is a Lipschitz constant of the full gradient V and 33 is Lipschitz constant of the sampled gradients Vl(.,) 2) for all z. Clearly, the inequality R < x holds and we will primarily be interested in settings where the two are comparable.
The primary reason why the PP. condition is useful for optimization is that it ensures linear con- vergence of gradient-type algorithms. Namely, it has been known since Polyak's seminal work
3

---

(a) ((w) satisfying the PL condition
(b) gradients v ((w) are far from zero
V (ww) converge at the linear rate O(epp(-t/He More recent papers have extended results of this type to wide variety of algorithms both for smooth and nonsmooth optimization (27 21 30 and to settings when the PL inequality
34 that the full-batch gradient descent iterates Wt++1 Wi B
holds only locally a ball [24]] 322
For stochastic optimization problems under the PL condition, the story is more subtle, since the rates achieved depend on moment bounds on the gradient estimator, such as:
V(W,, =)] v AC(w) B]
(1.1)
for A, B,C 2 In the setting where > 0 the classical regime- stochastic gradient methods converge sublinearly at best, due to well-known lower complexity bounds in stochastic optimiza- tion BIJ. On the other hand, the setting where
O-interpolatiin problems-stochent gradient methods converge linearly when equipped with appropriate stepsize, as shown in (2) Theorem 11. [2 Corollary 21, (38.. and [3] Theorem 4.6]. Although linear convergence assured, the rate of linear converge under the PL condition and interpolation is 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 the order of exp( (xp(-T/RK)). On the other hand, in the
C = 0, which is called the strong growth property, the paper (38) Theorem 4] yields the seemingly better rate exp(-t/Bit BR). The issue, however, is that B can be extremely large. As
illustrative example, consider the loss functions ((w, 2) 1 (w, Q=) where Q= are smooth manifolds. A quick computation shows that equality VV[(w, 2) = 26(w, 2) holds. Therefore. locally around the intersection of n2Q the estimate [I.] with A = C = 0 is exactly equivalent to the PE-condition with
1/a. As further illustration, Figure 3 shows the possible large value of the constant B along the SGD iterates for training neural network on MNIST. Another related paper is (36]: assuming so-called small gradient confusion and requiring stronger version of PL condition (i.e., each individual loss I is p-PL), the paper (36] showed a slow rate of convergence exp(-t/n21s), where the dataset size.
40k
SGC coefficient
30k
20k
10k
epochs (x100)
Figure 3: We train a fully-connected neural network on the MNIST dataset. The network has 4 hidden layers, each with 1024 neurons. We optimize the MSE loss using SGD with batch size 512 and learning rate 0.5. The training was run over 1 epochs, and the ratio [[[[V[( / [(w)) is evaluated every 100 epochs. The ratio grows almost linearly during training, suggesting that strong growth is practically not satisfied with constant coefficient B.

---

Quasar
Reference
Bound |||| V((w, =)1
Rate
Convex?
2
Theorem 1]
2/( ( - 2))
No
exp
Theorem 4]
B[[V[())]]
No
exp
BR
22
Corollary 2]
AC(w) + BII ((w)l?
No
exp
max ( (B,A/a)
Theorem 4.6]
2 - (* + [V[(w)]
No
exp
AR
Corollary 3.3]
0² + 2a w - w.ll
Yes
Sublinear
20
Theorem 4.4]
0² + [[(w)]
Yes
Sublinear
This work
23 (3([(() - L*)
Aiming
exp
Table I: Comparison to recent work on nonconvex stochastic gradient methods under the a-PE and smoothness conditions. We define 3/a and 3/a, where 33 is Lipschitz constant of the full gradient V and 3 is Lipschitz constant of the sampled gradients Vé(.,2) 2) for all
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) PE condition and interpolation. Unfortunately, the PL condition alone appears too weak to yield improved rates. Instead, we take inspiration from recent work accelerated deterministic nonconvex optimization, where the recently introduced quasar-convexity condition has led to improved rates [1]]. We note that quasar-convexity 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 accelerated sublinear rates of convergence for certain stochastic optimization problems 20 Theorem 4.4] (and the concurrent work [11, Corollary 3.31), 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 large stepsize. Moreover, contrast most available results, we only assume that regularity conditions hold on ball-the common setting in applications. The local nature of the assumptions requires us to bound the probability of the iterates escaping.
Main results
Throughout the paper, will consider the stochastic optimization problem
min C(w)
E BmP
E(w, 2),
where P is a probability distribution that is accessible only through sampling and f(., .) is differen- tiable function on Rdd. We let S denote the set of minimizers of L. We impose that satisfies the following assumptions on a set W. The two main examples are when W is a ball (ww)) and when W is tube around the solution set:
S, é {w R dist(w, S) < r).
Assumption 1 (Running assumptions). Suppose that there exist constants o, 0,3,0 2 and a set W C R. satisfying the following.
1. (Interpolation) The losses ((w, 2) are nonnegative, the minimal value of is zero, and the set of minimizers S argmin L is nonempty.
2. (Smoothness) For almost every 2 ~ P, the loss ((.22) differentiable and the gradient V((;2) is B-Lipschitz continuous on W.
3. (Quadratic growth) The estimate holds:
((w) g dist (ws S) Vw E W.
(2.1)
(Aiming) For all w W there exists point E proj (w S) such that
(VL(w), w क) 0 L(w).
(2.2)
5

---

We define the condition number K 3/a.
As explained in the introduction, the first three conditions ( are classical in the literature. In particular, both quadratic growth (3) ball W (ww) and existence of solutions in W follow from a local Pt.-inequality. In order to emphasize the local nature of the condition, following (24)
we say that C is a-PL" on W if the inequality (PL) holds for all w E W. We recall the proof of the following lemma in Section A.2
Lemma 2.1 (PL* condition implies quadratic growth). Suppose that C is differentiable and is a-PL" on ball Bz (w.). Then long as ((w)) 2ar², the intersection S B, (2)))) is nonempty and
[(w) 2 dist (w, S)
Vw E (wwo))
The aiming condition (4) is very closely related to quasar-convexity, which requires 2.2 hold for all w E R and distinguished point v E S that is independent of w. This distinction may seem mild, but in fact important because aiming holds for much wider class of problems. As concrete example, consider the function C(x, y) = {3 )² for any O. It is straightforward to see that C satisfies the aiming condition on some neighborhood of the origin. However, for any neighborhood U of the origin, the function C is not quasar- quasar-convex on U relative to any point (I, ax²) ) E U; see Section C We now show that (4) is valid locally for any C9-smooth function satisfying quadratic growth, and we may take 0 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 C2 -moothh and V L L-Lipschitz continuous on the tube Sr. Suppose moreover that C satisfies the quadratic growth condition (2.1 and < 5L'
60
Then the aiming condition 2.2 holds with parameter 0 = 2 5hr
An analogous statement holds if
3a
S, is replaced by a ball B, (w)) for some Wo E S.
Algorithm SGD (oo T)
Initialize: Initial up E Rd learning rate r 0, iteration counter T N. For t = I,...,T do:
Sample 24 ~ P
Set w1+1 Wt V
Return: WT.
2.1 SGD under regularity on tube S,
Convergence analysis for SGD (Algorithm I short and elementary in the case W S, and therefore this is where we begin. We note, however, that the setting W (ww)) is much more realistic, as we will see, but also more challenging.
The converge analysis of SGD proceeds by familiar one-step contraction argument.
Lemma 2.3 (One-step contraction on tube). Suppose that Assumption 7 Tholds on a tube W = S2r and fix point w E S,. Define the updated point w = w f/w, where z ~ P. Then for any stepsize 1i < 3, the estimate holds:
E dist? (ws S) (1-11-arm an(0 - Bn)) dist?(w,,S) znPP
(2.3)
Using the one step guarantee of Lemma 2.3 we can show that SGD iterates converge linearly to if Assumption [hools hollds on tube W S2r- The only complication is to argue that the iterates are
unlikely to leave the tube if we start in a slightly smaller tube S,, for some r' r. We do so with a simple stopping time argument.
Theorem 2.4 (Convergence a tube). Suppose that Eesumptionnhoole relative tube W Sar for some constant > 0. Fix stepsize 1) O satisfying 1) 3. Fix constant 8, > 0 and a point

---

wo E ST. Then with probability at least 1- 11 the SGD iterates (weten) remain in W. Moreover, 51 82. the estimate dist? 5) S e dist ddisst(wu (wo, S) holds after
with probability at least
1
log
iterations.
an(0 Bn)
Thus as long SGD is initialized at a point wo S/ar with probability at least 61 62. the iterates remain in S, and converge at linear rate O 22 xxp(-tt2a/3)), Note that the dependence on 82 is logarithmic, while the dependence on ó1 appears linearly in the initialization requirement Wg E S S. One simple way remove the dependence on ô1 is 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
2.2 SGD under regularity ball (wo)
Next, we describe convergence guarantees for SGD when Assumption ! holds on ball W = B, (wo). The key complication is the following. While wt are in the ball, the distance dist dist? S) shrinks in expectation. However, the iterates may in principle quickly escape the ball (ww)) ). after which point we lose control on their progress. Thus we must lower bound the probability that the iterates wa remain in the ball. To this end, will require the following additional assumption. Assumption (Uniform aiming). The estimate
( (VL(w), w u) 2 #L(w) - p dist(w,S)) S) holds for all w E (w)) and E (wow) n.s
(2.4)
The intuition underlying this assumption is as follows. We would like to replace w in the aiming condition (..2) by an arbitrary point e (wx) S. thereby having condition of the form (VL(w), w t) 2 [(w). The difficulty is that this condition may not be true for the main problem
interested in training wide neural networks. Instead, it suffices to lower bound the inner product by (C(w)- p dist (., 5) where p is a small constant. This weak condition provably holds for wide neural networks,
will see in the next section. The following is main result. Theorem 2.5 (Convergence on ball). Suppose that Assumptions / and 2 hold on a ball W ? (wo, S) S 67,2. Fix 6 and suppose p I (f - Br))r. Then with probability at least 1 - 56,, all the SGD
Bar (wo). Fix constants & E (0, 3)) and 82 E (0,1), 1), and assume dist
stepsize 17 < 3
iterates {we] J120 remain in B, (wo). Moreover, with probability at least 1 - 5Ø1 82, the estimate dist? (w,, S) < - diss?
8) holds after
1
t an BT) log
iterations.
Thus as long as / is sufficiently small and the initial distance satisfies dist ( S) S 87r2, with probability at least 56, 52, the iterates remain in (,(w)) and converge at fast línear rate O(2, cx)(-ttfa/3)). While the dependence on 82 is logarithmic, the constant 81 linearly impacts the initialization region. Section 4 discusses way to remove this dependence. As explained in Lemma 2.1 both quadratic growth and the initialization quality holds if C is a-PE" on the ball B, (wo), and is sufficiently big relative to 1/a.
3 Consequences for nonlinear least squares and wide neural networks
We next discuss the consequences of the results in the previous sections 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 nonlinear map's Jacobian V has a small Lipschitz constant in operator norm.
will make is that the
Theorem 3.1. Consider function L(w) = (w))(, where F: Rd R is -s-soothh. Suppose that there is point we satisfying dist(wn, (wn, S) S r and such that on the ball B2, (wo), the gradient VL is 3- B-Lipschitz, the Jacobian V is L-Lipschitz in the operator norm, and the quadratic growth
condition 2.1 holds. Then as long as L < TV/3. the aiming (.2) and uniform aiming (2.4g conditions
hold on B, (wo) with 0 = 2 rLVE and LBB
7

---

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 (wu; z). with parameters input and linear output layer is defined as follows:
( =
a a f(w;
( (vm=T W
-i= 1==1..... 1i=1...,1--
ymu - 1W00
Here, m, is the width (i..., number of neurons) of i-th layer, a R denotes the vector of i-th hidden layer neurons, w := {W()) W ...... W W(++) denotes the collection of the parameters (or weights) W E Rm of each layer, and a is the activation function, e.g., sigmoid, lanh, linear activation. We also denote the width of the neural network as m == min;e]y mhy i.e., the minimal width of the hidden layers. The neural network is usually randomly initialized, i.c., each individual parameter is initialized i.i.d. following (N0,,) 1). Henceforth, we assume that the activation functions a are twice differentiable, L--Lipschitz, and 8 9--smooth. In what follows, the order notation (2(.) and O()) will suppress multiplicative factors of polynomials (up to degree 1) of the constants C, Lo and Bo.
Given a dataset D = {1.; )= we fit the neural network by solving the least squares problem
1 n Σ (f(w, u²²
min L(w) [[))
where
We assume that all the the data inputs v; are bounded, i.e., ||<; S C for some constant C.
Our immediate goal is to verify the assumptions of Theorem 3.1 which quadratic growth and (uniform) aiming. We begin with the former. Quadratic growth is consequence of the PL-condition. Namely, define the Neural Tangent Kernel K K(wn) V FFFww (Wo) F at the random initial point Wg ~ (0, T) and let Ao be the minimal eigenvalue of K (wo). The value An has been shown to be positive with high probability in [10 Specifically, it was shown that, under mild non-degeneracy condition on the data set, the smallest eigenvalue Aac of NTK of an infinitely wide neural network is
positive (see Theorem 3.1 of (10). Moreover, if the network width satisfies m [(("
X log )
the estimate Åo > Ase holds [8] Remark E.7]. Of course, this is worst case bound and for our purposes we will only need to ensure that Ào is positive. It will also be important to know that || [[F(w))]] O(1), which indeed occurs with high probability shown in To simplify notation, let us lump these two probabilities together and define
then with probability at least
18
p p > 0, (wo)ll < C).
Next, we require the following theorem, which shows two fundamental properties on B, B,(wp) (wo) when the width m is sufficiently large: (1) the function + f(w, x) is nearly linear and (2) the function L satisfies the PL condition with parameter 2o/2.
Theorem 3.2 (Transition to linearity (25) and the Pł. condition (24] Given any radius > 0, with probubility 1 - 2 exp( ml)) (1/m)0(Im) m)) of initialization wo NO, T), it holds:
2 llop = -(())
Ve E (,()), ||>||| S
(3.1)
VT
In the same event, long as the width of the network satisfies m = Ω PL on B, (wo) with parameter 0o/2.
nr 0i+2 x
the function L. is
Note that directly implies that the Lipschitz constant of F'ss bounded by Ô and can therefore be made arbitrarily small. Quadratic growth is now a direct consequence of the PE. condition while (uniform) aiming follows from application of Theorem 3.1
(wo))
Vm
Theorem 3.3 (Aiming and quadratic growth condition for wide neural network). With probability at (0lm)((lnn)) m) with respect the initialization we N (0. I), as long as
least 1 p- exp( 2
(1/m) ((nn
61+2
m á
and
r ==
Ω
X
the following are true:
8

---

1. the quadratic growth condition
2.1 holds on B, with parameter Xu/2 and the intersection
B, (wo) S is nonempty,
aiming 2.2 and uniform aiming 2.4 conditions hold in B, (wg) with and = Ô
I .8++2
mm
the gradient of each functiin ((w)) A (f(w,z2)) - y is 3 3-Lipschitz on (w( with 3 O(1)
Please see Appendix D for numerical verification of an estimate of the aiming condition. It remains deduce convergence guarantees for SGD by applying Theorem 2.5
Corollary 3.4 (Convergence of SGD for wide neural network). Fix constants di E (0, g). ó2 E (0, 1), & >> and E N. There a stepsize di = ((1) such that the following true. With probability at least 1 - p 51 - 82 - 2 expp( mi - (1/m) (In m), as long as
nr 81+
m = Q
and T= Ω
61
all the SGD iterates [we)12 remain in B, (wo) and the estimate dist S) S E dists S) holds after t 2 x log
iterations.
Thus, the width requirements for SGD to converge at a fast linear rate are nearly identical those for gradient descent 24 with the exception being that the requirement 5 <
is strengthened to
r= Ω ( /o That is, the radius needs shrink by the probability of failure.
4 Boosting to high probability
possible unsatisfying feature of Theorems 2.4 and 2.5 and Corollary 3.4 that the size of the initialization region shrinks with the probability of failure 51. A natural question is whether this requirement may be dropped. Indeed, we will now see how boost the probability success to be independent of oj. A first reasonable idea is to simply rerun SGD few times from an initialization region corresponding to
1/2. Then by Hoeffding's inequality, after very trials, at least third of them will be successful. The difficulty is to determine which trial was indeed successful. The fact that the solution set is not singleton rules out strategies based on the geometric median of means 31 243], 229.. 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/=²) samples. The saving grace in the interpolation regime is that f(w,-) )) is a nonnegative function of the samples. While estimating the mean of nonnegative random variables still requires O(1/²²) samples, detecting that nonnegative random variable is large requires very few samples! This basic idea is often called the small ball principle and the basis for establishing generalization bounds with heavy tailed data (28). With this in mind, will require the following mild condition, stipulating that the empirical average >m l(w, 2;) to be lower bounded by L(w) with high probability over the iid samples 24.
Assumption (Detecting large values). Suppose that there exist constants er > 0 and C2 > 0 such that for any Rd, integer m E N, and iid samples 21. 2m ~ P, the estimate holds:
P l(w 2:) 2 (1((()) 2 exp(p
Importantly, this condition does not have anything to do with light tails. A standard sufficient condition for Assumption iis a small ball property.
Assumption (Small ball). There exist constants > and E (0, 1) satisfying
P(l(w, 2) 2 F L(w)) 2P Yw E Rd
The small ball property simply asserts that E(w: .) should not put too much mess small values rela- tive to mean L(w).. Bernstein's 's inequality directly shows that Assumption 4 implies Assumption We summarize this observation in the following theorem.

---

Lemma 4.1. Assumption 4 implies Assumption 3 with pr. and
valid bound for the small ball probabilities is furnished by the Paley-Zygmund inequality [33
[(w)2
P(e(w, () 2 IL(w)) 2 11--7 (1= (B((w,))3
Vr E [0,
Thus if the ratio _x) is bounded by some 0, then the small ball condition holds with [(w)
p (1 (1- T) where [0, 1] arbitrary.
The following lemma shows that under Assumption 3 we may turn any estimation procedure for finding minimizer of L that succeeds with constant probability into that succeeds with high probability. The procedure simply draws a small batch of samples .1..... 2m ~ P and rejects those trial points w; for which the empirical average دド=1 E(w;, 22) is too high.
...
Lemma 4.2 (Rejection sampling). Let W1....,, Wk be independent random variables satisfying P(C(w;)) I E) 1/2. For each i 1,...,K k draw samples 21...... :2m ltd P. For any A > 1, define admissible indices I = {i i [k] 1h2 Σ3= l(w,, 23) S Xe Then with probability k kexp(-c2m) X k/4 the set I is nonempty and [(m)) S for any E I.
1 exp(
We may 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 62 As we saw from the Paley-Zygmond inequality, we always have c2 s where
D upper bounds the ratios B((w.))² It remains an interesting open question investigate the scaling of small ball probabilities for overparametrized neural networks.
2(w)
Conclusion
Existing results ensuring convergence of SGD under interpolation and the PL condition require the method to use small stepsize, and therefore converge slowly. this work we isolated conditions that enable SGD to take a large stepsize and therefore have similar iteration complexity gradient descent. Consequently, our results align theory better with practice, where large stepsizes are routinely used. Moreover, argued that these conditions reasonable because they provably hold when training sufficiently wide feedforward neural networks with a linear output layer.
Acknowledgements
The work of Dmitriy Drusvyatskiy was supported by the NSF DMS 1651851 and CCF 1740551 awards. The work of Damek Davis is supported by an Alfred P. Sloan research fellowship and NSF DMS award 2047637. Yian Ma is supported by the NSF SCALE MoDL-2134209 and the CCF-2112665 (TILOS) awards, as well as the U.S. Department of Energy. Office of Science, and the Facebook Research award. Mikhail Belkin acknowledges support from National Science Foundation (NSF) and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning (https: / /ttp://eeepfffoo ouneepooonatiionss ai/) through awards DMS- 20SS-203188 and #8 44339 and the TILOS institute (NSF CCF-2112665) This work used the programs (1) XSEDE (Extreme science and engineering discovery environment) which is supported by NSF grant numbers ACI-1548562, and (2) ACCESS (Advanced cyberinfrastructure coordination ecosystem: services & support) which supported by NSF grants numbers #21 #8258259, #2138286, #2138307, #2137603, and #2138296. Specifically, used the resources from SDSC Expanse GPU compute nodes, and NCSA Delta system, via allocations TG-CIS220009.
References
[1] Hedy Attouch, Jérôme Bolte, and Benar Fux Svaiter. "Convergence of descent methods for semi-algebraic and tame problems: proximal algorithms, forward-bacc baradd-bbwwrrdd splitting, and regularized methods". In: Mathematical Progrumming 137.1-2 (2013), pp. 91-
129.
10

---

[2] Raef Bassily, Mikhail Belkin, and Siyuan Ma. "On exponential convergence of sgd in non- convex wver-paramerrized learning". In: arXiv preprint arXiv.1811.0256 (2018).
[3] Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. "Reconciling modern machine- learning practice and the classical aas-vnriance trade-off". In: Proceedings of the National Academy of Sciences 116.32 (2019), pp. 15849-15854.
[4] 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.
[5] 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.0 00005 (2021).
[6] 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.
[7] 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. 10.7/oo/oorr... 2017 0889 eprint: https: //do///o....o org/ 10. 1287 moor 2017 .889 URL:
attpp://doii /tw..//////. 1287/moor moor. 2017 0889
[8] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. "Gradient Descent Finds Global Minima of Deep Neural Networks". In: International Conference Machine Learning. 2019, pp. 1675-1685.
[9] Simon S Du, Wei Hu, and Jason D Lee. "Algorithmic regularization in learning deep homo- geneous models: Layers are automatically balanced". In: Advances in neural information processing systems 31 (2018).
[10] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. "Gradient Descent Provably Optimizes Over awer-parameterized Neural Networks". In: International Conference on Learning Representationss 2018.
[11] Qiang Fu, Dongchu Xu, and Ashia Wilson. "Accelerated Stochastic Optimization Methods under Quasar-convexity" In: arXiv preprint arXiv:2305.044: (2023).
[12] Guillaume Garrigos and Robert M Gower. "Handbook of convergence theorems for (stochastic) gradient methods", In: arXiv preprint arXiv:2301. 12335 (2023).
[13] Robert Gower, Othmane Sebbouh, and Nicolas Loizou. "Sgd for structured nonconvex func-
tions: Learning rates, minibatching and interpolation". In: International Conference on Artifi-
cial Intelligence and Statistics. PMLR. 2021, pp. 1315-1323.
[14] Moritz Hardt, Tengyu Ma, and Benjamin Recht. "Gradient descent learns linear dynamical systems". In: arXiv preprint arXiv: 1609.05 (2016).
[15] Oliver Hinder, Aaron Sidford, and Nimit Sohoni. Optiialp Methods for Minimizing CoaveCovex Functions and Beyond". In: Proceedings of Thirty Third Conference Learning Theory. Ed. by Jacob Abernethy and Shivani Agarwal. Vol. 125. Proceedings of Machine Learning Research. PMLR, Sept. 2020, pp. 1938. URL: htttps://proceeddnng/ mmlr Setta/v125/hinder222 hhmml
[16] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe 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).
[17] Aleksandr Davidovich Ioffe. "Metric regularity and subdifferential calculus". In: Russian Mathematical Surveys 55.3 (2000), p. 501.
[18] 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).
[19] Stanislaw Jastrzębski, 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).
[20] Jikai Jin. "On the convergence of first order methods for quasar-convex optimization". In: arXiv preprint arXiv:2010.04937 (2020).

---

[21] Hamed Karimi, Julie Nutini, and Mark Schmidt. "Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewiig condition". In: Machine Learning and Knowl- edge Discovery in Databases: European Conference, ECML. PKDD 2016, Riva del Garda, Ital, September 19-23, 2016, Proceedings, Part 16. Springer. 2016, pp. 795-811.
[22] Ahmed Khaled and Peter Richtárik. "Better theory for SGD in the nonconvex world". In: arXiv preprint www.v::002..03333..
(2020).
[23] 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-ECC! 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part V 16. Springer. 2020, pp. 491-507.
[24] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. *Loss landscapes and optimization in over- parameterized non-lineat systems and neural networks". In: Applied and Computational Harmonic Analysis 59 (2022), pp. 85-116.
[25] 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.
[26] 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.
[27] 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. DOT: 10. 1007 /F07/BF00096261
URL: https: :/dps:://oi.i/ org/ 10. 1007 L007/BP02096261
[28] Shahar Mendelson. "Learning without concentration" In: Journal of the ACM (JACM) 62.3
(2015), pp. 1-25.
[29] STANISLAV MINSKER. "Geometric median and robust estimation in Banach spaces". In: Bernoulli 21.4 (2015), pp. 2308-23003 2335. ISSN: 13507265, 15739759. URL: htt::http:///we jstor org/ stable/ 43590532 (visited on 05/14/2023).
[30] Ion Necoara, Yu Nesterov, and Francois Glineur. "Linear convergence of first order methods for non-strongly convex optimization". In: Mathematical Programming 175 (2019), pp. 69-
107.
[31] Arkadij Semenovič Nemirovskij and David Borisovich Yudin. "Problem complexity and method efficiency in optimization". In: (1983).
[32] 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.
[33] 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.
[34] B. T. Poljak. "Gradient methods for minimizing functionals". In: 2. Vyčisl. Mat i Mat. Fiz. 3 (1963), pp. 643-653. ISSN: 0044-4669.
[35] Quentin Rebjock and Nicolas Boumal. "Fast convergence to non-isolated minima: four equiva- lent conditions for C2 functions". In: arXiv preprini 1-Xxi:2303.0009: (2023).
[36] Karthik Abinav Sankararaman, Soham De, Zheng Xu, 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. 8469- 8479.
[37] Mingxing Tan and Quoc Le. "Efficientnet: Rethinking model scaling for convolutional neural networks". In: International conference on machine learning. PMLR. 2019, pp. 6105-6114 14.
[38] Sharan Vaswani, Francis Bach, and Mark Schmidt. "Fast and faster convergence of sgd for over- parameterized models and an accelerated perceptron". In: The 22nd international conference on artificial intelligence and statistics. PMLR. 2019, pp. 1195-1204.
[39] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. "Under- standing deep learning (still) requires rethinking generalization" " In: Communications of the ACM 64.3 (2021), pp. 107-115.
12

---

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 differentiable function. Fix two points w, w R satisfying V (() = and L(w) = Suppose that the Hessian 4 C exists and is L-Lipschitz continuous the segment [w, T.. Then we have
1 [(w) + 2 [VL(w,, it w) 1
5L
Proof. Define the function g(0) = L(w + ((w - (.)). The theorem is evidently equívalent to
5L
9(0) + 2 (0) (() - 12 I T||
In order to establish this estimate, we first note that g" is Lipschitz continuous with constant := Lil - 3|| as follows from quick computation. Taylor's theorem with remainder applied to g and g'. respectively, then gives
g(1) 9 + + 2 "0) + E1
0= = g0) + Eg,
f. 2. Combining the two estimates yields where EE s t and |Ea]
9(1) = (()) + 2 g'(() + 5 - (0)
= 9(0) + ((() (0) + E2) + 2 g0) (0) + E
Eg 2
= 9(0) + (0) + E
thereby completing the proof.
Turning to the proof of Theorem 2.2 an application of Lemma A.1
guarantees
5Lr
()=)), w w) 2 2[(w)
w|
6
for all w E S, and any w E proje (w). the quadratic growth condition
3 is satisfied on the tube
Sr, then we have the upper bound GLr ||w - -012 3a prove Lemma A.I
5Le
[(w). Theorem 2.2 follows. It remains to
A.2 Proof of Lemma 2.1
In this section, verify the classical result that the PE." condition ball implies quadratic growth. We begin with the following lemma estimating the distance of a single point to sublevel set of function; this result is special instance 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 differentiable function L: Rd - [0, oo) and ball (y and define = (w L(w) = 0}. Suppose that satisfies the Pt. condition on B, (w())) with parameter a. Then as long as L(wo) < Jr²a, a, the intersection S (yw)) is nonempty and the estimate holds:
((wo) 2 a
dist (wo. S).
2
Proof. Define the function [w) / Ll [(w) and observe that for any B,(w)) with L(w)
have [ [(
4[(w) g. Therefore application of the descent principle 6 Lemma
2.5] implies that the set S n B, is nonempty and the estimate dist(wo, S) < Va/2
(ww)
holds. Squaring both sides completes the proof.
13

---

We may now complete the proof of Lemma 2.1 by extending from single point we to neighborhood of Wo as follows. First, Lemma A.2 ensures that (2) intersects at some point wp and the inequality ((wb) e dist (wo, S) holds. Fix a point w (wo)) Then clearly L satisfies the 2 PL condition on (-(w) with parameter a. Let us now consider two cases: [(w) < 2 and [(w) 2 2r²a. In the former case, Lemma A.2 implies the claimed estimate L(w) 2 2 dist (w, S).
In the remaining case [(w) 2 2 compute
distš (w) < || - uoll v 4, 8f(w)/a.
Rearranging completes the proof of Lemma 2.1
A.3 Proof of Lemma
Fix point S, and let projs (w:) be point satisfying the aiming condition 2.2 that Lipschitz continuity of V 2) and interpolation ensures
Observe
[[V((w, =) = Ve(w, [) Vl(w, =)]l S 3 dist(w, S) S 3r.
Therefore for every 7 E [0, 1).. the point w == w T VV((w 2) satisfies
dist(w' S) S dist(w, S) V[[V
Therefore the gradient v)( B-Lipschitz on the entire line segment {w, 0 s 1/8). The descent lemma therefore guarantees VV(w, 2) < 2,3(ww,, z). Therefore upon taking expectations we obtain the second moment bound: E ( =)[ 2.C(w). Next, we compute
E dist (w+, s 11w
= E [ - w) V/ =)1?
||w w| - V L(w)(()) (w), w w)) 1 (Lw,, 2]
< dist dis S) 2np0 L(w) + 3/(w
(A.1)
dist (w, (w, S) 2n 31) ((w)
S (1 on n)3))dist 5),
(A.2)
where A.I
follows from 2.2 while A.2) follows from 2.1 The proof is complete.
A.4 Proof of Theorem
Define the stopping time = inf (t 2 1 we & Sr} and set U, = dist? S). Note that we may equivalently write t inf(t U, > r2). Now, multiplying (2.3) through by 7 yields the estimate
E 1> Wha] S (1 an Br) U,1+
(A.3)
An application Theorem B.2 with q = (1 - or7 Bn)) 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 Br(wo), we have Bl|e(w, =)1l² < 2,3 L(w).
Proof. Let Wa E projs(ux) be arbitrary. Observe that Lipschitz continuity of Ve(,,2) on (sr(w)) and interpolation ensure
[[V[(w, =)l| [[V[(w, V.(Wo. =) 3 || woll 23r.
Therefore for every 7 E [0, 1/). the point w" w 7 ~V((w, =) satisfies
||w" woll < || woll + T V((w, 2)1l < 3r.
Therefore the gradient V((,,)) is B-Lipschitz on the entire line segment {w, 0 S 1/B)). The descent lemma therefore guarantees [V =) < 2,8( 2). Taking the expectation both sides completes the proof.
D
14

---

Next we prove the following lemma that simultaneously estimates (1) one step progress of the iterates towards and (2) how far the iterates move away from the the center uin-
Lemma A.A. Fix a point w (r))(wo) and choose wo E projs (wo) n Brya(wog) (wo). Assume 11 S 0 and define w
N(w, z) where P. Then the following estimates hold:
Eldist (w+ {)] - (1 ny( Br))a) (w;
(A.4)
Ellw"* wo [1 s ||w wD 11 + 2mp dist S)
(A.5)
Proof. Fix any point v E S and observe that
I - v1l2 = || w1|2 ( ( 2), w - v) + V =)12
Taking the expectation with respect to = and using Lemma A.3 deduce
E I|w 01| - 1|w- u112 2ny( V (ww),, v) + 2.3))
(A.6)
We will use this estimate multiple times for different vectors v.
Choose any point projs(w) satisfying the aiming condition
2.2
Define dist dist?(w, S) and
U+ == dist (w+, 8). Then setting = in (A.6). we deduce
E[U+] < 2ny v ((w), w i) 2Br ((w).
Using the aiming condition
2
therefore deduce
E < U - 2n0 + 2,17 [(w)
U 2n (0 -Bn) ((w)
S (1 r7(0 - BB)))UU
where the last inequality follows from quadratic growth. This establishes (A.4).
Next, set e == Wp in (A.6). Defining V = || wo wol V+ = =ll woll2 and using Assumption 2 we therefore deduce
[ S V - -2n(0((w) - PVU) + 28n2² ((w).
= - -2- (0 1133 L(w) + 2npvU
= V + 21ppVU
where the last estimate follows from the inequality Bn S 6. This establishes (A5).
For each 2 0, let 1, denote the indicator of the event E. = {wo,.... we E B, (wa)]. Define
the random variables U, = dist (w;, S) and V/2 = w wo 2 We may now multiply A.4) by Ix. Noting that Ee41 1g, we may iterate the bound yielding
E[ S (1 7 Bn)a)
(A.7)
Similarly, multiplying A.4 by 1g, and using (A.7)
deduce
E 1441 V421l < E[1 Vi + 2mp /IU,
< 1 V. + 2mp(1 ni(0 31) Uo
Iterating the recursion gives
20
1LVV < Uo + 1- /(1 17(0 - Br)a)
4pVU.
< Uo + (0 - Bn)a
(A.8)
We now lower bound the probability of escaping from the ball. Note that within event E,, have
Vi < (w woll ll + 1Iwo woll)² (1 + 3r2.
15

---

Therefore,
P(E)) < PIV > (1 11+511 61)2,
S PIV (1 + 6,)
Edl P(E.)
E g. V. ()))
(A.9)
(1 8,)
[I Vil
),)
Up + ApVUg
(0-9n)a
(A.10)
(1 + 51)3T2
4p
1 + (0 - -Br))rr
81
where (A.9) follows from Markov's 's inequality and (A.10) follows from (A.))
Next, we estimate the probability that U, remains small within the event E,, To that end, let define the constant C, = (1 15--/p 1 ) 82. Then Markov's inequality yields
E Ed]
E U.] C,F(E,)
82
P(U, > C, E,)
S
Ct
P(E.)
Finally, we unconditionally bound the probability that U, remains small:
P(U S C.) 2 P(U, S C. E.).(E))
82
2
P(E.).
P(E)))
= ((E)) - 82.
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 the Jacobian V..
Lemma A.5. Consider a function L(w) (w))l where F: Rd R" Ca and the Jacobian F is L-Lipschitz on the ball (,o)). Then the following estimates
(VL(w), u v) < 8r L 22(w))
(A.11)
[V w - w) -[((w) < / L(w)/2 || /w hold for all w, u, (,o satisfying F(u) F(v) =
(A.
Proof. Fix any E (wp) and u, v E (wo) satisfying F(u) = F(v) = 0. We first prove AA.11.
To this end, we compute
(VL(w), w)1 (VF(w) T F(w), u v)
(F(w), (F(w) w))l
< < w)]b
where the last estimate follows from the Cauchy- inequaliit... Next, the fundamental theorem of calculus and Lipschitz continuity of F yields
0 F(v) - F(w)-F(u) I VF t(y u))(u u) dt
-V[(w)( w) +)+ E,
16

---

where
E I L u||( 11 w|l + ||e ull) _ 8r Thus we have proved A.11
In order to see
A.122
we compute
(VL(w), w (I ( = (F(w), VF(w)(w
u) -Fw)))
- 2 || ull
where the last inequality follows from Cauchy-Schwarz and Lipschtiz continuity of the Jacobian V.. Thus A.12
holds.
Lemma A.5
quickly yields the following corollary.
Corollary A.6. Consider a function ((w) = | [(w)ll where F; R + R is C -mooth and the
Jacobian V is I-LLpschitz on the ball (, Suppose moreover that
((w) 2 2
a
(w;
where S = (w: F(w) = 0}. Then the estimates
(V(w)), w) 2 2f(w)
LV2 L(w)3/2.
(A.13)
V [Vw)), v) 2 2L(ww Lv/ L(w))/2 8r 1 2L(w), hold for all w E (wo), v (wo) n S, and w E (wv) (projs(w).
(A.14)
Proof. We will apply Lemma A.5 Setting a =: v in Lemma A.12 and using quadratic growth yields the estimate:
and
(VL(w), w w) 2 2[(ww) I. (ww)/2 / || w
L ( [(w).
2 20((w)
which establishes A.13).
Adding this estimate
A.12 with w yields
Theorem 3.1
follows immediately from the corollary. Indeed, note that let w E projs(w) satisfies 2r. Therefore we may apply Corollary and and use the estimate L(w) S 2 [|l 1
||
A.7 Proof of Theorem 3.3
Theorem 3.2 implies that with probability at least 1 - p- expp(-=2 ml
2 (1/m) (lamm, the Jacobian (wo) and the PE-condition holds on (ww)
of V is Lipschitz with constant L = Vm
with parameter 22/ Using the assumption
may apply Lemma 2.1 to deduce order to deduce (uniform) auming with parameters 0 To this end, it suffices to ensure 2r Lr/ I 1, which follows from the
2 /
quadratic growth. Next, we will apply Theorem 3.1
prerequisite assumption m = 22(
6642
a²
To see the last claim, for any w E (wo) we compute the Hessian
v2((w) = v
I.) (f(w,x+) vi) f
Therefore,
(v)(((w)llop < V f x1 lf(w, Vel | Setting Wi E prooss (wp) B, (wo), observe that
(A.15)
(
/334+1
VV x1)1 f. x.) Vf [w x:)ll S
sup ve (we)
| x.) =
where the last inequality follows from transition to linearity 3.I.. Thus f(,,r;) 11)) is Lipschitz continu- ous on the ball (wo). Therefore, we may also bound
3L+2
[f xx)) ysl (f(w, x;;) f(w, x;)1 =Ō
m
Returning to (A.15). we thus have [ Ö (=ontt)
Taking into account that m
m
we deduce pLL2
O() thereby completing the proof.
17

---

A.8 Proof of Corollary 3.4
The goal 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 | and hold on ball (20) with >o/2, 2, 0 = 1, and
p=
33-2
We may then choose = 23 (1). In order to make sure that S (0 37)
VT
61+2 suffices to in the regime m 2
Finally remains to ensure that dist (ws, S) S 8p.2.
Ag
To do so, using quadratic growth, we have dist (wo, S) - 4 ((wo) ( Thus it suffices to let
Ao
O Agg ). An application of Theorem 2.5 completes the proof.
A.9
Proof of Lemma 4.2
Define the set J {i [k] : f(w))) e}. Then Hoeffding inequality ensures that the inequality [J] 2 holds with probability at least 1 - exp( exp(-k/16). Conditioned on this event, consider any index i E [k] satisfying [(w:) Ae, - exp(-e2m)) we have
Then from we know that with probability at least
c
Σ E(w;, 2,) 2 e [(w;) > xe,
and therefore & I. Let now estimate the probability that I is nonempty conditioned on | k To this end, by Markov's 's inequality for any E J, we have
1
[(ws) Xf(w;) X
P
Eww, z4) > de < P
E(w,, 23) 2 AL(w;)
Consequently, the probability that the set I is empty at most --4/4 4/4
B Auxiliary results on stopping times
Theorem B.1 (Stopping time argument). Let (Ur)e20 be sequence of nonnegative random variables and define the stopping time т = inf t 2 0: U2 > tn for some constant u > 0. Suppose that
ElUr-I U.d] S Uel Ce
Vt 2 0.
(B.1)
Blo+ E 4: holds for all >
Then the estimate P S t]
Proof. Observe that by Markov's inequality, we have
B[UeAr] u
P/r s 1] P(U u)
Let us therefore bound the expectation of the stopped random variable Vg = Unr Letting E,[]] denote the conditional expectation E[- Unr), we successively compute
E2, [V+4]]
Eg [U41AT]
E, [U + E,
E, + E, Unrlr
-E, + Unrl
< U,l,>t+ + Unstrse
= V22 +
Taking the expectation with respect to UL., applying the tower rule, and iterating the recursion we deduce [[]] [EU.) + E S, thereby completing the proof.
Theorem B.2 (Stopping time with contractions). Ler (Uth22 be a sequence of nonnegative random
variables and define the stopping time T inf 2 0 U, 4t for some constant u. > 0. Suppose that there exists q E (0, 1) such that
ElUH+lLee Una S U,L,s
V2 0.
(B.2)
18

---

Then as long as BUD S 61u, the event {T= =o) occurs with probability at least 1- with probability at least 1 - 51 - 82, the estimate U, S sUo holds after i 2 -- log ō2ē
51. Moreover,
iterations.
Proof. Define the stopping time 7 = in 2 U, > u). An application Lemma therefore
implies IP[T < t] I Wile < 8,. Taking the limit as t - oo, we deduce that the event { oc) occurs
with probability least 1 o1. Next taking the expectation with respect to U1:t in (B22) and applying the tower rule gives
ElU4+11r2l S q E[U,,,>i]-
Taking into account that 17st 2 lyse+1 we may iterate the recursion thereby yielding
B[U,lrsa] S d'Uo-
Now, setting e' ;= =Uo, Markov's inequality yields
E[
UOO E¹
Pl
S 82.
e¹
Finally observe
P(U, 2 , x0)
P(U1l 26)
PIU, 2 =
P| x0
P|r ocl
Therefore we deduce
PIU, < < P(U < < T= = oo]
P1-P 2 r= oo]) P[r ool
2 P/r oo] P r= -oo--Pp PPrr oo]
2 F/7 xo| -(U,l 2
2 1 - 61 62,
as claimed.
C A bad example
Lemma C.1. Consider the objective L(x, v) = (y 2x for any a > 0. Then is O near (0, 0) and the function C satisfies the Pt. inequality (PE. with constant a = 1. Therefore L satisfies the aiming condition on some neighborhood of the origin. However, for any neighborhood U of the origin, the function C is not quasar-conve on U relative to any point (I, ax²) E U.
Proof. To see the validity of the Pt.-condition, observe that
V Lxx x)ll
= V (Cr, = V
az²
C (x,y)
v2"
follows immediately from Lemma 2.2 that C satisfies the aiming condition 2.2) on some neigh- borhood of the origin. Next verify the failure of quasar convexity. Consider an arbitrary point (x. ax?) on the parabola. For any (u, v) E R²2, we compute
<VL(u, (w, (xz, ax²)) {( (2u, 1), (u x, ax²)) (u aa))
(-2aaulu x) + v ax²) au²).
In particular, setting 2, = (0,7) we obtain
(VL(2,). - (a, ax²)) = = -aax2
Note the right hand side is negative for any 0 < 7 < ax² Thus, letting > 0 tend to zero, deduce that is not quasar-convex relative to (r, ax²) E U with x + 0. Let us consider now the setting 0. Then compute
(VL(u, v), (u, (0, 0)) (v 2au an?)
The right side is negative if an² < v < 2an2 and therefore V.C. not quasar-convex relative to (0, 0)
on any neighborhood of 0).
19

---

MINIST
CIFAR-10
2.00
0.7
1.5
150-
0.5-
0.4
2.00
0.3
0.75
025
0.1
0.00
1030
1500
2000
2500
3500
0.0
500
2500
3000
25.00
4002
iterations
Iaration
Figure 4: Verifying aiming condition. Left: MNIST on fully-connected neural network; Right: CIFAR-10 CNN.
D Numerical verification of aiming condition on neural networks
In principle, the aiming condition difficult to verify, as it involves projs(w) - the nearest point of w to the solution set S, which we have no way of computing. Instead, we replace projs(w) (w) by the last iterate w of an SGD run, and compute an estimate of the aiming condition coefficient by the quotients 0g A (Ww),, W² w) ((w)) along the iterate path. Note that this quotient is not the same quotient would appear in the definition of quasar-convexity because w is random point that depends on the iterate path taken by SGD.
Settings: We conduct the experiments 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-1), and record the parameters of this trained network as the optimal solution W. For CIFAR-10, 10, we train a two-layer CNN for 4000 iterations, and record the final parameter setting w. Then, for each experiment, we took second run, and at each iteration Wt we compute the estimate the aiming condition coefficient using the equation mentioned above.
Results: Figure 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 nonnegative function L is called quasar-convex relative to a point w with L(w) if the incquality (VL(w), w w) 2 0 -(w) holds for all near Aiming, in contrast, stipulates the analogous inequality but with tú crucially replaced by projs(w), where is the set of zero loss solutions. We argue that in the interpolation setting, quasar connsrrcoovvxityy is 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 (44 D. To this, note that quasar-convexity implies that S is star-convex relative to w (Observation in (15). That is, there exists > O such that for any w E B, (w) NS, the line segment joining w and w is fully contained S. This is very strong conclusion. For example, a smooth-manifold S' is star convex near v if and only it coincides with a linear subspace around w. The PL-condition in turn implies by Theorem 2.16 in 35 that is 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. contrast, as we have proved, any C -mooth function satisfying PL automatically satisfies the aiming inequality near point S. In summary, quasar-convexity does not hold when training wide neural networks (since otherwise it would imply that interpolating solutions form linear subspace), while aiming provably holds one of our main results.
20


Reviews:
[{"review_id": "ZQGqspFGzf", "reviewer": "Reviewer_eCMo", "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.", "questions": "Please see the question above.", "limitations": "The authors have adequately addressed the limitations.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\nIt 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.\n2. It would be helpful if the authors can highlight the technical novelties.\n3. 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. \n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the constructive comments. We address the concerns as follows:\n\n**W1:** *The main application of the results is for wide enough neural networks in the NTK regime, which seems restrictive.*\n\n**A:** We would like to point out that the ”NTK regime“ setting is in line with much of the literature of optimization theories for neural networks. \nNearly  all convergence results for GD/SGD for neural networks so far are in the NTK regime (or more precisely, require sufficiently large network width) (see, e.g., [1-5]). \nTo the best of our knowledge, there is no general theory for convergence of GD/SGD outside of large width regimes.   The goal of our work is to improve the existing results on convergence of SGD by showing essentially the same rate of convergence as GD (something which has been empirically observed) and thus requires a setting where the existing analyses of GD exists. \nFurthermore, the large width regime should logically be the first to be explored before further progress can be made. \n\n[1] Allen-Zhu, Li, Song. A Convergence Theory for Deep Learning via\nOver-Parameterization. ICML. 2019.\n\n[2] Du, Zhai, Poczos, Singh. Gradient Descent Provably Optimizes\nOver-parameterized Neural Networks. ICLR, 2018.\n\n[3] Zou, Cao, Zhou, Gu. Gradient descent optimizes overparameterized deep ReLU networks. Machine Learning, 2020.\n\n[4] Oymak, Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. : IEEE Journal on Selected Areas in\nInformation Theory, 2020.\n\n[5] Liu, Zhu, Belkin, Loss landscapes and optimization in over-parameterized non-linear systems and neural networks. ACHA, 2022.\n\n**W2:**  *It would be helpful if the authors can highlight the technical novelties.*\n\n\n**A:** Thanks for the suggestion. We will include a more comprehensive summary of contributions in the revision.\nFrom a high level viewpoint, we (1) introduced the aiming condition and (2) proved that sufficiently wide neural networks satisfy the aiming condition. Note that the width requirement we have is the same as that in establishing convergence of GD in prior work, e.g., [1]; hence, we are not imposing stronger assumptions.\nSecond, we proved that aiming, interpolation, and quadratic growth conditions enable application of SGD with a large step size and fast linear rate.  Indeed, Section 1.2 explains in detail that the linear rate we obtain is much faster than existing convergence guarantees for SGD.\n\nFrom a technical viewpoint, we introduced a number of new analysis techniques. First, the proof that neural networks satisfy aiming and uniform aiming relied on some careful computations with nonlinear least squares and the transition to linearity phenomenon. Second, the proof of convergence of SGD relied on decomposing the iterate trajectories into    “tangent” and “orthogonal” components to the solution set $S$. The aiming condition ensured that the iterates make steady progress in orthogonal directions towards $S$ while uniform aiming controlled the tangent motion. Third, using this decomposition together with a novel martingale stopping time argument, we bounded the probability that the iterates escape the ball. We expect that such an orthogonal decomposition  with a stopping time argument to be useful more broadly. \n\n\n [1] Liu, Zhu, Belkin, Loss landscapes and optimization in over-parameterized non-linear systems and neural networks, ACHA, 2022.\n\n**W3：** *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.*\n\n\n**A:** Thank you. This comment, made in passing, may have caused some confusion.\nWe discuss generalization and the \"edge of stability\" as the motivations for understanding the training efficiency of SGD with large stepsize. We don't claim that our analysis implies better generalization or that it is in the  \"edge of stability\" regime. Instead, we focus on the optimization performance of SGD with large step size on the training loss only. As we have stated in Section 1.1 (outline of main results) and Section 1.2 (comparison to existing work), our goals are to improve the theoretically allowed stepsize range (especially to include the practically used large stepsizes), and to deduce faster convergence rate on the training loss. Apologies for the confusion; we will make this point clear in the revision.\n\n**W4** *Is it possible to verify the proposed conditions for common neural network models, at least empirically?*\n\n**A:** Thank you; this is a good question. The difficulty is that the aiming condition involves $proj_S(w)$---the nearest point of $w$ to the solution set $S$---and which we have no way of computing. Instead, we performed the following experiment on MNIST and CIFAR-10 datasets. We let SGD run and convergence to some point $\\bar w$. Then we plot the quotients $\\hat \\theta_t\\triangleq\\langle \\nabla L(w_t),w_t-\\bar w\\rangle/L(w_t)$ along the iterate path. Note that positivity of $\\hat \\theta_t$ suggests that aiming holds, but does not verify it exactly. Indeed, for MNIST we see that $\\hat\\theta_t>0.5$ while for CIFAR $\\hat\\theta_t>0.05$. See the results in the response to all reviewers above. Figures are shown in the pdf file attached to it.  We will include the empirical verification in the revision. We note that it may seem that the experiments verify the related quasar-convexity property, but this *not* the case because $\\bar w$ is a *random point* that depends on the sample path taken by SGD. "}, {"review_id": "enNwFm69sQ", "reviewer": "Reviewer_FvUq", "summary": "This 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\n", "questions": "There is no figure 2\n", "limitations": "The authors have described the limitations.\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive comments. We address your concerns below in detail.\n\n**W1:** *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.*\n\n**A:** Apologies; this statement was not carefully stated. In line 133-134, we mean to state our goal of obtaining *theoretically acceptable* stepsizes, which better match the large stepsizes used in practice. \nIn this paper, we successfully showed that, under certain conditions (including aiming and PL), our theory indeed allows larger stepsizes than prior works and leads to a significantly faster convergence rate, as explained in Section 1.2. Moreover, we proved that these conditions hold for sufficiently wide neural networks.\n\n**W2:**  *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.)*\n\n**A:** Thank you for pointing out this reference (We call it \"the reference\" later on). We agree that this reference is relevant and indeed has some discussion on large step size for SGD. We will be sure to cite it in our revision and include a comparison. \n\nWe summarize the differences as follows:\n\n1. (PL condition for each summand) The reference requires that every loss function $f_i$, $i\\in[N]$, satisfies the $\\mu$-PL condition (see (A2) in page 3 of the reference), which is {\\em much stronger} than the usual PL condition for the ERM loss $F=\\frac{1}{N}\\sum_{i=1}^N f_i$, which we use in our paper. Indeed, it is not clear what the limiting value of $\\mu$ in the reference would be as the width increases to infinity. In contrast, the PL constant that appears in our bounds is the minimal eigenvalue of the NTK of the infinitely wide neural network.\n\n\n2. (slower convergence rate) According to Theorem 3.1 in the reference,  the required stepsize is $\\alpha=1/N L$, which leads to a slow rate of convergence $O(\\frac{N^2L}{\\mu}\\log(1/\\epsilon))$ to a fixed constant loss value $\\epsilon\\approx \\eta N/\\mu$, which itself can be quite large. Here, $\\eta$ is a gradient confusion constant. In contrast, the number of samples $N$ does not influence the convergence rate we obtain, and we prove fast convergence to the global minimum. Note, that the reference proves that a small $\\eta<4$ is a valid gradient confusion constant with high probability only when the data are sampled from a unit sphere. \n\n**W3:** *... 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\n**A:** Thanks for the comment. We conducted some preliminary experiments on CIFAR-10 dataset for additional empirical evaluations. Please see the results and settings in the response to all reviewers above. Figures are shown in the pdf file attached to it. More extensive numerical evaluations are currently running and will be included in the revision. From what we can tell so far, the qualitative results we see on MNIST are also visible on CIFAR10, with the main difference that both GD and SGD are *both* much slower, as one expects.\n\n\n\n**Q1:** *There is no figure 2.*\n\n**A:** This is the figure on top of page 4; we will add a label to this figure in the revision."}, {"review_id": "Aq77VOqFl1", "reviewer": "Reviewer_vP9H", "summary": "This 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. ", "questions": "See points 1 and 2. ", "limitations": "The main limitation of this work is a strong assumption on the locality of the SGD analysis. Some further clarification would be helpful.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1) 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. ", "weaknesses": "1) 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback. We address the concerns below:\n\n**W1:** *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$\n 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\n**A:** Thank you for bringing this to our attention; we can see where the confusion arises in the informal statement of Theorem 1.2. \n In the statement of the theorem, one can replace $\\lambda$---the minimal eigenvalue of the NTK at initialization---by the the minimal eigenvalue $\\lambda_{\\infty}$ of the NTK of an infinitely wide neural network, which *does not* depend on the initialization. See the discussion on page 7 for a precise statement. The two theorem statements are equivalent  because as the width increases, the parameter $\\lambda$ approaches $\\lambda_{\\infty}$. More precisely, in the regime $m=\\Omega(nr^{6l+2}/\\lambda_{\\infty}^2)$, with high probability one has $\\lambda\\geq \\lambda_{\\infty}/2$. Therefore the requirement on the radius $r= \\Omega(1/\\lambda \\delta^2)$ is independent of the initialization, with high probability. Summarizing,  *for any fixed radius $r= \\Omega(1/\\lambda_{\\infty} \\delta^2)$, one may take the width sufficiently large $m=\\Omega(nr^{6l+2}/\\lambda_{\\infty}^2)$, so that with probability roughly $1-\\delta$ aiming holds on the entire ball $B_r(w_0)$ and the SGD iterates converge at the linear rate $O(\\exp(-t\\lambda_{\\infty}))$.*\n\nWe emphasize that in this parameter regime (i.e. when $m$ is sufficiently large), the ball $B_r(w_0)$ is in fact large enough to contain the whole optimization path of gradient descent (see for example Section 5 in [1]), and the initial point need *not* be close to the true minima. That means, even if \"initialization is made far away from the minima\", $B_r(w_0)$ still covers both the initialization and the global minima. Hence, the aiming condition would hold throughout the training process. \n\nWe will modify the statement of the theorem and the preceding discussion to make these points clear. We hope that this explanation convinces the reviewer that aiming is indeed a reasonable condition.\n\n[1] Liu, Zhu, Belkin, Loss landscapes and optimization in over-parameterized non-linear systems and neural networks, ACHA, 2022.\n\n**W2:** *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 escape $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\n**A:** Thank you for the question. Interpolation implies the basic bound $E[||\\nabla \\ell(w,z)||^2] \\le 2\\beta \\cdot L(w)$, where $\\beta$ is the smoothness constant of the loss functions (Line 408 of our appendix). Thus the second moment of the stochastic gradient (and its variance) tends to zero as the iterates approach the solution set $S$. The \"rate\" at which the variance tends to zero is controlled by $\\beta$. Now, if one wants to use this bound in order to ensure that the expected function value $E[L(w_t)]$ contracts in each iteration, one would need to use a tiny learning rate $\\eta=\\frac{\\alpha}{\\beta^2}$ leading to a slow rate of convergence (as simple examples show).  To overcome this difficulty, we focus on the contraction in the squared distance $E[dist^2(w_t,S)]$. The aiming condition together with the aforementioned bound on the variance ensures that $E[dist^2(w_t,S)]$ indeed contracts in each iteration when using a stepsize on the order $\\eta=\\frac{\\theta}{\\beta}$ and while the iterates remain in $B_r(w)$. That is, the the iterates makes very good progress moving \"orthogonally\" towards $S$.\n It remains to bound the probability that the iterates escape the ball $B_r(w)$. This bit of the argument is involved. We first establish an auxiliary condition called uniform aiming (page 6) for wide NNs, and use it to bound the \"tangent motion\" relative to $S$. A careful stopping time argument then bounds the probability of escaping the ball. We will add a more detailed discussion to the revision.\n\n**W3:** *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$\\ldots$*\n\n**A:** Thank you for pointing this out. We will cite this reference instead."}, {"review_id": "5G06kdu1jw", "reviewer": "Reviewer_APCi", "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.", "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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The aiming condition is impressive as a condition for SGD to achieve the same sample complexity as GD.\n- 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.\n- This work does not seem to be enough strong technically.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feeback. We address your concerns below.\n\n**W1:** *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.*\n\n**A:** Thank you for bringing attention to this point.  We will now argue  that in the interpolation setting, quasar-convexity is a very restrictive condition because it requires the set of zero-loss solutions to be a linear subspace, which is usually not the case (see for example [1]).\nTo be precise, quasar-convexity relative to a point $\\bar w$ with $L(\\bar w)=0$ means that $\\langle \\nabla L(w),w-\\bar w\\rangle\\geq \\theta\\cdot L(w)$ for all $w$ near $\\bar w$. Aiming, in contrast, stipulates the analogous inequality but with $\\bar w$ crucially replaced by $proj_S(w)$, where $S$ is the set of zero loss solutions. \nTo see the limitation of quasar-convexity, note that quasar-convexity implies that $S$ is star-convex relative to $\\bar w$ (Observation 3 in reference [13] in the paper).\n That is, there exists $\\epsilon>0$ such that for any $w\\in B_\\epsilon(\\bar w)\\cap S$, the line segment joining $w$ and $\\bar w$ is fully contained in $S$. This is a very restrictive assumption. For example, a smooth-manifold $S$ is star convex near $\\bar w$ if and only if it coincides with a linear subspace around $\\bar w$. The PL-condition in turn implies by Theorem 2.16 in [2] 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, on any fixed radius ball around the initial point, sufficiently wide neural networks satisfy the aiming condition. In parallel, we also showed that any $C^3$-smooth function satisfying PL automatically satisfies the aiming inequality near a point $\\bar w\\in 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. We will add a discussion along these lines to the revision.\n\n\n[1] Liu, Zhu, Belkin, Loss landscapes and optimization in over-parameterized non-linear systems and neural networks, ACHA, 2022.\n\n[2] Rebjock, Boumal, Fast convergence to non-isolated minima:\nfour equivalent conditions for $C^2$ functions, arxiv.org/pdf/2303.00096.pdf.\n\n**W2:** *This work does not seem to be enough strong technically.*\n\n**A:**  In this paper, we address a fundamental question of interest both for optimization and machine learning: why is that numerically the iteration complexity of SGD is comparable to GD when training wide neural networks. In order to answer this question, we introduced a number of new analysis techniques. First, the proof that neural networks satisfy aiming and uniform aiming---new conditions introduced in the paper---relied on some careful computations with nonlinear least squares and the transition to linearity phenomenon. Second, the proof of convergence of SGD relied on decomposing the iterate trajectories into \"tangent\" and \"orthogonal\" components to the solution set. The aiming condition ensured that the iterates make steady progress in orthogonal directions towards while uniform aiming controlled the tangent motion. Third, using a novel martingale stopping time argument together with this decomposition, we bounded the probability that the iterates escape the ball. We expect that such an orthogonal decomposition together with a stopping time argument to be useful more broadly.\n\nWith these new techniques, we were able to contribute the following novelties. We proved that sufficiently wide neural networks satisfy our proposed aiming condition. Note that the width requirement we have is the same as that in establishing convergence of GD in prior work; hence, we are not imposing stronger assumptions. Second, we proved that aiming, interpolation, and quadratic growth conditions enable application of SGD with a large step size and fast linear rate. Indeed, Section 1.2 explains in detail that the linear rate we obtain is much faster than existing convergence guarantees for SGD. We believe that this paper is strong technically, and the developed techniques and regularity conditions will be useful for understanding optimization algorithms and landscapes of loss functions in nonconvex optimization and learning.\n\n**W3:** *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.*\n\n**A:** Thank you for pointing this out. We will correct them in the revision and do a thorough pass through the paper."}]

Author rebuttal:
We provide additional experimental results here, requested by the reviewers. Please see the other rebuttals below, directly after each review.

**1: An estimate of aiming condition:** 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$---and which we have no way of computing. Instead, we replace $proj_S(w)$ by the last iterate $\bar{w}$ of an SGD run, and compute an estimate of the aiming condition coefficient by the quotients $\hat \theta_t\triangleq\langle \nabla L(w_t),w_t-\bar w\rangle/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 $\bar 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 $< 10^{-4}$), and record the parameters of this trained network as the optimal solution $\bar{w}$. 
For CIFAR-10, we train a two-layer CNN for $4000$ iterations $^\dagger$, and record the final parameter setting as $\bar{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 $\hat \theta_t$ using the equation mentioned above.


*Results:* Figure 1 and Figure 2 in the attached pdf file show the plots of $\hat \theta_t$ against the iteration number. In both cases, we observe that the estimate $\hat \theta_t$ stays positive, which suggests that aiming holds. 

We will include the empirical verification in the revision. 


**2: Mini-batch SGD on CIFAR-10:** In this experiment, we show that, with the same number of iterations, mini-batch SGD has almost the same convergence behavior as the full batch GD on the CIFAR-10 dataset (same as in Figure 1 for MNIST).

*Setting:* We run mini-batch SGD with different batch sizes $b$ ($b=60, 120, 250, 500, 1000$), as well as full batch GD, on a two-layer CNN with CIFAR-10 dataset for $4000$ iterations $^\dagger$.

*Results:* Figure 3 in the attached pdf file shows the training loss curves of the mini-batch SGD and full batch GD. We observe that with the same number of iterations, mini-batch SGD has almost the same convergence behavior as the full batch GD.

$\dagger$: We didn't train this CNN until convergence, as the convergence is slow for CIFAR-10 on this simple CNN. We are currently running experiments on deeper networks so that a near-zero training loss could be obtained. We will include it in the revision.

Meta-review: The authors have addressed the reviewers’ concerns. Please incorporate all the discussions and new experiments into the final version.
Decision: Accept (poster)