Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity

Introduction

Understanding the performance of neural networks is certainly one of the most thrilling challenges for the current machine learning community. From the theoretical point of view, progress has been made in several directions: we have a better functional analysis description of neural networks and we steadily understand the convergence of training algorithms as well as the role of initialisation . Yet there remain many unanswered questions. One of which is why do the currently used training algorithms converge to solutions which generalise well, and this with very little use of explicit regularisation .

To understand this phenomenon, the concept of implicit bias has emerged: if over-fitting is benign, it must be because the optimisation procedure converges towards some particular global minimum which enjoys good generalisation properties. Though no explicit regularisation is added, the algorithm is implicitly selecting a particular solution: this is referred to as the implicit bias of the training procedure. The implicit regularisation of several algorithms has been studied, the simplest and most emblematic being that of gradient descent and stochastic gradient descent in the least-squares framework: they both converge towards the global solution which has the lowest squared distance from the initialisation. For logistic regression on separable data, Soudry et al. show in the seminal paper that gradient descent selects the max-margin classifier. This type of result has then been extended to neural networks and to other frameworks. Overall, characterising the implicit bias of gradient methods has almost always come down to unveiling mirror-descent like structures which underlie the algorithms.

While mostly all of the results focus on gradient descent, it must be pointed out that this full batch algorithm is not used in practice for neural networks since it does not lead to solutions which generalise well . Instead, results on stochastic gradient descent, which is widely used and shows impressive results, are still missing or unsatisfactory. This has certainly to do with the fact that grasping the nature of the noise induced by the stochasticity of the algorithm is particularly hard: it mixes properties from the model’s architecture, the data’s distribution and the loss. In our work, by focusing on simplified neural networks, we answer to the following fundamental questions: do SGD’s and GD’s implicit bias differ? What is the role of SGD’s noise over the algorithm’s implicit bias?

The simplified neural networks which we consider are diagonal linear neural networks; despite their simplicity they have become popular since they already enable to grasp the complexity of more general networks. Indeed, they highlight important aspects of the theoretical concerns of modern machine learning: the neural tangent kernel regime, the roles of over-parametrisation, of the initialisation and of the step size. For a regression problem where we assume the existence of an interpolating solution, we study stochastic gradient descent through its continuous version, namely stochastic gradient flow (SGF). Though the continuous modelling of SGD has not yet led to many fruitful results compared to the well studied gradient flow, we believe it is because capturing the essence of the stochastic noise is particularly difficult. It has generally been done in a non realistic and over simplified manner, such as considering constant and isotropic noise. In our work, we attach peculiar attention to the adequate modelling of the noise. Tools from Itô calculus are then leveraged in order to derive exact formulas, quantitative bounds and interesting interpretations for our problem.

In Section 2, we start by introducing the setup of our problem as well as the continuous modelisation of stochastic gradient descent. Then, in Section 3, we state our main result on the implicit bias of the stochastic gradient flow. We informally formulate it here and illustrate it in Figure 1:

Stochastic gradient flow over diagonal linear networks converges with high probability to a zero-loss solution which enjoys better generalisation properties than the one obtained by gradient flow. Furthermore, the speed of convergence of the training loss controls the magnitude of the biasing effect: the slower the convergence, the better the bias.

Unlike previous works , in addition to characterising the implicit bias effect of SGF, we also prove the convergence of the iterates towards a zero-loss solution with high-probability. To accomplish this, we leverage in Section 4 the fact that the iterates follow a stochastic continuous mirror descent with a time-varying potential. We support our results experimentally and validate our model in Section 5.

2 Related work

While the links between SGD’s stochasticity and generalisation have been looked into in numerous works , no such explicit characterisation of implicit regularisation have ever been given. It has been empirically observed that SGD often outputs models which generalise better than GD . One suggested explanation is that SGD is prone to pick flatter solutions than GD and that bad generalisation solutions are correlated with sharp minima, i.e., with strong curvature, while good generalisation solutions are correlated with flat minima, i.e., with low curvature . This idea has been further investigated by adopting a random walk on random landscape modelling , by suggesting that SGD’s noise is smoothing the loss landscape, thus eliminating the sharp minima , by considering a dynamical stability perspective or by interpreting SGD as a diffusion process . Recently, label-noise has been shown to influence the implicit bias of SGD, by biasing the solution towards the origin for quadratically-parameterized models or by implicitly regularising the expected squared norm of the gradient of the model with respect to the weights . Thus, if the notion of implicit bias of GD is fairly well understood both in the cases of regression and classification, it remains unclear for SGD, and its explicit characterisation is missing.

The linear diagonal neural networks we consider have been studied in the case of gradient descent and stochastic gradient descent with label noise . In both cases the authors show that this model has the ability to implicitly bias the training procedure to help retrieve a sparse predictor. The link between gradient descent and mirror descent for this model has been initiated by and further exploited by the same author in for its sparse inducing property.

Contrary to the deterministic case, the modelling of stochastic gradient descent as a stochastic differential equation is quite recent, see . However, as highlighted by , early attempts often suffer from the drawback that they model the noise using a constant covariance matrix. On the contrary, state dependant noise has now become the legitimate manner for modelling SGD as a stochastic gradient flow and it is shown in that it can be done consistently. Yet, noise modelling still remains the principal issue as it influences largely the behaviour of the dynamics .

3 Notations

Setup and preliminaries

where by abuse of notation we use $L(w)=L(\beta_{w})$.

-layer diagonal linear network.

Stochastic Gradient Descent.

With this quadratic parametrisation, the loss now rewrites as: $L(w)=\frac{1}{4n}\sum_{i=1}^{n}\langle w_{+}^{2}-w_{-}^{2}-\beta^{*},x_{i}\rangle^{2}.$ Note that despite its simplicity, this loss is non convex and its minimisation is non trivial. The algorithm we shall consider is the well known SGD algorithm, where for a step size $\gamma>0$:

It is convenient to rewrite this recursion as

2 Stochastic gradient flow

Continuous time modelling of sequential processes offer a large set of tools, such as derivation, which come in helpful to understand the dynamics of the processes. This has led to a large part of the recent literature to consider continuous gradient flow in order and understand the behaviour of gradient descent on complicated architectures such as neural nets. However, the continuous time modelling of stochastic gradient descent is more challenging: it requires to add on top of the gradient flow a diffusion term whose covariance matches the one of SGD. Hence, it is fundamental to understand its structure and scale.

As seen in equation (2), evaluated at $w_{\pm}$, the stochastic noise $\gamma\mathop{\rm diag}(w_{\pm})X^{\top}\xi_{i_{t}}(w)$ has two main characteristics which we want to preserve:

Stochastic differentiable equation modelling.

Guided by the previous considerations, we study the following stochastic gradient flow:

The implicit bias of the stochastic gradient flow

Main result.

In the main theorem we show that, for an initialisation scale $\alpha$, the stochasticity of SGF biases the flow towards solutions which still minimise the hyperbolic entropy. However, what is remarkable is that it does so with an effective parameter $\alpha_{\infty}$ which is strictly smaller than $\alpha$. The recovered solution therefore minimises an optimisation problem which has better sparsity inducing properties than that of gradient flow.

$(\beta_{t})_{t\geq 0}$ converges towards a zero-training error solution $\beta_{\infty}^{\alpha}$

the solution $\beta^{\alpha}_{\infty}$ satisfies

The theorem is three-fold: with high probability and for an explicit choice of constant step size $\gamma$, (i) the flow $(\beta_{t})_{t\geq 0}$ converges, (ii) its limit $\beta_{\infty}^{\alpha}$ is an interpolating solution, i.e. $X\beta_{\infty}^{\alpha}=y$ , (iii) this solution minimises the hyperbolic entropy problem with a parameter that depends on the dynamics. We illustrate these results in Figure 2. Now let us comment further the theorem.

Beneficial implicit bias through effective initialisation.

Kernel regime.

Step size.

Convergence and proof sketch.

Let us put emphasis on the fact that since we deal with a non-convex problem, neither convergence nor convergence towards a global minimum are obvious. In most of similar works, convergence of the iterates is assumed . In fact, the hardest and most technical part of our result is to show the convergence of the flow with high probability: once the convergence is shown, describing the minimisation problem $\beta_{\infty}^{\alpha}$ verifies is straightforward. In the following section we give several properties which constitute the major keys of the theorem’s proof.

Links with mirror descent

We start by recalling known results on the link between implicit bias and mirror descent. We recall also convergence guarantees for mirror descent dynamics.

where $D_{\Psi}(\beta,\beta_{0})=\Psi(\beta)-\Psi(\beta_{0})-\langle\nabla\Psi(\beta_{0}),\beta-\beta_{0}\rangle$ is the Bregman divergence w.r.t. $\Psi$.

Link with our model.

Stochastic Mirror descent with a time varying potential.

To address the problem where $(w_{t})_{t}$ follows a stochastic gradient flow instead of a gradient flow, it is natural, as in the deterministic framework, to see what type of flow $(\beta_{t})_{t}$ follows. Because of the noise, we cannot hope to simply recover a classical mirror descent. However interestingly the next property shows that it follows a stochastic mirror-like descent with a geometry that depends on time.

Though it seems easy to derive the implicit minimisation problem (5) from the mirror-like structure of Eq.(7), it is necessary to ensure that the iterates converge towards an interpolator $\beta_{\infty}$. This is the purpose of the following proposition.

The convergence of the iterates is technical and requires several intermediate results. We start by considering an appropriate Bregman-type stochastic function with a time-varying potential and show that it converges with high probability. Leveraging the fact that we are able to bound the iterates $\beta_{t}$, we are able to show that the limit of the function is in fact . Owing to the fact that the function we consider also controls the distance of $\beta_{t}$ to a particular $\beta^{*}$ we finally get that the iterates converge.

Under the same setting as in Proposition 2 with initialisation $w_{0,\pm}=\alpha\mathbf{1}$, we have with probability at least $1-p$:

for some $\zeta>0$. Hence the smaller the initialisation scale $\alpha$ and the greater the benefit of SGD over GD in terms of implicit bias (see Section B.6 for more details).

Again, the proof of this proposition is technical and relies on considering appropriate Lyapunov functions which highly resemble to Bregman divergences, but which take into account the fact that the geometry changes over time. These overall decreasing Lyapunov’s enable to bound the iterates as well as lower and upper bound the integral of the loss. The stochastic integrals which naturally appear are controlled with high probability using time-uniform concentration of martingales .

Experiments

2 Validation of the SDE model

3 GD and SGD have the same implicit bias, but from different initialisations.

4 Doping the implicit bias with label noise

As in Section 2.2, we can derive its related stochastic gradient flow (see Section D.1 for more details):

Conclusion and Perspectives

In this paper, we have shown the benefit of using stochastic gradient descent over gradient descent for diagonal linear networks in terms of their implicit bias. Indeed, we prove that stochastic gradient flow acts as gradient flow but initialised at a smaller scale: this induces a sparser finale iterate. This effect is controlled by the speed of convergence of the loss. Moreover, we prove the convergence of the flow and exhibit an interesting link with mirror descent. Fully understanding this novel type of dynamics could help to grasp the implicit biasing properties of stochastic gradient descent in other frameworks. It is also natural to ask whether the integral of the loss also controls the difference of implicit regularisation for more general architectures. It would also be interesting to analyse how this property adapts to $\log$ losses known to lead to max-margin solutions in classification.

NF would like to thank Nathan Srebro for introducing him to the question of SGD’s implicit bias as well as for the stimulating discussions they had during his visit at EPFL.

References

Organisation of the Appendix.

The Appendix is structured as follows. In Section A, we give more precisions regarding the way we model stochastic gradient descent as a stochastic gradient flow. Section B.6 is the core of the Appendix as it provides the proof of the theorem in a self-contained fashion. For the sake of completeness, in Section 7 we gather the results on the link between mirror-descent and implicit bias as well as give convergence results in the deterministic case (gradient flow). In Section 5, we provide more experiments supporting our results. In Section E.2, we discuss some extensions of our results ; (E.1) regarding a more general stochastic gradient flow model and in (E.2) we extend our results to depths $p\geq 3$. Finally, Section F provides the technical material needed for the proofs of our results.

Appendix A Details on the SDE modelling

We recall that the SGD recursion writes for $t\geqslant 1$ as:

As we are interested in the stochastic differential model of the SGD recursion, let us now compute the covariance of the SGD noise. We first notice that

where $L_{i}(\beta)=\frac{1}{4}\langle\beta-\beta^{*},x_{i}\rangle^{2}$ is the individual loss of the observation $x_{i}$, such that $L(\beta)=\frac{1}{n}\sum_{i=1}^{n}L_{i}(\beta)$.

The overall SGD’s noise structure is then captured by

This leads us in considering the following SDE:

since its Euler discretisation with step size $\gamma$ is :

Appendix B Proofs of the main results

This section contains all the proofs of the main results. It is self contained as we recall each time the propositions we prove. In subsection B.1, we derive the mirror-descent-like flow which the iterates follow as in 1 of the main text. Then, we upper bound the loss integral in subsection B.2. This leads us in proving the convergence of the iterates towards an interpolator in subsection B.3. Equipped with these results we prove the main result of the paper (1) in subsection B.4. Finally, to complete the proof of 3 of the main text we derive a lower bound of the loss in subsection B.5.

For the sake of easy reading, we adopt the following notations in this section: we denote by $\bar{X}:=X/\sqrt{n}$, and $\lambda_{\max}:=\lambda_{\max}(\bar{X}^{\top}\bar{X})$.

In order to prove 1, we introduce the following lemma:

Note that this is not an explicit closed form for $\beta_{t}$ since the right hand side depends on $(\beta_{s})_{0\leq s\leq t}$.

Recall that the SDE we consider writes as:

It turns out that there is an implicit closed form solution to this SDE. Indeed deriving the Itô formula on $\ln(w_{t,\pm})$ gives the following integral expression:

Since $\beta=w_{+}^{2}-w_{-}^{2}$, we get:

For clarity we recall the statement of 1. See 1

The results immediately follows from Lemma 1. Indeed, inverting the implicit equation on $\beta_{t}$, Eq. (11), we have,

B.2 Upperbound of the integral of the loss

This section contains several technical arguments that permit us to derive the upperbound of the integral of the loss [3, right side]. Let us try to highlight the key features of this proof. First, as for classical mirror descent, we define a Lyapunov function that resembles a Bregman divergence plus a necessary control term [Eq. (12)]. Then, we fix a high-probability event on which we have a control of the Brownian diffusion term [Eq. (13)]. This gives an equation involving a weighted integral of the loss. After lower bounding this weight to access directly the loss integral [Lemma 4], we show that the iterates themselves are in fact bounded [Lemma 3]. We finally conclude the proof in 4.

A first Lyapunov function.

In this subsection we shall consider the following (stochastic) Lyapunov function:

For all $t>0$, $V_{t}$ verifies the following equation:

We are now equipped to apply the Itô formula on $V_{t}$. Indeed, it is clear that $\phi$ is a $C^{2}$ function of $(\beta,\alpha)$, hence,

The fifth term is explicit. Let us treat the first four terms separately:

First term. This term cancels with a compound of the fourth term.

Second term. We apply simply the chain rule for this term as $\alpha_{t}$ does not have any quadratic variation:

Fourth term. We apply Itô formula once again to get:

and thanks to Eq. (7), we have an expression for the first and last term, giving

Integrating this equation between and $t$ concludes the proof. ∎

Control of the martingale term and definition of 𝒜𝒜\mathcal{A}.

From now on and until the end of the Section, we place ourselves on the event $\mathcal{A}$, that is, all (in)equalities between random variables should be considered pointwise for any $\omega\in\mathcal{A}$. To make it clear, we will recall from time to time laconically this fact by writing, “on $\mathcal{A}$”.

From Lemma 2, we deduce the following inequalities,

Hence, we have the following control on $V_{t}$ with respect to a weighted loss integral:

Let us place ourselves on the event $\mathcal{A}$. Let $\tau>0$. Assume $(U_{t})_{0\leq t\leq\tau}$ is positive. Then for all $t\leqslant\tau$ we have the following explicit upper bound on both $\|\beta_{t}\|_{1}$ and $\|\xi_{t}\|_{1}$,

Since the last two terms cancel and for all $i$, $\sqrt{|\beta_{i}(t)|^{2}+4\alpha_{i}(t)^{4}}\leqslant\|\xi\|_{1}$, we have

where in the last inequality we plug in the value of $a$. This concludes the proof of the lemma. ∎

Let us define the stopping time $\tau=\inf\{t\geq 0\ \text{such that}\ U(t)\leq\frac{1}{2}\}$. Note that

where the last inequality comes from the upperbound on $\gamma$. Since $U_{t}$ is continuous we have that $\tau>0$. Assume that $\tau<+\infty$, by definition of the stopping time, for $t\leq\tau$: $U(t)\geq 0$ and we can apply Lemma 3 at time $\tau$:

where the last inequality comes from the choice of $\gamma$.

This is inconsistent since $U_{\tau}=\frac{1}{2}$. Hence $\tau=+\infty$ and thus $U_{t}\geq 1/2$ for all $t$. ∎

From the result of Lemma 4, with Equation 14, we obtain:

Hence it remains to lower bound $V_{t}$ in order to get the convergence of the integral of the loss.

On $\mathcal{A}$, let $\gamma$ be set as in Lemma 3, for all $t>0$, we have the following lower bound on $V_{t}$:

We follow exactly the same proof as for upperbounding the iterates.

Hence $(V_{t})_{t\geq 0}$ is lowerbounded and we can derive an upper bound on the loss integral to show the right part of 3. We recall it here in the following proposition.

On $\mathcal{A}$, let $\gamma$ be set as in Lemma 3, we have the following upper bound on the loss integral:

and thanks to the lower bound on $V_{t}$ from Lemma 5, it yields,

B.3 Proof of the convergence of the iterates: 2

In this subsection we prove the convergence of the iterates which corresponds to 2 of the main text. For the sake of completeness, we recall this fact in the following lemma.

Consider the following Bregman divergence style function for any interpolator $\beta^{*}$ :

where the first inequality is because $\alpha\mapsto\phi_{\alpha}(\beta)$ is decreasing and $\alpha_{t}\geq\alpha_{\infty}$. Therefore $D_{\phi_{\alpha_{t}}}(\beta_{\infty}^{\alpha},\beta_{t})\to 0$. Finally, since:

for some $\mu$ since the iterates are bounded. Therefore for all $t\geq 0$, $\phi_{\alpha_{t}}$ is $\mu$-strongly convex on some convex set in which the iterates $\beta_{s}$ stay in. Which means that: $D_{\phi_{\alpha_{t}}}(\beta_{\infty}^{\alpha},\beta_{t})\geq\frac{\mu}{2}\|\beta_{t}-\beta_{\infty}^{\alpha}\|_{2}^{2}$. Hence $\beta_{t}\to\beta_{\infty}^{\alpha}$. ∎

Lemma 6 along with the fact that the event $\mathcal{A}$ has probability at least $1-\frac{p}{2}$ (see Lemma 9 and paragraph around 13) concludes the proof of 2.

B.4 Proof of 1

We are now equipped to prove the main result of the paper. For clarity we recall the statement of the theorem here. See 1

Recall first that on $\mathcal{A}$, Lemma 6 implies that the iterates converge towards a zero-training error we denote by $\beta_{\infty}^{\alpha}$. From 1 we also have that:

Similarly to what has been done in subsection B.2, in order to lower bound the loss integral, we need a (different) control on the deviation of the local martingale $S_{t}$. We choose $\hat{a}:=W_{0}^{\alpha}/2$ and $\hat{b}:=\frac{1}{2}\ln(4/p)\hat{a}^{-1}$ so that once again $\hat{a}\hat{b}=\frac{1}{2}\ln(4/p)$. We refer to Lemma 7 for the definition of $W_{0}^{\alpha}$. Now that these parameters are fixed, consider the new event:

According to Lemma 6, the flow converges to an interpolator $\beta_{\infty}^{\alpha}$. We consider the same Lyapunov as before:

which is such that, following the same computations as in Lemma 2:

Now since we put ourselves on $\mathcal{B}$:

where $\beta^{*}_{\alpha}=\underset{\beta\ \text{s.t}\ X\beta=Y}{\text{argmin}}\phi(\beta,\alpha^{2})$. Therefore, we control the integral of the loss as

We now plug in the values $\hat{a}=\frac{W_{0}^{\alpha}}{2}$ and $\hat{b}=\frac{1}{W_{0}^{\alpha}}\ln(\frac{4}{p})$:

which along with Lemma 7 concludes the proof. ∎

Therefore through this lemma we see that by picking the biggest step-size which ensures convergence, we have a dependency of the integral of the loss as $\ln\frac{1}{\alpha}$.

Now we are equipped to prove 3. We recall it here to be self-contained.

In the final proposition of this subsection, we give the scale of $\alpha_{\infty}$ we obtain thanks to our analysis. Indeed though we know that in all case $\alpha_{\infty}<\alpha$, we would like to quantitatively know how much smaller the effective initialisation is in order to have an idea of the gain of SGD over GD (in terms of implicit bias).

This concludes the proof of the Proposition. ∎

This upperbound is significantly better than that of 5: the smaller the initialisation scale $\alpha$ and the greater the benefit of SGD over GD in terms of implicit bias. More precisely, Proposition 6 shows that the benefit scales as a power law with respect to the initialization $\alpha$.

Appendix C Deterministic framework

In this section we recall some known results concerning the implicit bias of deterministic mirror descent as well as give convergence guarantees. In the previous section, the stochasticity of the flow made the analysis much more involved. In contrast, the analysis is straightforward in the deterministic setting and we believe this simple case can serve as a warmup to gain further intuition. Note that even though these results are known independently, we did not find a clear reference gathering them. See for example for the convergence of the iterates towards an interpolator and for the associated implicit minimisation problem.

Then the iterates $(\beta_{t})_{t}$ converge to an interpolator $\beta_{\infty}$ which satisfies:

where $D_{\Psi}(\beta,\beta_{0})=\Psi(\beta)-\Psi(\beta_{0})-\langle\nabla\Psi(\beta_{0}),\beta-\beta_{0}\rangle$ is the Bregman divergence w.r.t. $\Psi$.

where the inequality is by convexity of the potential $\Psi$. Hence the loss is decreasing. Now consider the Bregman divergence between an arbitrary interpolator $\beta^{*}$ and $\beta_{t}$:

where the first inequality is by the convexity of the loss. Therefore:

Second step: the iterates converge towards an interpolator $\beta_{\infty}$.

and therefore $\beta_{t}$ converges towards the interpolator $\beta_{\infty}$.

Appendix D Experiments

Using the same notations and following the same derivations as in Appendix A, we can rewrite the recursion as:

Since $\Delta_{t}$ is zero-mean and independent of $i_{t}$ we get:

Now following the same reasoning as in Appendix A, it is natural to consider the following SDE:

We experimentally validate the advantage of adding label noise by choosing the sequence $\delta_{t}=1$ if $t\leq 10^{3}$ and $\delta_{t}=0$ if $t>10^{3}$. The results are illustrated Figure 5. Note that the training loss is heavily slowed down, however the recovered solution at iteration $t=10^{6}$ is much better than that of SGD, and it has not even converged yet. However, it must be kept in mind that adding too much label noise can significantly slow down the convergence of the validation loss or even prevent the iterates from converging.

Appendix E Extensions

We introduce two extensions of our results: subsection E.1 extends our results for a very general stochastic gradient flow model and subsection E.2 discuss them in the depth $p\geq 3$ case.

Recall from the SDE modelling of Appendix A that $\text{Cov}_{i_{t}}[\xi_{i_{t}}(\beta)]=\frac{4}{n}\mathop{\rm diag}(L_{i}(\beta))_{1\leq i\leq n}+O(\frac{1}{n^{2}})$. If we assume $n$ large enough we can neglect the second order term of order $1/n^{2}$:

Assume we do not consider that $L_{i}(\beta)\sim L(\beta)$, then the overall SGD noise structure is captured by

This leads us in considering the following SDE:

As previously, this SDE admits an implicit integral formulation (multiplication must be understood component-wise):

And we obtain the following mirror-type descent flow:

Assuming convergence of the iterates and of $\alpha_{t}$ (we do not show the convergence, though we think the proof could straightforwardly be adapted following Appendix B ), the corresponding minimisation problem is:

Note that the main result of the paper is very similar, the difference relies in:

E.2 Higher order models: the cases of depth p>2𝑝2p>2

Until now, we have focused on a $2$-homogeneous parametrisation of the estimator. A legitimate question is how the implicit bias changes as we go to a higher degree of homogeneity. In terms of networks architecture, this corresponds to increasing the depth of the neural networks. Let us fix $p\geqslant 3$ with the new parametrisation $\beta_{w}=w_{+}^{p}-w_{-}^{p}$, the loss of our new model writes: $L(w)=\frac{1}{4n}\sum_{i=1}^{n}\langle w_{+}^{p}-w_{-}^{p}-\beta^{*},x_{i}\rangle^{2}$. As previously, we want to consider the stochastic differential equation related to stochastic gradient descent on the above loss. With the same modelling as in Section 2.2, stochastic gradient flow writes:

We apply the Itô formula on $w_{t,+}^{2-p}$ and $w_{t,-}^{2-p}$ to get the following:

Appendix F Technical lemmas

In this section, we state and prove technical lemmas which we use to prove our main results.

Since $(S_{t})_{t\geq 0}$ is a is a locally square-integrable martingale with a.s. continuous paths, [19, Corollary 11] gives that

Let $A,B>0$ such that $\frac{A}{B}+\ln(B)\geq 2$. Assume that $x\leq A+B\ln x$, then

$x\leq A+B\ln x$ is equivalent to $x\leq\exp(-\frac{A}{B})\exp(\frac{x}{B})$. Standard analysis on the Lambert $W$ function shows that this leads to $x\leq-B\ W_{-1}(-\frac{1}{B}\exp(-\frac{A}{B}))$, where $W_{-1}$ is the lower branch see https://en.wikipedia.org/wiki/Lambert_W_function for more details. For $-\frac{1}{e}\leq z\leq 0$, the branch $W_{-1}$ can be lower bounded as: $W_{-1}(z)\geq-\sqrt{-2(1+\ln(-z))}+\ln(-z)$ (see Theorem 1 of ). Since $\ln(-z)=\ln(\frac{1}{B}\exp(-\frac{A}{B}))=-(\frac{A}{B}+\ln(B))$:

For the other term of the $\max$, let us define $g(\beta):=\frac{1}{4}\beta\ln\frac{\beta}{2\alpha^{2}}$, we have that

Hence, $f-g$ increases and as $f(0)-g(0)=0$, we have that $f>g$ which concludes the proof. ∎