(S)GD over Diagonal Linear Networks: Implicit bias, Large Stepsizes and Edge of Stability
Authors: Mathieu Even
Scott Pesme
Suriya Gunasekar
Nicolas Flammarion
Venue: NeurIPS 2023 poster (2023)
Keywords: SGD
GD
implicit bias
large stepsizes
edge of stability
diagonal linear networks
TL;DR: We study the convergence  and the implicit regularization of (S)GD with large stepsizes over diagonal linear networks.

Abstract: In this paper, we investigate the impact of stochasticity and large stepsizes on the implicit regularisation of gradient descent (GD) and stochastic gradient descent (SGD) over $2$-layer diagonal linear networks. We prove the convergence of GD and SGD with macroscopic stepsizes in an overparametrised regression setting and characterise their solutions through an implicit regularisation problem. Our crisp characterisation leads to qualitative insights about the impact of stochasticity and stepsizes on the recovered solution. Specifically, we show that large stepsizes consistently benefit SGD for sparse regression problems, while they can hinder the recovery of sparse solutions for GD. These effects are magnified for stepsizes in a tight window just below the divergence threshold, in the ``edge of stability'' regime. Our findings are supported by experimental results.

(S)GD over Diagonal Linear Networks: Implicit Bias, Large Stepsizes and Edge of Stability
Mathieu Even" Inria ENS Paris
Scott Pesme" EPFL
Suriya Gunasekar Microsoft Research
Nicolas Flammarion EPFL
Abstract
In this paper, we investigate the impact of stochasticity and large stepsizes on the implicit regularisation gradient descent (GD) and stochastic gradient descent (SGD) over 2-layer diagonal linear networks. We prove the convergence of GD and SGD with macroscopic stepsizes in an overparametrised regression setting and provide characterisation of their solution through an implicit regularisation prob- lem. Our characterisation provides insights on how the choice minibatch sizes and stepsizes lead to qualitatively distinct behaviors in the solutions. Specifically, we show that for sparse regression learned with 2-layer diagonal linear networks, large stepsizes consistently benefit SGD, whereas they hinder the recovery of sparse solutions for GD. These effects are amplified for stepsizes in a tight window just below the divergence threshold, known as the "edge of stability" regime.
Introduction
The stochastic gradient descent algorithm (SGD) [51] is the foundational algorithm for almost all neural network training. Though remarkably simple algorithm, it has led to many impressive empirical results and a key driver of deep learning. However the performances of SGD are quite puzzling from theoretical point of view (1) its convergence is highly non-trivial and (2) there exist many global minimums for the training objective which generalise very poorly [66].
To explain this second point, the concept of implicit regularisation has emerged: if overfitting is harmless many real-world prediction tasks, must because the optimisation process implicilly favoring solutions that have good generalisation properties for the task. The canonical example is overparametrised linear regression with more trainable parameters than number of samples: although there infinitely many solutions that fit the samples, GD and SGD explore only small subspace of all the possible parameters. As a result, it can be shown that they implicitly converge to the closest solution in terms of the distance, and this without explicit regularisation [66, 24).
Currently, most theoretical works on implicit regularisation have primarily focused on continuous time approximations of (S)GD where the impact of crucial hyperparameters such as the stepsize and the minibatch size are ignored. One such common simplification to analyse gradient flow, which is continuous time limit of GD and minibatch SGD with an infinitesimal stepsize. By definition, this analysis does not capture the effect of stepsize or stochasticity. Another approach is to approximate SGD by stochastic gradient flow [60, 48), which tries to capture the noise and the stepsize using an appropriate stochastic differential equation. However, there are no theoretical guarantees that these results can be transferred to minibatch SGD as used in practice. This is limitation our understanding since the performances of most deep learning models are often sensitive the choice of stepsize and minibatch size. The importance of stepsize and SGD minibatch size common knowledge practice and has also been systematically established in controlled experiments [36, 42, 20].
Denotes equal contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Test losses at convergence
Figure I: Noiseless sparse regression with diagonal linear network using SGD and GD, with parameters initialized at the scale of a = 0.1 (Section 2). The test losses at convergence for various stepsizes are plotted for GD and SGD. Small stepsizes correspond gradi- ent flow (GF) performance. We see that increasing the stepsize improves the generalisation properties of SGD, but deteriorates that of GD. The dashed vertical lines 5,50D
10-1
test luss
SGD
10
GD
at stepsizes "max and 7max denote the largest stepsizes for which SGD and GD, respectively, converge. See Section for the precise experimental setting.
10 &CCD
TIDX
7max
Stepsize
In this work, we aim to expand our understanding of the impact of stochasticity and stepsizes by analysing the (S)GD trajectory in 2-layer diagonal networks (DLNs). In Fig. 1, show that even in our simple network, there are significant differences between the nature of the solutions recovered by SGD and GD at macroscopic stepsizes. We discuss this behavior further in the later sections.
The 2-layer diagonal linear network which we consider is simplified neural network that has received significant attention lately [61, 57, 26, 50]. Despite its simplicity, it surprisingly reveals training characteristics which are observed in much more complex architectures, such as the role of the initialisation [61], the role of noise [48, 50), or the emergence of saddle-to-saddle dynamics [6, 49]. therefore serves as an ideal proxy model for gaining a deeper understanding of complex phenomenons such the roles of stepsizes and of stochasticity as highlighted in this paper. We also point out that implicit bias and convergence for more complex architectures such as 2-layer ReI.U networks, matrix multiplication are not yet fully understood, even for the simple gradient flow. Therefore studying the subtler effects of large stepsizes and stochasticity in these settings is currently out of reach.
1.1
Main results and paper organisation
The overparametrised regression setting and diagonal línear networks are introduced in Section 2. We formulate our theoretical results (Theorems and 2) in Section 3: we prove that for macroscopic stepsizes, gradient descent and stochastic gradient descent over 2-layer diagonal linear networks converge to a zero-training loss solution 8 We further provide a refined characterization of 3% through trajectory-dependen implicit regularisation problem, that captures the effects of hyperparameters of the algorithm, such as stepsizes and batchsizes, in useful and analysable ways. In Section 4 we then leverage this crisp characterisation to explain the influence of crucial parameters such as the stepsize and batch-size on the recovered solution. Importantly our analysis shows stark difference between the generalisation performances of GD and SGD for large stepsizes, hence explaining the numerical results seen in Fig. for the sparse regression setting. Finally, in Section 5, we use our results to shed new light on the Edge of Stability (FoS) phenomenon [14].
1.2 Related works
Implicit bias. The concept of implicit bias from optimization algorithm in neural networks has been studied extensively in the past few years, starting with early works of Telgarsky [55], Neyshabur et al. [45], Keskar et al. [36]. Soudry et al. [53]. The theoretical results on implicit regularisation have been extended to multiplicative parametrisations [23, 25], linear networks [34], and homoge- neous networks [40, 35, 13]. For regression loss on diagonal linear networks studied in this work, Woodworth et al. [61] demonstrate that the scale of the initialisation determines the type of solution obtained, with large initialisations yielding minimum & norm solutions the neural tangent kernel regime [30] and small initialisation resulting in minimum l norm solutions the rich regime [13]. The analysis relies the link between gradient descent and mirror descent established by Ghai et al. [21] and further explored by Vaskevicius et al. [56], Wu and Rebeschini [62]. These works focus on full batch gradient, and often in the inifitesimal stepsize limit (gradient flow), leading to general insights and results that do not take into account the effects of stochasticity and large stepsizes.
The effect of stochasticity in SGD on generalisation. The relationship between stochasticity in SGD and generalisation has been studied in various works [41, 29, 11, 38, 64). Empirically, models generated by SGD exhibit better generalisation performance than those generated by GD [37, 27].

---

Explanations related to the flatness of the minima picked by SGD have been proposed [28]. Label noise has been shown to influence the implicit bias of SGD [26, 8, 15, 50] by implicitly regularising the sharp minimisers. Recently, studying stochastic gradient flow that models the noise of SGD in continuous time with Brownian diffusion, Pesme et al. [48] characterised for diagonal linear networks the limit of their stochastic process as the solution of an implicit regularisation problem. However similar explicit characterisation of the implicit bias remains unclear for SGD with large stepsizes.
The effect of stepsizes in GD and SGD. Recent efforts to understand how the choice of stepsizes affects the learning process and the properties of the recovered solution suggest that larger stepsizes lead to the minimisation of some notion of flatness of the loss function [52, 37, 44, 33, 64, 43], backed by empirical evidences or stability analyses. Larger stepsizes have also been proven to be beneficial for specific architectures or problems: two-layer network [39], regression [63], kernel regression [7] or matrix factorisation [59]. For large stepsizes, it has been observed that GD enters an Edge of Stability (EoS) regime [32, 14], in which the iterates and the train loss oscillate before converging to a zero-training error solution; this phenomenon has then been studied on simple toy models [1. 67, 12, 16] for GD. Recently, [2] presented empirical evidence that large stepsizes can lead to loss stabilisation and towards simpler predictors.
2 Setup and preliminaries
Overparametrised linear regression. We consider a linear regression over inputs X (0i,-..,,f7 E (Rdd)" and outputs y = (.) Mn) E R.. We consider overparametrised problems where input dimension d is (much) larger than the number of samples n. In this case, there exists infinitely many linear predictors 3 TRd which perfectly fit the training set, i.e,, yi ( .na) for all 1 & 8 & n. We call such vectors interpolating predictors or interpolators and we denote by S the set of all interpolators S = {,8 E Idd s.t. (8* ,I,) = yi: Vi E [n]]] Note that S is an affine space of dimension greater than d - n and equal to 3 span(21 F1............. I..) for any 3 E S. We consider the following quadratic loss: L(3) 2n
Σ (3) )x for / E IR4.
2-layer linear diagonal network. We parametrise regression vectors 3 as functions 3w of trainable parameters w E IP, Although the final prediction function I + (sw.r) is linear in the input x, the choice of the parametrisation drastically changes the solution recovered by the optimisation algorithm [25]. In the case of the linear parametrisation B.. = w many first-order methods (SGD, GD, with or without momentum) converge towards the same solution and the choice of stepsize does not impact the recovered solution beyond convergence. In an effort to better understand the effects of stochasticity and large stepsize, we consider the next simple parametrisation, that of 2-layer diagonal linear neural network given by:
(1) This parametrisation can be viewed as a simple neural network x +> (u, (diag()))) where the output weights are represented by the inner weights is the diagonal matrix diag(v), and the activation a is the identity function. In this spirit, we refer to the entries of w (u, v) E Rdd as the weights and to B == u @ u E Rd as the prediction parameter. Despite the simplicity of the parametrisation (1), the loss function F over parameters = (u, v) IR22d is non-convex (and thus the corresponding optimization problem is challenging to analyse), and is given by:
B = u where w (u, v) E R2dd
F(w) = L(u =
1
I(: - (uo
(2)
2n
Mini-batch SGD. We minimise F' using mini-batch SGD: let wo
(uo, vo) and for l > 0,
Wk+1 = wa 7tk V FB, (wx), where FIs (w) = 2h
(yr (u vr,))))
(3)
where "/k are stepsizes, B& C [n]] are mini-batches of b E [n]] distinct samples sampled uniformly and independently, and V Fs, (wx) are minibatch gradients of partial loss over Br, Fi, (w) == Ls, (w@v) defined above. Classical SGD and full-batch GD are special cases with For k > 0, we consider the successive prediction parameters Bq. = up @ eg built from the weights
and b = n,, respectively.
3

---

In Appendix we show that (S)GD trajectory with this initialisation exactly matches that of another common parametrisation 8., w w2 with initialisation w4.0 w-.o = a. The second layer of our diagonal linear network is set to in order to obtain results that are easier to interpret. However. our proof wk= = (ug, 0k). We analyse SGD initialised at do = v2a E R and vo = 0 E R&, resulting in 80 = 0 E IR independently of the chosen weight initialisation o²
Experimental details. We consider the noiseless sparse regression setting where (Za)iE(n) (O, [2) and y/5 .z)) for some s-sparse vector S/. We perform (S)GD over the DLN with a uniform initialisation a setup (n, d, a) = (20, 30, 0.1), Fig. 2 (right) to (n, d, 8,0) = (50, 100, 4, 0.1) and Fig. 3 to (n, (n,d,s, s, a) = (50, 100, 0.1))
E Rd where > O. Fig. I and Fig. (left) correspond to the
denote the Hessian of L, and for batch B C [n] let HB == LB [8] Eien T. denote the Hessian of the partial loss over the batch B. Let L denote the "smoothness" such that VB, |HH & L|| 2. | < =|| for all batches B C [n]] of size b. A real function (e.g, log, exp) applied to vector must be understood as element-wise application, and for vectors u E R u² (uā) u/v = (u) We write 1, O for the constant vectors with coordinates and O respectively. The Bregman divergence [9] of differentiable convex function h Rd R is defined as Dr, 82) h(51) - (h(32) + (32), B1 32)).
Notations. Let H = AーC= 1 - x,,
(u,u))eed) and
3 Implicit bias of SGD and GD
We start by recalling some known results on the implicit bias of gradient flow on diagonal linear networks before presenting our main theorems on characterising the (stochastic) gradient descent solutions (Theorem 1) as well as proving the convergence of the iterates (Theorem 2).
3.1 Warmup: gradient flow
We first review prior findings on gradient flow on diagonal linear neural networks. Woodworth al. [61] show that the limit 8% of the gradient flow du; = F (wF(())d )dt initialised at (uo, no)
(VZa, 0) the solution of the minimal interpolation problem:
aa = argmin do (f"), where (3) = Σ (aarcsinh ES
8 + a + aZ.)
(4)
The convex potential vo is the hyperbolic entropy function (or hypentropy) [21]. Depending on the structure of the vector the generalisation properties of 3 highly vary. We point out the two main characteristics of a that affect the behaviour of vaa and therefore also the solution 9
1. The Scale of a. For an initialisation vector a we call the f--norm |allı the scale of the initialisation. It is an important quantity affecting the properties of the recovered solution 3 To see this let us consider uniform initialisation of the form a al for a scalar value a > 0. In this case the potential doa has the property of resembling the -nomr as the scale a vanishes: va ~ In(1 =n[11//)]] |l-]l as + O. Hence, small initialisation results in low f--norm solution which known induce sparse recovery guarantees [10]. This setting is often referred to as the "rich" regime [61]. In contrast, using large initialisation scale leads to solutions with low fo-mr:: vo ~ (1-I)/(2aa) as a + oo, a setting known as the "kernel" or "lazy" regime. Overall, to retrieve the minimum E -noom solution, one should use a uniform initialisation with small scale a, see Fig. 7 in Appendix for an illustration and [61, Theorem 2] for a precise characterisation.
2. The Shape of a. In addition to the scale of the initialisation a lesser studied aspect is its "shape", which
term we use refer the relative distribution of {a,}, along the d coordinates [3]. It is a crucial property because having -> 0 does not necessarily lead to the potential va being close to the normm. Indeed, we have that via (3) a=0 Σd_ In( 1 )B3/ (see Appendix D), therefore the vector In(1/ a/)) ) has entries changing at different rates, then tha (3) is a weighted é -nnor.. In words, if the entries of do not go zero "uniformly". then the resulting implicit bias minimizes
az
techniques be applied directly to general initialisation, at the cost of additional notations in Theorems.

---

weighed & nn-norm This phenomenon can lead to solutions with vastly different sparsity structure than the minimum & norm interpolator. See Fig. 7 and Example I in Appendix D.
3.2 Implicit bias of (stochastic) gradient descent
In Theorem I, we prove that for an initialisation v2ax E IR and for arbitrary stepsize sequences (2k) if the iterates converge to an interpolator, then this interpolator is the solution a constrained minimisation problem which involves the hyperbolic entropy aaa defined in (4), where ax E I&d is an effective initialisation which depends on the trajectory and on the stepsize sequence. Later, we prove the convergence of iterates for macroscopic step sizes in Theorem 2.
Theorem 1 (Implicit bias of (S)GD). Let (ug, vk)k30 follow the mini-batel SGD recursion (3) initialised at (wo, vo) (V2x, 0) and with stepsizes (nh) Let (3x)az0 (ux O uk)azo and assume that they converge to some interpolator 3%0 E S. Then, 8% satisfies:
B argmin Dejaze (0)) E*
where Dvan is the Bregman divergence with hyperentropy potential Vay of the effective initialisa- tion Oloo, and 9o is a small perturbation term. The effective initialisation On given by,
ax = O
- an VL, (3x)))x)))) (3x)) k=0
(6)
where ((6)) = - In((1 - 22) satisfies g(z)) > for |ae| < 22, with the convention q(1) = 400.
The perturbation term 3o Rd is explicitly given by 30 1 (a a²), where (= F2x
1n((1= F ()²), and a a² o exp (- - CL, (3x))).
The characterisation of 3% in Theorem holds for any stepsize schedule such that the iterates converge and goes beyond the continuous-time frameworks previously studied [61, 48). The result even holds for adaptive stepsize schedules which keep the stepsize scalar such as AdaDelta [65]. An important aspect of our result is that & and 3o depend on the iterates" trajectory. Nevertheless, ingredients for understanding the implicit regularisation effects of (SJGD for this problem compared to trivial characterisations (such as e.g., ming || 8 Il). Importantly, the key parameters ax.. 3 depend on crucial parameters such as the stepsize and noise in a useful and analysable manner: understanding how they affect Olso and 30 coincides with understanding how they affect the recovered solution &%x and its generalisation properties. This is precisely the object of Sections 4 and 5 where we discuss the qualitative and quantitative insights from Theorem in greater detail.
Trajectory-dependent characterisation.
argue that our formulation provides useful
The perturbation 30 can be ignored. We show in Proposition 16, under reasonable assumptions The magnitude of Bo is therefore
on the stepsizes, that |3ul & a² and ax & (component-wiss negligible in front of the magnitudes of 3* S and one can roughly ignore the term Bo. Hence, the implicit regularisation Eq. (5) can be thought of as 3* & argming-es Daa.. (3*, 0) = d'a (3*), and thus the solution 3 minimises the same potential function that the solution of gradient flow (see Eq. (4)), but with an effective initialisation axcc Also note that for 7 = y + 0 we have aeoc + a and 30 + 0 (Proposition 19), recovering the previously known result for gradient flow (4).
Deviation from gradient flow. The difference with gradient flow is directly associated with the quantity Σe ((TH VLs, (Bx)). Also, as the (stochastic) gradients converge to 0 and q(x) =20 x?, one should think of this sum as roughly being Σ VLs, (3&)2: the larger this sum, the more the recovered solution differs from that of gradient flow. The full pieture of how large stepsizes and stochasticity impact the generalisation properties of 3* and the recovery of minimum fn-mr solution is nuanced clearly seen in Fig. 1.
3.3 Convergence of the iterates
Theorem provides the implicit minimisation problem but says nothing about the convergence the iterates. Here we show under very reasonable assumptions the stepsizes that the iterates indeed
5

---

converge towards global optimum. Note that since the loss F is non-convex, such convergence result non-trivial and requires an involved analysis.
Theorem (Convergence of the iterates). Let (ua, us ux) follow the mini-batch SGD recursion (3) initialised at ug = /2ax E R and to = O, and let (Bt)k20 = (ux @ 0k kk) Recall the "smoothness" parameter L the minibatch loss defined in the notations. There exist B 0 verifying B Ominnges ||3 loc)) and numerical constant > 0 such that for stepsizes satisfying * LB* the iterates (3x) k220 converge almost surely to the interpolator 8x solution of Eq. (5).
In fact, we can be more precise by showing an exponential rate of convergence of the losses as well as characterise the rate of convergence of the iterates as follows.
Proposition 1 (Quantitative convergence rates). For uniform initialisation al and under the assumptions of Theorem 2, we have:
1 E [[C(B)]] < (1- ((8)) and E [1 - - 11] <
2a Ao )
11-2700 2
where db > is the largest value such that NoH 2B(a2 X )-1 (1 + (4BXmux) X (())) and Xmin Amax >
Eg Hg) C
respectively the smallest non-null and the largest eigenevalues of H, and Bo, is the interpolator that minimises the perturbed hypentropy h& of parameter ork, as defined in Eq. (7) in the next subsection.
The convergence of the losses is proved directly using the time-varying mirror structure that exhibit in the next subsection, the convergence of the iterates is proved by studying the curvature of the mirror maps small neighborhood around the affine interpolation space.
3.4 Sketch of proof through a time varying mirror descent
As in the continuous-ting framework, our results heavily rely on showing that the iterates (3a)s follow a mirror descent recursion with time-varying potentials on the convex loss L(3). To show this, we first define the following quantities:
(a2 a²
a == a+.k 2 ah.N
and
Φx = 2
aresinh
E
ERd,
20
-i a+ (70 Lo, (B6)) E Rd. Finally for h > 0, we define the
where 0=,k = a eexp
potentials (h Rdd - R) as:
hx(3) = May (0aa,(1) (3) (0k, 3).
Where via is the hyperbolic entropy function defined Eq. (4). Now that all the relevant quantities defined, we can state the following proposition which explicits the time-varying stochastic mirror descent.
The iterates (3x up C uk)xzo from Eq. (3) satisfy the Stochastic Mirror Descent recursion with varying potentials (ha)k:
Proposition
Vhgta (Bx+1) (Bk( 7h VLB((a))
where by : Rdd R for are defined Eq. (7). Since (ho 0 we have:
Vh span(zi.........
(8)
Theorem and 2 and Proposition follow from this key proposition: by suitably modifying classical convex optimization techniques to account for the time-varying potentials, we can prove the con- vergence of the iterates towards an interpolator 8 along with that of the relevant quantities 0..kk. and dhg. The implicit regularisation problem then directly follows from: (1) the limit condition Vhico (.9x0) E Span(a) (#I,..., F..) as seen from Eq. (8) and (2) the interpolation condition X X3 y. Indeed, these two conditions exactly correspond to the KKT conditions of the convex problem Eq. (5).

---

4 Analysis of the impact of the stepsize and stochasticity on axe
In this section, we analyse the effects of large stepsizes and stochasticity on the implicit bias of (S)GD. We focus on how these factors influence the effective initialisation axx, which plays a key role as shown in Theorem vector Σx 9(7% V LB, (5%)). We henceforth call this quantity the gain vector. For simplicity of the discussions, from now on, we consider constant stepsizes T = ? for all k > 0 and a uniform initialisation of the weights = al with > O. We can then write the gain vector as:
From its definition in Eq. (6), we see that a is a function of the
=
Gain, = In
Rd.
Following our discussion in section 3.1 on the scale and the shape of a we recall the link between the scale and shape of Gain, and the recovered solution:
1. The scale of Gain,, i.e. the magnitude of ||Gain/]|: indicates how much the implicit bias of (S)GD differs from that of gradient flow: |[Gain, ~ implies that ao ~ a and therefore the recovered solution is close to that of gradient flow. On the contrary, ||Gain,||a >> In(1/a) implies that aeoc has effective scale much smaller than a thereby changing the implicit regularisation Eq. (5).
2. The shape of Gain, indicates which coordinates of 33 in the associated minimum weighted & problem are most penalised. First recall from Section 3.1 that a uniformly large Gain, leads to &a being closer to the & norm. However, with small weight initialisation - 0, have,
Vaw (B) In( 1 11111 +
(i)
In this case, having a heterogeneously large vector Gain, leads to a weighted & norm as the effective implicit regularisation, where the coordinates of B corresponding to the largest entries of Gain, less likely to be recovered.
4.1
The scale of Gain, is increasing with the stepsize
The following proposition highlights the dependencies of the scale of the gain [|Gainy]]i in terms of various problem constants.
Proposition 3. Let An, As > 0 be the largest and smallest values, respectively, such that A,H <
Es [ < A,H. For any stepsize 7 > O satisfying 7 & BL (as in Theorem 2), initialisation ol and batch size E [m], the magnitude of the gain satisfies:
2o?? Σ (E((3k) < [[||/ain, 2A0 [BC(B))
(10)
where the expectation is over uniform and independent sampling of the batches (Ba)az0-
The slower the training, the larger the gain. Eq. (10) shows that the slower the training loss converges to 0, the larger the sum of the loss and therefore the larger the scale of Gain,. This means that the GD)GD trajectory deviates from that gradient flow if the stepsize and/or noise slows down the training. This supports observations previously made from stochastic gradient flow [48] analysis.
The bigger the stepsize, the larger the gain. The effect of the stepsize on the magnitude of the gain is not directly visible in Eq. (10) because a larger stepsize tends to speed up the training. For stepsize BL as in Theorem 2 we have that (see Appendix G.1):
0 7 & Tmax
~x²((B))= = (> in ( l10..1.)
(11)
Eq. (11) clearly shows that increasing the stepsize boosts the magnitude || |/aiin up until the limit of 7max- Therefore, the larger the stepsize the smaller is the effective scale of a In turn, larger gap between aeo and a leads to a larger deviation of (S)GD from the gradient flow.
"A,, Au, > are data-dependent constants; x (H) is the smallest non-null eigenvalue of H; for b 1. have min; ||0:1| d1 & A1 & max: ||a:
n, we have (An, An) (At (H), Amax Amax(()) where
7

---

Scale:
||Gain,|| ||
Shape of Gain,
800
SGD
50
GD
SGD:
100
0
GD:
0
10
2SGD
Tiax
GD Trmax
2
Coordinate
Stepsize Y
Support of Ssparse
Figure 2: Left: the scale of Gain, explodes as 7 Fmax for both GD and SGD. Right: Paparse is fixed, perform 100 runs GD and SGD with different feature matrices, and we plot d coordinates Gain, (for GD and SGD) the r-axis (which is in log scale for better visualisation). The shape of Gain, SGD homogeneous whereas that of GD is heterogeneous with much higher magnitude the support of Paparse: The shape of Gain.
GD is proportional to the expected gradient at initialisation which is (Ssparse)*
Large stepsizes and Edge of Stability. The previous paragraph holds for stepsizes smaller than Ymax for which we can theoretically prove convergence. But what if we use even bigger stepsizes? Let (BE)) denote the iterates generated with stepsize " and let us define Ímax sup 17 s.t. V/y (0,7), Σg ((02))
oo), which corresponds to the largest stepsize such that the iterates still converge for given problem (even if not provably so). From Proposition 3 we have that Ymua & mmax- As approach this upper bound convergence y + 7max the sum Σx ((52) diverges. For such large stepsizes, the iterates of gradient descent tend to "bounce" and this regime is commonly referred as the Edge of Stability. In this regime, the convergence of the loss can be made arbitrarily slow due to these bouncing effects. As a consequence, as seen through Eq. (10), the magnitude of Gain, be become arbitrarily big observed in Fig. 2 (left). In this regime, the recovered solution tends to dramatically differ from the gradient flow solution, as seen in Fig. 1.
Impact of stochasticity and linear scaling rule. Assuming inputs z; sampled from .N(0, o² La) with ²² > 0, we obtain E [|/ainn]la] = θ (f In (2) I1). w.h.p. over the dataset (see Ap- pendix G.3, Proposition 17). The scale of Gain, decreases with batch size and there exists a factor n between that of SGD and that of GD. Additionally, the magnitude of Gain,, depends z. resembling the linear scaling rule commonly used in deep learning [22].
By analysing the magnitude ||Gain, l1, we have explained the distinct behavior of (S)GD with large stepsizes compared gradient flow. However, our current analysis does not qualitatively distinguish the behavior between SGD and GD beyond the linear stepsize scaling rules, in contrast with Fig. 1. deeper understanding of the shape of Gainy is needed to explain this disparity.
4.2
The shape of Gain,, explains the differences between GD and SGD
In this section, we restrict our presentation to single batch SGD (b 1) and full batch GD = n). When visualising the typical shape of Gain,, for large stepsizes (see Fig. 2 right), we note that GD and SGD behave very differently. For GD, the magnitude of Gain,, is higher for coordinates in the support of 18, and thus these coordinates are adversely weighted in the asymptotic limit of va... (per (9)). This explains the distinction seed in Fig. 1, where GD in this regime has poor sparse recovery despite having small scale of ao as opposed to SGD that behaves well.
The shape of Gain, is determined by the sum of the squared gradients Σn VL8, 23 and in particular by the degree of heterogeneity among the coordinates of this sum. Precisely analysing the sum over the whole trajectory of the iterates (8& )a is technically out of reach. However, we empirically observe for the trajectories shown in Fig. 2 that the shape is largely determined within the first few iterates as formalized in the observation below.
Observation 1. Σe LL., (8 & LB,
8

---

In the simple case of a Gaussian noiseless sparse recovery problem (where y; (3 sparse Tii) for some sparse vector & we can control these gradients for GD and SGD (Appendix G.4) as:
V VL(6))2 (Ssparse) 1 +E, for some verifying || Egg VC (86) (1
<<
(12)
(13)
The gradient of GD is heterogeneous. Since Spparse is sparse by definition, we deduce from Eq. (25)
that VC(6)) is heterogeneous with larger values corresponding to the support of - sparse' Along
with observation I, this means that Gain, has much larger values on the support of 3 Ssparse- The corresponding weighted & -nom therefore penalises the coordinates belonging to the support of 33 Ssparse which hinders the recovery of Psparse (as explained in Example I, Appendix D).
The stochastic gradient of SGD is homogeneous. On the contrary, from Eq. (26), we have that the initial stochastic gradients are homogeneous, leading to a weighted -noor where the weights are roughly balanced. The corresponding weighted -norm is therefore close to the uniform & -norm and the classical & recovery guarantees are expected.
Overall summary of the joint effects of the scale and shape. In summary we have the following trichotomy which fully explains Fig. 1:
1. for small stepsizes, the scale is small, and (S)GD solutions are close to that of gradient flow;
2. for large stepsizes the scale is significant and the recovered solutions differ from GF:
for SGD the shape of &to is uniform, the associated norm is closer to the -norr and the recovered solution is closer the sparse solution;
for GD, the shape is heterogeneous, the associated norm is weighted such that it hinders the recovery of the sparse solution.
In this last section, relate heuristically these findings to the Edge of Stability phenomenon.
Edge of Stability: the neural point of view
In recent years has been noticed that when training neural networks with 'large' stepsizes at the limit of divergence, GD enters the Edge of Stability (EoS) regime. In this regime, as seen in Fig. 3, the iterates of GD 'bounce' 'oscillate". In this section, we come back to the point of view of the weights w/k (ug, v&) E 2d and make the connection between our previous results and the common understanding of the EoS phenomenon. The question we seek to answer is: in which case does GD enter the EoS regime, and if so, what are the consequences on the trajectory? Keep in mind that this section aims to provide insights rather than formal statements. We study the GD trajectory starting from small initialisation that gradient flow converges close to the sparse interpolator 3aparse = Prsppere corresponding to
al where < 1 such that can consider
the weights Waparse (v /ipprsel sigu( (Sspars))))) V [Biparsal
(see Lemma in [49] for the mapping from the predictors to weights for gradient flow). The trajectory of GD as seen in Fig. 3 (left) can be decomposed into up phases.
The stepsize is appropriate for the local curvature (as seen in Fig. 3, lower right) around initialisation and the iterates of GD remain close to the trajectory of gradient flow (in black in Fig. 3). If the stepsize is such that 7 < Amax (FFF(eipsra)) then is compatible with the local curvature and the iterates can converge: in this case GF and GD converge to the same point (as seen in Fig. 1 for small stepsizes). For larger >> Xmax (V#F( (as is the case for YGD in Fig. 3, lower right), the iterates cannot converge to 8 & and we enter the oscillating phase.
First phase: gradient flow.
Second phase: oscillations. The iterates start oscillating. The gradient of F writes V ((a,)) F(w) (V C(3) C v. VL(B) 3 u) and for w in the vicinity of Waparse we have that u, se RJ 0 for i E supp(();arse). Therefore for Waparse we have that V.F(w), ss V. V.F(), R 0 for i of supp(8)paree) and the gradients roughly belong to Span(ei, (Eup)))) This means

---

Train losses
Coordinates 8
GD at Eo5
Gradient descent
200
Grodient flow
supp(P coordinates
Other coordinates
101
Gradient flow
10 10"
Drift towards 0
10
10F
10? herasion
Anuox V F (1AA.)
Goadient Bow
100
10
10
10
Itrabon
10
7GD
10
102
102
11P
CD GF
Oscillations
Convergence
Iteration
Figure 3: GD at the EoS. Left: For GD. the coordinates on the support of Beparse oscillate and drift towards 0. Right, top: The GD train losses saturate before eventually converging. Bottom: GF converges towards solution that has a high hessian maximum eigenvalue. GD cannot converge towards this solution because of its large stepsize: it therefore drifts towards a solution that has a curvature just below 2/7-
that only the coordinates of the weights (u;, us) for i E supp( can oscillate and similarly for (B, (B)))esupp(Pppa..)) (as seen Fig. left).
Last phase: convergence. Due the oscillations, the iterates gradually drift towards a region of lower curvature (Fig. 3, lower right, the sharpness decreases) where they may (potentially) converge. Theorem enables us understand where they converge: the coordinates 3& that have oscillated significantly along the trajectory belong to the support of $ Spparse. and therefore Gain,, (i) becomes much larger for i E supp(3" (paaee)) than for the other coordinates. Thus, the coordinates of the solution recovered in the EoS regime are heavily penalised on the support of the sparse solution. This is observed in Fig. 3 (left): the oscillations of (3:) kesppp((tt))) lead to gradual shift of these coordinates towards O, hindering an accurate recovery of the solution Siparse.
SGD in the EoS regime. In contrast the behavior of GD where the oscillations primarily occur on the non-sparse coordinates of ground truth sparse model, for SGD we see different behavior in Fig. 6 (Appendix A). For stepsizes in the EoS regime, just below the non-convergence threshold: the fluctuation of the coordinates occurs evenly over all coordinates, leading to uniform ar These fluctuations are reminiscent of label-noise SGD [2], that have been shown to recover the sparse interpolator in diagonal linear networks [50].
6 Conclusion
We study the effect of stochasticity along with large stepsizes when training DLNs with (S)GD. We prove convergence of the iterates as well as explicitly characterise the recovered solution by exhibiting an implicit regularisation problem which depends on the iterates' trajectory. In essence the impact of stepsize and minibatch size captured by the effective initialisation parameter &xe that depends on these choices in an informative way. We then use our characterisation to explain key empirical differences between SGD and GD and provide further insights the role of stepsize and stochasticity. In particular, our characterisation explains the fundamentally different generalisation properties of SGD and GD solutions large stepsizes of large stepsizes can prevent the recovery of the sparse interpolator, even though the effective seale of the initialization decreases with larger stepsize for both SGD and GD. We also provide insights on the link between the Edge of Stability regime and our results.
in Fig. 1: without stochasticity, the
Aknowledgements
M. Even deeply thanks Laurent Massoulié for making it possible to visit Microsoft Research and the Washington state during an internship supervised by Suriya Gunasekar, the MSR Machine Learning Foundations group for hosting him, and Martin Jaggi for inviting him for week in Lausanne EPFL, making it possible to meet and discuss with Scott Pesme and Nicolas Flammarion.
10

---

References
[1] Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via the "edge of stability". arXiv preprint, 2022.
[2] M. Andriushchenko, Varre, Pillaud-Vivii and N. Flammarion. SGD with large step sizes learns sparse features. arXiv preprint, 2022.
[3] Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference on Machine Learning, pages 468-477 PMLR, 2021.
[4] R. Baraniuk, M. Davenport, R. DeVore, and M. Wakin. A simple proof of the restricted isometry property for random matrices. Constructive Approximation, 28(3):253-265 January 2008.
[5] H. H Bauschke, J. Bolte, and M. Teboulle. A descent lemma beyond Lipschitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research, 42(2):330-348, 2017.
[6] Raphaël Berthier. Incremental learning in diagonal linear networks. arXiv preprint arXiv:2208. 14673, 2022.
[7] G. Beugnot, J. Mairal, and A. Rudi. On the benefits of large learning rates for kernel methods. In Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 252-282. PMLR, 02-05 Jul 2022.
[8] G. Blanc, N. Gupta, G. Valiant, and P. Valiant. Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeek like process. In Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 483-513. PMLR, 09-12 Jul 2020.
[9] L.M. Bregman. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR Computational Mathematics and Mathematical Physics, 7(3):200-217, 1967. ISSN 0041-5553,
[10] E. Candès, J. Romberg, and T. Tao. Stable signal recovery from incomplete and inaccurate measurements. Communications on Pure and Applied Mathematics, 59(8): 1007-112, 2006.
[11) Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational infer- ence, converges to limit cycles for deep networks. In International Conference on Learning Representations, 2018
[12] Lei Chen and Joan Bruna. On gradient descent convergence beyond the edge of stability, 2022.
[13] Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Pro- gramming. 2019.
[14] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representations, 2021.
[15] Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global minimizers. In Advances in Neural Information Processing Systems, 2021.
[16] Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization:: The implicit bias of gradient descent at the edge of stability. In International Conference on Learning Representations, 2023.
[17] J. I., Doob. Stochastic Processes. John Wiley & Sons, 1990.
[18] Radu Alexandru Dragomir, Mathieu Even, and Hadrien Hendrikx. Fast stochastic Bregman gra- dient methods: Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 2815-2825. PMLR, 18-24 Jul 2021.
Sharp analysis and variance reduction. In Proceedings of the 38th International

---

[19] Mathieu Even and Laurent Massoulie. Concentration of non non-isotropic random tensors with applications to learning and empirical risk minimization. In Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 1847-1886. PMLR, 15-19 Aug 2021.
[20] Jonas Geiping, Micah Goldblum, Phillip E Pope, Michael Moeller, and Tom Goldstein. Stochas- tic training is not necessary for generalization. In International Conference Learning Representations, 2022.
[21] Udaya Ghai, Elad Hazan, and Yoram Singer. Exponentiated gradient meets gradient descent. In Proceedings of the s1 International Conference Algorithmic Learning Theory, volume 117 of Proceedings of Machine Learning Research, pages 386-407. PMLR, 08 Feb-11 11 Feb 2020.
[22] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training imagenet in hour. arXiv preprint arXiv: 1706.0077 2017.
[23] Suriya Gunasckar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
[24] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1832-1841. PMLR, 10-15 Jul 2018.
[25] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Systems, volume 31, 2018.
[26] Jeff Z. HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pages 2315-2357. PMLR, 15-19 Aug 2021.
[27] Fengxiang He, Tongliang Liu, and Dacheng Tao. Control batch size and learning rate to generalize well: Theoretical and empirical evidence. In Advances in Neural Information Processing Systems, volume 32, 2019.
[28] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, January 1997.
[29] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: Closing the generalization gap in large batch training of neural networks. In Proceedings of the 31st International Conference on Neural Information Processing Systems, page 1729-1739, 2017.
[30] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, page 8580-8589, 2018
[31] Stanislaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD, 2017.
[32] Stanislaw Jastrzebski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019.
[33] Stanislaw Jastrzkebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Width of minima reached by stochastic gradient descent influenced by learning rate to batch size ratio. In Artificial Neural Networks and Machine Learning ICANN 2018, pages 392-402, 2018
[34] Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations, 2019.
12

---

[35] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In Advances in Neural Information Processing Systems, volume 33, pages 77176-17186 2020.
[36] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[37] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
[38] Bobby Kleinberg, Yuanzhi Li, and Yang Yuan. An alternative view: When does SGD escape local minima? In Proceedings of the 35th International Conference on Machine Learning. volume 80 of Proceedings of Machine Learning Research, pages 2018.
PMLR, 10-15 Jul
[39] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, 2019.
[40] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint
2019.
[41] Stephan Mandt, Matthew D. Hoffman, and David M. Blei. A variational analysis of stochastic gradient algorithms. In Proceedings of the 33rd International Conference on International Conference on Machine Learning Volume 48, page 354-363, 2016.
[42] Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprini 1rXiv:1804.0000 1804.0661,,, 2018.
[43] Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: A view from function space. In Advances in Neural Information Processing Systems, 2021.
[44] Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry. Implicit bias of the step size in linear diagonal neural networks. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 16270-16295. PMLR, 17-23 Jul 2022.
[45] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint arXiv:1412.6614, 2014.
[46] Ryan O' O'Donnell Analysis of boolean functions, 2021.
[47] Francesco Orabona, Koby Crammer, and Nicolò Cesa-Bianchi. A generalized online mirror descent with applications to classification and regression. Mach. Learn., 99(3):111-433 jun 2015.
[48] S. Pesme, L. Pillaud-Vivien, and N. Flammarion. Implicit bias of SGD for diagonal linear networks: a provable benefit of stochasticity. In Advances in Neural Information Processing Systems, 2021.
[49] Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. arXiv preprint w:Xx::2304.004888, 2023.
[50] L. Pillaud-Vivie,, J. Reygner, and N. Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Proceedings Thirty Fifth Confer- ence on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pages 2127-2159. PMLR, 2022.
[51] H. Robbins and S. Monro. A stochastic approxiation method. Ann. Math. Statist, 2203):400-407, 1951.
[52] Samuel I. Smith and Quoe V. Le. A Bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018.
13

---

[53] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. J. Mach. Learn. Res., 1911):2822-287 2878, jan 2018.
[54] Terrence Tao. Concentration of measure. 254A, Notes I, Blogpost, 2010.
[55] Matus Telgarsky. Margins, shrinkage, and boosting. In International Conference Machine Learning, pages 307-315. PMLR, 2013.
[56] Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. The statistical complexity of early- stopped mirror descent. In Advances in Neural Information Processing Systems, volume 33, pages 253-264, 2020.
[57] Tomas Vaškevičius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, 2019.
[58] Roman Vershynin. High Sigh--imensional Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
[59] Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect. In International Conference on Learning Representations, 2022.
[60] Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part II: Continuous time analysis. arXiv preprint 1rXiv:2106.0220 02588888, 2021.
[61] Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Proceedings of Thirty Third Conference Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673. PMLR, 09-12 Jul 2020.
[62] Fan Wu and Patrick Rebeschini, A continuous-time mirror descent approach to sparse phase retrieval. In Advances in Neural Information Processing Systems, volume 33, pages 20192- 20203, 2020.
[63] Jingfeng Wu, Difan Zou, Vladimir Braverman, and Quanquan Gu. Direction matters: On the implicit bias of stochastic gradient descent with moderate learning rate. In International Conference on Learning Representations, 2021.
[64] Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized learning: dynamical stability perspective. In Advances in Neural Information Processing Systems, volume 31, 2018.
[65] Matthew D Zeiler. Adadelta: adaptive learning rate method. arXiv preprint 121.v.12 2012.
[66] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference Learning Representations, 2017.
[67] Xingyu Zhu, Zixuan Wang, Xiang Wang. Mo Zhou, and Rong Ge. Understanding edge-of- stability training dynamics with a minimalist example. International Conference on Learning Representations, 2023.
14

---

Organisation of the Appendix.
Appendix A, we provide additional experiments for uncentered data as well as on behaviour of the sharpness and trace of the Hessian along the trajectory of the iterates. We finally provide an experiment highlighting the EoS regime for SGD.
2. In Appendix B, we prove that (3x) follows Mirror descent recursion with varying potentials. We explicit these potentials and discuss some consequences.
3. In Appendix C we prove that (S)GD on the 1 {w w2 and u @ v parametrisations with suitable initialisations lead the same sequence (Bx))
4. In Appendix D, we show that the hypentropy va converges to weighted-(y-none when a converges to 0 non-uniformly. We then discuss the effects of this weighted -norr for sparse recovery.
5. In Appendix E, we provide our descent lemmas for mirror descent with varying potentials and prove the boundedness of the iterates.
6. In Appendix convergence (Proposition 1).
prove our main results: Theorem and Theorem 2, well quantitative
7. Appendix G, we prove the lemmas and propositions given in the main text.
8. In Appendix H, we provide technical lemmas used throughout the proof of Theorem 1 and Theorem 2.
9. In Appendix I, we provide concentration results for random matrices and random vectors, used to estimate with high probability (w.r.t the dataset) quantities related the data.
15

---

A Additional experiments and results
A.1 Uncentered data
When the data is uncentered, the discussion and the conclusion for GD are somewhat different. This paragraph is motivated by the observation of Nacson et al. [44] who notice that GD with large stepsizes helps to recover low 4 solutions for uncentered data (Fig. 4), We make the following assumptions on the uncentered inputs.
Assumption 1. There exist μ Rd and ö, Co,Ci C1, C2 such that for all s-sparse vectors verifying (1,, B) & co 31 |1 there exists & E Rd such that (X x)3 = and ei (3, p1) 12 < + 2, x (mi, & c2 c2(3,p) p)
where ||e|l2 & 8
Assumption 1 is not restrictive and holds with high probability for N (w1, 1a) inputs when >>> ol (see Lemma 9 in Appendix). The following lemma characterises the initial shape of SGD and GD gradients for uncentered data.
Proposition (Shape of the (stochastic) gradient at initialisation). Under Assumption and if (1s, (ppare) > col1 Iplloc the squared fill fill batch gradient and the expected stochastic gradient descent at initialisation satisfy, for some satisfying ||=1loo < =|pasparellll
VL(3)) = (3Sparse 11)21x² +
(14)
E [in) (V2, = + (8ipaarse
(15)
In this case the initial gradients SGD and of GD are both homogeneous, explaining the behaviours of gradient descent in Fig. 4 (App. A): large stepsizes help in the recovery of the sparse solution the presence of uncentered data, as opposed to centered data. Note that for decentered data with μ E IRd orthogonal to Spprrse there is no effect of decentering the recovered solution. If the support of μ is the same as that B Biparser the effect is detrimental and the same discussion as the centered data case applies.
Fig. 4: for uncentered data the solutions of GD and SGD have similar behaviours, corroborating Proposition 4.
Test losses || B 11
10
10 10-3
10 4
SGD
10 10-5
GD
100
10
Stepsize y
Figure 4: Noiseless sparse regression with a 2-layer DLN with uncentered data æ; (p11, Ia) where μ All the stepsizes lead to convergence to a global solution and the solutions of SGD and GD have similar behaviours, corroborating Proposition 4. The setup corresponds to (n, d, s, a) = (20, 30, 3, 0.1).
A.2 Behaviour of the maximal value and trace of the hessian
Here in Fig. 5, we provide some additional experiments on the behaviour of: (1) the maximum eigenvalue of the hessian 2 F²F((z)) the convergence of the iterates of SGD and GD (2) the trace
16

---

of hessian at the convergence of the iterates. As is clearly observed. increasing the stepsize for GD leads to 'flatter' minimum in of the maximum eigenvalue of the hessian, while increasing the stepsize for SGD leads to 'flatter' minimum in terms of its trace. These two solutions have very different structures. Indeed from the value of the hessian Eq. (22) at a global solution, and (very) roughly assuming that 'X" X Id and that "ax ~ O' (pushing the EoS phenomenon), one can see that minimising Amax((2 (V²F" under the constraints (w w²)
and w4 @ w_ is equivalent minimising [o under the constaint X y. On the other hand minimising the trace of the hessian is equivalent to minimising the -norm.
max (2F(w>))
Trace V2F(W Trace(V2F(wv
20.0
5.0
SGD
17.5
GD
4.5
15.0
4.0
12.5
SGD
10.0
3.5
GD
10
10
10
10
Stepsize Y
Stepsize y
Figure 5: Noiseless sparse regression setting. Diagonal linear network. Centered data. Behaviour 2 different types of flatness of the recovered solution by SGD and GD depending on the stepsize. The setup corresponds to (n, d, a)) = (20, 30, 3, 0.1).
A.3 Edge of Stability for SGD
Coordinates B( of SGD at EoS
200
supp(pipaase) coordinates other coordinates
150
Gradient flow
100
50
100
10a
10
10a
10a
10
Iteration k
Figure 6: SGD at the edge of stability: all coordinates fluctuate, and the sparse solution is recovered. As opposed to GD at the EoS, since all coordinates fluctuate, the coordinates to recover are not more penalised than the others.
17

---

B Main ingredients behind the proof of Theorem and Theorem 2
In this section, we show that the iterates (Be) follow a stochastic mirror descent with varying potentials. At the core of our analysis, this result enables to (i) prove convergence of the iterates to an interpolator and (ii) completely characterise the inductive bias of the algorithm (SGD or GD). Unveiling a mirror-descent like structure to characterise the implicit bias of a gradient method is classical. For gradient flow over diagonal linear networks [61], the iterates follow a mirror flow with respect to the hypentropy (4) with parameter the initialisation scale, while for stochastic gradient flow [48] the mirror flow has a continuously evolving potential.
B.1 Mirror descent and varying potentials
We recall that for strictly convex reference function h Rd - R. the (stochastic) mirror descent iterates algorithm write as [5, 18], where the minimum assumed be attained over IR and unique:
3 argmingeR4 {nk (gk=, 8) + Dí (B, B2))
(16)
for stochastic gradients yh, stepsize 7kk 2 0, and the Bregman divergence associated to h. Iteration (16) can also be cast
= h(33 h(B)) (Vh(8), 3*)
vh(3<+1) = (b(3,)) Tr:lk
(17)
Now, let (h&) ) be strictly convex reference functions Rे R. Whilst in continuous time, there is only one natural way to extend mirror flow to varying potentials, in discrete time the varying potentials can be incorporated in (16) (replacing h by hę and leading to (h+(3k+1) Wh (3x) 7a.9%). the mirror descent with varying potentials we study in this paper incorporates h&+1 and h2 in (17). The iterates are thus defined as through:
Bu+1 = argmin {nA B) + Dhz (8,3k) Bx)
where Dheti.he by (8,8") = hx+1(3) hx(33) ( (Vhn(3'), 3¹), recursion that can also be cast as:
Vhe+i (B Th+(BB) "(k-
To derive convergence of the iterates, we prove analogs to classical mirror descent lemmas, generalised to vare-vg potentials.
B.2 The iterates (Bx) follow a stochastic mirror descent with varying potential recursion
In this section we show and prove that the iterates (Ba)a follow stochastic mirror descent with varying potentials. Before stating the proposition, we recall the definition potentials. To do so we introduce several quantities.
Let g, -= R RU { be defined as:
a+ (x) = F2r ((1 F x)))
1
(④x) (a+ (x) 9- (x)) = In ((1 2²)²),
with the convention that ((1) co. Notice that q(x) > 0 for |x| < 2 and q(z) < 0 otherwise, For the iterates 34- O Up E Rd, we recall the definition of the following quantities:
L exexp
41 V (Be)) E R
Y=0
a a+ aa-k.
a
dy = arcsinh
E Rdd
20
Finally for l: > 0, we define the potentials (hr- Rdd + R)a30 as:
hx(33) Voy (3) (0k, B)
(18)
18

---

where Va, is the hyperbolic entropy defined in (4) of scale axp:
d
Bi
Va. (3) - Σ Brarcsinh
V B3 + a2, + oz.5)
where i corresponds to the 7th coordinate of the vector O'--
Now that all the relevant quantities are define, we can state the following proposition which explicits the time-varey varying stochastic mirror descent followed by (Bk)x
Proposition 5. The iterates (3g, up C oR)xzofrom Eq. (3) satisfy the Stochastic Mirror Descent recursion with varying potentials (ha)k:
Tha+1(3x+11) Vhu VA (3x))
(19)
where h2 Rd R for 2 are defined Eq. (18). Since (ho = 0 we have:
(ha( E sspan(z)......
Proof. Using Proposition 6, we study the {w] w2 parametrisation instead of the u v, indeed this is the natural parametrisation to consider when doing the calculations as it "separates" the recursions on w4 and w...
Let us focus on the recursion of w4:
W-,k+1 - 77 (Bx)) W4.6.
We have:
w k+1 = (1 11--0 VA VLB. (3x))² w = exx (In((1 (1 > VL. c (3x)) w
with the convention that exp(In(0)) = 0. This leads to: / +++++ exp (- 27t LL, (wxx ) + 27x V Lo, (3x) + In((1 7t VL (3x))²)) wh. = exp ( 27x V LB, (3x) 7+ (7% VLs (3x))) -
since Q+ (x) 2r In((1 x) )) Expanding the recursion and using that w+. is initialised
W+, k= a, we thus obtain:
k-
h-1 Σ Zm
w k = a²
expl
Σ
9 VLE, exp (-2
v VL, (B))
a² exp (-2 Σ
"/ V Lo, (3e)),
where we recall that a²,x = expppp get:
Σ au (- One can easily check that we similarly
--1 1
w² a² & exp (+2
Σ
V LB, (B))
t=
leading
1
3x
2 (wi w2,x)
k-1
k-1
1 -² v exp (+2 Σ 5
a + exp (-2 Σ VL. (3))
70 Lz, (8)))
2
t=
1=
Using Lemma 4, the previous equation be simplified into:
A-1
2 a a
-
Bx 04.60 & sinh (-2
Σ
78 V (Bc) + aresinh
20
1=0
19

---

which writes as:
aresinh 3x
k-1
Φk
vt VLB, E span( F.))
1=
where d2 2/2 aresinh ( -i.+00
a a+, o a- & and since the potentials hy are defined in
20
Eq. (18) as hę = Vax - (0x: with
(Barcsinh( B
a (3)
VBB
B + at +a;)
(20)
i=1
specifically such that Vhu (Bx) aresinh (6.7 Ok. Hence,
(ha( = Σ
Ne LB, (32)
8<k
so that:
Vhy+1(Bk+1) = (hx(3s) Vik VLB, (3x),
which corresponds Mirror Descent with varying potentials (hu)a.
C Equivalence of the @ and (w -w=) parametrisations
We here prove the equivalence between the [ w²) and u @ parametrisations, that we use throughout the proofs in the Appendix.
Proposition 6. Let (Bx)k20 and (8))k20 be respectively generated by stochastic gradient descent on the u C and 1 [w? - w2) parametrisations:
(u 0641) = (ug, 0k) mh V Vu,o (Ls, (u & u)) (uk, uk),
and
Wi,k+1 Wi.k 7h V wx (CB, ( ([( w2 ))) ( w)
initialised as ug = 20, 10 = and 0.0 = w w- = a. Then for all 3 O, we have Bx = 3k
Proof. We have:
W1.0 = a: Wi.++1 (1 VL. (
and
No 22a,
Up = =0, U1+1 Uk 7/ VLB, (3&)Kk, Tk+1 Uh 7 V (L..))u
Hence,
(u Be+1 (1 + 7 VL (M +
Ls,,
and
P1 (1 + 2 L.. 11 - mw + w² (2,xvV/l V (8%)
Then, let 24 = (ug b and 2 w A k. We have 20 = a² 2/ a² and:
2h+1 = (1 - 22 v LB, (8x) 2)2k
2 = (1 22 LB. 2
Using a² + 12 = / (2ab)) + (a2 12 for a, b RR we finally obtain that:
2 o (23x)² (2
w w + (233)2 +2=6))
We conclude by observing that (3k, 2x) and (8% 2() follow the exact same recursions, initialised at the same value (0, a²))
20

---

D Convergence of dh to a weighted é norm and harmful behaviour
We show that when taking the scale of the initialisation 0, one must be careful in the characterisation of the limiting norm, indeed if each entry does not go to zero "at the same speed", then the limit norm a weighted ---oorm rather than the classical norm.
R. Then we have Proposition 7. For 2 O and vector h E IRd, let á = weepppeppnn(((ee that for all 3 E Rd
)n( Σ (1+h)))811
(a(3) a-u
Proof. Recall that
d
va(3) = Σ (,arcsiinh
33. + ai, + a?
Using that arcsinh(())
sgu(x) (()x) and that In + = (1 + h::) In( +
we obtain that
10/
(B( a-0
Σ
sgn 3,) + 14) In
= h
+ k.)
a
The following Fig. 7 illustrates the effect of the non-uniform shape a on the corresponding potential daa
Figure 7: Left: Uniform al: smaller scale leads the potential vh., being closer the & -orm. Right: non uniform can to recovery solution which is very far from the minimum -oor solution. The affine line corresponds to the set of interpolators when n = 1, and = 1.
More generally, for such that + 0 for all i E [d]] at rates such that In(1/as) In(1/ max; ac). we retrieve a weighted & norm:
(B)
qi 11
ln(1/a2)
1=1
21

---

Hence, even for arbitrary small max; a, if the shape of a is 'bad', the interpolator Pa that minimizes % can be arbitrary far away from Sh the interpolator of minimal & norm.
We illustrate the importance of the previous proposition in the following example.
Example We illustrate how, even for arbitrary small max, an the interpolator 3% that minimizes Va can be far from the minimum & norm solution, due to the shape of a that is not uniform. The message of this example is that for a + O non-uniformly across coordinates, if the coordinates of that go slowly to O coincide with the non-null coordinates of the sparse interpolator we want to retrieve, then 3 will be far from the sparse solution.
A simple counterexample can be built: let PSparse = (1, (...... 1.O....,0 (with only the = o(d) first coordinates that are non-null), and let (1:), (1s) be generated as yi = (B) sparse: Fz) with xq ~ (0, 1)- For n large enough (n of order s =I(dd) where is the sparsity), the design matrix X is RIP [10], so that the minimum & norm interpolator Sp. exactly equal Siparse
However, if or is such that max; 04 - 0 with h; >> 1.for j & and h; = 1 2 + 1 (h: as in Proposition 7), 3a will be forced verify Ba, 0o i & and hence 11 B7. 2 8.
E
Main descent lemma and boundedness of the iterates
The goal of this section is to prove the following proposition, our main descent lemma: for well- chosen stepsizes, the Bregman divergences (Di. (3", (8',))))) (B3)) decrease. We then use this proposition to bound the iterates for both SGD and GD..
Proposition 8. There exist a constant > 0 and B >> 0 such that B Oinfg-es || Il)) for GD and = (In(1/a inf es || ) SGD, such that if " & LB for all k, then we have, for all k > 0 and any interpolator 8 E S::
Dis+i (8", Bx+1) & Dne (K*, Yu (3x)
To prove this result, we first provide a general descent lemma for time-varying mirror descent (Proposition 9, appendix E. I), before proving the proposition for fixed iteration & and bound B on the iterates infinity norm in Appendix E.2 (Proposition 10). We finally use this to prove a bound on the iterates infinity norm in appendix E.3.
E.1 Descent lemma for (stochastic) mirror descent with varying potentials
In the following we adapt a classical mirror descent equality but for time varying potentials, that differentiates from Orabona al. [47] in that it enables us to prove the decrease of the Bregman divergences of the iterates. Moreover, as for classical MD, it is an equality.
Proposition 9. For h, IR + R. functions, let Dr [B,8) 3* = h(B) -(0) (Vg(3",3 B"y for 3' E Rd. Let (hx) strictly convex functions defined IR L convex function defined on Rdd Let (Bx) defined recursively through So E Rd, and
Bx+1 E argmin (TH (V B -Ba) + DreiiAh ny (8, (b.)
where we assume that the minimum is unique and attained in IRd, Then, (Bx) satisfies
Vhu+ (34+1) = (hv(33) - 7A VL(x))
for any B E IR.d
Dresi (3,3k-1) = (B, (VC()))) -3) + Dreai (3k. Bx41) (h441 - 11x) (3x) + (1x++1 - (B).
Proof. Let 3 E 1Rd. Since we assume that the minimum through which Ak+1 is computed attained in Rd, the gradient of the function Vz) Be41 is null, leading to Vhg+1 (8x+1) Wh (Bx) 7/ (Bx2)
(V ((3)), BE) + Dhetishy (3, 3x) evaluated at
for h g. recover the classical Bregman divergence that we denote Dr Da.h
22

---

Then, since VV = 0, we have Dv, (8, = VV Using V = h2
we also have Dv, D Hence:
Diaia (8, 1k++1) - 7 (VL(3k); Bu+1) + Dis 11 (3, Din 1.ke (BR+1,
We write 7A (V(VL(3),
8x+11 7A (V (bh) + Th (VL(BB), 8n - Bk+1). We also
have 7a (V (VL)(B)) By Bk+1) (Bk) Vhx+1 (Bk+1), ,,,,,,,.0 3x Bk+1) Dat (Bk, 3k+1) + D be (5K+1, (") so that nu nh (VL(Bk). .xx Bb+1) D Diattbb .b (Bx+1, 1. 18h) Dietieats (3k, 3x+ 1). Thus,
DrL (k+1) = D Dhzt (3,8%) - 7 (D, (8. + (,(3,, B)) + Dha,hets (3kr Bk+1)s and writing Dr.g (3, 3*) (g( lx) ((3 concludes the proof.
E.2 Proof of Proposition 10
In next proposition,
Proposition 9 to prove our main descent lemma. To that end, we bound
the error terms that appear in Proposition 9 as functions of LB, (8x) and norms of 3x, Bk41, so that for explicit stepsizes, the error terms can be cancelled by half of the negative quantity -2 (3x).
Additional notation:
let L2, Loo > 0 such that v3, |HB3112 < L|18||2. |1
< _l for all
batches B C [n|] of size b.
Proposition 10. Let k R 0 and B > O. Provided that [|3x || 11 2mx & B and 7 < LBB
where c 0 is some numerical constant, we have;
Diuen 8x+11) & Di, (8" (k) VkLB. (3x)
(21)
Proof. Let 8 E S be any interpolator. From Proposition 9: Dian (18,,,+))
Dr. (3*, BBx) 27 Lss, (Bx) + Dan (Bk+1, Bx) (h++1 hx) + (h++1 - -ha)
We want to bound the last three terms of this equality. First, to bound the last two we apply Lemma 7 assuming that || : 1| < B:
- -hax11 - (3k) + (ha+1 - hw) < 24BL (8x)
We now bound Diati 1 (3k, 3k-1). ). Classical Bregman manipulations provide that Day (Bk, Bx+1) = Dai_, (Vh++1 ()++)
- (Vh+(5*) - > VLs, (8x), Vh++1
From Lemma 6 we have that hx+11 is m 1/(4B) strongly convex on the ec-tceneeeed ball of radius B therefore hi is max(4a 4B) = 4B (for a small enough or B big enough) smooth on this ball, leading to:
Dine 11 (3k, Bk+1)
2B (3t) 7t VLs. (3x) Vha+i (Bx))l
< ( - Vha+i (3x)
1|2 2 + Th V LLs (8x) 112)
Using (3) Vhe+1(3) < 26 where dz = ((7 LL. (3t,),, we get that:
DiLIs (Be, 3k+1) < <8B11b + 4BL (9x).
Now, 110
| | and using Lemma 5, || loc 21/ LL. (Bx) since 1h VL, (Bx)
41|78 VL2, (Bx) v LB. 11x)
E & 2h 2LB 1/2 is verified for 2 & 1/(4B)). Thus,
MhLoo 11 Booll
Dny, 1 (Bk, Bk+1) < 40B Ls, (.x).
Hence, provided that ||Bxll x0 < B, 1| loo < B and 2h < 1/(4LB), we have: Dists (8", Bk+1)) & Dry (8*, B&6) 2% (3x) + 64 (8k),
and thus
- ( Draan (8+, 8x41) < (B* Bx)
(Bx)
The & e BL
where == 64
23

---

E.3
Bound on the iterates
We now bound the iterates (3x) by an explicit constant B that depends on || (for any fixed 8* E S).
The first bound we prove holds for both SGD and GD, and is of the form 0(1| In(1/o² while the second bound, that holds only for GD (b = ni) is of order (1, II1) (independent of o). While a bound independent of is only proved for GD, we believe that such result also holds for SGD, and in both cases B should be thought of order (11))*1),)
E.3.1 Bound that depends on a for GD and SGD
consequence of Proposition 10 the boundedness of the iterates, as shown in next corollary, Hence, Proposition 10 can be applied using B uniform bound on the iterates 8 norm.
Corollary 1. Let B = =3113*11 In (1 11 For stepsizes "Ik & BL k +>0 O.
we have 11 & B for all
Proof. We proceed by induction. Let k > 0 such that | & B for some B 0 and
DNe (8*,8e) & Di, 1 So) (note that these two properties are verified for k = 0, since 3n = 0).
For 7 sufficiently small (ie., that satisfies Yik & By where B' > || pl || lo: || aoo)) using Proposition 10, we have Drats (B", Bk+1) & Dny (3 3Bc) so that Dhats ( ,x+11) & (3 , ,0)) which can be rewritten as:
Σ at aut 111-
Shttii 1h1,
&
Σ 1=1
B aresinh
SktLi a²
Hence, 1|3x+1ll1 < || 10 [10atall)). We then notice that for x. x x 3y (n + y): i > yym(1+ + y) and x > y, we have that y + y) < y yln(1 + x), so that 1 1+ y < + x, which contradicts our assumption. Hence, max(y. In(1l111111111111 + v)). In our case,
y. > 0, yln(1 +x)
3 || /a2 y = 11 so that for small alpha, In(1 + v) > 1.
Hence, we deduce that 1|3 & B, where B = || In(1 (1 + -
This is true as long as Yt is tuned using B' bound on max
of Bk+1 as a function of yx (3g being fixed), we show that 76 < || llxe) Using the continuity Indeed, let Φ IR+ >> Idd be the function that takes as entry 7 > 0 and outputs the corresponding || & is continuous. Let 7 2 x rf. for set is upper-bounded, if is is empty, we do not need what follows since it means that any stepsize leads ||,3+11l
BE can be used using this B.
c
0 and
sup fr 2 0: B < <(%+) (the
B). By continuity of é, p(7x) B. Furthermore, for all r that satisfies r b max((6(7)) B) 2 max(
| Iloes |1.3 we have, using what is proved just above, that 11 & B and thus ((7)) & B for such
Lemma 1. For > O such that max(@(?))) B). we have d(x.) < B.
Now, if > B, by definition of and by continuity of de since ((F) B, there exists some B < < F such that x(7+) > B (definition of the supremum) and d(7,)) & 2B (continuity of Φ). This particular choice of r thus satisfies r > B and and ((7,) < 2B < 2r, leading to @(T)) < B, using Lemma 1, hence contradiction: we thus have 7 < B.
This concludes the induction: for all r > B, we have r > so that p(y,) < B and thus for all stepsizes < 2 we have || < B.
E.3.2 Bound independent of a
We here assume in this subsection that 6 = We prove that for gradient descent, the iterates bounded by a constant that does not depend on a.
Proposition 11. Assume that n (full batch setting). There exists some B 0(118+11, such that for stepsizes 7kk BL- have || & B for all k 2 0.
24

---

Proof. We first begin by proving the following proposition: for sufficiently small stepsizes, the loss values decrease. In the following lemma we provide a bound on the gradient descent iterates (w+,k; (..))) k)) which will be useful to show that the loss is decreasing.
Proposition 12, For Tok LB where B b max(]]
|| |l we have L(BK+1) & ((3x)
Proof. Oddly, using the time-varying mirror descent recursion is not the easiest way to show the decrease of the loss, due to the error terms which come up. Therefore show that the loss is decreasing we use the gradient descent recursion. Recall that the iterates wi = (W+,bs w-,k) E Rdd follow a gradient descent on the non convex loss F F(w) = ~ lly x w²)ll2
For & 0, using the Taylor formula we have that F F(Wx+1) < F(wx) 7h(1 2 F(V))()))]
with the local smoothness L3 suPwe[[w..ze Amax [V)) Hence if 7/ 1/ La for all & we get that the loss is non-increasing. We now bound Lxk Computing the hessian ot F, we obtain that:
F²F(wx)= [diag( (LLk))
0
0
(iag)))L(Bk))
(22)
diag(u (+,, x)H diag(w (w+,a) +- -diag((w+ (w+ H diag (w-.))
diag(w [ (lag(w+,+)) diag(w-,x:) diag(w-.k)
Let us denote by M ( M+
M+ M+.-)
E R2dx2d the second matrix in the previous equality. With M M=-)
M+,-
(VV(())) + M|| (where the norm corresponds to the Schatten 2-norm which the largest eigenvalue for symmetric matrices). Now, notice that:
this notation ((w))( )I &
||M||
sup wERed ||||==1
|Mull
aa W M
112
sup u4 " ll ||=11=111 n- M-ERd,11 11 ||=1 (a.b)ER* ,0²
b u
We have:
M
a a.u+) u4 b u
1|- =
(aM+u_ + BM4 u
(a -- 44 + bM u
|laM + bM 112 ||aM u4 + bM 2 _11
2 2 u4 112 1
1+ u_ 1|2 a²l - 12 12 u- 12132
< 2 |M- 2 + 11 |lM 112)
Since || & Amaz
Imll 113c and || < Amaz w |0 | 1l,o we finally get that
|
& 612 max [ 12
|| 1
< 612 2mp (1|w 1 1l
||=2 )l
< 1212 [ 2 I13
We now upper bound this quantity in the following lemma.
Lemma 2. For all > the following inequality holds component-wise:
w + w²
/ 4aa + 18.
Proof. Notice from the definition of w+. and w-,k given in the proof of Proposition 5 that:
1up,ullue -|llwa_ Ml 0 oz.
(23)
25

---

And = =²² Now since ar is decreasing coordinate-win (under our assumptions on the stepsizes, 2h (L(x)) (1/2)² 1), we get that.:
2 +,6 + w =
2/ + B 22 + S
leading to w w² 440 B
2vor From Lemma 2, w + w² & is bounded by 2v ar + B Putting things together we finally get that
[[2)( < [VC(3) dc + 8Amar / B². Hence,
La < sup =
|V 82 max v /1 B LB 8A VIlael B < 10LBB
for B ||a.
We finally prove the bound on || independent of a for a uniform initialisation a a1, using the monotonic property of
Proposition 13. Assume that b = (full batch setting). There exists some B = O(1 such that for stepsizes T/k c we have 1/ ll & B for all k 2 O.
Proof. In this proof, we first let B be bound on the iterates. Tuning stepsizes using this bound, we prove that the iterates are bounded by some B' = 011)8*11). Finally, we conclude by using the continuity of the iterates (at a finite horizon) that this explicit bound can be used to tune the stepsizes.
Writing the mirror descent with varying potentials, we have, since Vho( 0.
(Ba( = Σ VL(B))
eck
leading to, by convexity of hy:
hx(32) - hx((8)) < (Ven, - -B)) Σ 3x
Esk
We then write, using VC(3)= = H(3 - -3-) for for H XXT, that Eech VL(()), V Br - 8*)
Eecr xe(x (Bg 8x),,' (Bx 102--1)) & Ve Te ((Bx) leading to:
< Σ (e(()3)
Σ (e(())) < Σ (e(()) & 2D 30))
tcc
tsc
tck
where the last inequality holds provided that 74 < ULB Thus,
vo (3x) < Va. (8") + 21h (3") + (ches 3r - 3²)
Then, (0x, Br - B*) < @ 11 - 1l and 1 < C.euax 2 (BC( Ymaxx?maxxio (3*). Then, using
(B) In(1/a2) < [1+ In([1]]] + +a²) a²)) Im(1/a?) (
113 - In(1/a?) <
we have:
|| to < In(1/a2) + 11 In([ In(113*111 + a²)
+ 11 In || + a²)
In(1/a2)
In(1/(2
+ CahC''mmxTaaaaa ho In lm/1/22)
< R + CmCAmax ( (3*)) in(1/n²)
where R= -(1,331 ||, is independent of a. Hence, since Bg = supseco 11 < 00, we have:
Bo (1 CAmax Ymaxx ho
(13)/// R Bo < 2R.
26

---

provided that Ymax < 1/(20) / ( (note that ho(33)//n(1/6) In(1/a²) ( is independent of a²).
Hence, if for all k we have "he & LB where B bounds all | we have 11 & 2R for all k
where R (11)3*11, is independent of and stepsizes "
Let K be fixed, and
inf 7-20
sup kgk
| > 2R
For y 2 a constant stepsize, let
(0(7) = sup
||
1o0
kgK
which is continuous function of For > 0, let
C
An important feature to notice is that if y < 7r and bounds all |Bxll k & K, then (p()) & R. as shown above. We will show that we have 7 //2 72R- Reasoning by contradiction, if < 72R by continuity of yp, we have (77) & R and thus, there exists some small 0 < < 728 such that for all E [7, 7 + El, we have (7) < 2R.
However, such y's verify both p()) & 2R (since 7 e) and by definition of E) and y & 72n (by definition of E), and hence (77) < R. This contradicts the infimum of 7, and hence 72R- 2C"RR. we have 13 < R.
Thus, for 7 & 72R
F Proof of Theorem and 2, and of Proposition 1
F.1 Proof of Theorem and 2
We are now equipped to prove Theorem and Theorem 2, condensed in the following Theorem. Theorem 3. Let (ug, va) follow the mini-batch SGD recursion (3) initialised at V2a e RS and un = 0, and let (3k)sz0 (ux O vk)xzo There exists and explicit B > 0 and a numerical constant e such that:
1. For stepsizes satisfying "Ik & LB the iterates satisfy Ih VLs (Be)ll < 1 and 11 & for all
2. For stepsizes satisfying " & LB (Bx)a30 converges almost surely some 3 S,
3. If (Bu)x and the neurons (uN, Uk)a respectively converge model 3% and neurons (Woo, Voc) satisfying 3% E S (and 3 ux O voch then for almost all stepsizes (with respect to the Lebesgue measure), the limit 3% satisfies:
8 = argmin Den (8*, 18), 3.
for O0 E ER and Bo E Rdd satisfying
a = =a 3 exp
Σ
( VL. (B)))
k=0
where q(x) = - [nn(1 x2)2) > 0 for
|x| < 22, and Bo is a perturbation term equal to:
3, = a a²),
where, (z) F2x (n((1 ()) and af = a² c exp (- De a (7h VL8, (Bx)))
Proof. Point The first point the Theorem a direct consequence of Corollary and the bounds proved in appendix E.3.
27

---

Point 2. Then, for stepsizes
< c LB using Proposition 8 for any interpolator 5 E S:
D (8t, (B8+1)= Dia (3*, (Bk) Tw (3x)
(24)
Hence, summing:
TBLBB(B2) < (8", Σ &
so that the series converges.
Under our stepsize rule, 1h VL. (3x < leading to Vll
31/78 VLs, (B+) by Lemma 5, Using VLB, (3x)]l & 2L2CB, (Bu), we have that In(a+,)),
In(as) all converge.
We now show that En (x((3k) < co. We have:
((k)) = InCB. (3x) + Mr
8ck
tCk
where M₁ = =eck (x((())) LB. (3x)). We have that (M&) is martingale with respect to the filtration (Fx) defined as Fx = 0 (Be, l & k). Using our upper-bound on Eeck Tk (3x), we
have:
M& >
Σ TkC((B) - - (3x) > -Dho (3", ,0),
Eec
Ecc
and hence (M&) is lower bounded martingale. Using Doob's first martingale convergence theorem (a lower bounded super-martingal converges almost surely, Doob [17]). (Mx) converges almost surely. Consequently, since Deck (xx) Σec# Tn LL, + we have that > Rcch 7n. (5x) converges almost surely (the first term upper bounded, the second converges almost surely).
We now prove the convergence of (3x). Since it is bounded sequence, let (a(()) be convergent sebuse and let 18 denote its limit: Ba(k)) + 8%
Almost surely, En N/s ((3x) < oo and so ((3k) - O, leading to ((xx) - 0 since the stepsizes are lower bounded, so that (((a(k) + O, and hence ((3²) O: this means that 3 is an interpolator.
Since the quantities (cc)), (a+,))x and (Φk)x converge almost surely to a a+ and d.c. get that the potentials h& uniformly converge to l1o0 = vaw - (#xx -) on all compact sets. Now notice that we can decompose Wh (1²²) as:
Vh (B%%) (hh
Vho (Bo())) (Bok)() ((o(3)))) + Whocky (k)()))
The first two terms converge to O: the first is a direct consequence of the convergence of the extracted subsequence, the second is a consequence of the uniform convergence of (l to hos on compact sets. Finally the last term is always in Span(ar...... ...,,,,) .n) due Proposition 5, leading to Vhoc ( E (ppan(11....,,,,, Rn)) Consequently, Wh (B*)) Span(21..... Notice that from the definition Vva (Bt)) dxc. Now since dxo arcsinh 20
of hoc. we have that Vh (B
one notice that 3g is preciscly defined such that V (Bo) 2x Therefore Via (Bx) Vrh.... (3)) Span(x1, Zin) This condition along with the fact that 8 is an interpolator exactly the optimality conditions of the convex minimisation problem:
min
Denan (B*, Bo)
BLES
Therefore 8% must be equal to the unique minimiser of this problem. Since this is true for any sub-sequence we get that 3r converges almost surely to:
8
argmin D Dvan (8*,8) BES
Point 3. From what we just proved, note that it is sufficient to prove that Orky Ot dh converge to limits Otoo, Ott,oo &oo satisfying aoo, a E IR (with positive and non-null coordinates) and &x E IRd. Indeed, if this holds and since we assume that the iterates converge some interpolator, we proved just above that this interpolator is uniquely defined through the desired implicit regularization problem. We thus prove the convergence of Cii, Ot Øk.
28

---

Note that the convergence of uk. Up is equivalent to the convergence of W in the w - w² parameterisation used our proofs, that we use there too. We have:
W+,++1 = (1 79 V LL. (3x)) @ w
so that
In(w 1n(w2,x)
Σ VL Bcck
We now assume that stepsizes are such that for all L f 0 and i E [d], stepsizes are such that we have Io V, LB, (Be) * 1: this is true for all stepsizes except countable number of stepsizes, and this is true for almost all stepsizes. Since we assume that the iterates 3ą converge to some interpolator, this leads to 7/e VLs, (3t) - 0 we assume that stepsizes do not diverge.
Taking the limit, we have
In(w7,od )= = E 7L
EX000
This limit is in ({-oo) U IR) (since Wi E IR.), and coordinate of the limit is equal to -oo if and only the sum on the RHS diverges -0o (note that from our assumption just above, no term of the sum can be equal to -00).
We have
F2% VLB (Bt) as l + oo, so that if for some coordinate we have Ee 7/ (, Foc, then the coordinate i of the limit satisfies In(w) = +20, which is impossible. Hence, the sum Σ. 7e VL, (Be) is in Rd (and is thus converging); consequently, Ee Ls, (Be)² converges and thus E. g V LB, (Be)) and Ee (Be) all converge: the sequences akk a thus converge to limits in R and Øe converges, concluding our proof.
VL,,
F.2 Proof of Proposition 1
We begin with the following Lemma, that explicits the curvature of Da around the set of interpolators.
Lemma For all is > O, if ((3k) 2B 2B 2 LL(L).
<
1 (a ( ). we have | 8z 1|
2Awas
Proof. Recall that the sequence z (3** satisfies z 0 and z - so that
we have that z E = Im for all k > O. Then, let (% be the unique minimizer of h4 over S the space of interpolators: S6 is exactly characterized by X ph = Y and (Bg) V. We define z E V as z Vh, (8g).
Now, fix x = z and h = he, and let us define v/2 2 E V - Da- (z. 2*) and Φ 2 VEV - L(Vh' (z)). We next show that for all z E V. there exists μe such that V²Ø(z) > Ps (z) and that μ4 > a for z in an open convex set of around z", for some > 0. For A E R Rdxd an V.
operator/matrix on Rdd let us denote Ay its restrictioncco-restriiioon
First, for z E V. we have (z)(2) = (²( (z) (²(2)) - (2 2d) ( (VV
Then, Vø(z) (VVh())) so that v2ó(z)
Since h is 1/(2a²) smooth (on Rd and thus on V), he is 20² strongly con- vex (on V and on Rd). Using V have (( (z)) = V2h*(z)v V (z) 2 2x i V2h"(z)v.
lm and V = xx we and thus
(PP
we compute V h(()))))) 1==z= sinh(2;). leading to: (I( (z),, = diag(2a² sinh(z) @ (XX (2a² sinh(z) - ))))) Thus, writing B2 20 sinh(z) (z"(z) the primal surrogate of
For the other term of V namely (V
29

---

we have:
VVV(0VVZ(VV --vv diag(2a sinh(z) : (XX T (82 - - 1))
= xX (5e - 5%) [lag(2afz o sinh(z)])v
= xx (Be - [8) diag(2a cosh(z))v
xx (62 - (E)1 Pzb(z)..
Wrapping things together,
V (20(Z) = (2a² x min |X T (32 - - 1F
Let 2
{ V |XX (32 11) < a²>x X that satisfies
{BEV V ((22) < 2dmar
1 (a min 12 C 2. 2 is open convex set of V containing 2n On 2,
V2oz 2 V and 2x(2)) = 6(z²)( 0. that for all z E Z, we have d(z) 2 Xtin (z). Hence, for all z 2, we have Drc (3g 32) < Dr, 2") & (a²) ((22). (((2)) and using the fact
that Di. is 4n strongly convex, we obtain, for 82 Be (since z E V): ((3x) & 2A (a² (
have 1 Bxl &
Proposition 14. As assume L is L.-relatively smooth with respect to all the h2 's. Then for all B we have the following inequality.
(e(((()+ ((3)) & (3, - Disis (3, Bx-1) 1) - (1 (3u+1, 3s) + (hx+1 - h3) - (hx-1 - h+)(34+1).
Proof. For any BK Ba+1, the following holds (three points identity for time varying potentials, Proposition 9):
Dbs (8,3k) D Do (8,3k+1) = [hv(3) (hx(3)) (Vh) Bx))]
- (x)+1 (ha-1(8k+1) +
- Bx+1))]
(B) - h (8) + (Vh++ Vh2 (5x), BKt1)
+ h4+ (3x+1) - (Bxe (Vhn(3a), 3k+1 - -3a)]
hb4 h (B) + (Vh++I (Bx+1) (3 3 BK+1))
+ h (3++1) - (B + Dn. 1x))
Rearranging and plugging in our mirror update we obtain that for all f:: "ik (VL(3)), Bx41
33 Din (B, 3x) D (B,
- DD, (BK+1, 3x) (he+1 - ha (8e+1) + (ha)1-ha)(3)) From the convexity of C and its -ellatiee smoothness we also have that: L(3++1) & L(B) + (V (3x)))) Bu+1 38) + L,, Dro (8K41, Bk),
Finally:
(c(33k)1) (()) & (B, Dises (8, 3/-1) (1- -1, (3k+1, 3a) + (hx+1 (3) - (ha-1 - (+
Note that in our setting, for any 8, k + hb(3) is increasing. We can therefore write that: (c((kk))) L(B)) & Dru (B, 3k) Drn.i Bk+1) (1--(br) Drs (8x+1, BB) + (b441-h4)CC (3) In particular, for 3 = -3:: 7hk ((B+2) < Dine ,3x) Dren (k+1) - -1-7x L) (3x+1: 3x) + (hx+1 - hx)(3*)
(h++1 (8k+1)
< Di. (3",,, Drum
- (1 7th I, Dr, (3k41. + (he41 kax)
and in 3 = =Bx:
7h (xx+1))1) & Tx Diu (Bkr Bk+z) (1 L, Di,, (8k+1, Bk) + (ho+1 (3k) (he+1 (Bk41)
< (BBB)) -ru-i (5k, Bk+1) (1 - 7h I., Dr., (Be+1, Bk) + (14+1 (Bke
30

---

Proof of Proposition 1. We apply Proposition 14 for B = 5k. with L, = 4BL (using Lemma 6) and
replacing by LB., obtain:
Nnu (CB, (5x+1) LBu (3x) -Dhas (3k, (k+1) (1 111-20 -> (0 (8++1.Bk) (h++1 hxa(3a) (h++1 (Bk+1),
and thus, taking the mean wrt BBs
Yo ( ((3++1) - [(Bx)) < - -Bg, Disfi (5k, 3R+1) - (1 - L, Dne (Be+1,
+ 1B (hx+1 - (Bx) - EB, (1x+1 h2)
Th Tx )B, DNu (3k+1, Bkx
+ Es, (hx+1 - (Bx) EB, (140+11 - (xx(1)+1)
First, as in the proof of Proposition 10, using the fact that h4 is 1n(1/oz)) smooth,
1
Da, (x)
>
(Bx) - 7h VLB, (8a) - V/u + Vhg+ (86+1(9941)+++) Vheg (3k+ 1)
2 (n(1/ak)
Vhu (Ba) Vha 1 (3x)112 4 41(1//x)) 17 VL5. (3w))2
2 l
and thus
E.D (3K+1, 3x) E
2 2l((1/))) [Bh) Vh&+1
1
2 2I((1/ow) (c
A
Now, we apply Lemma assuming that || 1| & (which is satisfied since we are under the assumption of Theorem 2):
(h++1 (3,) - (ha+1 - hx) <
Using (hx Vh++1 (3)] & 26 where 6y = 4(7h VL. (3x)) as in Proposition 10, we have:
E V(++)) Vha+ (3x)112 < 16B (B=))l < (3211
Wrapping everything together,
A E (1186+1)-((B22)) & -(1 -244BI)
21n(1 BLkk)))
32BL
+ - (r(1-4%aBB)2BBB + In(1/ow) )2 )2
Thus, for "i < LB (m(1/ "))
we have the first part of Proposition 1.
Using Lemma 3, we then have:
|18b 8%
191
= E 1 ((B)) IAmax (a2At Ama 13 |135 - 3/a 11
E (C(BW) TAEmA a }2 11:b- -Baa 112)
EE 1 CCBBES ZA (1 12] 2B 2B(a² 2 ) ((Bn))
P [(BB)> 2Amax
(a² )²) 4B2
2B(a2 Attnn) -'E [C(3a)]
E [((b)
x 4B2
(a22 x )2
2Amax
4BAmax
= (ttma
E [[C(3x)
a axx
31

---

G Proof of miscellaneous results mentioned in the main text
In this section, we provide proofs for results mentioned in the main text and that not directly directed to the proof of Theorem 3.
G.1 Proof Proposition 3 and the sum of the losses
We start by proving the following proposition, present as is in the first 9 pages of this paper. We then continue with upper and lower bounds (of similar magnitude) on the sum of the losses.
Proposition 3. Let As, As > 0$ be the largest and smallest values, respectively, such that A I EB H < A, For any stepsize y > satisfying Y & BL (as in Theorem 2), initialisation al and batch size [n], the magnitude of the gain satisfies:
An7 Σ
EE((K) [|||Gainy/la] 2A07 EBC(B))
(10)
where the expectation is over uniform and independent sampling of the batches (Ba)
Proof. From Lemma 5, for all -1/2 & x < 1/2, it holds that x² < (()) & 2x.. We have, using Ih v LB, (Bx) < 1/2 (which holds under the stepsize assumption):
E E|| II, = -E Σ In Oroctt
Σ Σ
Eq(
1:000
E VLL, (B)))
1000
V LB. (Be)1
1000
< 4A,
Σ
re ELB (9e)
E:C00
since V CB. (Be)11 & (Be). For the left handside we use ((x) > x² for |x| < 1/2 and E Ls, (Be)l1 > 2A,LB, (Be). Finally, since Be independent freom Be, have ELs, (Be) (B:(
Proposition 15. For stepsizes 7K Ey E LB (as in Theorem 2), we have:
(B)( = θ (111,8*11, In(1/o))
k
Proof. We first lower bound Excoo 7t LB, (Bx)- We have the following equality, that holds for any
k:
Dinas ;x+11 = (3* ( 3xx) 2YLB, (Ba) - D Diees (Bk, Bx+1)
+ (hi- - h h++1) (Ba)) - (h+ hx+1)
leading to, by summing for k E N:
Σ 27LB (8x) = Dro (3* (8)
lim (3* Σ
Diven (3kx Bx+2)+
Σ hho-ho+11) (hu- (o0-hh+1)00 (8*).
Rico
1:000
1:000
First, since ha - 11 Bk - Bro, we have lima-sos Di,, (B*, 3x) = Then, Phy-1BBk Bk+1) > 0.
Finally, (hi - hx++)) (3x) - (hu - ho+1) < 16BL2 LB, (Bx).
Hence:
Σ 10000
(Bk) > Dn,
"A,, Av, are data-dependent constants; nt, we have (Am, An) (A (H), Amax (H)) where
X (H) is the smallest non-null eigenvalue of H; for b = 1. have min; ||2:11 2 d1 & A1 & max; ||a:l:.
32

---

and thus Ensoo BLB > (3", ( for & cBL) (with > 16). This gives the RHS inequality. The LHS direct consequence of bounds proved previous subsections.
Hence, we have that
Σ ((3x) = 0 (7Dr (3*,))))
Noting that Dno (B", ho( (In(1/a))) concludes the proof.
G.2 3o is negligible
In the following proposition we show that 3g is close to 0 and therefore one should think of the implicit regularization problem 8 argming&es Va (3*)
Proposition 16. Under the assumptions of Theorem 2,
1301 a
where the inequality must be understood coordinale-wise. Proof.
1301 = 2 1l - 2
= na² eep(( [<(())(((B)) exp(- (x-(TVV((83)
< a
where the inequality is because 9+ V7x > 0, 0. 0- (xx)) > for all k.
G.3 Impact of stochasticity and linear scaling rule
Proposition 17. With probability dn1- - 3/n² over the x, mia NN(0, c do² (1
o(1)) & A < As & do² (1+)(1)).
so that under these assumptions,
Σ (EE(()))= = - (1
Proof.
The bound on λs Ag a direct consequence of the concentration bound provided in Lemma 13.
D
G.4 (Stochastic) gradients the initialisation
To understand the behaviour and the effects of the stochasticity and the stepsize on the shape of Gainy, we analyse noiseless sparse recovery problem under the following standard assumption [10] and as common in the sparse recovery literature, we make the following assumption on the inputs.
Assumption 2. There exists s-sparse ground truth vector Psparso where verifies R(s Ind))) such that y; (Bāparsoy Izi) for all [m).
Assumption 3. There exists 8, C1, C2 > 0 such that for all s-sparse vectors 8, there exists Idd such that (X T x) 3 + where || < 8|]B8|2 and e₁ 111111111 1 2f B3² < 0z 31
The first part of Assumption 3 closely resembles the classical restricted isometry property (RIP) and relevant for GD while the second part relevant for SGD. Such an assumption restrictive and holds with high probability for Gaussian inputs (0,, (2)a) (see Lemma 10 in Appendix).
33

---

Based on the claim above, we analyse the shape of the (stochastic) gradient at initialisation. For GD and SGD, it respectly writes, where 9o V. (Bo)2, io Unif([n)::
VL(3))) x Eio [oo] = 1 Σ of (",,
The following lemma then shows that while the initial stochastic gradients of SGD are homogeneous, it is not the case for that of GD.
Proposition 18. Under Assumption 3, the squared full batch gradient and the expected stochastic gradient at initialisation satisfy. for some E verifying ||
<< ||
V (Separse) tE,
(25)
10 Vi (9o) = θ (11
(26)
Proof of Proposition 18. Under Assumption 3, have using:
VL(3)) (X X Paparse)
(Siparse =)²
Saparso + == + 25 2ppapes
We have
|E2
25 Deraparrnlll ||Ell + 21|= s1pareel and we conclude by using ||
<
5
5|.5 ( sparss 2
Then,
x. (z., 8 Spparee)
and we conclude using Assumption 3.
Proof of Proposition 4. The proof proceeds as that of Proposition 18.
G.5 Convergence aes and Bo for y
Proposition 19. Let 80(7), (s))) (7) be defined in Theorem 1, for constant stepsizes 7/ y We have:
B0(7) +0, ax - al,
when 7 0.
Proof. We have, as proved previoulsy, that
Σ
2 V VLB, (Bx)
k
2 VL., (Bx)ll
< 2T Σ Li, (Bx)
< 2Ly (3*,
for & BT. Thus, Σ,,> LLLL, (8x)2 + 0 as > 0 (note that Br implicitly depends on so that this result is not immediate).
Then, for 7 < LB*
|In(22//
< Σ Ig(r
< < VLB,
VLB, (3x)
34

---

which tends to - 0. Similarly, | - 0 and (anaa200//))]] + as + 0, leading to B0(7) + as -
H
Technical lemmas
In this section we present few technical lemmas, used and referred throughout the proof of ??.
Lemma 4, Let a+; a- > 0 and I E P. and B = =ft - -a
We have:
arcsinh (
3
I + In a4+ = x + aresinh a a²
224a
2040
Proof. First,
3
1 1/
6²
204
2
e"tinlay/a - 5 in(a+/aaa
2
sinh(x + In(a+/a_))
hence the result by taking the arcsinh of both sides. Note also that we have In(a+/a_) af -a²
arcsinh("
20+0-
Lemma 5. [x 1/2 then x.2 & q(x)) S 2x
[x]]
Lemma 6. On the foo ball of radius B, the quadratic loss function 3 + L(3) is 4Amax max B, a²) relatively smooth w.rl all the hy?
Proof. We have:
V h3) diag 2 + 82
= diag 2 2/ 32
since Og < a component-wise.
Thus, (3)x >
1 min (minigigd 2| 2a
1 L
2
max(4|))l A Ta, and hy max(4B8,a -Stronnnyyyyyyy -strongly convex on the goo norm of radius B. Since Amax- smooth over Rd, we have our result.
Lemma 7. For k > 0 and for all B E Rd:
(B)+2 ha(B)l (L
0+h and a² k+1 a² = 5 - for 84, 7(7x V LB, (3x)) and Proof. We have of 141 02 o2
8_ & = - 71-7x VL.. (3x)). And 0x+1 an exppexppo 6b) where oz = 8+ 6 5 4(7u V Lo, (3x)).
To prove the result we will use that for B E Ied we have 1(h+11 bx) <
EL_ So pl V V (z)ldd.
First, using that arcsinh(a) arcsinh(b)) < In(a/b)) for ab > O. We have that
aresinh 2
az
arcsinh
In
+++1
since & > 0 due to our stepsize condition.
We now prove that l - 0x & 16+.6-8.kkl We have Φe = arcsinh ( and hence,
102 a= aresinl / 2 k+1
[ 0xl arcsinh
200 1
20 20 ++10-+11- 20-
35

---

Then, assuming that a+, 2 a kiú we have:
a 0,,11, - a² kt1, = ch../2 44 k.
6.A. b - a k, 0b,
20
20
44,k, k. 2 K, - 0 ki
if 8+,kit 3 d
20 00- Kb,
11 but a k, a² ki
e
if 8_ 8_ 3 8+
20 k,
6 kit 0 k. 0² k.
e
8 2 ni.
20+ k -0,5
bi - di a ki - a²,k. kit
if 8_ 2 8+,k.
20+, k aa- ,,,3
We thus have a sti. e
- kid
0 4. kii
t, -a_ and this holds
20
similarly if a a a.i.
Then, using |arcsinh(a)
arcsinh (b)
< |In(a/b) obtain that:
aresinh ( a²x²)
a7, - a² -
1 0x1
204.aa_
arcsinh
2a+ A+1
b++ d_1kl
2
Wrapping things up, we have:
16+, - -1
(B) - Vha+i (3) & 6x
$ 28k
2
This leads the following bound:
(B)+ h3)( < (125x)
131)
< 2||6 ||
Recall that &2 97% ( LL, (Bx), hence from Lemma 5 if yA V LL, (Bx) & 1/2, we get that
|| & L (3x) < 4L Ls. (3x).
Putting things together we obtain that
(bx+1 - hx)) < 1128x1.18)
< 8T22R 11&)
I Concentration inequalities for matrices
In this last section of the appendix, we provide and prove several concentration bounds for random vectors and matrices, with (possibly uncentered) isotropic gaussian inputs. These inequalities can easily be generalized to subgaussian random variables via more refined concentration bounds, and to non-isotropic subgaussian random variables [19], leading to dependence on an effective dimension and on the subgaussian matrix Σ. We present these lemmas before proving them in a row
The next two lemmas closely ressemble the RIP assumption, for centered and then for uncentered gaussians.
Lemma 8. Let Z1...., In E Rd be i.idd random variables of law (0, Ia) and H = = Σ TA Then, denoting by the set of all s-sparse vector 3 E IRd satisfying || < 1, there exist C4, C > 0 such that for any > O, if n 2 Cas (d)en
sup BES
H3
B1lo0 >E Con
36

---

Lemma 9. Let 21,... 21...., R be ii.d. random variables of law NH, I2) and H 1 Ei= T.0." Then, denoting by C the set of all 8-sparse vector 3 E Re satisfying || & 1, there exist C4, Cg > 0 such that for any ε > 0, in d C48 ?n(d))
F
sup BES
HB - plp, 8) o² B 3 e
< Can
We then provide two lemmas that estimate the mean Hessian of SGD.
Lemma 10. Let ,I be i.i.d. random variables of law NO, Ia). Then, there exist cı 02 > 0 such that with prohability 1 d and if n = 2//5/4 In(d)), we have for all 8-sparse vectors f3:
1 Σ n
a 1
o (2,, By² & 021131121,
where the inequality meant component-wise.
Lemma 11. Let In be i.i.d. random variables of law (N,, o² Then, there exist
Co2 C1, C2 > 0 such that with probability - o and if n 22(55/4 In(d)) and 40 v I we have for all s-sparse vectors 3:
n²2
+ 11811
x (Mi, (xi, < 41² (12, + 20 11111111))
< n
where the inequality meant component-wise.
Finally, next two lemmas are used to estimate X6, As in our paper.
Lemma 12. Let 21,.... In E Idd be i..d. random variables of law N(p1, Np11,22aa) Let H 1 Ei ziae and = E || ( x... There exist numerical constants Cz C3 > 0 such that
P P(C2(12= + + od 3 3 Csl + +o²)dH) 1 2ne 11-2ne-d/1i -d/16
Lemma 13. Let 3t, In E IRd be i.idd random variables of law N N o² for some μ E IR.
Let H = E.= T and for 1 & b < n let Ho = Eg [t: ErEB 2,,27 ) where B C [n] is
sampled uniformly at random in {B c [m] s.n. [B] = b). With probability
2ne d/16 3/n2,
have, for some numerical constants C1, C2 C3, C >>0 0:
d(p² +2)
(0² + In(n)
pdd
c(duu++++
+0+)
+ (a2 + 122) In(n)
H Ho 1 =C
+ v²d
C1
C2
Vd
c3 n
b
dd
Proof of Lemma 8. For j E [d], we have:
(13), = 1 Σ
Tig (2. 3)
1=1
Fif Bz
1
2/
+
I tif By
71
71
f'
8
20
+
Izj By
We thus notice that IE [H3] 3, and
8
(13),
=
Σ
1)
71
37

---

-3 where 2, = Tij Σÿ Tyy and Σソ Bgy ~ (O, || - 53) and || - 5] < 1. Hence, 2j + zij is centered subexponential random variables (with subexponential parameter of order 1). Thus, for / < 1:
P
1 1) +
Z4 ==1
26 1
$
2e
71
72
==1
Hence, using an c-net of C = { Rd: 18112
& 1. ] (of cardinality less than d x (C/e)²,
and for e of order 1), we have, using the classical g-net trick explained in [Chapt. 9, [58] or [App. C. Even and Massoulie [19]]]
(H3), - -B112) & <dxd(C/e) x2e ant = exp (-cIn(2)nt + (s + 1)dn(d) + sln(//s))
sup JECd Consequently, for and if 3 C4sln(d)/5² we have:
P sup JEC,
(HH3) - B31 21 < <ex <exp (-Conl))
Proof of Lemma 9, We write I; =2; + ps where 2; ~ ~N 12). We have:
1 a
X x3
Σ n1
(1 + ou + ozs.
=,(p,
0²
2) (z (8)
n
1 11 Σ (1,
0²
oiH B)
=1p, B)
2 (1 B3 + ohp
21,8) B) +
1=1
i=1
i=1
The first term is deterministic and is be kept. The second one is of order o² 3 whp using Lemma 8. Then, 1 - 24 ~ (0, Ta/n), so that
P
Zia
B)
>
1
< 2c nt² /2)
i=1
and
2e -ot22/11
P
7 t
<
1=1
Hence,
IP sup DEC
Zzj
t, sup SEC
1 Σ
24s B)
>
< 4e (n( -nt?/2
20
i=1
Thus, with probability 1
Ce nn- and under the assumptions of Lemma 8, have
| X3 - p(p, B) - 02
< E
Proof of Lemma 10. To ease notations, assume that 1. We remind (O'Donnel [46], Chapter and Tao [54]) that for i..dd real random variables a1,-.., an that satisfy a tail inequality of the form
P(ln - Eail > 1) < Ce
(27)
for p < 1, then for all s > 0 there exists C², e' such that for all t,
n
PP(1 =
04 Eaj > E d'et
11
1=1
38

---

We now expand 1 Ei= 2 (2,, 8)2:
2 (T 33 = 1
Σ
1
Σ
Σ
zfauaaBBBe.
ie(m]
Thus, for j E [d].
- a B)²
2
Y2
kE|d
sEin]
keEE|d
We notice that for all indices, all aa. and 23 2 satisfy the tail inequality Eq. (27) for C 8, 1/2 and p = 1/2, so that for E = 1/4:
P(1 = 12 o mal > 1) & Cee e'ne
Σ
P(
Σ
> 1) & C'e C'e-d
For # k we have E [ 1 while for j = k, we have E [
E (2 3. Hence,
P
3j
Σ
a
o aaad > 1,
_ n
Σ x
E |
12
& CC'de-dnen
i=1
Thus, with probability
C'd -cnt" for all j E [d].
Σ
o (z,, 8)2
-233 - 11:111
t A,E
| $|/81
Using the classical technique of Baraniuk al. [4], to make union bound on all s-sparse vectors, consider an e-net of the set of s-sparse vectors of f2-norm smaller than 1. This e-net is of cardinality less than (Co/s)'d", and we only need take of order obtain the result for all s-sparse vectors. This leads to:
P 3B E IRd s-sparse and ||31| 1 Aj E TRd
o (mi, 3/²
- 2,8 - -113112
> 11,11
==
< & -
This probability is equal to /²/d& for t = (o+4) In(d))ess) e'n
4
We conclude that with probability
1 - C² C²/d?, all s-sparse vectors B satisfy:
618) 4 2 (z,, 8)2 - 23² - -1,311 < (8+ +018)
4 || < ( (8 (6++4) (n)d((())) + 61s
s $11812.
c'n
c'n
and the RHS is smaller than ||012/2 for 1(55 Indd))
Proof of Lemma 11. We write x; = + σz, where ~ NO, 1). We have:
P(Vi [n], E [d], |zyl > t) & In(nd) 12/2 1
nd"
for /I(nd)). Thus, if > 40 / In(nd) we have $ Ix & 2u, so that
< 4
piy B)2 E n
Σ
,,(,,,
2n
n
39

---

Then, () ~ (H, B), ) For now, we assume that || = 1. We have P( 8) (n, o3 t) & Ce oo/o and for & 1, using concentration of subexponential random variables [58]:
Σ
Ca 3) - () o² 111312 31) CC'e-no'ioo
n1
and using the E-net trick of Baraniuk et al. [4],
Σ (),,3)2 - (4, - o 02 || st < CCeelln()-nooe1/ n
1 =
C¹
P
sup BEC
d
for = 02113118V
2(0s 21 In(d)
Consequently, we have, with probability
C¹
1
d2
nd
12²
1 0011111 (u, 3)² +
< n Σ
a (2,,
< 44²2 (11, B)² + 22020 1081)
2
2
Proof of Lemma 12. First, we write 2₁ p1l + 022, where 2, ~N I), leading to:
(o21==:112 + du² + 201(11, zi) p,a,
sEin
sEin]
We use concentration of X2 random variables around d:
P(x² > d 2t 2/tt) =t) and F(x. > d d 2V < E
so that for all i [n].
P([2 12 $ dd - 11-2Vdt d 2t + 2Vdt)) < 2
Thus,
P(Vi E [n]. 1 E [d - 2/dt, d + 2t 2vdt)) > 2ne
Taking d/16.
P(V [n], | 2 8) > 2ne d/16
Then, for all i, (1,21) is of law /N(0,)), so that P( (1, 2| t) < 2e (22))) and
< 2nc z/
P(Vi E [m]. [1, 21) > t)
Taking v203/4
P(Vi E [n], (1,
> d < 2ne 2//2
Thus, with probability 1 - 2n(e d/16 + e Vd we have Vi [n], ((1, 21)l > d3/4 and |
v-a
[g. 13d/8], so that
2 + dp² 2/2 ) = H 1
d
13d
du² + 2/ord/+) H.
8
leading to the desired result.
Proof of Lemma 13. We have:
A = E
(
12
A B
= E
Σ
(1e, 23
12
62
iEB
i,jEB, itj
Σ
P(i
B)
I
Pli, E (11,
82
1 it
40

---

Then, since P(i E B) = f/2 and P(i.j E B)
b(b-1) n(n-I) for i # j, we get that:
(b 1)
H
|
Σ
bn
bn(n 1)
Using Lemma 12, the first term satisfies:
d(p² + 0²)
d(u² + 0²)
CaHl
Σ
C2fl 2 1 -nne dd/16
b
bn iE(n]
We now show that the second term is of smaller order. Writing Tj tdl ozz where 2, (O, we have:
(b - 1) bn(n 1)
Σ it1
(6-1) bn(n 1) itj
Σ [y))4
For + j. (ra; 2y) Ei=1 B. Σ1 k= dg where ag Fat-Bikk satisfies Ea; 0, Eag 1/22 Hence, aş centered subexponential random variables. Using concentration of subexponential random variables [58], for t < 1,
and P(ae > t) < 2P(
v <
P
7 (on الارةة | {(
< 2e -ods"
Thus,
P
Vi + j, " [(( <t) 1 n(n 1)e- odt"
Then, taking t = d In (n)/cc c, we have:
1 2 (oo, Vi x j.
< 41n(n)
> 1 n²
cvd
Going back our second term,
(b - 1) bn(n - 1) if1
(b 1b-11 1)
Σ [,,
(z) 2)) + xạoر))
bmin - 1)
(b 1)
(1i))z))(( (1:2) + tz,2,),
bn(n 1)
ijj
where we used 150, + m.ao 2 x + z.y.. Thus,
(b 10-1) 1) bm(n 1)
Σ 21),
sup ify
x (b 1)
(2aa]] + x5277)
bn(n 1)
b b-1 1
sup if3
[
n1
b b- 1
sup ifi
|| 2لح
x
H
Similarly, we have
(b - 1) Σ
b-1 1 H.
(11, 2 sup
bn(n - 1) jg1
if1
Hence, with probability 1 - 1/ni²
4ln(n) b- 1 11 1
(b - 1)
Σ
(I (1)
4ln(n)
b b-1 n
x
c/d
bn(n 1)
cvd
b
41

---

Wrapping things up, with probability 1 - 1/n² - -2ne 4/116
n (=ovd cvd b n -
4 (n((n) b-in--a c22) x H 3 Hl 3
(4lm(n) b b-1 + +c, d) x H.
cvd n-1
Thus, provided that 4lm(n) Cad & 2b
and d > 48 In (n), we have with probability 1 - 33/n²
c x 11 3 Ho 3 C55 c 11.
CITx
Proof of Lemma 13. We have:
HB = E
(a 2,00
b
i.j
E
+
E
(Te, 2,)oaa3
12
62
IE
i,jEB, itj
1
Σ
P(i 11:
Σ
P E (
82
8²
iE(n]
Then, since Pii E B) and P(i.j E B)
6(b-1) m(n-1) for i + j. we get that:
I + (b 10-1] 1)
H bn VEiN
Σ
(cis 2,
bnin 1)
Using Lemma 12, the first term satisfies:
P d(u² + + +
C
1 bn vE(m)
Σ
1 d(u² + 0²) d(u2+
> 1 1-2ne d/16 CgH)
b
We now show that the second term is of smaller order. Writing x; jll + az; where 2; NOO, Ia). we have:
(b 10-1 1)
(b- 1)
Σ 20
Σ
(o2 23) + tpl 2, + 2,) + p= x
bn(n- - 1)
bm (n - 1)
inj
inf
(b-1)
(0² 23) + op(1, + 23)) [, T + (b - 1)
11² 2d
2 T
bm(n - 1)
bn(n 1)
ifj
(
For * j. (2,, = ER= 2ト = Ei= as where = 24 satisfies Ea;- = 0, Eag = 1 and 4e -t/2, Hence, as is centered subexponential random variables. Using concentration of subexponential random variables [58], for t & 1,
P(ax > t) & 21P(|=uR| >
Vt)
IP Kam 2t)
< 2e nodt
Thus,
P
Vi + j, 2(o0 st) > 21 - n(n 1)e- cddt?
Then, taking / - -d-2 (n (n(n// cc have:
1 IP Vi + j. 7 | 2لحه < 4ln(n)
> 1 n2
cvd
For E [n], (1) (0, d) that P(K 2, > t) -c-1" (/2d))
and 2ne e(2d))
P(V [n], (1, 2 < t)
n2²
42

---

for = 3/d(n(n)). Hence, with probability 1- - 3/n², for all i # we have (e) 23) + 00u(1,
zy)l & (o² ton)Cv/nn//vd
Now,
(b 10-1] 1)
(b- 16-1)
(0 (2 2n) + #a(11 21 + 2))) bn(n 1)
(2) + to + (x,a) + xzza))
bn(n- ifj
(6 1)
Σ
( + tow(1, + ))))) + tzz
bn(n - 1)
isj
where we used 13a3 + x.po 3 nxo nza]. Thus,
2 (b 1) Σ (2,, + tou(1, 2 + xz)) 3 sup |0 23) + + +2))1 x bnin- 1) ifs
(b - 1)
(x + xzaa]]
bn(n 1)
(f
i<j
- 1 1
sup it}
(2, + 0 24 + 2ل)1
b n 1
i=1
sup
(2 a 24 + +z)) if
b-1
n
H
b
Similarly, we have
bn(n Σ
(b 1)
- - sup -ssplp/)ss b- n H izj
Hence, with probability 1 - 3//n, (o² (o²+ b-1 H 1 (b (b- 1)
(o 23) + oo((1, + 23)) 25a)
vd
bn(n 1)
(o" oo) In(n)
b-1 n b n
x
H
Va
We thus have shown that this term (the one in the middle of the above inequality) is of smaller order.
We are hence left with
bn(n 1J 4² Eizz I Denoting ã = n Ei L4, we have # Σitj TTI =
(h-1)
n Ens var] 3 E, x.r. that:
(b 10-1) 1)
(b - 1)n 112d dd TT. H
1
1 d
2
FiJ
bn(n 1)
b(n 1)
$$1
We note that we have H = 1 Σ;
T
1 Σ
Eiej I + T2,z,
T
- #/2 Exs 2 +2zI7 F2
12
using x. + x _ 3 IJ Thus, H
33 0, and:
(b 1)m
(b- 1)
p F 1 (b- 1)n 12 dd 1/n)H.
bn 13 11 dd
H 1 bn(n 1)
b(n - 1)
(
We are now able to wrap everything together. With probability 1 - 2ne dd/16 -3/n², we have, for some numerical constants c1, C2, Ca. C > 0:
9 dyp² +02) 02 (02 + + In(n)
1d
+02)
b + (02 + (2) pdd)
c3
H = Ho C
Vd
Vd
43


Reviews:
[{"review_id": "efXxweKK3P", "reviewer": "Reviewer_uxxP", "summary": "The paper analyzes 2-layer diagonal linear networks and explains the different characteristics of the solutions received through GF, GD, and mini-batch SGD.", "questions": "1) What conditions were used to determine if the training converges, i.e., the training reached \\\\( w^\\gamma_\\infty \\\\)? In particular, in Figure 5, it seems like \\\\( \\lambda_\\max\\left( \\nabla^2F \\left(  w^\\gamma_\\infty \\right) \\right) \\\\) is above the value of \\\\( 2/\\gamma \\\\), which is impossible at convergence.\n\n2) What is the effect of learning rate schedules? In particular, about Observation 1 and Figure 3, what would be the effect on coordinates not in \\\\( \\text{supp}\\left( \\beta^*_{\\text{ sparse }} \\right) \\\\) if the step increases slowly so that the Oscillation phase is arbitrarily long?\n\n3) In lines 290 and 295, did you perhaps mean to quote Eq.(10) and Eq.(11)?", "limitations": "The limitations are adequately addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper analyzes the impact of the different step sizes and batch sizes and explains the differences between the solution received from training 2-layer diagonal linear networks using GF, GD, and mini-batch SGD. The paper shows that some of the differences are explainable using the linear scaling rule of \\\\( \\frac{\\text{step size}}{\\text{batch size}} \\\\). However, the paper also shows other differences which are not explainable using the linear scaling rule.", "weaknesses": "Part of the theory, including theorems 2 and Proposition 1, only apply when the step size is small enough, and the edge of stability phenomenon does not occur. However, in this case, both GD and SGD behave approximately like GF. Thus any difference between the solutions received via the different optimization algorithms is minimal.\n\nFor larger step sizes, the theoretical analisis can not explain the full picture by itself. Rather, additional empirical observation are first being made, and then assumed to be true in order to complete the full analisis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough reviewing, the helpful remarks, and the noticed typos.\n\nAs rightfully noticed and as acknowledged in our paper (lines 41-42), part of our results do not apply for large step sizes at the edge of stability. Nonetheless our main result (the implicit bias result in Theorem 1) holds **for any stepsize schedule such that the iterates converge**: we use this result in sections 4.2 and 5 together with empirical observations to provide informations on the behavior of GD at the edge of stability. Rigorously proving the observations made remains however a hard problem which we currently believe to be out of reach.\n\nThe condition for the convergence is $\\mathcal L (w_k)\\leq \\varepsilon$ for a small $\\varepsilon$ ($\\varepsilon=10^{-20}$).\nIndeed as accurately noticed there was a mild mistake in the code used to create Figure 5 (left and right): in this figure the y-axis should be rescaled by a factor 1/8 that comes from wrong normalizations (a 1/2 factor in front of the loss, and a 1/2 factor in the argument of the loss that leads to a 1/4 factor for second order derivatives, giving us this $1/8=1/2*1/4$ normalization error), we apologize for this mistake and will correct it. The trend of the figure is however unchanged (monotonicity of sharpness for GD and SGD). \n\n**Learning rate schedules.** The proposed learning rate schedule (slowly increasing the stepsize to keep oscillations) *will magnify the effects of large stepsizes*: for SGD, this will mimic label-noise SGD if we keep the magnitude of the oscillations of the same order, and lead to sparse recovery through l1 norm minimization; however, for GD, this will lead to recovery of the wrong support due to a weighted l1 norm minimization which is adversarial, as explained in the paper.\n\nIn lines 290-295, we indeed meant to quote Eq. 10-11."}, {"review_id": "heBCuJ61Ok", "reviewer": "Reviewer_Y9wy", "summary": "This paper studies the implicit bias of 2-layer diagonal linear networkss. \nThe authors show the convergence of GD and SGD with macroscopic stepsizes and characterise their solutions through an implicit regularization problem. \nMoreover, the theoretical results reveal the difference between the generalization performances of GD and SGD for large stepsizes.", "questions": "1. While the authors have provided explanations of the Gain in Section 4 from various perspectives, I am still confused by its inisightful meaning. To address this, the authors can consider providing additional geometric insights or visual aids to help clarify the intuitive meaning of Gain.\n\n2. The authors assert that the Gain of SGD is homogeneous in Equation (11) and Figure 2. However, it is important to note that Equation (11) corresponds to the expectation of the stochastic gradient rather than the realistic stochastic gradient. Previous studies, such as Zhu et al. (2018), have argued that the realistic stochastic gradient is highly anisotropic. This raises a potential contradiction, and it would be helpful if the authors could address this concern in the paper.\n\n\nZhu et al (2018). The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects.", "limitations": "The results in the paper is limited to 2-layer diagonal linear network and specific initialization, but it should not be the reason for rejection. \nThis paper indeed offers a solid theoretical perspective on the implicit regularization of SGD.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is clear and well-written. \nWnderstanding what solutions do GD/SGD converges to is an important theoretical issue.\nThis paper achieves a clear characterization by deriving the solution found GD/SGD through an implicit regularization problem.\nFrom a technical standpoint, this paper cleverly establishes an equivalent relationship between gradient descent and mirror descent of another problem, showcasing its potential for broader applications.\nFurthermore, the authors analyze the gain in the implicit regularization problem, offering insightful explanations for the differences in generalization abilities between GD and SGD with varying step sizes.\n", "weaknesses": "This paper generalizes the existing results of the implicit regularization on GF to SGD, and the results are relatively similar, although I know this is indeed technically difficult.\nFuthermore, 2-layer diagonal linear network and specific initialization restrict the influnce of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time spent reviewing, the thoughtful comments and reference.\n\nWe first answer the ‘weakness’ part of the review.\nOur paper indeed generalizes existing results: references [48,50,61]  provide the implicit bias of gradient flow and stochastic gradient flow over DLNs, while we prove convergence of GD and SGD for macroscopic stepsizes, as well as determine the implicit bias for any arbitrary stepsizes that lead to convergence. Our results are thus much stronger since they directly study the algorithms which are used in practice and enables to gain insights on the effects of stepsizes and stochasticity. Similar characterisations were  not possible with previous works. We thus believe that our results are not just simple technical and incremental contributions. \n\nAs we acknowledge in our paper, DLNs are indeed a very simple model of neural nets; however, we believe this model to be (currently) the only one for which a non-trivial implicit bias can be rigorously proved for large stepsize GD and SGD, as rightfully noted by Reviewer JT8J.\n\nFinally, we use a specific initialization of the neurons that lead to $\\beta_0=0$, in order to have results that are clearer to understand. Our analysis can directly be extended to general initializations: non-centered initializations would only modify the term $\\tilde \\beta_0$ in Theorem 1; we can add a more thorough discussion about the initialization in a revised version.\n\n**Meaning of Gain.**\nThe quantity Gain naturally appears in the analysis as being the key quantity which quantifies how much the solution found by (S)GD with given stepsizes deviates from that found by GF with same initialization. This quantity is insightful because it is  tractable and analyzable: the goal of sections 4 and 5 is precisely to leverage its definition in order to understand the impact of the stepsize and stochasticity on the recovered solution.\nA geometric insight could however be: for all iteration $k$, the interpolating vector that minimizes the potential $h_k$ (eq. (16)) related to the effective initialization $\\alpha_k$ (line 621 in the appendix) is the solution found by GF started at $w_k$; hence the effective initializations $\\alpha_k$ track the solutions found by GF starting from the iterates of (S)GD. The Gain and the effective initialization $\\alpha_\\infty$ are limits of these quantities for $k\\to\\infty$.\n\n**Anisotropic noise.**\nWe thank the reviewer for this very interesting paper.\nEq. (10) is indeed the expected squared stochastic gradient: it is thus related to both the true gradient and the noise, but the latter dominates. It is however important to note that these are the stochastic gradients of the **convex** loss $\\mathcal L$, and not of the **non-convex** loss $F$ (that takes neurons as input). We thus do not believe Eq. (10) to be in contradiction with Zhu et al.: they argue that the noise of stochastic gradients of the **non-convex** loss is anisotropic, which we do not contradict. A thorough study of the SGD noise of the non-convex loss $F$ would show that for $s$-numerically sparse (i.e., almost or approximately sparse) neurons, the noise covariance and the non-convex Hessian (see Eq. (19)) are highly non-isotropic.\nWe thank again the reviewer for this reference, and suggest to add it with such a discussion in a revised version."}, {"review_id": "DrwjxIixxL", "reviewer": "Reviewer_1MhA", "summary": "This paper studies the impact of stochasticity and step size on the implicit regularization of GD and SGD, in the setting of two-layer diagonal linear network. It is show that large step size can benefit SGD, but hinder the performance of GD. Both the implicit bias and convergence are proved for (S)GD, providing a complete trajectory-dependent result. The results are further used to provide insights for explaining the \"edge of stability\" regime.", "questions": "Theorem 2 gives the convergence guarantee. Is it possible to obtain an explicit convergence rate? If not, what's the difficulty here?\n\nIn Figure 1, the largest step sizes for GD and SGD are different. Why this difference doesn't appear in Theorem 2?", "limitations": "The authors have adequately addressed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The writing is very clear and easy to follow, and highlighting of the important parts helps the reading a lot.\n\nTheorem 1&2 together provides a complete set of results for trajectory-dependent analysis for GD and SGD for two-layer diagonal linear network, and the subsequent discussion on the impact of stochasticity and step size is clear and insightful.", "weaknesses": "It seems that the technical tools used in this paper have already been developed in previous related papers, especially [48, 50, 61] (corresponding to the reference number in the paper). Can the authors discuss the technical novelty of the current paper?\n\nThe gain vector does not have a closed-form expression, and it seems unclear what exact property does the minimization in (5) lead to. More specifically, does it give exact sparse recovery in some case?\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time spent reviewing, the questions, and for the positive feedback.\n\n**Technical novelty.**\nOur technical tools are in fact very different from the cited references: our underlying process (SGD on the neurons) is **discrete** while theirs are **continuous-time** processes: gradient flow or stochastic GF.  Our analysis is thus much more challenging, and studies the real algorithms used in practice, as opposed to continuous time approximations.\nOn the technical side,  [48, 50, 61] leverage a continuous mirror flow structure; our analysis is based on leveraging a (discrete) time varying mirror descent structure, indeed showing some similarity with  [48, 50, 61]. Analyzing the implicit bias in sections 4-5 is however completely new both technically speaking and in terms of insights on the effects of stepsizes and noise.\n\nIn the general case, the minimization in (5) is hard to grasp; particular cases that lead to exact sparse recovery are GD or SGD with small initialization and stepsizes that are not in the EoS regime, or SGD with any initialization, but with an EoS stepsize.\n\n**Convergence rate.** From Proposition 6, we directly have $1/T \\sum_{t<T} L(\\beta_t)=O(1/T)$. However, for a quantitative bound on $||\\beta_k-\\beta_\\infty||$, a more thorough analysis is required. In fact, we can show that if the loss becomes small (as can be proved with Prop. 6 in our paper), the loss becomes locally relatively strongly convex wrt the hypentropy. This property leads to a linear convergence (on the quantity $\\Vert \\beta_k-\\beta_\\infty \\Vert_2$) to 0 with a rate of the form $O(\\exp(-\\alpha^2\\mu\\gamma k))$ where $\\mu$ is the smallest non-null eigenvalue of $H$.\nWe are open to incorporating this result into a revised version of the paper if the reviewer believes it would be beneficial. We did not include it due to space constraints and our primary focus on the implicit bias aspect of the problem.\n\nQuestion 2. : In Figure 1, the largest step sizes for GD and SGD are different, but this is also the case in Theorem 2, where the constant $L$ depends on the batchsize (see def in line 131)."}, {"review_id": "bqfuIyhX7u", "reviewer": "Reviewer_JT8J", "summary": "This is a theoretical work on implicit bias in diagonal linear networks. It proves the empirical observation by Pesme et al that SGD with large learning rates can recover the sparse signal while GD or small learning rates cannot. Technically it utilizes mirror descent with time-vary potentials. ", "questions": "See above in Weakness.\n\nFor now, I would like to recommend an acceptance for its theoretical proof of an open problem.", "limitations": "N/A\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The motivation and setting are great and non-trivial.\n\nIt is a popular and significant trend to understand the implicit bias with large learning rates, which is practically meaningful and much harder than previous works on small learning rates. For instance, the small-learning-rate version of this work is Woodworth et al for gradient flow.\n\nAlthough the main empirical observation has been reported by Pesme et al last year, this work successfully proves it, completing the story. Meanwhile, for now, the setting of diagonal linear networks might be the most complicated one we can expect for SGD-large-learning-rate bias to rigorously prove.\n\n2. The presentation and discussion are comprehensive.\n\nThe target setting is with several factors, including initialization scale, initialization shape, stochasticity, learning rates. This work does a good job to clearly its position in terms of these factors. In Section 3, it connects the main theorem 1 with previous gradient-flow result, with an emphasize on effective initialization $\\alpha_{\\infty}$. Then in Section 4, it carefully discusses how these four factors affect $\\alpha_{\\infty}$.", "weaknesses": "1. Is such a phenomena robust to different learning rates? In Figure 1, the x-axis region of significant drop of test loss for SGD seems to be not broad, and it looks hard to find the optimal $\\gamma$ since it is too close to $\\tilde{\\gamma}_\\max$ to explode. Meanwhile, the definition of $\\tilde{\\gamma}_\\max$ around 257 is not much helpful to find the exact value. Moreover, my guess is that $\\tilde{\\gamma}_\\max$ is not stable due to stochasticity in data generation and sgd sampling.\n\nIn other words, how can we choose a good lr for good sparse recovery?\n\n2. In terms of practical computation, I am wondering how much SGD+large lr is better than GD+small initialization or SGD+label noise. It is mentioned in Section 4.1 that ''the slower the training, the larger the gain''. Hence it seems to bring more computation cost in practice for SGD to achieve good test error, while GD + small initialization may be faster in steps. And SGD+label noise might be another strong candidate.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the overall very positive feedback, all the helpful remarks and questions, that could lead to additional valuable discussions in our paper.\n\n**1.** The quantity $\\tilde \\gamma_{max}$ is not robust to the sampling of the random inputs $x_i$, as it is defined conditionally on them (for different samples of $(x_i)$'s, we empirically observe that $\\tilde \\gamma_{max}$ varies a lot, making it hard to estimate). However, it is empirically robust to SGD noise (to the noise induced by the sampling of mini-batches at each iteration): if SGD converges for some stepsize, we empirically observe that SGD converges for any smaller stepsize. Proving such a phenomenon seems however out of reach (even for GD).\n\nFor sparse recovery, a good learning rate for SGD would be to choose a learning rate that produces oscillation in the initial phase, but still converges. As noticed, such a learning rate is hard to estimate, hence the interest of *learning rate schedules*, that start with a (too) large learning rate, but decreases it slowly during different phases.  These schedules are in fact also used in practice to increase generalization on more complex architectures; DLNs here offer a provable benefit for such schedules. \nWe propose adding this discussion in a revised version, as we believe it will provide additional valuable insights into the effects of learning rates for SGD by discussing lr schedules.\n\n**2.** There is indeed a tradeoff between generalization power and computational efficiency here: empirically, taking into account its reduced complexity per iteration, SGD with large stepsizes performs better for large dimensions $d$, which can be explained by the fact that $\\alpha$ needs to be taken too small for GD to converge quickly. Finally, SGD with large stepsizes or label noise SGD perform equally well. If the reviewer believes that illustrating this discussion with experiments could be a valuable addition, we propose to do so."}, {"review_id": "VwgyVRo3n1", "reviewer": "Reviewer_ZVE8", "summary": "This paper analyzes the implicit biases of GD and SGD on two-layer diagonal linear networks, specifically with large step sizes. This paper shows that SGD converges to a limit that is determined by the trajectory, and specifically by a certain effective initialization. Moreover, while both GD and SGD are close to gradient flow when the step size is small, it is shown that SGD may recover a sparse solution with larger step sizes. This can help explain the generation gap between GD and SGD with large step sizes.", "questions": "N/A", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper analyzes the implicit bias of SGD for a wide range of step sizes, which is an important contribution. In particular, it is pointed out that while larger batch sizes generally hurt the performance of GD, they may actually improve the accuracy of SGD before divergence, which is interesting.", "weaknesses": "My main concern is that the two-layer diagonal linear network model might be too simple, and more evidence is needed to show that observations in this setting can be transferred to more practical settings. Also the presentation is a bit technical; for example, the definition of the effective initialization is a little hard to follow. It could be better to consider a simple case to simplify the presentation while highlighting the interesting results, such as the different implicit biases of GD and SGD.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments and the feedback. We would first like to point out a misunderstanding in the review (this is most certainly only a typo between batchsize and stepsize): *it is pointed out that while larger batch sizes generally hurt the performance of GD, they may actually improve the accuracy of SGD before divergence, which is interesting* should be replaced by *it is pointed out that while larger **stepizes** generally hurt the performance of GD, they may actually improve the accuracy of SGD before divergence, which is interesting*.\n\n**Limitations of DLNs.** DLNs are indeed a simplistic models of neural networks, a limitation acknowledged in our work. However, as we explain in lines 41-50, we believe that the rigorous study of complex phenomenons such as the intricate effect of noise and stepsizes has to start on such simple models. As put by Reviewer JT9J, *for now, the setting of diagonal linear networks might be the most complicated one we can expect for SGD-large-learning-rate bias to rigorously prove*.\n\nThe definition of the effective initialization is indeed quite technical, but there is no simple case for which it has a simpler expression. However, its effect on the solution recovered is summarized as follows for sparse regression and large stepsizes: the effective initialization has heterogeneous coordinates that hinder the recovery of sparse vectors for GD, while for SGD its coordinates are of same order and lead to good sparse recovery."}]

Author rebuttal:

Meta-review: There is consensus among the reviewers on the validity of the problem setting and soundness of the results, but also on limitations of the theoretical characterization proposed in the paper. Nevertheless, on the balance the reviewers and I find the paper a worthy contribution, and consequently I recommend acceptance.
Decision: Accept (poster)