Title: (S)GD over Diagonal Linear Networks: Implicit bias, Large Stepsizes and Edge of Stability
Venue: NeurIPS 2023
Accepted: yes
Keywords: SGD

GD

implicit bias

large stepsizes

edge of stability

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

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 provide a characterisation of their solution through an implicit regularisation problem. Our characterisation provides insights on how the choice of 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 can 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.

1 Introduction

The stochastic gradient descent algorithm (SGD) [51] is the foundational algorithm for almost all neural network training. Though a remarkably simple algorithm, it has led to many impressive empirical results and is a key driver of deep learning. However the performances of SGD are quite puzzling from a theoretical point of view as (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 in many real-world prediction tasks, it must be because the optimisation process is implicitly 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 are infinitely many solutions that fit the samples, GD and SGD explore only a 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 ℓ 2 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 is to analyse gradient flow, which is a 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 a 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 a limitation in our understanding since the performances of most deep learning models are often sensitive to the choice of stepsize and minibatch size. The importance of stepsize and SGD minibatch size is common knowledge in practice and has also been systematically established in controlled experiments [36, 42, 20].

Test losses

β

β

∗

/lscript

1

2

10

-

1

˜

γ

SGD

˜

γ

Stepsize

γ

10

-

5

10

-

2

Test losses

||

β

∞

-

β

/lscript

1

||

2

SGD

Test losses at convergence

10

-

1

˜

γ

SGD

˜

γ

Stepsize

γ

10

-

5

10

-

2

Test losses

||

β

γ

∞

-

β

∗

/lscript

1

||

2

2

SGD

10

-

1

˜

γ

SGD

˜

γ

Stepsize

γ

10

-

5

10

-

2

SGD

10

-

1

GF

test loss

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, we 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 a 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]. It therefore serves as an ideal proxy model for gaining a deeper understanding of complex phenomenons such as 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 ReLU 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 linear networks are introduced in Section 2. We formulate our theoretical results (Theorems 1 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 β ⋆ ∞ . We further provide a refined characterization of β ⋆ ∞ through a trajectory-dependent 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 a stark difference between the generalisation performances of GD and SGD for large stepsizes , hence explaining the numerical results seen in Fig. 1 for the sparse regression setting. Finally, in Section 5, we use our results to shed new light on the Edge of Stability ( EoS ) 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 homogeneous 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 ℓ 2 norm solutions-the neural tangent kernel regime [30] and small initialisation resulting in minimum ℓ 1 norm solutions-the rich regime [13]. The analysis relies on 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, 31, 27].

<latexi sh

1_b

64="OJ7kS

qVI9

+

E

z

j

5

8

d

>A

B

n

c

N

3

W

X

Q

rHo

MF

w

L

v

p

2

YT

Cf

D

P

g

/

m

y

u

G

U

K

0

R

Z

||

γ

γ

∞

-

∗

||

2

2

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 a 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 = ( x 1 , . . . , x n ) ∈ ( R d ) n and outputs y = ( y 1 , . . . , y n ) ∈ R n . 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 β ⋆ ∈ R d which perfectly fit the training set, i.e. , y i = ⟨ β ⋆ , x i ⟩ for all 1 ⩽ i ⩽ n . We call such vectors interpolating predictors or interpolators and we denote by S the set of all interpolators S = { β ⋆ ∈ R d s . t . ⟨ β ⋆ , x i ⟩ = y i , ∀ i ∈ [ n ] } . Note that S is an affine space of dimension greater than d -n and equal to β ⋆ +span( x 1 , . . . , x n ) ⊥ for any β ⋆ ∈ S . We consider the following quadratic loss: L ( β ) = 1 2 n ∑ n i =1 ( ⟨ β, x i ⟩ -y i ) 2 , for β ∈ R d .

2-layer linear diagonal network. We parametrise regression vectors β as functions β w of trainable parameters w ∈ R p . Although the final prediction function x ↦→ ⟨ β w , x ⟩ 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 β w = 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 a 2 -layer diagonal linear neural network given by:

This parametrisation can be viewed as a simple neural network x ↦→ ⟨ u, σ (diag( v ) x ) ⟩ where the output weights are represented by u , the inner weights is the diagonal matrix diag( v ) , and the activation σ is the identity function. In this spirit, we refer to the entries of w = ( u, v ) ∈ R 2 d as the weights and to β := u ⊙ v ∈ R d as the prediction parameter . Despite the simplicity of the parametrisation (1), the loss function F over parameters w = ( u, v ) ∈ R 2 d is non-convex (and thus the corresponding optimization problem is challenging to analyse), and is given by:

Mini-batch SGD. We minimise F using mini-batch SGD: let w 0 = ( u 0 , v 0 ) and for k ⩾ 0 ,

where γ k are stepsizes, B k ⊂ [ n ] are mini-batches of b ∈ [ n ] distinct samples sampled uniformly and independently, and ∇ F B k ( w k ) are minibatch gradients of partial loss over B k , F B k ( w ) := L B k ( u ⊙ v ) defined above. Classical SGD and full-batch GD are special cases with b = 1 and b = n , respectively. For k ⩾ 0 , we consider the successive prediction parameters β k := u k ⊙ v k built from the weights

w k = ( u k , v k ) . We analyse SGD initialised at u 0 = √ 2 α ∈ R d > 0 and v 0 = 0 ∈ R d , resulting in β 0 = 0 ∈ R d independently of the chosen weight initialisation α 2 .

Experimental details. We consider the noiseless sparse regression setting where ( x i ) i ∈ [ n ] ∼ N (0 , I d ) and y i = ⟨ β ⋆ ℓ 1 , x i ⟩ for some s -sparse vector β ⋆ ℓ 1 . We perform (S)GD over the DLN with a uniform initialisation α = α 1 ∈ R d where α > 0 . Fig. 1 and Fig. 2 (left) correspond to the setup ( n, d, s, α ) = (20 , 30 , 3 , 0 . 1) , Fig. 2 (right) to ( n, d, s, α ) = (50 , 100 , 4 , 0 . 1) and Fig. 3 to ( n, d, s, α ) = (50 , 100 , 2 , 0 . 1) .

Notations. Let H := ∇ 2 L = 1 n ∑ i x i x ⊤ i denote the Hessian of L , and for a batch B ⊂ [ n ] let H B := ∇ 2 L B = 1 |B| ∑ i ∈B x i x ⊤ i denote the Hessian of the partial loss over the batch B . Let L denote the 'smoothness' such that ∀ β , ∥ H B β ∥ 2 ⩽ L ∥ β ∥ 2 , ∥ H B β ∥ ∞ ⩽ L ∥ β ∥ ∞ for all batches B ⊂ [ n ] of size b . A real function (e.g, log , exp ) applied to a vector must be understood as element-wise application, and for vectors u, v ∈ R d , u 2 = ( u 2 i ) i ∈ [ d ] , u ⊙ v = ( u i v i ) i ∈ [ d ] and u/v = ( u i /v i ) i ∈ [ d ] . We write 1 , 0 for the constant vectors with coordinates 1 and 0 respectively. The Bregman divergence [9] of a differentiable convex function h : R d → R is defined as D h ( β 1 , β 2 ) = h ( β 1 ) -( h ( β 2 ) + ⟨∇ h ( β 2 ) , β 1 -β 2 ⟩ ) .

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 et al. [61] show that the limit β ∗ α of the gradient flow d w t = -∇ F ( w t )d t initialised at ( u 0 , v 0 ) = ( √ 2 α , 0 ) is the solution of the minimal interpolation problem:

The convex potential ψ α is the hyperbolic entropy function (or hypentropy ) [21]. Depending on the structure of the vector α , the generalisation properties of β ⋆ α highly vary. We point out the two main characteristics of α that affect the behaviour of ψ α and therefore also the solution β ⋆ α .

The Scale of α . For an initialisation vector α we call the ℓ 1 -norm ∥ α ∥ 1 the scale of the initialisation. It is an important quantity affecting the properties of the recovered solution β ⋆ α . To see this let us consider a uniform initialisation of the form α = α 1 for a scalar value α > 0 . In this case the potential ψ α has the property of resembling the ℓ 1 -norm as the scale α vanishes: ψ α ∼ ln(1 /α ) ∥ . ∥ 1 as α → 0 . Hence, a small initialisation results in a low ℓ 1 -norm solution which is known to induce sparse recovery guarantees [10]. This setting is often referred to as the 'rich' regime [61]. In contrast, using a large initialisation scale leads to solutions with low ℓ 2 -norm: ψ α ∼ ∥ . ∥ 2 2 / (2 α 2 ) as α →∞ , a setting known as the 'kernel' or 'lazy' regime. Overall, to retrieve the minimum ℓ 1 -norm solution, one should use a uniform initialisation with small scale α , see Fig. 7 in Appendix D for an illustration and [61, Theorem 2] for a precise characterisation.

The Shape of α . In addition to the scale of the initialisation α , a lesser studied aspect is its 'shape', which is a term we use to refer to the relative distribution of { α i } i along the d coordinates [3]. It is a crucial property because having α → 0 does not necessarily lead to the potential ψ α being close to the ℓ 1 -norm. Indeed, we have that ψ α ( β ) α → 0 ∼ ∑ d i =1 ln( 1 α i ) | β i | (see Appendix D), therefore if the vector ln(1 / α ) has entries changing at different rates, then ψ α ( β ) is a weighted ℓ 1 -norm. In words, if the entries of α do not go to zero 'uniformly" , then the resulting implicit bias minimizes a

weighed ℓ 1 -norm. This phenomenon can lead to solutions with vastly different sparsity structure than the minimum ℓ 1 -norm interpolator. See Fig. 7 and Example 1 in Appendix D.

3.2 Implicit bias of (stochastic) gradient descent

In Theorem 1, we prove that for an initialisation √ 2 α ∈ R d and for arbitrary stepsize sequences ( γ k ) k ⩾ 0 if the iterates converge to an interpolator , then this interpolator is the solution of a constrained minimisation problem which involves the hyperbolic entropy ψ α ∞ defined in (4), where α ∞ ∈ R 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 ( u k , v k ) k ⩾ 0 follow the mini-batch SGD recursion (3) initialised at ( u 0 , v 0 ) = ( √ 2 α , 0 ) and with stepsizes ( γ k ) k ⩾ 0 . Let ( β k ) k ⩾ 0 = ( u k ⊙ v k ) k ⩾ 0 and assume that they converge to some interpolator β ⋆ ∞ ∈ S . Then, β ⋆ ∞ satisfies:

where D ψ α ∞ is the Bregman divergence with hyperentropy potential ψ α ∞ of the effective initialisation α ∞ , and ˜ β 0 is a small perturbation term . The effective initialisation α ∞ is given by, where q ( x ) = -1 2 ln((1 -x 2 ) 2 ) satisfies q ( x ) ⩾ 0 for | x | ⩽ √ 2 , with the convention q (1) = + ∞ .

The perturbation term ˜ β 0 ∈ R d is explicitly given by ˜ β 0 = 1 2 ( α 2 + -α 2 -) , where q ± ( x ) = ∓ 2 x -ln((1 ∓ x ) 2 ) , and α 2 ± = α 2 ⊙ exp( -∑ ∞ k =0 q ± ( γ k ∇L B k ( β k ))) .

Trajectory-dependent characterisation. The characterisation of β ⋆ ∞ in Theorem 1 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 ˜ β 0 depend on the iterates' trajectory. Nevertheless, we argue that our formulation provides useful ingredients for understanding the implicit regularisation effects of (S)GD for this problem compared to trivial characterisations (such as e.g. , min β ∥ β -β ⋆ ∞ ∥ ). Importantly, the key parameters α ∞ , ˜ β 0 depend on crucial parameters such as the stepsize and noise in a useful and analysable manner : understanding how they affect α ∞ and ˜ β 0 coincides with understanding how they affect the recovered solution β ⋆ ∞ and its generalisation properties. This is precisely the object of Sections 4 and 5 where we discuss the qualitative and quantitative insights from Theorem 1 in greater detail.

The perturbation ˜ β 0 can be ignored. We show in Proposition 16, under reasonable assumptions on the stepsizes, that | ˜ β 0 | ⩽ α 2 and α ∞ ⩽ α (component-wise). The magnitude of ˜ β 0 is therefore negligible in front of the magnitudes of β ⋆ ∈ S and one can roughly ignore the term ˜ β 0 . Hence, the implicit regularisation Eq. (5) can be thought of as β ⋆ ∞ ≈ argmin β ⋆ ∈ S D ψ α ∞ ( β ⋆ , 0) = ψ α ∞ ( β ⋆ ) , and thus the solution β ⋆ ∞ minimises the same potential function that the solution of gradient flow (see Eq. (4) ), but with an effective initialisation α ∞ . Also note that for γ k ≡ γ → 0 we have α ∞ → α and ˜ β 0 → 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 ∑ k q ( γ k ∇L B k ( β k )) . Also, as the (stochastic) gradients converge to 0 and q ( x ) x → 0 ∼ x 2 , one should think of this sum as roughly being ∑ k ∇L B k ( β k ) 2 : the larger this sum, the more the recovered solution differs from that of gradient flow. The full picture of how large stepsizes and stochasticity impact the generalisation properties of β ⋆ ∞ and the recovery of minimum ℓ 1 -norm solution is nuanced as clearly seen in Fig. 1.

3.3 Convergence of the iterates

Theorem 1 provides the implicit minimisation problem but says nothing about the convergence of the iterates. Here we show under very reasonable assumptions on the stepsizes that the iterates indeed

converge towards a global optimum. Note that since the loss F is non-convex, such a convergence result is non-trivial and requires an involved analysis.

Theorem 2 (Convergence of the iterates) . Let ( u k , v k ) k ⩾ 0 follow the mini-batch SGD recursion (3) initialised at u 0 = √ 2 α ∈ R d > 0 and v 0 = 0 , and let ( β k ) k ⩾ 0 = ( u k ⊙ v k ) k ⩾ 0 . Recall the 'smoothness' parameter L on the minibatch loss defined in the notations. There exist B > 0 verifying B = ˜ O (min β ⋆ ∈S ∥ β ⋆ ∥ ∞ ) and a numerical constant c > 0 such that for stepsizes satisfying γ k ⩽ c LB , the iterates ( β k ) k ⩾ 0 converge almost surely to the interpolator β ⋆ ∞ 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 a uniform initialisation α = α 1 and under the assumptions of Theorem 2, we have:

where λ b > 0 is the largest value such that λ b H ⪯ E B [ H B ] , C = 2 B ( α 2 λ + min ) -1 ( 1 + (4 Bλ max )( α 2 λ + min ) -1 ) L ( β 0 ) and λ + min , λ max > 0 are respectively the smallest non-null and the largest eigenevalues of H , and β ⋆ α k is the interpolator that minimises the perturbed hypentropy h k of parameter α k , as defined in Eq. (7) in the next subsection.

The convergence of the losses is proved directly using the time-varying mirror structure that we exhibit in the next subsection, the convergence of the iterates is proved by studying the curvature of the mirror maps on a small neighborhood around the affine interpolation space.

3.4 Sketch of proof through a time varying mirror descent

As in the continuous-time framework, our results heavily rely on showing that the iterates ( β k ) k follow a mirror descent recursion with time-varying potentials on the convex loss L ( β ) . To show this, we first define the following quantities:

where α ± ,k := α exp ( -1 2 ∑ k -1 i =0 q ± ( γ ℓ ∇L B ℓ ( β ℓ ) ) ) ∈ R d . Finally for k ⩾ 0 , we define the potentials ( h k : R d → R ) k ⩾ 0 as:

Where ψ α k is the hyperbolic entropy function defined Eq. (4). Now that all the relevant quantities are defined, we can state the following proposition which explicits the time-varying stochastic mirror descent.

Proposition 2. The iterates ( β k = u k ⊙ v k ) k ⩾ 0 from Eq. (3) satisfy the Stochastic Mirror Descent recursion with varying potentials ( h k ) k :

where h k : R d → R for k ⩾ 0 are defined Eq. (7) . Since ∇ h 0 ( β 0 ) = 0 we have:

Theorem 1 and 2 and Proposition 1 follow from this key proposition: by suitably modifying classical convex optimization techniques to account for the time-varying potentials, we can prove the convergence of the iterates towards an interpolator β ⋆ ∞ along with that of the relevant quantities α ± ,k , α k and ϕ k . The implicit regularisation problem then directly follows from: (1) the limit condition ∇ h ∞ ( β ∞ ) ∈ Span( x 1 , . . . , x n ) as seen from Eq. (8) and (2) the interpolation condition Xβ ⋆ ∞ = 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 α ∞

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 α ∞ , which plays a key role as shown in Theorem 1. From its definition in Eq. (6), we see that α ∞ is a function of the vector ∑ k q ( γ k ∇L B k ( β k )) . We henceforth call this quantity the gain vector . For simplicity of the discussions, from now on, we consider constant stepsizes γ k = γ for all k ⩾ 0 and a uniform initialisation of the weights α = α 1 with α > 0 . We can then write the gain vector as:

Following our discussion in section 3.1 on the scale and the shape of α ∞ , we recall the link between the scale and shape of Gain γ and the recovered solution:

The scale of Gain γ , i.e. the magnitude of ∥ Gain γ ∥ 1 indicates how much the implicit bias of (S)GD differs from that of gradient flow: ∥ Gain γ ∥ 1 ∼ 0 implies that α ∞ ∼ α and therefore the recovered solution is close to that of gradient flow. On the contrary, ∥ Gain γ ∥ 1 > > ln(1 /α ) implies that α ∞ has effective scale much smaller than α thereby changing the implicit regularisation Eq. (5).

The shape of Gain γ indicates which coordinates of β in the associated minimum weighted ℓ 1 problem are most penalised. First recall from Section 3.1 that a uniformly large Gain γ leads to ψ α ∞ being closer to the ℓ 1 -norm. However, with small weight initialisation α → 0 , we have,

In this case, having a heterogeneously large vector Gain γ leads to a weighted ℓ 1 norm as the effective implicit regularisation, where the coordinates of β corresponding to the largest entries of Gain γ are 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 ∥ Gain γ ∥ 1 in terms of various problem constants.

Proposition 3. Let Λ b , λ b > 0 3 be the largest and smallest values, respectively, such that λ b H ⪯ E B [ H 2 B ] ⪯ Λ b H . For any stepsize γ > 0 satisfying γ ⩽ c BL (as in Theorem 2), initialisation α 1 and batch size b ∈ [ n ] , the magnitude of the gain satisfies:

where the expectation is over a uniform and independent sampling of the batches ( B k ) k ⩾ 0 .

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 (S)GD trajectory deviates from that of 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 0 < γ ⩽ γ max = c BL as in Theorem 2 we have that (see Appendix G.1):

Eq. (11) clearly shows that increasing the stepsize boosts the magnitude ∥ Gain γ ∥ 1 up until the limit of γ max. Therefore, the larger the stepsize the smaller is the effective scale of α ∞ . In turn, larger gap between α ∞ and α leads to a larger deviation of (S)GD from the gradient flow.

<latexi sh

1_b

64="DfTXJMY

Qw ord3z

W

n

uE/q

I

>A

C

c

V

7

FH

v

K

L

g

B

P

2

p

O

R

G

Nj

+

9

S

y

mkU

5

8

0

Z

<latexi sh

1_b

64="VS3

y

HAX20

Z

o

nR

7

+w

>

C

c

DJ

gNBEO

5

L

K

jMq

G

W

TQ

k

p

u

I

/

F

PUv

d

9

z

r

Y

f

8

m

<latexi sh

1_b

64="oUg

8

wEn0Q

Y

yXOf

HKZ

>A

c

V

NS

3

r

q/

9L

GPR

2

p

D

k

T

W

v

5

7

d

M

C

J

j

m

I

F

u

z

<latexi sh

1_b

64="0L/

3

+

G5C

I

gcE

y

DV

Pk

>A

BX

JS

N

O

j

R

Y

m

qM

F

8

T

U

H

r

o

f2

n

9K

p

z

Z

7

W

1

2

3

4

d

0

50

Shape of

Gain

γ

for GD

∝

β

∗

sparse

Shape of

1

2

3

4

d

Coordinate

i

0

50

Shape of

Gain

γ

for SGD

︸

︷︷

︸

Support of

β

/star

sparse

1

2

3 4

d

Coordinate

i

0

50

SGD:

Gain

γ

10

-

1

˜

γ

SGD

˜

γ

Stepsize

γ

0

200

400

600

800

Scale:

‖

Gain

γ

‖

1

SGD

Averaged

Averaged

(

β

/star

sparse

)

2

Large stepsizes and Edge of Stability. The previous paragraph holds for stepsizes smaller than γ max for which we can theoretically prove convergence. But what if we use even bigger stepsizes? Let ( β γ k ) k denote the iterates generated with stepsize γ and let us define ˜ γ max := sup γ ⩾ 0 { γ s.t. ∀ γ ′ ∈ (0 , γ ) , ∑ k L ( β γ ′ k ) < ∞} , which corresponds to the largest stepsize such that the iterates still converge for a given problem (even if not provably so). From Proposition 3 we have that γ max ⩽ ˜ γ max . As we approach this upper bound on convergence γ → ˜ γ max , the sum ∑ k L ( β γ k ) diverges. For such large stepsizes, the iterates of gradient descent tend to 'bounce' and this regime is commonly referred to 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 γ can be become arbitrarily big as 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 x i sampled from N (0 , σ 2 I d ) with σ 2 > 0 , we obtain E [ ∥ Gain γ ∥ 1 ] = Θ ( γ σ 2 d b ln ( 1 α ) ∥ β ⋆ ℓ 1 ∥ 1 ) , w.h.p. over the dataset (see Appendix 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 on γ b , resembling the linear scaling rule commonly used in deep learning [22].

By analysing the magnitude ∥ Gain γ ∥ 1 , we have explained the distinct behavior of (S)GD with large stepsizes compared to 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. A deeper understanding of the shape of Gain γ 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 ( b = 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 β ⋆ ℓ 1 and thus these coordinates are adversely weighted in the asymptotic limit of ψ α ∞ (per (9)). This explains the distinction seed in Fig. 1, where GD in this regime has poor sparse recovery despite having a small scale of α ∞ , as opposed to SGD that behaves well.

The shape of Gain γ is determined by the sum of the squared gradients ∑ k ∇L B k ( β k ) 2 , 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 ( β k ) k 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.

In the simple case of a Gaussian noiseless sparse recovery problem (where y i = ⟨ β ⋆ sparse , x i ⟩ for some sparse vector β ⋆ sparse ), we can control these gradients for GD and SGD (Appendix G.4) as:

The gradient of GD is heterogeneous. Since β ⋆ sparse is sparse by definition, we deduce from Eq. (25) that ∇L ( β 0 ) is heterogeneous with larger values corresponding to the support of β ⋆ sparse . Along with observation 1, this means that Gain γ has much larger values on the support of β ⋆ sparse . The corresponding weighted ℓ 1 -norm therefore penalises the coordinates belonging to the support of β ⋆ sparse , which hinders the recovery of β ⋆ sparse (as explained in Example 1, 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 ℓ 1 -norm where the weights are roughly balanced. The corresponding weighted ℓ 1 -norm is therefore close to the uniform ℓ 1 -norm and the classical ℓ 1 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:

for small stepsizes, the scale is small, and (S)GD solutions are close to that of gradient flow;

for large stepsizes the scale is significant and the recovered solutions differ from GF:

for SGD the shape of α ∞ is uniform, the associated norm is closer to the ℓ 1 -norm and the recovered solution is closer to 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, we relate heuristically these findings to the Edge of Stability phenomenon.

5 Edge of Stability: the neural point of view

In recent years it 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 = ( u k , v k ) ∈ R 2 d 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 a small initialisation α = α 1 where α < < 1 such that we can consider that gradient flow converges close to the sparse interpolator β ⋆ sparse = β w ⋆ sparse corresponding to the weights w ⋆ sparse = ( √ | β ⋆ sparse | , sign( β ⋆ sparse ) √ | β ⋆ sparse | ) (see Lemma 1 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 to 3 phases.

First phase: gradient flow. 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 γ < 2 λ max ( ∇ 2 F ( w ⋆ sparse )) , then it 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 γ > 2 λ max ( ∇ 2 F ( w ⋆ sparse )) (as is the case for γ GD in Fig. 3, lower right), the iterates cannot converge to β ⋆ sparse and we enter the oscillating phase.

Second phase: oscillations. The iterates start oscillating. The gradient of F writes ∇ ( u,v ) F ( w ) ∼ ( ∇L ( β ) ⊙ v, ∇L ( β ) ⊙ u ) and for w in the vicinity of w ⋆ sparse we have that u i ≈ v i ≈ 0 for i / ∈ supp( β ⋆ sparse ) . Therefore for w ∼ w ⋆ sparse we have that ∇ u F ( w ) i ≈ ∇ v F ( w ) i ≈ 0 for i / ∈ supp( β ⋆ sparse ) and the gradients roughly belong to Span( e i , e i + d ) i ∈ supp( β ⋆ sparse ) . This means

<latexi sh

1_b

64="9RgYH

N

GEC3

LS

nO

p+

Q

>A

c

VD

M

F

U

2vqj dBI

r

0W

/

om

u

Z

TJ

X

7

8

w

z

5

k

K

f

y

P

︸

Iteration

Iteration

k

k

(

i

)

k

10

0

10

1

10

2

10

3

10

4

Iteration

k

-

100

0

100

200

Coordinates

β

(

i

)

k

of GD at EoS

supp(

β

∗

sparse

)

coordinates

Other coordinates

Gradient flow

10

0

10

1

10

2

10

3

λ

(

∇

2

F

(

w

k

))

Gradient descent

Gradient flow

2

γ

GD ~ GF

Oscillations

Convergence

︸

︷︷

︸

︸

︷︷

︸

10

0

10

1

10

2

10

3

Iteration

k

10

-

1

10

1

10

3

Train losses

Gradient descent

Gradient flow

10

0

10

1

10

2

10

3

10

4

Iteration

k

-

100

0

100

200

Coordinates

β

(

i

)

k

of GD at EoS

supp(

β

∗

sparse

)

coordinates

Other coordinates

Gradient flow

10

0

10

1

10

2

10

3

10

4

-

100

0

100

200

Coordinates

β

(

i

)

k

of GD at EoS

supp(

β

∗

sparse

)

coordinates

Other coordinates

Gradient flow

10

0

10

1

10

2

Iteration

k

-

100

0

100

200

Coordinates

β

k

supp(

β

Drift towards 0

0

100

200

supp(

β

∗

sparse

)

coordinates

Other coordinates

Gradient flow

︷︷

Coordinates

︸

10

10

10

2

of GD at EoS

10

3

10

4

.

Iteration k that only the coordinates of the weights ( u i , v i ) for i ∈ supp( β ⋆ sparse ) can oscillate and similarly for ( β i ) i ∈ supp( β ⋆ sparse ) (as seen Fig. 3 left).

Last phase: convergence. Due to 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 1 enables us to understand where they converge: the coordinates of β k that have oscillated significantly along the trajectory belong to the support of β ⋆ sparse , and therefore Gain γ (i) becomes much larger for i ∈ supp( β ⋆ sparse ) 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 ( β i ) i ∈ supp( β ⋆ sparse ) lead to a gradual shift of these coordinates towards 0 , hindering an accurate recovery of the solution β ⋆ sparse .

SGD in the EoS regime. In contrast to the behavior of GD where the oscillations primarily occur on the non-sparse coordinates of ground truth sparse model, for SGD we see a 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 a uniform α ∞ . 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 are captured by the effective initialisation parameter α ∞ 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 on the role of stepsize and stochasticity. In particular, our characterisation explains the fundamentally different generalisation properties of SGD and GD solutions at large stepsizes as seen in Fig. 1: without stochasticity, the use of large stepsizes can prevent the recovery of the sparse interpolator, even though the effective scale 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.

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 a week in Lausanne at EPFL, making it possible to meet and discuss with Scott Pesme and Nicolas Flammarion.

<latexi sh

1_b

64="RvM8

YU0oTdzCW

un5S

f

Q

PA

>

c

VDL

g

F

3

2

J

E

m

y

Z

B

p

G

k

K

7

H

r

j

IO

w

N

X

/

9

q

+

β

<latexi sh

1_b

64="p nV

My/rBNuC

A

w0Im9

F

R8

>

3

c

DLSg

P

d

Y

T

W

o

k

Q

GZ

z

K25

O

j

H

v

X

q

7+

E

f

U

J

(

i

)

of GD at EoS

∗

sparse

)

coordinates

Other coordinates

Gradient flow

10

3

<latexi sh

1_b

64="

V

0J3dDB

Q

w

p

CXR

W

E

>A

n

c

LS

N

F

rPU

m

Ku

q

f

T

y

Y

O

G

I

j

7

o

+

8

M

H

v

2

z

5

/

k

Z

g

9

10

4

References

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.

M. Andriushchenko, A. Varre, L. Pillaud-Vivien, and N. Flammarion. SGD with large step sizes learns sparse features. arXiv preprint , 2022.

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.

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-263, January 2008.

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.

Raphaël Berthier. Incremental learning in diagonal linear networks. arXiv preprint arXiv:2208.14673 , 2022.

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 254-282. PMLR, 02-05 Jul 2022.

G. Blanc, N. Gupta, G. Valiant, and P. Valiant. Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeck 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.

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.

E. Candès, J. Romberg, and T. Tao. Stable signal recovery from incomplete and inaccurate measurements. Communications on Pure and Applied Mathematics , 59(8):1207-1223, 2006.

Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks. In International Conference on Learning Representations , 2018.

Lei Chen and Joan Bruna. On gradient descent convergence beyond the edge of stability, 2022.

Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On Lazy Training in Differentiable Programming . 2019.

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 on Learning Representations , 2021.

Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global minimizers. In Advances in Neural Information Processing Systems , 2021.

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.

J. L. Doob. Stochastic Processes . John Wiley & Sons, 1990.

Radu Alexandru Dragomir, Mathieu Even, and Hadrien Hendrikx. Fast stochastic Bregman gradient methods: Sharp analysis and variance reduction. In Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 2815-2825. PMLR, 18-24 Jul 2021.

Mathieu Even and Laurent Massoulie. Concentration of 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.

Jonas Geiping, Micah Goldblum, Phillip E Pope, Michael Moeller, and Tom Goldstein. Stochastic training is not necessary for generalization. In International Conference on Learning Representations , 2022.

Udaya Ghai, Elad Hazan, and Yoram Singer. Exponentiated gradient meets gradient descent. In Proceedings of the 31st International Conference on Algorithmic Learning Theory , volume 117 of Proceedings of Machine Learning Research , pages 386-407. PMLR, 08 Feb-11 Feb 2020.

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 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems , 30, 2017.

Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 1832-1841. PMLR, 10-15 Jul 2018.

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.

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.

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.

Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation , 9(1):1-42, January 1997.

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.

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.

Stanisław Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD, 2017.

Stanisław 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.

Stanislaw Jastrzkebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Width of minima reached by stochastic gradient descent is influenced by learning rate to batch size ratio. In Artificial Neural Networks and Machine Learning ICANN 2018 , pages 392-402, 2018.

Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representations , 2019.

Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In Advances in Neural Information Processing Systems , volume 33, pages 17176-17186, 2020.

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.

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.

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 2698-2707. PMLR, 10-15 Jul 2018.

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.

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890 , 2019.

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.

Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprint arXiv:1804.07612 , 2018.

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.

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.

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.

Ryan O'Donnell. Analysis of boolean functions, 2021.

Francesco Orabona, Koby Crammer, and Nicolò Cesa-Bianchi. A generalized online mirror descent with applications to classification and regression. Mach. Learn. , 99(3):411-435, jun 2015.

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.

Scott Pesme and Nicolas Flammarion. Saddle-to-saddle dynamics in diagonal linear networks. arXiv preprint arXiv:2304.00488 , 2023.

L. Pillaud-Vivien, J. Reygner, and N. Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 2127-2159. PMLR, 2022.

H. Robbins and S. Monro. A stochastic approxiation method. Ann. Math. Statist , 22(3):400-407, 1951.

Samuel L. Smith and Quoc V. Le. A Bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations , 2018.

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. , 19(1):2822-2878, jan 2018.

Terrence Tao. Concentration of measure. 254A, Notes 1, Blogpost , 2010.

Matus Telgarsky. Margins, shrinkage, and boosting. In International Conference on Machine Learning , pages 307-315. PMLR, 2013.

Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. The statistical complexity of earlystopped mirror descent. In Advances in Neural Information Processing Systems , volume 33, pages 253-264, 2020.

Tomas Vaškeviˇ cius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. In Proceedings of the 33rd International Conference on Neural Information Processing Systems , 2019.

Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science . Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.

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.

Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part II: Continuous time analysis. arXiv preprint arXiv:2106.02588 , 2021.

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 on Learning Theory , volume 125 of Proceedings of Machine Learning Research , pages 3635-3673. PMLR, 09-12 Jul 2020.

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 2019220203, 2020.

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.

Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized learning: A dynamical stability perspective. In Advances in Neural Information Processing Systems , volume 31, 2018.

Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701 , 2012.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations , 2017.

Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-ofstability training dynamics with a minimalist example. International Conference on Learning Representations , 2023.

Organisation of the Appendix.

In Appendix A, we provide additional experiments for uncentered data as well as on the 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.

In Appendix B, we prove that ( β k ) follows a Mirror descent recursion with varying potentials. We explicit these potentials and discuss some consequences.

In Appendix C we prove that (S)GD on the 1 2 ( w 2 + -w 2 -) and u ⊙ v parametrisations with suitable initialisations lead to the same sequence ( β k ) .

In Appendix D, we show that the hypentropy ψ α converges to a weighted -ℓ 1 -norm when α converges to 0 non-uniformly. We then discuss the effects of this weighted ℓ 1 -norm for sparse recovery.

In Appendix E, we provide our descent lemmas for mirror descent with varying potentials and prove the boundedness of the iterates.

In Appendix F, we prove our main results: Theorem 1 and Theorem 2, as well as quantitative convergence (Proposition 1).

In Appendix G, we prove the lemmas and propositions given in the main text.

In Appendix H, we provide technical lemmas used throughout the proof of Theorem 1 and Theorem 2.

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 to the data.

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 ℓ 1 solutions for uncentered data (Fig. 4). We make the following assumptions on the uncentered inputs.

Assumption 1. There exist µ ∈ R d and δ, c 0 , c 1 , c 2 > 0 such that for all s -sparse vectors β verifying ⟨ µ, β ⟩ ⩾ c 0 ∥ β ∥ ∞ ∥ µ ∥ ∞ , there exists ε ∈ R d such that ( X ⊤ X ) β = ⟨ β, µ ⟩ µ + ε where ∥ ε ∥ 2 ⩽ δ ∥ β ∥ 2 and c 1 ⟨ β, µ ⟩ 2 µ 2 ⩽ 1 n ∑ i x 2 i ⟨ x i , β ⟩ 2 ⩽ c 2 ⟨ β, µ ⟩ 2 µ 2 .

Assumption 1 is not restrictive and holds with high probability for N ( µ 1 , σ 2 I d ) inputs when µ > > σ 1 (see Lemma 9 in Appendix). The following lemma characterises the initial shape of SGD and GD gradients for uncentered data.

Proposition 4 (Shape of the (stochastic) gradient at initialisation) . Under Assumption 1 and if ⟨ µ, β ⋆ sparse ⟩ ⩾ c 0 ∥ β ∥ ∞ ∥ µ ∥ ∞ , the squared full batch gradient and the expected stochastic gradient descent at initialisation satisfy, for some ε satisfying ∥ ε ∥ ∞ < < ∥ β sparse ∥ 2 :

In this case the initial gradients of 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 in the presence of uncentered data, as opposed to centered data. Note that for decentered data with a µ ∈ R d orthogonal to β ⋆ sparse , there is no effect of decentering on the recovered solution. If the support of µ is the same as that of β ⋆ sparse , the effect is detrimental and the same discussion as in the centered data case applies.

Fig. 4: for uncentered data the solutions of GD and SGD have similar behaviours, corroborating Proposition 4.

10

0

10

1

Stepsize

γ

10

-

5

10

-

4

10

-

3

10

-

2

Test losses

||

β

γ

∞

-

β

∗

/lscript

1

||

2

2

SGD

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 ( w γ ∞ ) at the convergence of the iterates of SGD and GD (2) the trace

200

150

100

50

0

50

-

-

100

10

0

Coordinates

10

1

of hessian at the convergence of the iterates. As is clearly observed, increasing the stepsize for GD leads to a 'flatter' minimum in terms of the maximum eigenvalue of the hessian, while increasing the stepsize for SGD leads to a '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 = I d ' and that ' α ∼ 0 ' (pushing the EoS phenomenon), one can see that minimising λ max ( ∇ 2 F ( w )) under the constraints X ( w 2 + -w 2 -) = y and w + ⊙ w -= 0 is equivalent to minimising ∥ β ∥ ∞ under the constaint Xβ = y . On the other hand minimising the trace of the hessian is equivalent to minimising the ℓ 1 -norm.

10

-

1

10

0

Stepsize

γ

3

.

5

4

.

0

4

.

5

5

.

0

λ

(

∇

2

F

(

w

γ

∞

))

SGD

10

-

1

10

0

Stepsize

γ

10

.

0

12

.

5

15

.

0

17

.

5

20

.

0

Trace(

∇

2

F(w

γ

∞

))

SGD

A.3 Edge of Stability for SGD

β

(

i

)

k

of GD at EoS

supp(

β

∗

sparse

)

coordinates other coordinates

Gradient flow

10

2

10

Iteration

k

10

0

10

1

10

2

10

3

10

4

Iteration

k

-

100

-

50

0

50

100

150

200

Coordinates

β

(

i

)

k

of SGD at EoS

supp(

β

∗

sparse

)

coordinates

other coordinates

Gradient flow

3

10

4

B Main ingredients behind the proof of Theorem 1 and Theorem 2

In this section, we show that the iterates ( β k ) k ⩾ 0 follow a stochastic mirror descent with varying potentials . At the core of our analysis, this result enables us 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 a strictly convex reference function h : R d → R , the (stochastic) mirror descent iterates algorithm write as [5, 18], where the minimum is assumed to be attained over R d and unique:

for stochastic gradients g k , stepsize γ k ⩾ 0 , and D h ( β, β ′ ) = h ( β ) -h ( β ′ ) -⟨∇ h ( β ′ ) , β -β ′ ⟩ is the Bregman divergence associated to h . Iteration (16) can also be cast as

Now, let ( h k ) be strictly convex reference functions R d → 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 k and leading to ∇ h k ( β k +1 ) = ∇ h k ( β k ) -γ k g k ), the mirror descent with varying potentials we study in this paper incorporates h k +1 and h k in (17). The iterates are thus defined as through:

where D h k +1 ,h k ( β, β ′ ) = h k +1 ( β ) -h k ( β ′ ) -⟨∇ h k ( β ′ ) , β -β ′ ⟩ , a recursion that can also be cast as:

To derive convergence of the iterates, we prove analogs to classical mirror descent lemmas, generalised to time-varying potentials.

B.2 The iterates ( β k ) follow a stochastic mirror descent with varying potential recursion

In this section we show and prove that the iterates ( β k ) k follow a stochastic mirror descent with varying potentials. Before stating the proposition, we recall the definition of the potentials. To do so we introduce several quantities.

Let q, q ± : R → R ∪ {∞} be defined as:

with the convention that q (1) = ∞ . Notice that q ( x ) ⩾ 0 for | x | ⩽ √ 2 and q ( x ) < 0 otherwise. For the iterates β k = u k ⊙ v k ∈ R d , we recall the definition of the following quantities:

Finally for k ⩾ 0 , we define the potentials ( h k : R d → R ) k ⩾ 0 as:

where ψ α k is the hyperbolic entropy defined in (4) of scale α k :

where α k,i corresponds to the i th coordinate of the vector α k .

Now that all the relevant quantities are define, we can state the following proposition which explicits the time-varying stochastic mirror descent followed by ( β k ) k

Proposition 5. The iterates ( β k = u k ⊙ v k ) k ⩾ 0 from Eq. (3) satisfy the Stochastic Mirror Descent recursion with varying potentials ( h k ) k :

where h k : R d → R for k ⩾ 0 are defined Eq. (18) . Since ∇ h 0 ( β 0 ) = 0 we have:

Proof. Using Proposition 6, we study the 1 2 ( w 2 + -w 2 -) parametrisation instead of the u ⊙ v , indeed this is the natural parametrisation to consider when doing the calculations as it 'separates" the recursions on w + and w -.

Let us focus on the recursion of w + :

with the convention that exp(ln(0)) = 0 . This leads to:

since q + ( x ) = -2 x -ln((1 -x ) 2 ) . Expanding the recursion and using that w + ,k =0 is initialised at w + ,k =0 = α , we thus obtain:

where we recall that α 2 ± ,k = α 2 exp( -∑ k -1 ℓ =0 q ± ( γ ℓ g ℓ )) . One can easily check that we similarly get:

We have:

leading to:

Using Lemma 4, the previous equation can be simplified into:

which writes as:

where ϕ k = 1 2 arcsinh ( α 2 + ,k -α 2 -,k 2 α 2 k ) , α 2 k = α + ,k ⊙ α -,k and since the potentials h k are defined in Eq. (18) as h k = ψ α k -⟨ ϕ k , ·⟩ with

specifically such that ∇ h k ( β k ) = 1 2 arcsinh ( β k α 2 k ) -ϕ k . Hence,

which corresponds to a Mirror Descent with varying potentials ( h k ) k .

C Equivalence of the u ⊙ v and 1 2 ( w 2 + -w 2 -) parametrisations

We here prove the equivalence between the 1 2 ( w 2 + -w 2 -) and u ⊙ v parametrisations, that we use throughout the proofs in the Appendix.

Proposition 6. Let ( β k ) k ⩾ 0 and ( β ′ k ) k ⩾ 0 be respectively generated by stochastic gradient descent on the u ⊙ v and 1 2 ( w 2 + -w 2 -) parametrisations:

and initialised as u 0 = √ 2 α , v 0 = 0 and w + , 0 = w -, 0 = α . Then for all k ⩾ 0 , we have β k = β ′ k .

Proof. We have:

and

Hence, and

Then, let z k = 1 2 ( u 2 k -v 2 k ) and z ′ k = w + ,k w -k . We have z 0 = α 2 , z ′ 0 = α 2 and:

Using a 2 + b 2 = √ (2 ab ) 2 +( a 2 b 2 ) 2 for a, b R

We conclude by observing that ( β k , z k ) and ( β ′ k , z ′ k ) follow the exact same recursions, initialised at the same value (0 , α 2 ) .

so that:

D Convergence of ψ α to a weighted ℓ 1 norm and harmful behaviour

Weshow that when taking the scale of the initialisation to 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 is a weighted ℓ 1 -norm rather than the classical ℓ 1 norm.

Proposition 7. For α ⩾ 0 and a vector h ∈ R d , let ˜ α = α exp( -h ln(1 /α )) ∈ R d . Then we have that for all β ∈ R d

Proof. Recall that

Using that arcsinh( x ) ∼ | x |→∞ sgn( x ) ln( | x | ) , and that ln( 1 ˜ α 2 i ) = (1 + h i ) ln( 1 α 2 ) we obtain that

The following Fig. 7 illustrates the effect of the non-uniform shape α on the corresponding potential ψ α .

More generally, for α such that α i → 0 for all i ∈ [ d ] at rates such that ln(1 /α i ) ∼ q i ln(1 / max i α i ) , we retrieve a weighted ℓ 1 norm:

Hence, even for arbitrary small max i α i , if the shape of α is 'bad', the interpolator β α that minimizes ψ α can be arbitrary far away from β ⋆ ℓ 1 the interpolator of minimal ℓ 1 norm.

We illustrate the importance of the previous proposition in the following example.

Example 1. We illustrate how, even for arbitrary small max i α i , the interpolator β ⋆ α that minimizes ψ α can be far from the minimum ℓ 1 norm solution, due to the shape of α that is not uniform. The message of this example is that for α → 0 non-uniformly across coordinates, if the coordinates of α that go slowly to 0 coincide with the non-null coordinates of the sparse interpolator we want to retrieve, then β ⋆ α will be far from the sparse solution.

A simple counterexample can be built: let β ⋆ sparse = (1 , . . . , 1 , 0 , . . . , 0) (with only the s = o ( d ) first coordinates that are non-null), and let ( x i ) , ( y i ) be generated as y i = ⟨ β ⋆ sparse , x i ⟩ with x i ∼ N (0 , 1) . For n large enough ( n of order s ln( d ) where s is the sparsity), the design matrix X is RIP [10], so that the minimum ℓ 1 norm interpolator β ⋆ ℓ 1 is exactly equal to β ⋆ sparse .

However, if α is such that max i α i → 0 with h i >> 1 for j ⩽ s and h i = 1 for i ⩾ s +1 ( h i as in Proposition 7), β ⋆ α will be forced to verify β ⋆ α,i = 0 for i ⩽ s and hence ∥ β ⋆ α, 1 -β ⋆ ℓ 1 ∥ 1 ⩾ s .

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 wellchosen stepsizes, the Bregman divergences ( D h k ( β ⋆ , β k )) k ⩾ 0 decrease. We then use this proposition to bound the iterates for both SGD and GD.

Proposition 8. There exist a constant c > 0 and B > 0 such that B = O (inf β ⋆ ∈S ∥ β ⋆ ∥ ∞ ) for GD and B = O (ln(1 /α ) inf β ⋆ ∈S ∥ β ⋆ ∥ ∞ ) for SGD, such that if γ k ⩽ c LB for all k , then we have, for all k ⩾ 0 and any interpolator β ⋆ ∈ S :

To prove this result, we first provide a general descent lemma for time-varying mirror descent (Proposition 9, appendix E.1), before proving the proposition for fixed iteration k and bound B > 0 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 et 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, g : R d → R functions, let D h,g ( β, β ′ ) = h ( β ) -g ( β ′ ) -⟨∇ g ( β ′ ) , β -β ′ ⟩ 4 for β, β ′ ∈ R d . Let ( h k ) strictly convex functions defined R d L a convex function defined on R d . Let ( β k ) defined recursively through β 0 ∈ R d , and

where we assume that the minimum is unique and attained in R d . Then, ( β k ) satisfies

and for any β ∈ R d ,

Proof. Let β ∈ R d . Since we assume that the minimum through which β k +1 is computed is attained in R d , the gradient of the function V k ( β ) = γ k ⟨∇L ( β k ) , β -β k ⟩ + D h k +1 ,h k ( β, β k ) evaluated at β k +1 is null, leading to ∇ h k +1 ( β k +1 ) = ∇ h k ( β k ) -γ k ∇L ( β k ) .

Then, since ∇ V k ( β k +1 ) = 0 , we have D V k ( β, β k +1 ) = V k ( β ) -V k ( β k +1 ) . Using ∇ 2 V k = ∇ 2 h k +1 , we also have D V k = D h k +1 . Hence:

We write γ k ⟨∇L ( β k ) , β -β k +1 ⟩ = γ k ⟨∇L ( β k ) , β -β k ⟩ + γ k ⟨∇L ( β k ) , β k -β k +1 ⟩ . We also have γ k ⟨∇L ( β k ) , β k -β k +1 ⟩ = ⟨∇ h k ( β k ) -∇ h k +1 ( β k +1 ) , β k -β k +1 ⟩ = D h k ,h k +1 ( β k , β k +1 ) + D h k +1 ,h k ( β k +1 , β k ) , so that γ k ⟨∇L ( β k ) , β k -β k +1 ⟩-D h k +1 ,h k ( β k +1 , β k ) = D h k ,h k +1 ( β k , β k +1 ) . Thus,

E.2 Proof of Proposition 10

In next proposition, we use Proposition 9 to prove our main descent lemma. To that end, we bound the error terms that appear in Proposition 9 as functions of L B k ( β k ) and norms of β k , β k +1 , so that for explicit stepsizes, the error terms can be cancelled by half of the negative quantity -2 L B k ( β k ) .

Proposition 10. Let k ⩾ 0 and B > 0 . Provided that ∥ β k ∥ ∞ , ∥ β k +1 ∥ ∞ , ∥ β ⋆ ∥ ∞ ⩽ B and γ k ⩽ c LB where c > 0 is some numerical constant, we have:

Additional notation: let L 2 , L ∞ > 0 such that ∀ β , ∥ H B β ∥ 2 ⩽ L ∥ β ∥ 2 , ∥ H B β ∥ ∞ ⩽ L ∥ β ∥ ∞ for all batches B ⊂ [ n ] of size b .

Proof. Let β ⋆ ∈ S be any interpolator. From Proposition 9:

We want to bound the last three terms of this equality. First, to bound the last two we apply Lemma 7 assuming that ∥ β ⋆ ∥ ∞ , ∥ β k +1 ∥ ∞ ⩽ B :

We now bound D h k +1 ( β k , β k +1 ) . Classical Bregman manipulations provide that

From Lemma 6 we have that h k +1 is min(1 / (4 α 2 k +1 ) , 1 / (4 B )) strongly convex on the ℓ ∞ -centered ball of radius B therefore h ∗ k +1 is max(4 α 2 k +1 , 4 B ) = 4 B (for α small enough or B big enough) smooth on this ball, leading to:

Using |∇ h k ( β ) -∇ h k +1 ( β ) | ⩽ 2 δ k where δ k = q ( γ k ∇L B k ( β k )) , we get that:

Now, ∥ δ k ∥ 2 2 ⩽ ∥ δ k ∥ 1 ∥ δ k ∥ ∞ and using Lemma 5, ∥ δ k ∥ 1 ∥ δ k ∥ ∞ ⩽ 4 ∥ γ k ∇L B k ( β k ) ∥ 2 2 ∥ γ k ∇L B k ( β k ) ∥ 2 ∞ ⩽ 2 ∥ γ k ∇L B k ( β k ) ∥ 2 2 since ∥ γ k ∇L B k ( β k ) ∥ ∞ ⩽ γ k L ∞ ∥ β k -β ∞ ∥ ⩽ γ k × 2 LB ⩽ 1 / 2 is verified for γ k ⩽ 1 / (4 LB ) . Thus,

Hence, provided that ∥ β k ∥ ∞ ⩽ B , ∥ β k +1 ∥ ∞ ⩽ B and γ k ⩽ 1 / (4 LB ) , we have:

and thus

if γ k ⩽ c BL , where c = 1 64 .

E.3 Bound on the iterates

We now bound the iterates ( β k ) by an explicit constant B that depends on ∥ β ⋆ ∥ 1 (for any fixed β ⋆ ∈ S ).

The first bound we prove holds for both SGD and GD, and is of the form O ( ∥ β ⋆ ∥ 1 ln(1 /α 2 ) while the second bound, that holds only for GD ( b = n ) is of order O ( ∥ β ⋆ ∥ 1 ) (independent of α ). While a bound independent of α is only proved for GD, we believe that such a result also holds for SGD, and in both cases B should be thought of order O ( ∥ β ⋆ ∥ 1 ) .

E.3.1 Bound that depends on α for GD and SGD

Aconsequence of Proposition 10 is the boundedness of the iterates, as shown in next corollary. Hence, Proposition 10 can be applied using B a uniform bound on the iterates ℓ ∞ norm.

Proof. We proceed by induction. Let k ⩾ 0 such that ∥ β k ∥ ∞ ⩽ B for some B > 0 and D h k ( β ⋆ , β k ) ⩽ D h 0 ( β ⋆ , β 0 ) (note that these two properties are verified for k = 0 , since β 0 = 0 ). For γ k sufficiently small ( i.e. , that satisfies γ k ⩽ c B ′ L where B ′ ⩾ ∥ β k +1 ∥ ∞ , ∥ β k ∥ ∞ , ∥ β ⋆ ∥ ∞ ), using Proposition 10, we have D h k +1 ( β ⋆ , β k +1 ) ⩽ D h k ( β ⋆ , β k ) so that D h k +1 ( β ⋆ , β k +1 ) ⩽ D h 0 ( β ⋆ , β 0 ) , which can be rewritten as:

Hence, ∥ β k +1 ∥ 1 ⩽ ∥ β ⋆ ∥ 1 ln(1 + ∥ β k +1 ∥ 1 α 2 ) . We then notice that for x, y > 0 , x ⩽ y ln(1 + x ) = ⇒ x ⩽ 3 y ln(1 + y ) : if x > y ln(1 + y ) and x > y , we have that y ln(1 + y ) < y ln(1 + x ) , so that 1 + y < 1 + x , which contradicts our assumption. Hence, x ⩽ max( y, y ln(1 + y )) . In our case, x = ∥ ∥ β k +1 ∥ ∥ 1 /α 2 , y = ∥ β ⋆ ∥ 1 /α 2 so that for small alpha, ln(1 + y ) ⩾ 1 .

This is true as long as γ k is tuned using B ′ a bound on max( ∥ β k ∥ ∞ , ∥ β k +1 ∥ ∞ ) . Using the continuity of β k +1 as a function of γ k ( β k being fixed), we show that γ k ⩽ 1 2 × c BL can be used using this B . Indeed, let ϕ : R + → R d be the function that takes as entry γ k ⩾ 0 and outputs the corresponding ∥ β k +1 ∥ ∞ : ϕ is continuous. Let γ r = 1 2 × c rL for r > 0 and ¯ r = sup { r ⩾ 0 : B < ϕ ( γ r ) } (the set is upper-bounded; if is is empty, we do not need what follows since it means that any stepsize leads to ∥ β k +1 ∥ ∞ ⩽ B ). By continuity of ϕ , ϕ ( γ ¯ r ) = B . Furthermore, for all r that satisfies r ⩾ max( ϕ ( γ r ) , B ) ⩾ max( ϕ ( γ r ) , ∥ β k ∥ ∞ , ∥ β ⋆ ∥ ∞ ) , we have, using what is proved just above, that ∥ β k +1 ∥ ∞ ⩽ B and thus ϕ ( γ r ) ⩽ B for such a r :

Hence, we deduce that ∥ β k +1 ∥ 1 ⩽ B , where B = ∥ β ⋆ ∥ 1 ln(1 + ∥ β ⋆ ∥ 1 α 2 ) .

Lemma 1. For r > 0 such that r ⩾ max( ϕ ( γ r ) , B ) , we have ϕ ( γ r ) ⩽ B .

Now, if ¯ r > B , by definition of ¯ r and by continuity of ϕ , since ϕ (¯ r ) = B , there exists some B < r < ¯ r such that ϕ ( γ r ) > B (definition of the supremum) and ϕ ( γ r ) ⩽ 2 B (continuity of ϕ ). This particular choice of r thus satisfies r > B and and ϕ ( γ r ) ⩽ 2 B ⩽ 2 r , leading to ϕ ( γ r ) ⩽ B , using Lemma 1, hence a contradiction: we thus have ¯ r ⩽ B .

This concludes the induction: for all r ⩾ B , we have r ⩾ ¯ r so that ϕ ( γ r ) ⩽ B and thus for all stepsizes γ ⩽ c 2 LB , we have ∥ β k +1 ∥ ∞ ⩽ B .

E.3.2 Bound independent of α

We here assume in this subsection that b = n . We prove that for gradient descent, the iterates are bounded by a constant that does not depend on α .

Proposition 11. Assume that b = n (full batch setting). There exists some B = O ( ∥ β ⋆ ∥ 1 ) such that for stepsizes γ k ⩽ c BL , we have ∥ β k ∥ ∞ ⩽ B for all k ⩾ 0 .

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 , w -,k ) which will be useful to show that the loss is decreasing.

Proposition 12. For γ k ⩽ c LB where B ⩾ max( ∥ β k ∥ ∞ , ∥ β k +1 ∥ ∞ ) , we have L ( β k +1 ) ⩽ L ( β k )

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 to show that the loss is decreasing we use the gradient descent recursion. Recall that the iterates w k = ( w + ,k , w -,k ) ∈ R 2 d follow a gradient descent on the non convex loss F ( w ) = 1 2 ∥ y -1 2 X ( w 2 + -w 2 -) ∥ 2 .

For k ⩾ 0 , using the Taylor formula we have that F ( w k +1 ) ⩽ F ( w k ) -γ k (1 -γ k L k 2 ) ∥∇ F ( w k ) ∥ 2 with the local smoothness L k = sup w ∈ [ w k ,w k +1 ] λ max ( ∇ 2 F ( w )) . Hence if γ k ⩽ 1 /L k for all k we get that the loss is non-increasing. We now bound L k . Computing the hessian ot F , we obtain that:

Let us denote by M = ( M + M + , -M + , -M -) ∈ R 2 d × 2 d the second matrix in the previous equality. With this notation ∥∇ 2 F ( w k ) ∥ ⩽ ∥∇L ( β k ) ∥ ∞ +2 ∥ M ∥ (where the norm corresponds to the Schatten 2 -norm which is the largest eigenvalue for symmetric matrices). Now, notice that:

We have:

Since ∥ M ± ∥ ⩽ λ max · ∥ w ± ∥ 2 ∞ and ∥ M + -∥ ⩽ λ max ∥ w + ∥ ∞ ∥ w -∥ ∞ we finally get that

We now upper bound this quantity in the following lemma.

Lemma 2. For all k ⩾ 0 , the following inequality holds component-wise:

Proof. Notice from the definition of w + ,k and w -,k given in the proof of Proposition 5 that:

And α 0 = α 2 . Now since α k is decreasing coordinate-wise (under our assumptions on the stepsizes, γ 2 k ∇L ( β k ) 2 ⩽ (1 / 2) 2 < 1 ), we get that.:

leading to w 2 + ,k + w 2 -,k ⩽ 4 α 4 + B 2 .

From Lemma 2, w 2 + ,k + w 2 -,k is bounded by 2 √ α 4 + B 2 . Putting things together we finally get that ∥∇ 2 F ( w ) ∥ ⩽ ∥∇L ( β ) ∥ ∞ +8 λ max √ 4 ∥ α ∥ 4 ∞ + B 2 . Hence,

We finally prove the bound on ∥ β k ∥ ∞ independent of α for a uniform initialisation α = α 1 , using the monotonic property of L .

Proposition 13. Assume that b = n (full batch setting). There exists some B = O ( ∥ β ⋆ ∥ 1 ) such that for stepsizes γ k ⩽ c BL , we have ∥ β k ∥ ∞ ⩽ B for all k ⩾ 0 .

Proof. In this proof, we first let B be a bound on the iterates. Tuning stepsizes using this bound, we prove that the iterates are bounded by a some B ′ = O ( ∥ β ⋆ ∥ 1 ) . 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 ∇ h 0 ( β 0 ) = 0 ,

leading to, by convexity of h k :

We then write, using ∇L ( β ) = H ( β -β ⋆ ) for H = XX ⊤ , that -∑ ℓ<k ⟨ γ ℓ ∇L ( β ℓ ) , β k -β ⋆ ⟩ = -∑ ℓ<k γ ℓ ⟨ X ⊤ ( ¯ β k -β ⋆ ) , X ⊤ ( β k -β ⋆ ) ⟩ ⩽ ∑ ℓ<k γ ℓ √ L ( ¯ β k ) L ( β k ) , leading to:

where the last inequality holds provided that γ k ⩽ 1 CLB . Thus,

Then, ⟨ ϕ k , β k -β ⋆ ⟩ ⩽ ∥ ϕ k ∥ 1 ∥ β k -β ⋆ ∥ ∞ and ∥ ϕ k ∥ 1 ⩽ Cλ max ∑ k<K γ 2 k L ( β k ) ⩽ Cλ max γ max h 0 ( β ⋆ ) . Then, using

we have:

where R = O ( ∥ β ⋆ ∥ 1 ) is independent of α . Hence, since B 0 = sup k< ∞ ∥ β k ∥ ∞ < ∞ , we have:

provided that γ max ⩽ 1 / (2 Cλ max h 0 ( β ⋆ ) / ln(1 /α 2 )) (note that h 0 ( β ⋆ ) / ln(1 /α 2 ) is independent of α 2 ).

Hence, if for all k we have γ k ⩽ 1 C ′ LB where B bounds all ∥ β k ∥ ∞ , we have ∥ β k ∥ ∞ ⩽ 2 R for all k , where R = O ( ∥ β ⋆ ∥ 1 ) is independent of α and stepsizes γ k .

For γ ⩾ 0 a constant stepsize, let

Let K > 0 be fixed, and

which is a continuous function of γ . For r > 0 , let γ r = 1 C ′ Lr .

An important feature to notice is that if γ < γ r and r bounds all ∥ β k ∥ ∞ , k ⩽ K , then φ ( γ ) ⩽ R , as shown above. We will show that we have ¯ γ ⩾ γ 2 R . Reasoning by contradiction, if ¯ γ < γ 2 R : by continuity of φ , we have φ (¯ γ ) ⩽ R and thus, there exists some small 0 < ε < γ 2 R -¯ γ such that for all γ ∈ [¯ γ, ¯ γ + ε ] , we have φ (¯ γ ) ⩽ 2 R .

However, such γ 's verify both φ ( γ ) ⩽ 2 R (since γ ∈ [¯ γ, ¯ γ + ε ] and by definition of ε ) and γ ⩽ γ 2 R (by definition of ε ), and hence φ ( γ ) ⩽ R . This contradicts the infimum of ¯ γ , and hence ¯ γ ⩾ γ 2 R . Thus, for γ ⩽ γ 2 R = 1 2 C ′ LR , we have ∥ β k ∥ ∞ ⩽ R .

F Proof of Theorem 1 and 2, and of Proposition 1

F.1 Proof of Theorem 1 and 2

We are now equipped to prove Theorem 1 and Theorem 2, condensed in the following Theorem.

Theorem 3. Let ( u k , v k ) k ⩾ 0 follow the mini-batch SGD recursion (3) initialised at u 0 = √ 2 α ∈ R d > 0 and v 0 = 0 , and let ( β k ) k ⩾ 0 = ( u k ⊙ v k ) k ⩾ 0 . There exists and explicit B > 0 and a numerical constant c > 0 such that:

For stepsizes satisfying γ k ⩽ c LB , the iterates satisfy ∥ γ k ∇L B k ( β k ) ∥ ∞ ⩽ 1 and ∥ β k ∥ ∞ ⩽ B for all k ;

For stepsizes satisfying γ k ⩽ c LB , ( β k ) k ⩾ 0 converges almost surely to some β ⋆ ∞ ∈ S ,

If ( β k ) k and the neurons ( u k , v k ) k respectively converge to a model β ⋆ ∞ and neurons ( u ∞ , v ∞ ) satisfying β ⋆ ∞ ∈ S (and β ⋆ ∞ = u ∞ ⊙ v ∞ ), then for almost all stepsizes (with respect to the Lebesgue measure), the limit β ⋆ ∞ satisfies:

for α ∞ ∈ R d > 0 and ˜ β 0 ∈ R d satisfying

where q ( x ) = -1 2 ln((1 -x 2 ) 2 ) ⩾ 0 for | x | ⩽ √ 2 , and ˜ β 0 is a perturbation term equal to:

where, q ± ( x ) = ∓ 2 x -ln((1 ∓ x ) 2 ) , and α 2 ± = α 2 ⊙ exp( -∑ ∞ k =0 q ± ( γ k ∇L B k ( β k ))) .

Proof. Point 1. The first point of the Theorem is a direct consequence of Corollary 1 and the bounds proved in appendix E.3.

Point 2. Then, for stepsizes γ k ⩽ c LB , using Proposition 8 for any interpolator β ⋆ ∈ S :

Hence, summing:

so that the series converges.

Under our stepsize rule, ∥ γ k ∇L B k ( β k ) ∥ ∞ ⩽ 1 2 , leading to ∥ q ( γ k ∇L B k ( β k ) ∥ ∞ ⩽ 3 ∥ γ k ∇L B k ( β k ) ∥ 2 ∞ by Lemma 5. Using ∥∇L B k ( β k ) ∥ 2 ⩽ 2 L 2 L B k ( β k ) , we have that ln( α ± ,k ) , ln( α k ) all converge.

We now show that ∑ k γ k L ( β k ) < ∞ . We have:

where M k = ∑ ℓ<k γ k ( L ( β k ) -L B k ( β k )) . We have that ( M k ) is a martingale with respect to the filtration ( F k ) defined as F k = σ ( β ℓ , ℓ ⩽ k ) . Using our upper-bound on ∑ ℓ<k γ k L B k ( β k ) , we have:

and hence ( M k ) is a lower bounded martingale. Using Doob's first martingale convergence theorem (a lower bounded super-martingale converges almost surely, Doob [17]), ( M k ) converges almost surely. Consequently, since ∑ ℓ<k γ k L ( β k ) = ∑ ℓ<k γ k L B k ( β k )+ M k , we have that ∑ ℓ<k γ k L ( β k ) converges almost surely (the first term is upper bounded, the second converges almost surely).

We now prove the convergence of ( β k ) . Since it is a bounded sequence, let β σ ( k ) be a convergent sub-sequence and let β ⋆ ∞ denote its limit: β σ ( k ) → β ⋆ ∞ .

Almost surely, ∑ k γ k L ( β k ) < ∞ and so γ k L ( β k ) → 0 , leading to L ( β k ) → 0 since the stepsizes are lower bounded, so that L ( β σ ( k ) ) → 0 , and hence L ( β ⋆ ∞ ) = 0 : this means that β ⋆ ∞ is an interpolator.

Since the quantities ( α k ) k , ( α ± ,k ) k and ( ϕ k ) k converge almost surely to α ∞ , α ± and ϕ ∞ , we get that the potentials h k uniformly converge to h ∞ = ψ α ∞ -⟨ ϕ ∞ , ·⟩ on all compact sets. Now notice that we can decompose ∇ h ∞ ( β ⋆ ∞ ) as:

The first two terms converge to 0: the first is a direct consequence of the convergence of the extracted subsequence, the second is a consequence of the uniform convergence of h σ ( k ) to h ∞ on compact sets. Finally the last term is always in Span( x 1 , . . . , x n ) due to Proposition 5, leading to ∇ h ∞ ( β ⋆ ∞ ) ∈ Span( x 1 , . . . , x n ) . Consequently, ∇ h ∞ ( β ⋆ ∞ ) ∈ Span( x 1 , . . . , x n ) . Notice that from the definition of h ∞ , we have that ∇ h ∞ ( β ⋆ ∞ ) = ∇ ψ α ∞ ( β ⋆ ∞ ) -ϕ ∞ . Now since ϕ ∞ = 1 2 arcsinh( α 2 + -α 2 -2 α 2 ∞ ) , one can notice that ˜ β 0 is precisely defined such that ∇ ψ α ∞ ( ˜ β 0 ) = ϕ ∞ . Therefore ∇ ψ α ∞ ( β ⋆ ∞ ) -∇ ψ α ∞ ( ˜ β 0 ) ∈ Span( x 1 , . . . , x n ) . This condition along with the fact that β ⋆ ∞ is an interpolator are exactly the optimality conditions of the convex minimisation problem:

Therefore β ⋆ ∞ must be equal to the unique minimiser of this problem. Since this is true for any sub-sequence we get that β k converges almost surely to:

Point 3. From what we just proved, note that it is sufficient to prove that α k , α ± ,k , ϕ k converge to limits α ∞ , α ± , ∞ , ϕ ∞ satisfying α ∞ , α ± , ∞ ∈ R d > 0 (with positive and non-null coordinates) and ϕ ∞ ∈ R d . Indeed, if this holds and since we assume that the iterates converge to some interpolator, we proved just above that this interpolator is uniquely defined through the desired implicit regularization problem. We thus prove the convergence of α k , α ± ,k , ϕ k .

Note that the convergence of u k , v k is equivalent to the convergence of w ± ,k in the w 2 + -w 2 -parameterisation used in our proofs, that we use there too. We have:

We now assume that stepsizes are such that for all ℓ ⩾ 0 and i ∈ [ d ] , stepsizes are such that we have | γ ℓ ∇ i L B ℓ ( β ℓ ) | = 1 : this is true for all stepsizes except a countable number of stepsizes, and so this is true for almost all stepsizes. Since we assume that the iterates β k converge to some interpolator, this leads to γ ℓ ∇L B ℓ ( β ℓ ) → 0 if we assume that stepsizes do not diverge.

so that

̸

Taking the limit, we have

This limit is in ( {-∞} ∪ R ) d (since w ± , ∞ ∈ R d ), and a coordinate of the limit is equal to -∞ if and only if the sum on the RHS diverges to -∞ (note that from our assumption just above, no term of the sum can be equal to -∞ ).

We have ln((1 ∓ γ ℓ ∇L B ℓ ( β ℓ )) 2 ) ∼ ∓ 2 γ ℓ ∇L B ℓ ( β ℓ ) as ℓ →∞ , so that if for some coordinate i we have ∑ ℓ γ ℓ ∇ i L B ℓ ( β ℓ ) = ∓∞ , then the coordinate i of the limit satisfies ln( w 2 i, ± , ∞ ) = + ∞ , which is impossible. Hence, the sum ∑ ℓ γ ℓ ∇L B ℓ ( β ℓ ) is in R d (and is thus converging); consequently, ∑ ℓ γ 2 ℓ ∇L B ℓ ( β ℓ ) 2 converges and thus ∑ ℓ q ( γ ℓ ∇L B ℓ ( β ℓ )) and ∑ ℓ q ± ( γ ℓ ∇L B ℓ ( β ℓ )) all converge: the sequences α k , α ± ,k thus converge to limits in R d > 0 , and ϕ k converges, concluding our proof.

F.2 Proof of Proposition 1

We begin with the following Lemma, that explicits the curvature of D h around the set of interpolators.

Lemma 3. For all k ⩾ 0 , if L ( β k ) ⩽ 1 2 λ max ( α 2 λ + min ) 2 , we have ∥ ∥ β k -β ⋆ α k ∥ ∥ 2 ⩽ 2 B ( α 2 λ + min ) -1 L ( β k ) .

Proof. Recall that the sequence z k = ∇ h k ( β k ) satisfies z 0 = 0 and z k +1 = z k -γ k L ( β k ) , so that we have that z k ∈ V = Im( XX ⊤ ) for all k ⩾ 0 . Then, let β α k be the unique minimizer of h k over S the space of interpolators: β α k is exactly characterized by X ⊤ β α k = Y and ∇ h k ( β α k ) ∈ V . We define z α k ∈ V as z α k = ∇ h k ( β α k ) .

Now, fix z α = z α k and h = h k , and let us define ψ : z ∈ V → D h ∗ ( z , z α ) and ϕ : z ∈ V → L ( ∇ h ∗ ( z )) . We next show that for all z ∈ V , there exists µ z such that ∇ 2 ϕ ( z ) ⩾ µ z ∇ 2 ψ ( z ) , and that µ z ⩾ µ for z in an open convex set of V around z α , for some µ > 0 . For A ∈ R d × d an operator/matrix on R d , let us denote A V its restriction/co-restriction to V .

First, for z ∈ V , we have ∇ 2 ψ ( z ) = ∇ 2 ( h ∗ ( z ) -h ∗ ( z ) -⟨∇ h ∗ ( z α ) , z -z α ⟩ )( z ) = ∇ 2 h ∗ ( z ) V . Then, ∇ ϕ ( z ) = ∇ 2 h ∗ ( z ) ∇L ( ∇ h ∗ ( z )) , so that ∇ 2 ϕ ( z ) = ( ∇ 2 h ∗ ( z ) ∇ 2 L ( ∇ h ∗ ( z )) ∇ 2 h ∗ ( z ) ) V + ∇ 3 h ∗ ( z )( ∇L ( ∇ h ∗ ( z )) , · , · ) V .

For the other term of ∇ 2 ϕ , namely ∇ 3 h ∗ ( z )( ∇L ( ∇ h ∗ ( z )) , · , · ) V , we compute ∇ 3 ijk h ∗ ( z ) = 1 i = j = k 2 α 2 i,k sinh( z i ) , leading to: ∇ 3 h ∗ ( z )( ∇L ( ∇ h ∗ ( z )) , · , · ) V = diag(2 α 2 sinh( z ) ⊙ ( XX ⊤ (2 α 2 sinh( z ) -β α ))) V . Thus, writing β z = 2 α 2 i,k sinh( z ) = ∇ h ∗ ( z ) the primal surrogate of

Since h is 1 / (2 α 2 ) smooth (on R d and thus on V ), h ∗ is 2 α 2 strongly convex (on V and on R d ). Using V = Im( XX ⊤ ) and ∇ 2 L ≡ XX ⊤ , we have ( ∇ 2 h ∗ ( z ) ∇ 2 L ( ∇ h ∗ ( z )) ∇ 2 h ∗ ( z ) ) V = ∇ 2 h ∗ ( z ) V ∇ 2 L ( ∇ h ∗ ( z )) V ∇ 2 h ∗ ( z ) V , and thus ( ∇ 2 h ∗ ( z ) ∇ 2 L ( ∇ h ∗ ( z )) ∇ 2 h ∗ ( z ) ) V ⪰ 2 α 2 λ + min ∇ 2 h ∗ ( z ) V .

z , we have:

Wrapping things together,

Let Z = { z ∈ V : ∥ ∥ XX ⊤ ( β z -β α k ) ∥ ∥ ∞ < α 2 λ + min } that satisfies { β ∈ V : L ( β z ) < 1 2 λ max ( α 2 λ + min ) 2 } ⊂ Z . Z is an open convex set of V containing z α . On Z , ∇ 2 ϕ ⪰ α 2 λ + min ∇ 2 ψ , and ψ ( z α ) = ϕ ( z α ) = 0 , so that for all z ∈ Z , we have ϕ ( z ) ⩾ α 2 λ + min ψ ( z ) . Hence, for all z ∈ Z , we have D h k ( β α k , β z ) ⩽ D h ⋆ ( z , z α ) ⩽ ( α 2 λ + min ) -1 L ( β z ) , and using the fact that D h k is 1 4 B strongly convex, we obtain, for β z = β k (since z k ∈ V ): if L ( β k ) ⩽ 1 2 λ max ( α 2 λ + min ) 2 , we have ∥ β α k -β k ∥ 2 2 ⩽ ( α 2 λ + min ) -1 L ( β k ) .

Proposition 14. As assume L is L r -relatively smooth with respect to all the h k 's. Then for all β we have the following inequality.

Proof. For any β, β k , β k +1 , the following holds (three points identity for time varying potentials, Proposition 9):

Rearranging and plugging in our mirror update we obtain that for all β :

From the convexity of L and its L r -relative smoothness we also have that:

Finally:

γ k ( L ( β k +1 ) -L ( β )) ⩽ D h k ( β, β k ) -D h k +1 ( β, β k +1 ) -(1 -γ k L r ) D h k ( β k +1 , β k ) + ( h k +1 -h k )( β ) . In particular, for β = β ∗ :

Note that in our setting, for any β , k ↦→ h k ( β ) is increasing . We can therefore write that:

and in β = β k :

Proof of Proposition 1. We apply Proposition 14 for β = β k , with L r = 4 BL (using Lemma 6) and replacing L by L B k , to obtain:

and thus, taking the mean wrt B k ,

First, as in the proof of Proposition 10, using the fact that h k is ln(1 /α k ) smooth,

and thus

Now, we apply Lemma 7 assuming that ∥ β ⋆ ∥ ∞ , ∥ β k +1 ∥ ∞ ⩽ B (which is satisfied since we are under the assumption of Theorem 2):

Using |∇ h k ( β ) -∇ h k +1 ( β ) | ⩽ 2 δ k where δ k = q ( γ k ∇L B k ( β k )) as in Proposition 10, we have:

Wrapping everything together,

Thus, for γ k ⩽ c ′ LB ln(1 / (min i α k,i )) , we have the first part of Proposition 1.

Using Lemma 3, we then have:

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 are not directly directed to the proof of Theorem 3.

G.1 Proof of 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 Λ b , λ b > 0 5 be the largest and smallest values, respectively, such that λ b H ⪯ E B [ H 2 B ] ⪯ Λ b H . For any stepsize γ > 0 satisfying γ ⩽ c BL (as in Theorem 2), initialisation α 1 and batch size b ∈ [ n ] , the magnitude of the gain satisfies:

where the expectation is over a uniform and independent sampling of the batches ( B k ) k ⩾ 0 .

Proof. From Lemma 5, for all -1 / 2 ⩽ x ⩽ 1 / 2 , it holds that x 2 ⩽ q ( x ) ⩽ 2 x 2 . We have, using ∥ γ k ∇L B k ( β k ) ∥ ∞ ⩽ 1 / 2 (which holds under the stepsize assumption):

since E ∥∇L B ℓ ( β ℓ ) ∥ 2 2 ⩽ 2Λ b L B ℓ ( β ℓ ) . For the left handside we use q ( x ) ⩾ x 2 for | x | ⩽ 1 / 2 and E ∥∇L B k ( β ℓ ) ∥ 2 2 ⩾ 2 λ b L B k ( β ℓ ) . Finally, since B ℓ independent freom β ℓ , we have E L B ℓ ( β ℓ ) = E L ( β ℓ ) .

Proposition 15. For stepsizes γ k ≡ γ ⩽ c LB (as in Theorem 2), we have:

Proof. We first lower bound ∑ k< ∞ γ 2 k L B k ( β k ) . We have the following equality, that holds for any k :

leading to, by summing for k ∈ N :

First, since h k → h ∞ , β k → β ∞ , we have lim k →∞ D h k ( β ⋆ , β k ) = 0 . Then, D h k +1 ( β k , β k +1 ) ⩾ 0 . Finally, ( h k h k +1 ) ( β k ) ( h k h k +1 ) ( β ⋆ ) ⩽ 16 BL 2 γ 2 k ( β k ) . Hence :

| ---| L B

and thus ∑ k< ∞ γ L B k ( β k ) ⩾ D h 0 ( β ⋆ , β 0 ) / 4 for γ ⩽ c/ ( BL ) (with c ⩾ 16 ). This gives the RHS inequality. The LHS is a direct consequence of bounds proved in previous subsections.

Hence, we have that

Noting that D h 0 ( β ⋆ , β 0 ) = h 0 ( β ⋆ ) = Θ ( ln(1 /α ) ∥ β ⋆ ∥ 1 ) concludes the proof.

G.2 ˜ β 0 is negligible

In the following proposition we show that ˜ β 0 is close to 0 and therefore one should think of the implicit regularization problem as β ⋆ ∞ = argmin β ⋆ ∈ S ψ α ∞ ( β ⋆ )

Proposition 16. Under the assumptions of Theorem 2,

where the inequality must be understood coordinate-wise.

Proof.

where the inequality is because q + ( γ k ∇L ( β k )) ⩾ 0 , q -( γ k ∇L ( β k )) ⩾ 0 for all k .

G.3 Impact of stochasticity and linear scaling rule

Proposition 17. With probability 1 -2 ne -d/ 16 -3 /n 2 over the x i ∼ iid N (0 , σ 2 I d ) , c 1 dσ 2 b (1 + o (1)) ⩽ λ b ⩽ Λ b ⩽ c 2 dσ 2 b (1 + o (1)) , so that under these assumptions,

Proof. The bound on λ b , Λ b is a direct consequence of the concentration bound provided in Lemma 13.

G.4 (Stochastic) gradients at the initialisation

To understand the behaviour and the effects of the stochasticity and the stepsize on the shape of Gain γ , we analyse a noiseless sparse recovery problem under the following standard assumption 2 [10] and as common in the sparse recovery literature, we make the following assumption 3 on the inputs.

Assumption 2. There exists an s -sparse ground truth vector β ⋆ sparse where s verifies n = Ω( s ln( d )) , such that y i = ⟨ β ⋆ sparse , x i ⟩ for all i ∈ [ n ] .

Assumption 3. There exists δ, c 1 , c 2 > 0 such that for all s -sparse vectors β , there exists ε ∈ R d such that ( X ⊤ X ) β = β + ε where ∥ ε ∥ ∞ ⩽ δ ∥ β ∥ 2 and c 1 ∥ β ∥ 2 2 1 ⩽ 1 n ∑ i x 2 i ⟨ x i , β ⟩ 2 ⩽ c 2 ∥ β ∥ 2 2 1 .

The first part of Assumption 3 closely resembles the classical restricted isometry property (RIP) and is relevant for GD while the second part is relevant for SGD. Such an assumption is not restrictive and holds with high probability for Gaussian inputs N (0 , σ 2 I d ) (see Lemma 10 in Appendix).

Based on the claim above, we analyse the shape of the (stochastic) gradient at initialisation. For GD and SGD, it respectly writes, where g 0 = ∇L i 0 ( β 0 ) 2 , i 0 ∼ Unif([n]) :

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 ε verifying ∥ ε ∥ ∞ << ∥ ∥ β ⋆ sparse ∥ ∥ 2 ∞ :

Proof of Proposition 18. Under Assumption 3, we have using:

We have ∥ ∥ ε 2 +2 εβ ⋆ sparse ∥ ∥ ∞ ⩽ ∥ ε ∥ 2 ∞ + 2 ∥ ε ∥ ∞ ∥ ∥ β ⋆ sparse ∥ ∥ ∞ , and we conclude by using ∥ ε ∥ ∞ ⩽ δ ∥ ∥ β ⋆ sparse ∥ ∥ 2 .

Then,

and we conclude using Assumption 3.

Proof of Proposition 4. The proof proceeds as that of Proposition 18.

G.5 Convergence of α ∞ and ˜ β 0 for γ → 0

Proposition 19. Let ˜ β 0 ( γ ) , α ∞ ( γ ) be as defined in Theorem 1, for constant stepsizes γ k ≡ γ . We have:

when γ → 0 .

Proof. We have, as proved previoulsy, that

for γ ⩽ c BL . Thus, ∑ k γ 2 ∇L B k ( β k ) 2 → 0 as γ → 0 (note that β k implicitly depends on γ , so that this result is not immediate).

Then, for γ ⩽ c LB ,

which tends to 0 as γ → 0 . Similarly, ∥ ∥ ln( α 2 + , ∞ /α 2 ) ∥ ∥ 1 → 0 and ∥ ∥ ln( α 2 -, ∞ /α 2 ) ∥ ∥ 1 → 0 as γ → 0 , leading to ˜ β 0 ( γ ) → 0 as γ → 0 .

H Technical lemmas

In this section we present a few technical lemmas, used and referred to throughout the proof of ?? . Lemma 4. Let α , α > 0 and x R , and β = α 2 e x α 2 e -x . We have:

+ -∈ + --

Proof. First,

hence the result by taking the arcsinh of both sides. Note also that we have ln( α + /α -) = arcsinh( α 2 + -α 2 -2 α + α -) .

Lemma 5. If | x | ⩽ 1 / 2 then x 2 ⩽ q ( x ) ⩽ 2 x 2

Lemma 6. On the ℓ ∞ ball of radius B , the quadratic loss function β ↦→L ( β ) is 4 λ max max( B,α 2 ) -relatively smooth w.r.t all the h k 's.

Proof. We have:

since α k ⩽ α component-wise. Thus, ∇ 2 h k ( β ) ⪰ 1 2 min ( min 1 ⩽ i ⩽ d 1 2 | β i | , 1 2 α 2 ) I d = 1 max(4 ∥ β ∥ ∞ , 4 α 2 ) I d , and h k is 1 max(4 B, 4 α 2 ) -strongly convex on the ℓ ∞ norm of radius B . Since L is λ max -smooth over R d , we have our result.

Lemma 7. For k ⩾ 0 and for all β ∈ R d :

Proof. We have α 2 + ,k +1 = α 2 + ,k e -δ + ,k and α 2 -,k +1 = α 2 -,k e -δ -,k , for δ + ,k = ˜ q ( γ k ∇L B k ( β k )) and δ -,k = ˜ q ( -γ k ∇L B k ( β k )) . And α k +1 = α k exp( -δ k ) where δ k := δ + ,k + δ -,k = q ( γ k ∇L B k ( β k )) .

To prove the result we will use that for β ∈ R d , we have | ( h k +1 -h k )( β ) | ⩽ ∑ d i =1 ∫ | β i | 0 |∇ i h k +1 ( x ) -∇ i h k ( x ) | d x .

First, using that | arcsinh( a ) -arcsinh( b ) | ⩽ | ln( a/b ) | for ab > 0 . We have that

since δ k ⩾ 0 due to our stepsize condition.

We now prove that | ϕ k +1 -ϕ k | ⩽ | δ + ,k -δ -,k | 2 . We have ϕ k = arcsinh ( α 2 + ,k -α 2 -,k 2 α + ,k α -,k ) and hence,

Then, assuming that α + ,k,i ⩾ α -,k,i , we have:

We thus have α 2 + ,k +1 ,i -α 2 -,k +1 ,i 2 α + ,k +1 ,i α -,k +1 ,i ∈ [ e -| δ + ,k,i -δ -,k,i | 2 , e | δ + ,k,i -δ -,k,i | 2 ] × α 2 + ,k,i -α 2 -,k,i 2 α + ,k,i α -,k,i , and this holds similarly if α + ,k,i ⩽ α ,k,i . Then, using arcsinh( a ) arcsinh( b ) ⩽ ln( a/b ) we obtain that:

Wrapping things up, we have:

This leads to the following bound:

Recall that δ k = q ( γ k ∇L B k ( β k ) , hence from Lemma 5 if γ k ∥∇L B k ( β k ) ∥ ∞ ⩽ 1 / 2 , we get that ∥ δ k ∥ 1 ⩽ 2 γ 2 k ∥∇L B k ( β k ) ∥ 2 2 ⩽ 4 L 2 γ 2 k L B k ( β k ) .

Putting things together we obtain that

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 a 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 x 1 , . . . , x n ∈ R d be i.i.d. random variables of law N (0 , I d ) and H = 1 n ∑ n i =1 x i x ⊤ i . Then, denoting by C the set of all s -sparse vector β ∈ R d satisfying ∥ β ∥ 2 ⩽ 1 , there exist C 4 , C 5 > 0 such that for any ε > 0 , if n ⩾ C 4 s ln( d ) ε -2 ,

Lemma 9. Let x 1 , . . . , x n ∈ R d be i.i.d. random variables of law N ( µ, σ 2 I d ) and H = 1 n ∑ n i =1 x i x ⊤ i . Then, denoting by C the set of all s -sparse vector β ∈ R d satisfying ∥ β ∥ 2 ⩽ 1 , there exist C 4 , C 5 > 0 such that for any ε > 0 , if n ⩾ C 4 s ln( d ) ε -2 ,

We then provide two lemmas that estimate the mean Hessian of SGD.

Lemma 10. Let x 1 , . . . , x n be i.i.d. random variables of law N (0 , I d ) . Then, there exist c 1 , c 2 > 0 such that with probability 1 -1 d 2 and if n = Ω( s 5 / 4 ln( d )) , we have for all s -sparse vectors β :

where the inequality is meant component-wise.

Lemma 11. Let x 1 , . . . , x n be i.i.d. random variables of law N ( µ, σ 2 I d ) . Then, there exist c 0 , c 1 , c 2 > 0 such that with probability 1 -c 0 d 2 -1 nd and if n = Ω( s 5 / 4 ln( d )) and µ ⩾ 4 σ √ ln( d ) 1 , we have for all s -sparse vectors β :

where the inequality is meant component-wise.

Finally, next two lemmas are used to estimate λ b , Λ b in our paper.

Lemma 12. Let x 1 , . . . , x n ∈ R d be i.i.d. random variables of law N ( µ 1 , σ 2 I d ) . Let H = 1 n ∑ n i =1 x i x ⊤ i and ˜ H = 1 n ∑ n i =1 ∥ x i ∥ 2 x i x ⊤ i . There exist numerical constants C 2 , C 3 > 0 such that

Lemma 13. Let x 1 , . . . , x n ∈ R d be i.i.d. random variables of law N ( µ 1 , σ 2 I d ) for some µ ∈ R . Let H = 1 n ∑ n i =1 x i x ⊤ i and for 1 ⩽ b ⩽ n let ˜ H b = E B [ ( 1 b ∑ i ∈B x i x ⊤ i ) 2 ] where B ⊂ [ n ] is sampled uniformly at random in {B ⊂ [ n ] s.t. |B| = b } . With probability 1 -2 ne -d/ 16 -3 /n 2 , we have, for some numerical constants c 1 , c 2 , c 3 , C > 0 :

Proof of Lemma 8. For j ∈ [ d ] , we have:

We thus notice that E [ Hβ ] = β , and

̸

̸

̸

̸

where z i = x ij ∑ j ′ = j x ij ′ β j ′ , and ∑ j ′ = j x ij ′ β j ′ ∼ N (0 , ∥ β ∥ 2 -β 2 j ) and ∥ β ∥ 2 -β 2 j ⩽ 1 . Hence, z j + x 2 ij -1 is a centered subexponential random variables (with a subexponential parameter of order 1). Thus, for t ⩽ 1 :

Hence, using an ε -net of C = { β ∈ R d : ∥ β ∥ 2 ⩽ 1 , ∥ β ∥ 0 } (of cardinality less than d s × ( C/ε ) s , and for ε of order 1), we have, using the classical ε -net trick explained in [Chapt. 9, [58] or [App. C, Even and Massoulie [19]]:

Consequently, for t = ε and if n ⩾ C 4 s ln( d ) /ε 2 , we have:

Proof of Lemma 9. We write x i = σz i + µ where z i ∼ N (0 , I d ) . We have:

The first term is deterministic and is to be kept. The second one is of order σ 2 β whp using Lemma 8. Then, 1 n ∑ n i =1 z i ∼ N (0 , I d /n ) , so that and

Hence,

Thus, with probability 1 -Ce -nε 2 and under the assumptions of Lemma 8, we have ∥ ∥ X ⊤ Xβ -µ ⟨ µ, β ⟩ -σ 2 β ∥ ∥ ∞ ⩽ ε

Proof of Lemma 10. To ease notations, we assume that σ = 1 . Weremind (O'Donnell [46], Chapter 9 and Tao [54]) that for i.i.d. real random variables a 1 , . . . , a n that satisfy a tail inequality of the form

for p < 1 , then for all ε > 0 there exists C ′ , c ′ such that for all t ,

We now expand 1 n ∑ n i =1 x 2 i ⟨ x i , β ⟩ 2 :

Thus, for j ∈ [ d ] ,

̸

We notice that for all indices, all x 2 ij x ik x iℓ and x 2 ij x 2 ik satisfy the tail inequality Eq. (27) for C = 8 , c = 1 / 2 and p = 1 / 2 , so that for ε = 1 / 4 :

For j = k , we have E [ x 2 ij x 2 ik ] = 1 while for j = k , we have E [ x 2 ij x 2 ik ] = E [ x 4 ij ] = 3 . Hence,

̸

̸

Thus, with probability 1 -C ′ d 2 e -c ′ nt 1 / 4 , for all j ∈ [ d ] ,

Using the classical technique of Baraniuk et al. [4], to make a union bound on all s -sparse vectors, we consider an ε -net of the set of s -sparse vectors of ℓ 2 -norm smaller than 1. This ε -net is of cardinality less than ( C 0 /ε ) s d s , and we only need to take ε of order 1 to obtain the result for all s -sparse vectors. This leads to:

This probability is equal to C ′ /d 2 for t = ( ( s +4)ln( d )+ c 1 s c ′ n ) 4 . We conclude that with probability 1 -C ′ /d 2 , all s -sparse vectors β satisfy:

and the RHS is smaller than ∥ β ∥ 2 / 2 for n Ω( s ln( d )) .

Proof of Lemma 11. We write x i = µ + σz i where x i ∼ N (0 , 1) . We have:

̸

Then, ⟨ x i , β ⟩ ∼ N ( ⟨ µ, β ⟩ , σ 2 ∥ β ∥ 2 2 ) . For now, we assume that ∥ β ∥ 2 = 1 . We have P ( |⟨ x i , β ⟩ 2 -⟨ µ, β ⟩ 2 -σ 2 ∥ β ∥ 2 2 | ⩾ t ) ⩽ Ce -ct/σ 2 , and for t ⩽ 1 , using concentration of subexponential random variables [58]:

and using the ε -net trick of Baraniuk et al. [4],

Proof of Lemma 12. First, we write x i = µ 1 + σz i , where z i ∼ N (0 , I ) , leading to:

We use concentration of χ 2 d random variables around d :

P ( χ 2 d > d +2 t +2 √ dt ) ⩾ t ) ⩽ e -t and P ( χ 2 d > d -2 √ dt ) ⩽ t ) ⩽ e -t , so that for all i ∈ [ n ] ,

Thus,

Taking t = d/ 16 ,

∀

∈

[

n

]

,

∥

∥

∈

-

-

Then, for all i , ⟨ 1 , z i ⟩ is of law N (0 , d ) , so that P ( |⟨ 1 , z i ⟩| ⩾ t ) ⩽ 2 e -t 2 / (2 d ) and

Taking t = √ 2 d 3 / 4 ,

Thus, with probability 1 -2 n ( e -d/ 16 + e - √ d , we have ∀ i ∈ [ n ] , |⟨ 1 , z i ⟩| ⩾ d 3 / 4 and ∥ z i ∥ 2 2 ∈ [ d 2 , 13 d/ 8] , so that

leading to the desired result.

Proof of Lemma 13. We have:

z

i

2

2

[

d

2

√

dt, d

+2

t

+2

√

dt

])

⩾

1

P

(

i

̸

̸

2

ne

-

t

.

Thus,

̸

Then, taking t = d -1 / 2 4 ln( n ) /c , we have:

̸

Going back to our second term,

̸

where we used x i x ⊤ j + x j x ⊤ i ⪯ x i x ⊤ i + x j x ⊤ j . Thus,

̸

̸

Similarly, we have

̸

Hence, with probability 1 -1 /n 2 ,

̸

̸

̸

Then, since P ( i ∈ B ) = b n and P ( i, j ∈ B ) = b ( b -1) n ( n -1) for i = j , we get that:

̸

Using Lemma 12, the first term satisfies:

We now show that the second term is of smaller order. Writing x i = µ 1 + σz i where z i ∼ N (0 , I d ) , we have:

̸

̸

̸

For i = j , ⟨ x i , x j ⟩ = ∑ d k =1 x ik x jk = ∑ d k =1 a k where a k = x ik x jk satisfies E a k = 0 , E a 2 k = 1 and P ( a k ⩾ t ) ⩽ 2 P ( | x ik | ⩾ √ t ) ⩽ 4 e -t/ 2 . Hence, a k is a centered subexponential random variables. Using concentration of subexponential random variables [58], for t ⩽ 1 ,

̸

̸

Wrapping things up, with probability 1 -1 /n 2 -2 ne -d/ 16 ,

Thus, provided that 4 ln( n ) c √ d ⩽ C 2 d 2 b and d ⩾ 48 ln( n ) , we have with probability 1 -3 /n 2 :

̸

̸

Proof of Lemma 13. We have:

̸

Then, since P ( i ∈ B ) = b n and P ( i, j ∈ B ) = b ( b -1) n ( n -1) for i = j , we get that:

̸

Using Lemma 12, the first term satisfies:

We now show that the second term is of smaller order. Writing x i = µ 1 + σz i where z i ∼ N (0 , I d ) , we have:

̸

̸

̸

̸

̸

For i = j , ⟨ z i , z j ⟩ = ∑ d k =1 z ik z jk = ∑ d k =1 a k where a k = z ik z jk satisfies E a k = 0 , E a 2 k = 1 and P ( a k ⩾ t ) ⩽ 2 P ( | z ik | ⩾ √ t ) ⩽ 4 e -t/ 2 . Hence, a k is a centered subexponential random variables. Using concentration of subexponential random variables [58], for t ⩽ 1 ,

Thus,

̸

Then, taking t = d -1 / 2 4 ln( n ) /c , we have:

̸

For i ∈ [ n ] , ⟨ 1 , z i ⟩ ∼ N (0 , d ) so that P ( |⟨ 1 , z i ⟩| ⩾ t ) ⩽ 2 e -t 2 / (2 d ) , and

̸

̸

̸

̸

for t = 3 √ d ln( n ) . Hence, with probability 1 -3 /n 2 , for all i = j we have | σ 2 ⟨ z i , z j ⟩ + σµ ⟨ 1 , z i + z j ⟩| ⩽ ( σ 2 + σµ ) C ln( n ) / √ d . Now,

̸

̸

where we used x i x ⊤ j + x j x ⊤ i ⪯ x i x ⊤ i + x j x ⊤ j . Thus,

̸

Similarly, we have

̸

Hence, with probability 1 -3 /n 2 ,

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 ( b -1) bn ( n -1) µ 2 d ∑ i = j x i x ⊤ j . Denoting ¯ x = 1 n ∑ i x i , we have 1 n 2 ∑ i = j x i x ⊤ j = 1 n 2 ∑ i,j x i x ⊤ j -1 n 2 ∑ i x i x ⊤ i , so that:

̸

̸

We note that we have H = 1 n ∑ i x i x ⊤ i = 1 n 2 ∑ i<j x i x ⊤ i + x j x ⊤ j ⪰ 1 n 2 ∑ i<j x i x ⊤ j + x j x ⊤ i = ¯ x ¯ x ⊤ using x i x ⊤ i + x j x ⊤ j ⪰ x i x ⊤ j + x j x ⊤ i . Thus, H ⪰ ¯ x ¯ x ⊤ ⪰ 0 , and:

̸

We are now able to wrap everything together. With probability 1 -2 ne -d/ 16 -3 /n 2 , we have, for some numerical constants c 1 , c 2 , c 3 , C > 0 :

̸

Peer reviews:
comments: 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.

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.

QUESTIONS:
N/A

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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}$.\n\nWEAKNESSES:\n1. 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\nQUESTIONS:\nSee above in Weakness.\n\nFor now, I would like to recommend an acceptance for its theoretical proof of an open problem.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

STRENGTHS:
The writing is very clear and easy to follow, and highlighting of the important parts helps the reading a lot.

Theorem 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?

The 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?

QUESTIONS:
Theorem 2 gives the convergence guarantee. Is it possible to obtain an explicit convergence rate? If not, what's the difficulty here?

In 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThis 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper analyzes 2-layer diagonal linear networks and explains the different characteristics of the solutions received through GF, GD, and mini-batch SGD.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nPart 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.\n\nQUESTIONS:\n1) 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)?\n\nLIMITATIONS:\nThe limitations are adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}