SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We study the effect of stochasticity along with large step sizes 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  $\alpha_{\infty}$  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 step sizes as seen in Fig. 1: without stochasticity, the use of large step sizes 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.
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.
[53] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. J. Mach. Learn. Res., 19(1):2822-2878, jan 2018.
[54] Terrence Tao. Concentration of measure. 254A, Notes 1, Blogpost, 2010.
[55] Matus Telgarsky. Margins, shrinkage, and boosting. In International Conference on Machine Learning, pages 307-315. PMLR, 2013.
[56] Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. The statistical complexity of early-stopped mirror descent. In Advances in Neural Information Processing Systems, volume 33, pages 253-264, 2020.
[57] Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, 2019.
[58] Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018.
[59] Yuqing Wang, Minshuo Chen, Tuo Zhao, and Molei Tao. Large learning rate tames homogeneity: Convergence and balancing effect. In International Conference on Learning Representations, 2022.
[60] Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part II: Continuous time analysis. arXiv preprint arXiv:2106.02588, 2021.
[61] Blake Woodworth, Suriya Gunasekar, Jason D. Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 3635-3673. PMLR, 09-12 Jul 2020.
[62] Fan Wu and Patrick Rebeschini. A continuous-time mirror descent approach to sparse phase retrieval. In Advances in Neural Information Processing Systems, volume 33, pages 20192-20203, 2020.
[63] Jingfeng Wu, Difan Zou, Vladimir Braverman, and Quanquan Gu. Direction matters: On the implicit bias of stochastic gradient descent with moderate learning rate. In International Conference on Learning Representations, 2021.
[64] Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized learning: A dynamical stability perspective. In Advances in Neural Information Processing Systems, volume 31, 2018.
[65] Matthew D Zeiler. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701, 2012.
[66] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
[67] Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edge-of-stability training dynamics with a minimalist example. International Conference on Learning Representations, 2023.
1. 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.
2. In Appendix B, we prove that  $(\beta_{k})$  follows a Mirror descent recursion with varying potentials. We explicit these potentials and discuss some consequences.
3. In Appendix C we prove that (S)GD on the  $\frac{1}{2} (w_{+}^{2} - w_{-}^{2})$  and  $u \odot v$  parametrisations with suitable initialisations lead to the same sequence  $(\beta_k)$ .
4. In Appendix D, we show that the hypentropy  $\psi_{\alpha}$  converges to a weighted- $\ell_1$ -norm when  $\alpha$  converges to 0 non-uniformly. We then discuss the effects of this weighted  $\ell_1$ -norm for sparse recovery.
5. In Appendix E, we provide our descent lemmas for mirror descent with varying potentials and prove the boundedness of the iterates.
6. In Appendix F, we prove our main results: Theorem 1 and Theorem 2, as well as quantitative convergence (Proposition 1).
7. In Appendix G, we prove the lemmas and propositions given in the main text.
8. In Appendix H, we provide technical lemmas used throughout the proof of Theorem 1 and Theorem 2.
9. In Appendix I, we provide concentration results for random matrices and random vectors, used to estimate with high probability (w.r.t. the dataset) quantities related to the data.
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  $\ell_{1}$  solutions for uncentered data (Fig. 4). We make the following assumptions on the uncentered inputs.
Assumption 1. There exist  $\mu \in \mathbb{R}^d$  and  $\delta, c_0, c_1, c_2 > 0$  such that for all  $s$ -sparse vectors  $\beta$  verifying  $\langle \mu, \beta \rangle \geqslant c_0 \| \beta \|_{\infty} \| \mu \|_{\infty}$ , there exists  $\varepsilon \in \mathbb{R}^d$  such that  $(X^\top X)\beta = \langle \beta, \mu \rangle \mu + \varepsilon$  where  $\| \varepsilon \|_2 \leqslant \delta \| \beta \|_2$  and  $c_1 \langle \beta, \mu \rangle^2 \mu^2 \leqslant \frac{1}{n} \sum_i x_i^2 \langle x_i, \beta \rangle^2 \leqslant c_2 \langle \beta, \mu \rangle^2 \mu^2$ .
Assumption 1 is not restrictive and holds with high probability for  $\mathcal{N}(\mu \mathbf{1},\sigma^2 I_d)$  inputs when  $\mu \gg \sigma \mathbf{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  $\langle \mu, \beta_{\mathrm{sparse}}^{\star} \rangle \geqslant c_0 \| \beta \|_{\infty} \| \mu \|_{\infty}$ , the squared full batch gradient and the expected stochastic gradient descent at initialisation satisfy, for some  $\varepsilon$  satisfying  $\| \varepsilon \|_{\infty} \ll \| \beta_{\mathrm{sparse}} \|_2$ :
$$
\nabla \mathcal {L} \left(\beta_ {0}\right) = \left\langle \beta_ {\text {s p a r s e}} ^ {\star}, \mu \right\rangle^ {2} \mu^ {2} + \varepsilon , \tag {14}
$$
$$
\mathbb {E} _ {i \sim \operatorname {U n i f} ([ \mathrm {n} ])} \left[ \nabla \mathcal {L} _ {i} \left(\beta_ {0}\right) ^ {2} \right] = \Theta \left(\left\langle \beta_ {\text {s p a r s e}} ^ {\star}, \mu \right\rangle^ {2} \mu^ {2}\right). \tag {15}
$$
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  $\mu \in \mathbb{R}^d$  orthogonal to  $\beta_{\mathrm{sparse}}^{\star}$ , there is no effect of decentering on the recovered solution. If the support of  $\mu$  is the same as that of  $\beta_{\mathrm{sparse}}^{\star}$ , 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.
Figure 4: Noiseless sparse regression with a 2-layer DLN with uncentered data  $x_{i} \sim \mathcal{N}(\mu \mathbf{1}, I_{d})$  where  $\mu = 5$ . All the stepsizes lead to convergence to a global solution and the solutions of SGD and GD have similar behaviours, corroborating Proposition 4. The setup corresponds to  $(n, d, s, \alpha) = (20, 30, 3, 0.1)$ .
A.2 Behaviour of the maximal value and trace of the hessian
Here in Fig. 5, we provide some additional experiments on the behaviour of: (1) the maximum eigenvalue of the hessian  $\nabla^2 F(w_\infty^\gamma)$  at the convergence of the iterates of SGD and GD (2) the trace
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^{\top}X = I_{d}$  and that  $\alpha \sim 0$  (pushing the EoS phenomenon), one can see that minimising  $\lambda_{\mathrm{max}}(\nabla^{2}F(w))$  under the constraints  $X(w_{+}^{2} - w_{-}^{2}) = y$  and  $w_{+}\odot w_{-} = 0$  is equivalent to minimising  $\| \beta \|_{\infty}$  under the constraint  $X\beta = y$ . On the other hand minimising the trace of the hessian is equivalent to minimising the  $\ell_1$ -norm.
Figure 5: Noiseless sparse regression setting. Diagonal linear network. Centered data. Behaviour of 2 different types of flatness of the recovered solution by SGD and GD depending on the stepsize. The setup corresponds to  $(n,d,s,\alpha) = (20,30,3,0.1)$ .
A.3 Edge of Stability for SGD
Figure 6: SGD at the edge of stability: all coordinates fluctuate, and the sparse solution is recovered. As opposed to GD at the EoS, since all coordinates fluctuate, the coordinates to recover are not more penalised than the others.
B Main ingredients behind the proof of Theorem 1 and Theorem 2
In this section, we show that the iterates  $(\beta_{k})_{k\geqslant 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  $\alpha$  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: \mathbb{R}^d \to \mathbb{R}$ , the (stochastic) mirror descent iterates algorithm write as [5, 18], where the minimum is assumed to be attained over  $\mathbb{R}^d$  and unique:
$$
\beta_ {k + 1} = \operatorname {a r g m i n} _ {\beta \in \mathbb {R} ^ {d}} \left\{\eta_ {k} \langle g _ {k}, \beta \rangle + D _ {h} (\beta , \beta_ {k}) \right\}, \tag {16}
$$
for stochastic gradients  $g_{k}$ , stepsize  $\gamma_{k} \geqslant 0$ , and  $D_{h}(\beta, \beta') = h(\beta) - h(\beta') - \langle \nabla h(\beta'), \beta - \beta' \rangle$  is the Bregman divergence associated to  $h$ . Iteration (16) can also be cast as
$$
\nabla h \left(\beta_ {k + 1}\right) = \nabla h \left(\beta_ {k}\right) - \gamma_ {k} g _ {k}. \tag {17}
$$
Now, let  $(h_k)$  be strictly convex reference functions  $\mathbb{R}^d\to \mathbb{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  $\nabla h_{k}(\beta_{k + 1}) = \nabla h_{k}(\beta_{k}) - \gamma_{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:
$$
\beta_ {k + 1} = \operatorname {a r g m i n} _ {\beta \in \mathbb {R} ^ {d}} \left\{\eta_ {k} \langle g _ {k}, \beta \rangle + D _ {h _ {k + 1}, h _ {k}} (\beta , \beta_ {k}) \right\},
$$
where  $D_{h_{k + 1},h_k}(\beta ,\beta ') = h_{k + 1}(\beta) - h_k(\beta ') - \langle \nabla h_k(\beta '),\beta -\beta '\rangle$  , a recursion that can also be cast as:
$$
\nabla h _ {k + 1} \left(\beta_ {k + 1}\right) = \nabla h _ {k} \left(\beta_ {k}\right) - \gamma_ {k} g _ {k}.
$$
To derive convergence of the iterates, we prove analogs to classical mirror descent lemmas, generalised to time-varying potentials.
B.2 The iterates  $(\beta_{k})$  follow a stochastic mirror descent with varying potential recursion
In this section we show and prove that the iterates  $(\beta_{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_{\pm} : \mathbb{R} \to \mathbb{R} \cup \{\infty\}$  be defined as:
$$
q _ {\pm} (x) = \mp 2 x - \ln \left((1 \mp x) ^ {2}\right),
$$
$$
q (x) = \frac {1}{2} (q _ {+} (x) + q _ {-} (x)) = - \frac {1}{2} \ln \left((1 - x ^ {2}) ^ {2}\right),
$$
with the convention that  $q(1) = \infty$ . Notice that  $q(x) \geqslant 0$  for  $|x| \leqslant \sqrt{2}$  and  $q(x) < 0$  otherwise. For the iterates  $\beta_{k} = u_{k} \odot v_{k} \in \mathbb{R}^{d}$ , we recall the definition of the following quantities:
$$
\boldsymbol {\alpha} _ {\pm , k} = \boldsymbol {\alpha} \exp \left(- \frac {1}{2} \sum_ {i = 0} ^ {k - 1} q _ {\pm} \left(\gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} \left(\beta_ {\ell}\right)\right)\right) \in \mathbb {R} _ {> 0} ^ {d},
$$
$$
\boldsymbol {\alpha} _ {k} ^ {2} = \boldsymbol {\alpha} _ {+, k} \odot \boldsymbol {\alpha} _ {-, k},
$$
$$
\phi_ {k} = \frac {1}{2} \operatorname {a r c s i n h} \left(\frac {\boldsymbol {\alpha} _ {+, k} ^ {2} - \boldsymbol {\alpha} _ {- , k} ^ {2}}{2 \boldsymbol {\alpha} _ {k} ^ {2}}\right) \in \mathbb {R} ^ {d}.
$$
Finally for  $k\geqslant 0$  , we define the potentials  $(h_k:\mathbb{R}^d\to \mathbb{R})_{k\geqslant 0}$  as:
$$
h _ {k} (\beta) = \psi_ {\boldsymbol {\alpha} _ {k}} (\beta) - \left\langle \phi_ {k}, \beta \right\rangle , \tag {18}
$$
where  $\psi_{\alpha_k}$  is the hyperbolic entropy defined in (4) of scale  $\alpha_{k}$ :
$$
\psi_ {\boldsymbol {\alpha} _ {k}} (\beta) = \frac {1}{2} \sum_ {i = 1} ^ {d} \left(\beta_ {i} \mathrm {a r c s i n h} \left(\frac {\beta_ {i}}{\alpha_ {k , i} ^ {2}}\right) - \sqrt {\beta_ {i} ^ {2} + \alpha_ {k , i} ^ {4}} + \alpha_ {k, i} ^ {2}\right)
$$
where  $\alpha_{k,i}$  corresponds to the  $i^{th}$  coordinate of the vector  $\alpha_{k}$ .
Now that all the relevant quantities are defined, we can state the following proposition which explicits the time-varying stochastic mirror descent followed by  $(\beta_{k})_{k}$
Proposition 5. The iterates  $(\beta_{k} = u_{k}\odot v_{k})_{k\geqslant 0}$  from Eq. (3) satisfy the Stochastic Mirror Descent recursion with varying potentials  $(h_k)_k$  ..
$$
\nabla h _ {k + 1} \left(\beta_ {k + 1}\right) = \nabla h _ {k} \left(\beta_ {k}\right) - \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} \left(\beta_ {k}\right), \tag {19}
$$
where  $h_k: \mathbb{R}^d \to \mathbb{R}$  for  $k \geqslant 0$  are defined Eq. (18). Since  $\nabla h_0(\beta_0) = 0$  we have:
$$
\nabla h _ {k} \left(\beta_ {k}\right) \in \operatorname {s p a n} \left(x _ {1}, \dots , x _ {n}\right)
$$
Proof. Using Proposition 6, we study the  $\frac{1}{2} (w_{+}^{2} - w_{-}^{2})$  parametrisation instead of the  $u \odot 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_{+}$ :
$$
w _ {+, k + 1} = \left(1 - \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k})\right) \cdot w _ {+, k}.
$$
$$
\begin{array}{l} w _ {+, k + 1} ^ {2} = \left(1 - \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k})\right) ^ {2} \cdot w _ {+, k} ^ {2} \\ = \exp \left(\ln ((1 - \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k})) ^ {2})\right) \cdot w _ {+, k} ^ {2}, \\ \end{array}
$$
with the convention that  $\exp (\ln (0)) = 0$ . This leads to:
$$
\begin{array}{l} w _ {+, k + 1} ^ {2} = \exp \big (- 2 \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (w _ {k}) + 2 \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k}) + \ln ((1 - \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k})) ^ {2}) \big) \cdot w _ {+, k} ^ {2} \\ = \exp \big (- 2 \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k}) - q _ {+} (\gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k})) \big) \cdot w _ {+, k} ^ {2}, \\ \end{array}
$$
since  $q_{+}(x) = -2x - \ln ((1 - x)^{2})$ . Expanding the recursion and using that  $w_{+,k=0}$  is initialised at  $w_{+,k=0} = \alpha$ , we thus obtain:
$$
\begin{array}{l} w _ {+, k} ^ {2} = \boldsymbol {\alpha} ^ {2} \exp \left(- \sum_ {\ell = 0} ^ {k - 1} q _ {+} \left(\gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell})\right)\right) \exp \left(- 2 \sum_ {\ell = 0} ^ {k - 1} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell})\right) \\ = \boldsymbol {\alpha} _ {+, k} ^ {2} \exp \left(- 2 \sum_ {\ell = 0} ^ {k - 1} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell})\right), \\ \end{array}
$$
where we recall that  $\alpha_{\pm ,k}^2 = \alpha^2\exp (-\sum_{\ell = 0}^{k - 1}q_{\pm}(\gamma_\ell g_\ell))$ . One can easily check that we similarly get:
$$
w _ {-, k} ^ {2} = \boldsymbol {\alpha} _ {-, k} ^ {2} \exp \left(+ 2 \sum_ {\ell = 0} ^ {k - 1} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell})\right),
$$
leading to:
$$
\begin{array}{l} \beta_ {k} = \frac {1}{2} \left(w _ {+, k} ^ {2} - w _ {-, k} ^ {2}\right) \\ = \frac {1}{2} \boldsymbol {\alpha} _ {+, k} ^ {2} \exp \left(- 2 \sum_ {\ell = 0} ^ {k - 1} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell})\right) - \frac {1}{2} \boldsymbol {\alpha} _ {-, k} ^ {2} \exp \left(+ 2 \sum_ {\ell = 0} ^ {k - 1} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell})\right). \\ \end{array}
$$
Using Lemma 4, the previous equation can be simplified into:
$$
\beta_ {k} = \boldsymbol {\alpha} _ {+, k} \boldsymbol {\alpha} _ {-, k} \sinh \Big (- 2 \sum_ {\ell = 0} ^ {k - 1} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell}) + \arcsinh \left(\frac {\boldsymbol {\alpha} _ {+, k} ^ {2} - \boldsymbol {\alpha} _ {- , k} ^ {2}}{2 \boldsymbol {\alpha} _ {+, k} \boldsymbol {\alpha} _ {- , k}}\right) \Big),
$$
which writes as:
$$
\frac {1}{2} \operatorname {a r c s i n h} \left(\frac {\beta_ {k}}{\boldsymbol {\alpha} _ {k} ^ {2}}\right) - \phi_ {k} = - \sum_ {\ell = 0} ^ {k - 1} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} (\beta_ {\ell}) \in \operatorname {s p a n} (x _ {1}, \dots , x _ {n}),
$$
where  $\phi_{k} = \frac{1}{2}\arcsinh\left(\frac{\alpha_{+,k}^{2} - \alpha_{-,k}^{2}}{2\alpha_{k}^{2}}\right)$ ,  $\pmb{\alpha}_{k}^{2} = \pmb{\alpha}_{+,k} \odot \pmb{\alpha}_{-,k}$  and since the potentials  $h_k$  are defined in Eq. (18) as  $h_k = \psi_{\alpha_k} - \langle \phi_k, \cdot \rangle$  with
$$
\psi_ {\boldsymbol {\alpha}} (\beta) = \frac {1}{2} \sum_ {i = 1} ^ {d} \left(\beta_ {i} \operatorname {a r c s i n h} \left(\frac {\beta_ {i}}{\boldsymbol {\alpha} _ {i} ^ {2}}\right) - \sqrt {\beta_ {i} ^ {2} + \boldsymbol {\alpha} _ {i} ^ {4}} + \boldsymbol {\alpha} _ {i} ^ {2}\right) \tag {20}
$$
specifically such that  $\nabla h_k(\beta_k) = \frac{1}{2}\operatorname {arcsinh}\left(\frac{\beta_k}{\alpha_k^2}\right) - \phi_k$ . Hence,
$$
\nabla h _ {k} \left(\beta_ {k}\right) = \sum_ {\ell <   k} \gamma_ {\ell} \nabla \mathcal {L} _ {\mathcal {B} _ {\ell}} \left(\beta_ {\ell}\right),
$$
which corresponds to a Mirror Descent with varying potentials  $(h_k)_k$ .
C Equivalence of the  $u \odot v$  and  $\frac{1}{2} (w_{+}^{2} - w_{-}^{2})$  parametrisations
We here prove the equivalence between the  $\frac{1}{2} (w_{+}^{2} - w_{-}^{2})$  and  $u \odot v$  parametrisations, that we use throughout the proofs in the Appendix.
Proposition 6. Let  $(\beta_k)_{k\geqslant 0}$  and  $(\beta_k')_{k\geqslant 0}$  be respectively generated by stochastic gradient descent on the  $u\odot v$  and  $\frac{1}{2} (w_{+}^{2} - w_{-}^{2})$  parametrisations:
$$
(u _ {k + 1}, v _ {k + 1}) = (u _ {k}, v _ {k}) - \gamma_ {k} \nabla_ {u, v} \left(\mathcal {L} _ {\mathcal {B} _ {k}} (u \odot v)\right) (u _ {k}, v _ {k}),
$$
$$
w _ {\pm , k + 1} = w _ {\pm , k} - \gamma_ {k} \nabla_ {w _ {\pm}} \left(\mathcal {L} _ {\mathcal {B} _ {k}} \left(\frac {1}{2} \left(w _ {+} ^ {2} - w _ {-} ^ {2}\right)\right)\right) \left(w _ {+, k}, w _ {-, k}\right),
$$
initialised as  $u_0 = \sqrt{2}\pmb {\alpha},v_0 = 0$  and  $w_{+,0} = w_{-,0} = \pmb{\alpha}$  . Then for all  $k\geqslant 0$  , we have  $\beta_{k} = \beta_{k}^{\prime}$
$$
w _ {\pm , 0} = \boldsymbol {\alpha}, \quad w _ {\pm , k + 1} = (1 \mp \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k} ^ {\prime})) w _ {\pm , k},
$$
$$
u _ {0} = \sqrt {2} \alpha , \quad v _ {0} = 0, \quad u _ {k + 1} = u _ {k} - \gamma_ {k} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k}) v _ {k}, \quad v _ {k + 1} = v _ {k} - \gamma_ {k} \nabla \mathcal {L} (\beta_ {k}) u _ {k}.
$$
$$
\beta_ {k + 1} = \left(1 + \gamma_ {k} ^ {2} \nabla \mathcal {L} \left(\beta_ {k}\right) ^ {2}\right) \beta_ {k} - \gamma_ {k} \left(u _ {k} ^ {2} + v _ {k} ^ {2}\right) \nabla \mathcal {L} _ {\mathcal {B} _ {k}} \left(\beta_ {k}\right),
$$
$$
\beta_ {k + 1} ^ {\prime} = \left(1 + \gamma_ {k} ^ {2} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} \left(\beta_ {k} ^ {\prime}\right) ^ {2}\right) \beta_ {k} ^ {\prime} - \gamma_ {k} \left(w _ {+, k} ^ {2} + w _ {-, k} ^ {2}\right) \nabla \mathcal {L} _ {\mathcal {B} _ {k}} \left(\beta_ {k} ^ {\prime}\right).
$$
Then, let  $z_{k} = \frac{1}{2}\bigl (u_{k}^{2} - v_{k}^{2}\bigr)$  and  $z_k^\prime = w_{+,k}w_{-k}$ . We have  $z_0 = \alpha^2$ ,  $z_0' = \alpha^2$  and:
$$
z _ {k + 1} = \left(1 - \gamma_ {k} ^ {2} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k}) ^ {2}\right) z _ {k}, \quad z _ {k + 1} ^ {\prime} = \left(1 - \gamma_ {k} ^ {2} \nabla \mathcal {L} _ {\mathcal {B} _ {k}} (\beta_ {k} ^ {\prime}) ^ {2}\right) z _ {k} ^ {\prime}.
$$
Using  $a^2 + b^2 = \sqrt{(2ab)^2 + (a^2 - b^2)^2}$  for  $a, b \in \mathbb{R}$ , we finally obtain that:
$$
u _ {k} ^ {2} + v _ {k} ^ {2} = \sqrt {(2 \beta_ {k}) ^ {2} + (2 z _ {k}) ^ {2}}, w _ {+, k} ^ {2} + w _ {-, k} ^ {2} = \sqrt {(2 \beta_ {k} ^ {\prime}) ^ {2} + (2 z _ {k} ^ {\prime}) ^ {2}}.
$$
We conclude by observing that  $(\beta_{k},z_{k})$  and  $(\beta_k',z_k')$  follow the exact same recursions, initialised at the same value  $(0,\alpha^{2})$ .
D Convergence of  $\psi_{\alpha}$  to a weighted  $\ell_1$  norm and harmful behaviour
We show 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  $\ell_1$ -norm rather than the classical  $\ell_1$  norm.
Proposition 7. For  $\alpha \geqslant 0$  and a vector  $h\in \mathbb{R}^d$ , let  $\tilde{\alpha} = \alpha \exp (-h\ln (1 / \alpha))\in \mathbb{R}^d$ . Then we have that for all  $\beta \in \mathbb{R}^d$