Leveraging Continuous Time to Understand Momentum When Training Diagonal
Linear Networks

Hristo Papazov∗ Scott Pesme∗ Nicolas Flammarion

EPFL<EMAIL_ADDRESS>EPFL<EMAIL_ADDRESS>EPFL
<EMAIL_ADDRESS>
###### Abstract

In this work, we investigate the effect of momentum on the optimisation
trajectory of gradient descent. We leverage a continuous-time approach in the
analysis of momentum gradient descent with step size $\gamma$ and momentum
parameter $\beta$ that allows us to identify an intrinsic quantity
$\lambda=\frac{\gamma}{(1-\beta)^{2}}$ which uniquely defines the optimisation
path and provides a simple acceleration rule. When training a $2$-layer
diagonal linear network in an overparametrised regression setting, we
characterise the recovered solution through an implicit regularisation
problem. We then prove that small values of $\lambda$ help to recover sparse
solutions. Finally, we give similar but weaker results for stochastic momentum
gradient descent. We provide numerical experiments which support our claims.

##  1 Introduction

Momentum methods (Sutskever et al., 2013) have now become a staple of optimal
neural network training due to the provided gains in both optimisation
efficiency and generalisation performance. This pivotal role is underscored by
the widespread use of momentum in the successful training of most state-of-
the-art deep networks, including CLIP (Radford et al., 2021), Chinchilla
(Hoffmann et al., 2022), GPT-3 (Brown et al., 2020), and PaLM (Chowdhery et
al., 2022).

Originating in the work of Polyak (1964), momentum first featured in the
heavy-ball method devised to accelerate convergence in convex optimisation.
However, when applied to neural network training, momentum exhibits a distinct
and complementary characteristic: a steering towards models with superior
generalisation performance compared to networks trained with gradient descent.
We note that while the effect of momentum on optimisation has been researched
extensively (Defazio, 2020; Sun et al., 2019), the generalisation aspect of
momentum has been left relatively underexplored.

The performance of gradient descent methods presents intriguing challenges
from a theoretical perspective. First, establishing convergence is highly non-
trivial. Second, the existence of numerous global minima for the training
objective, some of which generalise poorly, adds to the puzzle (Zhang et al.,
2017). To elucidate this second point, the notion of implicit regularisation
has come to the forefront. It posits that the optimisation process implicitly
favors solutions with strong generalisation properties, even in the absence of
explicit regularisation. The canonical example is overparametrised linear
regression with more trainable parameters than the number of samples. While
there exist infinitely many solutions that fit the data, gradient methods
navigate in a restricted parameter subspace and converge towards the solution
closest in terms of the $\ell_{2}$ distance (Lemaire, 1996).

In this work, we aim to expand our understanding of the implicit bias of
momentum by analysing its impact on the optimisation trajectory in $2$-layer
diagonal linear networks. The $2$-layer diagonal linear network has garnered
significant attention recently (Woodworth et al., 2020; Vaškevičius et al.,
2019; HaoChen et al., 2021; Pesme et al., 2021; Pillaud-Vivien et al., 2022).
Despite its apparent simplicity, this network has surprisingly shed light on
training behaviours typically associated with much more complex architectures.
Some of these insights include the influence of initialisation (Woodworth et
al., 2020), the impact of noise (Pesme et al., 2021), and the role of the step
size (Even et al., 2023). Consequently, this architecture serves as an
excellent surrogate model for gaining a deeper understanding of intricate
phenomena such as the role of momentum in the generalisation performance.

###  1.1 Main Contributions

In this paper, we investigate the influence of momentum on the optimisation
trajectory of neural networks trained with momentum gradient descent (MGD).
Leveraging the continuous-time approximation of MGD – momentum gradient flow
(MGF), we show that the optimisation trajectory strongly depends on the key
quantity ${\lambda}=\frac{{\gamma}}{(1-{\beta})^{2}}$, where ${\gamma}$ and
${\beta}$ denote the step size and momentum parameter of MGD, respectively.
Surprisingly, this continuous-time framework experimentally proves to be a
good approximation of the discrete trajectory even for large values of
$\gamma$.

We proceed to list our main contributions.

  * •

First, using the key quantity ${\lambda}$, we derive a straightforward
acceleration rule that maintains the optimisation path while accelerating the
optimisation speed.

  * •

Then, focusing on MGF on 2-layer diagonal linear networks, we precisely
characterise the recovered solution and prove that for suitably small values
of ${\lambda}$, MGF recovers solutions which generalise better than the ones
selected by gradient flow (GF) in a sparse regression setting.

  * •

Finally, we provide similar but slightly weaker results for stochastic MGD.

###  1.2 Related Works

Momentum and Acceleration. Momentum algorithms have their roots in
acceleration methods, and many studies have investigated their convergence
speed when optimising both convex and non-convex functions: (Ghadimi et al.,
2015; Flammarion and Bach, 2015; Kidambi et al., 2018; Can et al., 2019;
Sebbouh et al., 2021; Mai and Johansson, 2020; Liu et al., 2020; Cutkosky and
Mehta, 2020; Defazio, 2020; Orvieto et al., 2020; Sebbouh et al., 2021).
Moreover, apart from accelerating training, heavy-ball methods come with the
additional advantage of always escaping saddle points (Jin et al., 2018; Sun
et al., 2019).

Momentum and Continuous-Time Models. Building upon the foundational work of
Alvarez (2000); Attouch et al. (2000), researchers have analysed accelerated
gradient methods using second-order differential equations. Su et al. (2014)
extended the previous ODE to encompass the Nesterov accelerated method,
demonstrating convergence rates similar to the discrete case. Wibisono et al.
(2016) adopted a variational perspective to scrutinise the mechanics of
acceleration. A significant advancement emerged with the introduction of
Lyapunov analysis, undertaken by Wilson et al. (2021); Sanz Serna and
Zygalakis (2021); Moucer et al. (2023). This analytical approach sheds light
on the stability and convergence properties of these methods. Further
refinement has been achieved by Shi et al. (2021), who developed high-
resolution ODEs tailored to various momentum-based acceleration techniques and
able to distinguish between Nesterov’s Accelerated Gradient and Polyak’s Heavy
Ball methods. Finally, error bounds for the discretisation of MGF have been
developed by Kovachki and Stuart (2021).

Momentum and Implicit Bias. Sutskever et al. (2013); Leclerc and Madry (2020)
have empirically shown significant generalisation improvements in
architectures trained with momentum on common vision tasks. Building on these
empirical observations, Jelassi and Li (2022) designed a synthetic binary
classification problem where a 2-layer convolutional neural network trained
with MGD provably generalises better than gradient descent (GD). Recently,
Ghosh et al. (2023) reveal that the MGD trajectory closely resembles the
gradient flow trajectory of a regularised loss. Through the specific
regularisation, the authors argue that the MGD trajectory favors flatter
minima than the GD trajectory. The study’s findings apply to any reasonable
loss, but due to the finite time horizon restriction, cannot characterise the
solution to which MGD converges. Additionally, Wang et al. (2023) show that in
deep diagonal linear networks with identical weights across layers, increasing
the depth biases the optimisation towards sparse solutions.

##  2 From Discrete to Continuous

Momentum Gradient Descent. We consider minimising a differentiable function
$F:{\mathbb{R}}^{d}\to{\mathbb{R}}$ using momentum gradient descent (MGD) with
step size $\gamma>0$ and momentum parameter $\beta\in[0,1)$. Initialised at
two points $(w_{0},w_{1})\in{\mathbb{R}}^{2d}$, the iterates follow the
discrete recursion for $k\geq 1$:

$w_{k+1}=w_{k}-{\gamma}\nabla F(w_{k})+\beta(w_{k}-w_{k-1}).$

($\mathrm{MGD}({\gamma},{\beta})$)

Momentum Gradient Flow. Directly analysing the discrete recursion
$\mathrm{MGD}({\gamma},{\beta})$ appears intractable in many settings. To
overcome this difficulty, we follow the classical approach of considering a
second order differential equation of the form

$\displaystyle a\ddot{w}_{t}+b\dot{w}_{t}+\nabla F(w_{t})=0$ (1)

with leading coefficient $a\geq 0$ and damping coefficient $b>0$. In fact,
without loss of generality, the previous differential equation can be reduced
to a new one which depends on a single parameter $\lambda$. Indeed, assume
that $w_{t}$ follows ODE (1) with initialisation
$(w_{t=0},\dot{w}_{t=0})=(w_{0},\dot{w}_{0})$, then a simple chain rule shows
that $\tilde{w}_{t}=w_{bt}$ follows

$\displaystyle\frac{a}{b^{2}}\ddot{\tilde{w}}_{t}+\dot{\tilde{w}}_{t}+\nabla
F(\tilde{w}_{t})=0,$

with initialisation
$(\tilde{w}_{t=0},\dot{\tilde{w}}_{t=0})=(w_{0},b\dot{w}_{0})$. Hence, up to a
time reparametrisation, it is sufficient to consider the following
differential equation which depends on a unique parameter $\lambda\geq 0$:

$\displaystyle\lambda\ddot{w}_{t}+\dot{w}_{t}+\nabla F(w_{t})=0.$
($\mathrm{MGF}(\lambda)$)

We call the differential equation $\mathrm{MGF}(\lambda)$ momentum gradient
flow (MGF) with parameter $\lambda$. To show the link with the
$\mathrm{MGD}({\gamma},{\beta})$ recursion, we discretise
$\mathrm{MGF}(\lambda)$ with a second-order central difference, first-order
backward difference, and discretisation step $\varepsilon>0$ as carried out by
Kovachki and Stuart (2021):

(2)

Rewriting, we obtain

$\displaystyle w_{k+1}=w_{k}-\frac{\varepsilon^{2}}{{\lambda}}\nabla
F(w_{k})+(1-\frac{\varepsilon}{\lambda})(w_{k}-w_{k-1}),$

which corresponds to momentum gradient descent with parameters
$\gamma=\frac{\varepsilon^{2}}{\lambda}$ and
$\beta=1-\frac{\varepsilon}{\lambda}$. Solving for $\varepsilon$ and $\lambda$
leads to the following proposition:

######  Proposition 1.

For $(w_{0},w_{1})\in{\mathbb{R}}^{2d}$, consider momentum gradient flow
$\mathrm{MGF}(\lambda)$ with

$\lambda=\frac{\gamma}{(1-\beta)^{2}}$

and initialisation $w_{t=0}=w_{0}$,
$\dot{w}_{t=0}=(w_{1}-w_{0})/\sqrt{\lambda\gamma}$. Then, discretising as (2)
with discretisation step
$\varepsilon=\sqrt{\lambda\gamma}={\gamma}/(1-{\beta})$ leads to the momentum
gradient descent recursion $\mathrm{MGD}({\gamma},{\beta})$ with step size
$\gamma$, momentum parameter $\beta$, and initialisation $(w_{0},w_{1})$.

Proposition 1 motivates studying $\mathrm{MGF}(\lambda)$ as a continuous proxy
for $\mathrm{MGD}({\gamma},{\beta})$ assuming that the discretisation (2)
closely approximates the continuous path.

Discretisation Error Bounds. Unfortunately, applying known discretisation
error bounds to our setting leads to very pessimistic bounds. Indeed, for step
size $\gamma$ and momentum parameter $\beta$, consider the iterates $w_{k}$
from $\mathrm{MGD}({\gamma},{\beta})$ initialised at $(w_{0},w_{1})$. Now, let
$w(t)$ be the solution of $\mathrm{MGF}(\lambda)$ with
$\lambda=\gamma/(1-\beta)^{2}$ and the appropriate initialisation from
Proposition 1. Then, for a finite horizon $K>0$, classical discretisation
error bounds (see Kovachki and Stuart (2021), Theorem 4) lead to a
catastrophic

$\sup_{k\leq K}\|w_{k}-w(k\varepsilon)\|\leq\exp(CK)\varepsilon,$

where the constant $C$ depends on $\lambda$ and $F$. Such an exponential
dependence in the time horizon $K$ questions the suitability of momentum
gradient flow as a good proxy for momentum gradient descent. However,
empirically, the above bound appears excessively pessimistic (see Figure 1:
Left and Middle). The MGF and MGD trajectories behave similarly in various
settings, even with non-convex losses $F$ and relatively large step sizes
${\gamma}$ (see Appendix F for additional experiments).

Figure 1: (M)GD over a $2$D quadratic. Left and Middle: The (M)GD trajectories
closely follow the continuous trajectories of (M)GF as suggested by
Proposition 1. Right: MGD$(4\gamma,\beta^{2})$ follows the same trajectory as
MGD$(\gamma,\beta)$ but twice as fast as suggested by Corollary 1. In
contrast, GD$(4\gamma)$ runs four times faster than GD$(\gamma)$.

Intertwined Roles of $\gamma$ and $\beta$. When the discretisation accurately
follows the continuous path, Proposition 1 implies that the trajectory of
$\mathrm{MGD}({\gamma},{\beta})$ is solely determined by a single parameter
$\lambda=\gamma/(1-\beta)^{2}$, intertwining step size and momentum as
observed in Figures 1 and 2. Consequently, ${\gamma}$ and ${\beta}$ serve
interchangeable roles in influencing the trajectory of
$\mathrm{MGD}({\gamma},{\beta})$. Note that this single-parameter dependence
aligns with empirical results from Leclerc and Madry (2020) where
generalisation performance with large step sizes can be replicated with
momentum and smaller step sizes. Though the quantity $\gamma/(1-\beta)^{2}$
spontaneously appears in works studying MGD (Ghosh et al., 2023), to the best
of our knowledge, its natural presence was never clearly explained and
motivated.

MGD Acceleration Rule. Though all couples $(\gamma,\beta)$ with the same same
value of $\lambda$ yield the same trajectory, the iterates do not follow this
path at the same speed.

######  Corollary 1 (Acceleration rule).

Let $\mathrm{MGD}({\gamma},{\beta})$ initialised at
$w_{0}=w_{1}\in{\mathbb{R}}^{d}$ correspond to the discretisation of
$\mathrm{MGF}(\lambda)$ with discretisation step $\varepsilon$. Now, for
$\rho\in{\mathbb{R}}_{>0}$, consider the different parameter couple

$\hat{{\gamma}}=\rho^{2}\gamma\quad\text{and}\quad\hat{\beta}=1-\rho(1-\beta)\approx_{\beta\to
1}\beta^{\rho}.$

Then, since $\hat{{\gamma}}/(1-\hat{{\beta}})^{2}=\lambda$,
$\mathrm{MGD}(\hat{{\gamma}},\hat{{\beta}})$ initialised at $w_{0}=w_{1}$
becomes the discretisation of the same $\mathrm{MGF}(\lambda)$ but with
discretisation step $\hat{\varepsilon}=\rho\cdot\varepsilon$.

Following the notations of the previous corollary for an integer $\rho\geq 2$
and letting $w_{k}$ and $\hat{w}_{k}$ denote the iterates of
MGD$(\gamma,\beta)$ and MGD$(\hat{{\gamma}},\hat{\beta})$, respectively,
Corollary 1 implies that we expect $w_{\rho\cdot k}$ and $\hat{w}_{k}$ to be
close. This is in contrast with gradient descent, where scaling the step size
by a factor $\rho^{2}$ leads to a speedup of $\rho^{2}$. This acceleration
rule is illustrated in Figure 1 with $\rho=2$.

Optimisation Regimes. The link between $\lambda$, $\gamma$, and $\beta$
highlights several regimes:

  * •

$\beta$ large – the iterates converge arbitrarily slow. Taking $\beta$ close
to $1$ while keeping $\gamma$ constant leads to $\lambda\gg 1$. As explained
previously, a chain rule shows that $\tilde{w}_{t}=w_{\sqrt{\lambda}t}$
follows the ODE
$\ddot{\tilde{w}}_{t}+\lambda^{-1/2}\cdot\dot{\tilde{w}}_{t}+\nabla
F(\tilde{w}_{t})=0$. Consequently, the damping parameter $\lambda^{-1/2}$ goes
to $0$, and we expect the iterates to heavily oscillate and converge
arbitrarily slowly.

  * •

$\gamma$ small – the iterates follow GF. Taking $\gamma\to 0$ while keeping
$\beta$ fixed leads to $\lambda\ll 1$, and $\mathrm{MGF}(\lambda)$ boils down
to gradient flow. We expect the $\mathrm{MGD}({\gamma},{\beta})$ iterates to
be close to the discretisation of GF with discretisation step
$\varepsilon=\gamma/(1-\beta)$. That is, $\mathrm{MGD}({\gamma},{\beta})$ will
approximate GD with step size $\gamma/(1-\beta)$. Hence, MGD gains a speed-up
of $1/(1-\beta)$ over GD without a change of trajectory.

  * •

The “momentum” regime. In this regime, $\gamma$ and $\beta$ are such that
$\lambda$ is non-degenerate, and gradient flow cannot capture the trajectory
of $\mathrm{MGD}({\gamma},{\beta})$. Hence, $\beta$ has an impact on the
optimisation path, and the iterates can still converge in reasonable time.

Figure 2: Teacher-student framework with a fully-connected $1$-hidden layer
ReLU network. The level lines of the test loss after training with
$\mathrm{MGD}({\gamma},{\beta})$ correspond to values of $\gamma,\beta$ which
have a fixed value $\lambda=\gamma/(1-\beta)^{2}$, as predicted by Proposition
1.

##  3 Momentum Gradient Flow over Diagonal Linear Networks

Overparametrised Linear Regression. We consider a linear regression over $n$
samples $(x_{i},y_{i})_{i=1}^{n}$ with inputs $x_{i}$ living in
${\mathbb{R}}^{d}$ and scalar outputs $y_{i}\in{\mathbb{R}}$. We assume the
dimension $d$ to be larger than the number of samples $n$, in which case there
exists an infinite number of vectors ${\theta}^{\star}$ which perfectly fit
the dataset with $y_{i}=\langle{\theta}^{\star},x_{i}\rangle$ for all $1\leq
i\leq n$. We call these vectors interpolators and we denote by $\mathcal{S}$
the set of such vectors:
$\mathcal{S}=\\{{\theta}^{\star}\in{\mathbb{R}}^{d}:y_{i}=\langle{\theta}^{\star},x_{i}\rangle,\
\forall i\in[n]\\}$. Note that $\mathcal{S}$ is an affine space of dimension
at least $(d-n)$ equal to
${\theta}^{\star}+\mathrm{span}(x_{1},\dots,x_{n})^{\perp}$ for any
interpolator ${\theta}^{\star}$. We consider the quadratic loss:

${\mathcal{L}}({\theta})=\frac{1}{2n}\sum_{i=1}^{n}(y_{i}-\langle
x_{i},{\theta}\rangle)^{2}.$ (3)

MGF over Least Squares. A classical result found in Lemaire (1996) and
Gunasekar et al. (2018) shows that when initialised at $\theta_{0}$, gradient
flow over the quadratic loss (3) converges to the orthogonal projection of the
initialisation on $\mathcal{S}$:
$\arg\min_{{\theta}^{\star}\in\mathcal{S}}\|{\theta}^{\star}-\theta_{0}\|_{2}$.
This next proposition from Alvarez (2000) shows that momentum does not
fundamentally change the implicit bias.

######  Proposition 2 (Alvarez (2000)).

Initialised at $\theta_{0}$ with initial speed $\dot{\theta}_{0}$, momentum
gradient flow $\mathrm{MGF}(\lambda)$ over the least squares loss (3)
converges towards

$\displaystyle\operatorname*{argmin}_{\theta^{\star}\in\mathcal{S}}\|\theta^{\star}-(\theta_{0}+\lambda\dot{\theta}_{0})\|_{2}.$

$\mathrm{MGF}(\lambda)$ recovers the same solution as gradient flow but with
an effective initialisation $\theta_{0}+\lambda\dot{\theta}_{0}$ which takes
into account the drift along $\mathrm{span}(x_{1},\cdots,x_{n})^{\perp}$ due
to the initial speed $\dot{\theta}_{0}$. Note that in practice,
$\dot{\theta}_{0}$ is chosen equal to $0$, in which case the presence of
momentum has no effect on the recovered solution.

To better understand momentum’s effect on neural networks, we move beyond
simple linear parametrization.

$2$-Layer Diagonal Linear Network. We consider a toy neural network, which
corresponds to reparametrising the regression vector as $\theta=u\odot v$ for
weights $(u,v)\in{\mathbb{R}}^{2d}$. This parametrisation can be viewed as a
simple neural network $x\mapsto\langle u,\sigma(\mathrm{diag}(v)x)\rangle$,
where the output weights are $u$, the inner weights are the diagonal matrix
$\mathrm{diag}(v)$, and where the activation function $\sigma$ is the
identity. The loss function over the trainable weights
$w=(u,v)\in{\mathbb{R}}^{2d}$ now writes

$F(w)={\mathcal{L}}(u\odot v)=\frac{1}{2n}\sum_{i=1}^{n}(y_{i}-\langle
x_{i},u\odot v\rangle)^{2},$

where $\odot$ denotes the Hadamard product. Despite the simplicity of this
reparametrisation, the loss function $F$ is non-convex and challenging to
analyse.

Momentum Gradient Flow. We consider momentum gradient flow
$\mathrm{MGF}(\lambda)$ with parameter $\lambda\geq 0$ over the diagonal-
linear-network loss $F$:

$\displaystyle\begin{split}&\lambda\ddot{u}_{t}+\dot{u}_{t}+\nabla{\mathcal{L}}(\theta_{t})\odot
v_{t}=0\\\
&\lambda\ddot{v}_{t}+\dot{v}_{t}+\nabla{\mathcal{L}}(\theta_{t})\odot
u_{t}=0.\end{split}$ (4)

We initialise the flow with zero speed $\dot{u}_{0}=\dot{v}_{0}=0$, and apart
from requiring the quantity $|u_{0}^{2}-v_{0}^{2}|$ to have non-zero
coordinates222If initially $u_{i,0}=\pm v_{i,0}$ for some coordinate
$i\in[d]$, then $u_{i,t}=\pm v_{i,t},\ \forall t\geq 0$. Hence, imposing
$|u_{0}^{2}-v_{0}^{2}|\neq 0$ becomes equivalent to working with $2d$ distinct
weights. See Section C.3.2 for the full argument from uniqueness., we impose
no further constraints on the weight initialisations $(u_{0},v_{0})$. In what
follows, we often rely on the reparametrisation
$(w_{+,t},w_{-,t})\coloneqq(u_{t}+v_{t},u_{t}-v_{t})$ which makes our formulas
more succinct. We will also make use of the initialisation scale ${\alpha}$,
which we define as ${\alpha}\coloneqq\max(\lVert u_{0}\rVert_{\infty},\lVert
v_{0}\rVert_{\infty})$ and consider as a small quantity.

Balancedness. In our results, the balancedness of the weights plays a key
role. We recall its definition here.

######  Definition (Balancedness).

The balancedness333The absolute value in the definition must be understood
coordinate-wise. of the weights of the diagonal linear network corresponds to
the quantity ${\Delta}_{t}\coloneqq|u_{t}^{2}-v_{t}^{2}|\in{\mathbb{R}}_{\geq
0}^{d}$. We define ${\Delta}_{\infty}\coloneqq\lim_{t\to\infty}{\Delta}_{t}$
as the asymptotic balancedness.

Notice that with the above definition we simply adapted the classical notion
of balancedness for general linear neural networks (see Du et al., 2018; Arora
et al., 2019) to our toy setting. In the case of gradient flow, a simple
derivation shows that balancedness is a conserved quantity: i.e.,
${\Delta}_{t}={\Delta}_{0}$ for all $t\geq 0$. However, the evolution of
${\Delta}_{t}$ becomes more complicated as soon as ${\lambda}>0$, and our
findings emphasise that the asymptotic balancedness ${\Delta}_{\infty}$ plays
a crucial role in the generalisation properties of the recovered solution.

Experimental Details. In our numerical experiments, we explore the effects of
momentum in the noiseless sparse regression setting with uncentered data as in
(Nacson et al., 2022). Specifically, we choose
$(x_{i})_{i=1}^{n}\stackrel{{\scriptstyle\scriptscriptstyle\text{i.i.d.}}}{{\sim}}\mathcal{N}(\mu\mathbf{1},{\sigma}^{2}I_{d})$
and $y_{i}=\langle x_{i},{\theta}_{s}^{\star}\rangle$ for $i\in[n]$, where
${\theta}_{s}^{\star}$ is $s$-sparse with nonzero entries equal to
$1/\sqrt{s}$. The use of uncentered data is necessary in order to
experimentally observe a clear impact of momentum over the training trajectory
(see Figure 9 for experiments with centered data). We train a 2-layer diagonal
linear network with (M)GD and (M)GF with a uniform initialisation
$u_{0}={\alpha}\cdot\mathbf{1}$, $v_{0}=0$, where ${\alpha}=0.01$. For the
plots presented in the main part of our paper, we fixed $(n,d,s)=(20,30,5)$,
$(\mu,{\sigma})=(1,1)$. We show results averaged over 5 replications. We refer
the reader to Appendix F for additional experiments where we vary the
parameters of the data distribution (e.g., centered data), change the
architecture of the trained model, and give further details on the
implementation of the (M)GF simulation.

Notations. We let $X=(x_{1},\dots,x_{n})^{\top}\in{\mathbb{R}}^{n\times d}$
denote the feature matrix and $y=(y_{1},\dots,y_{n})\in{\mathbb{R}}^{n}$ – the
output vector. For a vector $z\in{\mathbb{R}}^{d}$ and a scalar function
$f:{\mathbb{R}}\to{\mathbb{R}}$, the action of $f$ on $z$ must be understood
element-wise: $f(z)\in{\mathbb{R}}^{d}$ represents the vector
$(f(z_{k}))_{k=1}^{d}$. Inequalities between vectors will also be interpreted
as holding coordinate-wise. Additionally, when we write $q_{\pm}$ for some
place-holder quantity $q$, we mean that we refer to both $q_{+}$ and $q_{-}$.
For example: $w_{\pm,t}=(u_{t}\pm v_{t})$. Finally, for a strictly convex
function $\Phi:{\mathbb{R}}^{d}\to{\mathbb{R}}$, which we call a potential,
the Bregman divergence is defined as the nonnegative quantity
$D_{\Phi}(\theta_{1},\theta_{2})=\Phi(\theta_{1})-\Phi(\theta_{2})-\langle\nabla\Phi(\theta_{2}),\theta_{1}-\theta_{2}\rangle,\
\forall\theta_{1},\theta_{2}\in{\mathbb{R}}^{d}$.

###  3.1 Implicit Bias of Gradient Flow

Before analysing the effect of momentum, we start by recalling the known
results for gradient flow on diagonal linear networks, which corresponds to
taking $\lambda=0$ in eq. 4. Woodworth et al. (2020) show that the predictors
$\theta_{t}=u_{t}\odot v_{t}$ converge towards an interpolator
${\theta}^{\texttt{GF}}$ uniquely defined by the following constrained
minimisation problem:

${\theta}^{\texttt{GF}}=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}D_{\psi_{{\Delta}_{0}}}({\theta}^{\star},{\theta}_{0}),$
(5)

where for ${\Delta}\in{\mathbb{R}}_{>0}^{d}$,
$\psi_{{\Delta}}:{\mathbb{R}}^{d}\to{\mathbb{R}}$ denotes the hyperbolic
entropy function (Ghai et al., 2020) at scale ${\Delta}$:

$\psi_{\Delta}({\theta})=\frac{1}{4}\sum_{i=1}^{d}\left(2{\theta}_{i}\operatorname{arcsinh}\left(\frac{2{\theta}_{i}}{{\Delta}_{i}}\right)-\sqrt{4{\theta}_{i}^{2}+{\Delta}_{i}^{2}}+{\Delta}_{i}\right),$

(6)

and $D_{\psi_{\Delta}}$ is the Bregman divergence. Note that through eq. 5,
$\theta^{\texttt{GF}}$ corresponds to the Bregman-projection of the
initialisation on the set of interpolators.

Effect of the Initialisation Scale. For a small initialisation scale
${\alpha}$, ${\theta}_{0}=O({\alpha}^{2})$ becomes much smaller than any
interpolator ${\theta}^{\star}\in\mathcal{S}$. Hence,
$D_{\psi_{{\Delta}_{0}}}({\theta}^{\star},{\theta}_{0})$ roughly equals
$D_{\psi_{{\Delta}_{0}}}({\theta}^{\star},0)$, and eq. 5 should be thought of
as

$\displaystyle{\theta}^{\texttt{GF}}\approx\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\psi_{{\Delta}_{0}}({\theta}^{\star}).$
(7)

This last equation highlights the fact that the recovered solution simply
depends on the initial balancedness $\Delta_{0}$, making it a key quantity.
Importantly, the hyperbolic entropy is a convex function which interpolates
between the $\ell_{1}$ and $\ell_{2}$ norms as the magnitude of ${\Delta}_{0}$
goes from 0 to $+\infty$ (see Woodworth et al. (2020), Theorem 2). So, as
${\Delta}_{0}=O({\alpha}^{2})$ goes to 0, $\psi_{{\Delta}_{0}}$ becomes
asymptotically identical to the $\ell_{1}$-norm (see Appendix E). Hence, as
seen through eq. 7, a small initialisation scale ${\alpha}$ leads to the
recovery of a solution with a small $\ell_{1}$-norm, which facilitates sparse
recovery and explains why this setting is referred to as the “rich” or
“feature-learning” regime. On the other hand, larger initialisation scales
lead to the so-called “kernel” or “lazy” regime, where gradient flow selects
small-$\ell_{2}$-norm solutions. Overall, the smaller the initialisation
scale, the closer the retrieved solution will be to the
minimum-$\ell_{1}$-norm solution. We refer the reader to the work of Wind et
al. (2023) for precise recovery bounds. However, as noted in Even et al.
(2023), the picture remains incomplete if we do not take into account the
homogeneity of ${\Delta}_{0}$. Indeed, initialisations with entries of
different magnitudes can hinder the recovery of a sparse vector. However, in
our case, our experiments (for uncentered data) verify that the overall
magnitudes of ${\Delta}_{0}$ and ${\Delta}_{\infty}$ are sufficient to explain
the effects of momentum. We therefore put aside potential homogeneity
considerations.

###  3.2 Implicit Bias of Momentum Gradient Flow

We now move on to describe the impact of momentum on the solution recovered by
$\mathrm{MGF}(\lambda)$. Our work proceeds under the following assumption.

######  Assumption 1 (Boundedness).

The optimisation trajectory $(u_{t},v_{t})_{t\geq 0}$ of MGF (4) is bounded.

Unfortunately, even though 1 holds true in all our experiments, the
boundedness of the trajectory of a second-order gradient flow has only been
established under stronger assumption on the loss function (Alvarez, 2000;
Goudou and Munier, 2009; Apidopoulos et al., 2022). We defer further details
to Section C.1. Crucially, the boundedness assumption allows us to prove the
convergence of the iterates, and we let
$(u_{\infty},v_{\infty})\coloneqq\lim_{t\to\infty}(u_{t},v_{t})$. Our goal now
becomes to characterise the recovered predictor which we denote with
${\theta}^{\texttt{MGF}}\coloneqq u_{\infty}\odot v_{\infty}$. For our proofs,
we make the following additional assumption.

######  Assumption 2 (Balancedness).

The asymptotic balancedness ${\Delta}_{\infty}$ has non-zero coordinates:
${\Delta}_{\infty,i}>0$ for all $i\in[d]$.

Again, 2 holds true empirically in all our experiments, and in Proposition 3,
we prove that small values of ${\lambda}$ lead to nonzero asymptotic
balancedness. Positing 2 allows us to prove that the recovered solution
${\theta}^{\texttt{MGF}}$ interpolates the dataset.

####  3.2.1 General Characterisation of MGF Bias

In our main result for MGF, we prove that the iterates converge towards an
interpolator characterised as the solution of a constrained minimisation
problem which involves the hyperbolic entropy (6) scaled at the asymptotic
balancedness ${\Delta}_{\infty}$. Moreover, we derive an insightful
description of the asymptotic balancedness in terms of the full optimisation
trajectory which allows us to compare the generalisation properties of MGF and
GF for small values of ${\lambda}$. Before stating our main continuous-time
theorem, we define two integral quantities which appear in our results.

######  Lemma 1.

The following integral quantities ${\Omega}_{+}$ and ${\Omega}_{-}$ are well-
defined and finite:

$\displaystyle{\Omega}_{\pm}\coloneqq\int_{0}^{\infty}\mathrm{m.p.v.}\int_{0}^{t}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t$

where $\operatorname{sgn}(\cdot)$ denotes the sign function,
$w_{\pm,t}=u_{t}\pm v_{t}$, and $\mathrm{m.p.v.}$ denotes a modified Cauchy
principal value defined in Appendix A.

The fact that the weights $w_{\pm,t}$ can cross zero necessitates the use of
the modified Cauchy principal value since otherwise the integrals would
diverge. Now, for succinctness, let us introduce the integral quantities

$I_{\pm}\coloneqq{\Omega}_{\pm}+\Lambda_{\pm},$

where the terms $\Lambda_{\pm}$ vanish whenever the balancedness
${\Delta}_{t}$ remains strictly positive for all $t\in[0,\infty]$. The precise
form of $\Lambda_{\pm}$ is uninformative and can be found in Equation 19,
Section C.3.2. We now proceed to characterise the recovered solution
${\theta}^{\texttt{MGF}}$.

######  Theorem 1.

The solution ${\theta}^{\texttt{MGF}}$ of MGF (4) interpolates the dataset and
satisfies the following implicit regularisation:

$\displaystyle{\theta}^{\texttt{MGF}}=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\
D_{\psi_{{\Delta}_{\infty}}}({\theta}^{\star},\tilde{{\theta}}_{0}).$

In the above expression, $D_{\psi_{{\Delta}_{\infty}}}$ denotes the Bregman
divergence with potential $\psi_{{\Delta}_{\infty}}$, where the asymptotic
balancedness equals

$\displaystyle{\Delta}_{\infty}={\Delta}_{0}\odot\exp\big{(}-(I_{+}+I_{-})\big{)}$

and
$\tilde{{\theta}}_{0}=\frac{1}{4}(w_{+,0}^{2}\odot\exp\left(-2I_{+}\right)-w_{-,0}^{2}\odot\exp\left(-2I_{-}\right))$
denotes a perturbed initialisation term.

The proof of Theorem 1 appears in Section C.3 as well as explicit formulas for
${\Delta}_{\infty}$ and $\tilde{{\theta}}_{0}$. We explain the significance
and shed more light on the different parts of Theorem 1 below.

Perturbed Initialisation $\tilde{\theta}_{0}$. In all our experiments, we
observe that the perturbed initialisation $\tilde{\theta}_{0}$ remains
negligible in the sense that for any interpolator
${\theta}^{\star}\in\mathcal{S}$,
$\lVert\tilde{\theta}_{0}\rVert_{2}\ll\lVert{\theta}^{\star}\rVert_{2}$.
Moreover, in the next section, we prove that whenever the balancedness remains
nonzero during training, $\tilde{\theta}_{0}$ becomes smaller than
${\alpha}^{2}$, where ${\alpha}$ stands for the initialisation scale. Hence,
exactly for the same reasons as for gradient flow, the implicit regularisation
problem from Theorem 1 should be though of as

${\theta}^{\texttt{MGF}}\approx\operatorname*{argmin}_{\theta^{\star}\in
S}\psi_{\Delta_{\infty}}(\theta^{\star}).$ (8)

Section C.3.3 provides more details. Thus, the asymptotic balancedness
${\Delta}_{\infty}$ becomes the key quantity governing the properties of the
recovered solution.

Key Role of ${\Delta}_{\infty}$. If during optimisation the weights become
more balanced, i.e. ${\Delta}_{\infty}<{\Delta}_{0}$, then as discussed
previously, based on the properties of $\psi_{\Delta_{\infty}}$, the recovered
solution will enjoy better sparsity guarantees than the solution of gradient
flow. Figure 3 illustrates this point: the smaller the magnitude of
$\Delta_{\infty}$, the better the generalisation. Finally note that by eqs. 5
and 8, ${\theta}^{\texttt{MGF}}$ approximately equals the solution recovered
from gradient flow initialised at $u_{0}=\sqrt{{\Delta}_{\infty}},v_{0}=0$,
which we denote by ${\theta}_{\Delta_{\infty}}^{\texttt{GF}}$. We observe
$||{\theta}^{\texttt{MGF}}-{\theta}_{\Delta_{\infty}}^{\texttt{GF}}||_{2}/||{\theta}_{\Delta_{\infty}}^{\texttt{GF}}||_{2}<0.01$
in all our experiments, which validates the approximation in eq. 8.

Figure 3:  Test loss (in blue) and magnitude of balancedness (in red) at
convergence of $\mathrm{MGF}(\lambda)$ over a diagonal linear network in a
sparse regression setting with uncentered data. As predicted by Theorem 1, a
more balanced solution generalises better. The shaded zone corresponds to
values of $\lambda$ for which the balancedness never hits zero during training
and for which Corollary 2 therefore holds.

Path-Dependent Quantity. Unfortunately, the asymptotic balancedness depends on
the whole optimisation trajectory in a very intricate way, and we cannot
compare $\lVert{\Delta}_{\infty}\rVert$ and $\lVert{\Delta}_{0}\rVert$. Thus,
in general, we cannot meaningfully compare the recovered interpolators
${\theta}^{\texttt{MGF}}$ and ${\theta}^{\texttt{GF}}$. However, in the
following section we prove that with the additional assumption that the
balancedness remains nonzero, we have ${\Delta}_{\infty}<{\Delta}_{0}$.

###  3.3 Provable Benefits of Momentum for Small Values of $\lambda$

In this subsection, we prove that for small values of the momentum flow
parameter $\lambda$, the recovered solution becomes more balanced (and
therefore sparser) than the solution of gradient flow. As a starting point for
our argument, notice that if the balancedness
${\Delta}_{t}=|u_{t}^{2}-v_{t}^{2}|=|w_{+,t}w_{-,t}|$ remains strictly
positive throughout training, then the weights $w_{\pm,t}$ never change sign.
Hence, the integral quantities $\Lambda_{\pm}$ become 0, and
${\Omega}_{\pm}>0$. Thus, $I_{\pm}>0$, which combined with Theorem 1 implies
the following corollary.

######  Corollary 2.

For $\lambda>0$, if the balancedness ${\Delta}_{t}$ remains strictly positive
during training (i.e. $\Delta_{t}\neq 0$ for $t\in[0,+\infty]$), then the
perturbed initialisation satisfies $|\tilde{{\theta}}_{0}|<\alpha^{2}$ and

$\displaystyle{\Delta}_{\infty}$
$\displaystyle={\Delta}_{0}\odot\exp\Big{(}-\lambda\int_{0}^{\infty}\Big{(}\frac{\dot{w}_{+,t}}{w_{+,t}}\Big{)}^{2}+\Big{(}\frac{\dot{w}_{-,t}}{w_{-,t}}\Big{)}^{2}\mathrm{d}t\Big{)}.$

Importantly, ${\Delta}_{\infty}<{\Delta}_{0}$.

In words, the above corollary (proved in Section C.4) implies that if the
balancedness $\Delta_{t}$ does not hit zero during training, then (i) the
perturbation term $\tilde{{\theta}}_{0}$ is provably negligible, (ii) the
asymptotic balancedness is coordinate-wise smaller than initial balancedness
$\Delta_{0}$ which translates into a solution with better sparsity properties
than the gradient flow interpolator. This regime corresponds to the gray zone
in Figure 3. The following proposition proved in Appendix E demonstrates that
for small values of ${\lambda}$, the balancedness remains strictly positive.

######  Proposition 3.

For $\lambda\leq\frac{n}{\|y\|_{2}^{2}}\cdot(\min_{i\leq d}{\Delta}_{0,i})$,
the balancedness ${\Delta}_{t}$ never vanishes: ${\Delta}_{t}\neq 0,\ \forall
t\in[0,+\infty]$.

Hence, through Proposition 3 and Corollary 2, we show that small values of
$\lambda$ lead to solutions with better sparse recovery guarantees.

Limitations of Our Analysis. In Section C.3.2, we prove that ${\Delta}_{t}$
can vanish at most a finite number of times. Experimentally, ${\Delta}_{t}$
never hits $0$ for much larger values of ${\lambda}$ than
$\frac{n}{\|y\|_{2}^{2}}\cdot(\min_{i\leq d}{\Delta}_{0,i})$, making the bound
from Proposition 3 relatively loose. In Figure 3, we empirically observe an
interval $(0,\lambda_{max})$ in which MGF($\lambda$) outperforms GF in terms
of generalisation. Moreover, there exists an optimal value $\lambda^{\star}$
(roughly corresponding to the smallest ${\Delta}_{\infty}$) which brings about
the most improvement compared to gradient flow. Unfortunately, as observed
Figure 3, the balancedness vanishes for $\lambda=\lambda^{\star}$, and
therefore Corollary 2 does not cover the optimal value. Also note that
$(0,\lambda_{max})$ and $\lambda^{\star}$ depend on the data.

Behaviour of $\Delta_{\infty}$ for Small Values of $\lambda$. Unfortunately,
determining the precise effect of $\lambda$ on $\Delta_{\infty}$ is
challenging. Nonetheless, for small $\lambda$, we informally show in Section
C.5 that

$\displaystyle{\Delta}_{\infty}^{2}\underset{\lambda\to
0}{\approx}{\Delta}_{0}^{2}\odot\exp\big{(}-2\lambda\int_{0}^{\infty}\nabla{\mathcal{L}}(\theta_{s})^{2}\mathrm{d}s\big{)}.$

This approximate equivalence for small $\lambda$ echoes the implicit bias of
SGD (Even et al., 2023; Pesme et al., 2021), which involves a similar
formulation for the effective initialisation where the step size $\gamma$
appears instead of $\lambda$. Note that the above approximation suggests that
for small values of $\lambda$, ${\Delta}_{\infty}$ monotonically decreases
with $\lambda$ as experimentally confirmed by Figure 3.

###  3.4 Sketch of Proof

Implicit Bias through a Second-Order Time-Varying Mirror Flow. A natural way
of showing the implicit regularisation (5) of gradient flow on a 2-layer
diagonal linear network goes through proving that the predictors
$\theta^{\texttt{GF}}_{t}$ follow the mirror flow
$\mathrm{d}\nabla\psi_{\Delta_{0}}(\theta^{\texttt{GF}}_{t})=-\nabla{\mathcal{L}}(\theta^{\texttt{GF}}_{t})\mathrm{d}t$.
In our setting, we prove that the predictors $\theta^{\texttt{MGF}}_{t}$
follow a second-order time-varying mirror flow. Specifically, we define a
family of potentials $(\Phi_{t})_{t\geq 0}$ with
$\Phi_{t}(\theta)\coloneqq\psi_{{\Delta}_{t}}(\theta)-\langle\phi_{t},\theta\rangle$
where $\psi_{{\Delta}_{t}}$ corresponds to the hyperbolic entropy (6)
depending on the balancedness ${\Delta}_{t}$ and a perturbation function
$\phi_{t}$. We then prove the following proposition.

######  Proposition 4.

The predictors $\theta_{t}^{\texttt{MGF}}$ follow a momentum mirror flow with
time-varying potentials $\Phi_{t}$:

$\displaystyle\lambda\frac{\mathrm{d}^{2}\nabla\Phi_{t}(\theta_{t}^{\texttt{MGF}})}{\mathrm{d}t^{2}}+\frac{\mathrm{d}\nabla\Phi_{t}(\theta_{t}^{\texttt{MGF}})}{\mathrm{d}t}+\nabla{\mathcal{L}}(\theta_{t}^{\texttt{MGF}})=0.$

The implicit regularisation follows from integrating the ODE:
$\nabla\Phi_{\infty}(\theta^{\texttt{MGF}})=-\int_{0}^{\infty}\nabla{\mathcal{L}}(\theta_{t}^{\texttt{MGF}})\mathrm{d}t\in\texttt{span}(x_{1},\dots,x_{n})$
which exactly corresponds to the KKT conditions of the constrained
minimisation from Theorem 1. Assuming that $w_{\pm,t}$ do not change sign, the
proof of Proposition 4 comes naturally and relies on the writing
$w_{\pm,t}=\operatorname{sgn}(w_{\pm,0})\exp(\rho_{\pm,t})$. When the iterates
cross $0$, this reparametrisation does not hold anymore. The analysis can
still be carried out by decomposing ${\mathbb{R}}_{\geq 0}$ into intervals on
which the iterates have constant sign and appropriately sticking the intervals
using a modified Cauchy principal value.

Figure 4: (Non-stochastic) MGD over a diagonal linear network in a sparse
regression setting with uncentered data. As predicted by Proposition 1, the
three quantities at convergence only depend on the single parameter
$\lambda\coloneqq\gamma/(1-\beta)^{2}$. As predicted by Theorem 2, a more
balanced solution (center plot) leads to a solution with a smaller
$\ell_{1}$-norm (right plot), which in turn translates into better
generalisation (left plot). Finally, as predicted by Corollary 3, the
trajectories for which the iterates do not cross zero satisfy
$\Delta_{\infty}<\Delta_{0}$, where $\Delta_{0}$ (approximately) corresponds
to the asymptotic balancedness for $\beta=0$ and $\gamma=10^{-3}$.

##  4 Momentum SGD over Diagonal Linear Networks

In this section, we move from continuous to discrete time and focus on the
original $\mathrm{MGD}({\gamma},{\beta})$ recursion for which we can prove
similar but slightly weaker results than the ones for MGF. In fact, our
results hold for stochastic momentum gradient descent (SMGD) with any batch
size $B\in[n]$. For step size ${\gamma}>0$ and momentum parameter
${\beta}\in[0,1)$, the SMGD recursion writes as follows:

$\displaystyle\begin{split}&u_{k+1}=u_{k}-{\gamma}\nabla{\mathcal{L}}_{\mathcal{B}_{k}}({\theta}_{k})\odot
v_{k}+{\beta}(u_{k}-u_{k-1})\\\
&v_{k+1}=v_{k}-{\gamma}\nabla{\mathcal{L}}_{\mathcal{B}_{k}}({\theta}_{k})\odot
u_{k}+{\beta}(v_{k}-v_{k-1}),\end{split}$ (9)

where
$L_{\mathcal{B}_{k}}({\theta})\coloneqq\frac{1}{2B}\sum_{i\in\mathcal{B}_{k}}(y_{i}-\langle
u\odot v,x_{i}\rangle)^{2}$ corresponds to the partial loss over the batch
$\mathcal{B}_{k}\subset[n]$ of size $B$. The batches could be sampled with or
without replacement. As for continuous time, we let ${\theta}_{k}=u_{k}\odot
v_{k}$ correspond to the regression predictor. We initialise at $u_{1}=u_{0}$
and $v_{1}=v_{0}$, and we again consider the balancedness of the weights
${\Delta}_{k}\coloneqq|u_{k}^{2}-v_{k}^{2}|$ for $k\geq 0$, the reparametrised
iterates $w_{\pm,k}\coloneqq u_{k}\pm v_{k}$, and the initialisation scale
${\alpha}\coloneqq\max(\lVert u_{0}\rVert_{\infty},\lVert
v_{0}\rVert_{\infty})$. In contrast to our continuous-time prerequisites where
we only assumed boundedness of the optimisation trajectory, here we assume
that the iterates converge:

######  Assumption 3 (Convergence).

The iterates $(u_{k},v_{k})$ converge towards the limiting weights
$(u_{\infty},v_{\infty})$. We denote by $\theta^{\texttt{SMGD}}\coloneqq
u_{\infty}\odot v_{\infty}$ the recovered predictor.

As in continuous time, we again assume that the asymptotic balancedness is
nonzero.

######  Assumption 4 (Balancedness).

The asymptotic balancedness
$\Delta_{\infty}\coloneqq|u_{\infty}^{2}-v_{\infty}^{2}|$ has non-zero
coordinates.

Similar to Lemma 1, we define two discrete infinite sums which depend on the
entire trajectory and appear in our discrete-time result.

######  Lemma 2.

The following two sums $S_{+}$ and $S_{-}$ converge to finite vectors:

$S_{\pm}=\frac{1}{1-\beta}\sum_{k=1}^{\infty}\left[r\Big{(}\frac{w_{\pm,k+1}}{w_{\pm,k}}\Big{)}+{\beta}r\Big{(}\frac{w_{\pm,k}}{w_{\pm,k+1}}\Big{)}\right],$

where $r(z)=(z-1)-\ln(|z|)$ for $z\neq 0$.

Importantly, the function $r(z)$ from Lemma 2 is positive for $z>0$. Contrary
to the continuous-time case, in discrete time, the iterates $w_{\pm,k}$ never
exactly equal zero. Indeed, since $\nabla L$ is linear, we have that for all
$k\geq 0$, $w_{\pm,k}({\gamma},{\beta})$ is a polynomial in
$({\gamma},{\beta})$. Therefore, the set of pairs $({\gamma},{\beta})$ for
which there exists $k\geq 0$ such that $w_{\pm,k}({\gamma},{\beta})=0$ is a
negligible set in ${\mathbb{R}}^{2}$. The iterates therefore ‘jump’ over zero,
making the sums from Lemma 2 well-defined.

###  4.1 General Characterisation of SMGD Bias

The following theorem represents the discrete counterpart of Theorem 1 and
generalises (Even et al., 2023, Theorem 1) which considers SGD without
momentum.

######  Theorem 2.

The solution ${\theta}^{\texttt{SMGD}}$ of SMGD (9) interpolates the dataset
and satisfies the following implicit regularisation:

$\displaystyle{\theta}^{\texttt{SMGD}}=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\
D_{\psi_{{\Delta}_{\infty}}}({\theta}^{\star},\tilde{{\theta}}_{0}).$

In the above expression, $D_{\psi_{{\Delta}_{\infty}}}$ denotes the Bregman
divergence with potential $\psi_{{\Delta}_{\infty}}$, where the asymptotic
balancedness equals

${\Delta}_{\infty}={\Delta}_{0}\odot\exp\big{(}-(S_{+}+S_{-})\big{)}$

and
$\tilde{{\theta}}_{0}=\frac{1}{4}(w_{+,0}^{2}\odot\exp(-2S_{+}))-w_{-,0}^{2}\odot\exp(-2S_{-}))$
denotes a perturbed initialisation term.

Due to the strong similarities with Theorem 1, we proceed by making similar
comments. In our experiments, the norm of the perturbed initialisation
$\tilde{\theta}_{0}$ remains much smaller than that of any interpolator
$\theta^{\star}$. Hence, arguing as before, the implicit regularisation
problem from Theorem 2 should be though of as

${\theta}^{\texttt{SMGD}}\approx\operatorname*{argmin}_{\theta^{\star}\in
S}\psi_{\Delta_{\infty}}(\theta^{\star}).$ (10)

Again, the asymptotic balancedness ${\Delta}_{\infty}$ controls the
generalisation properties of the recovered solution. Thus, if
$\lVert{\Delta}_{\infty}({\gamma},{\beta})\rVert_{2}<\lVert{\Delta}_{\infty}({\gamma}^{\prime},{\beta}^{\prime})\rVert_{2}$,
we expect the interpolator ${\theta}^{\texttt{SMGD}}({\gamma},{\beta})$ to be
sparser than ${\theta}^{\texttt{SMGD}}({\gamma}^{\prime},{\beta}^{\prime})$.
Figure 4 illustrates this point: the smaller the magnitude of
$\Delta_{\infty}$ (center plot), the better the sparsity of the interpolator
(right plot), which translates into better generalisation (left plot).
Unfortunately, as for MGF, the asymptotic balancedness ${\Delta}_{\infty}$
depends on the whole optimisation trajectory in an intricate way, which
prevents us from extracting an insightful formula for ${\Delta}_{\infty}$ in
terms of ${\gamma}$ and ${\beta}$. However, Figure 4 indicates that
${\Delta}_{\infty}$ effectively depends on the single parameter
$\lambda=\gamma/(1-\beta)^{2}$. As in Figure 2, $\lambda$ again clearly
appears to be the relevant quantity which governs the performance of MGD, and
not $\gamma$ and $\beta$ considered individually. These empirical observations
support the idea that even for ‘practical’ step sizes ${\gamma}$ and momentum
parameters ${\beta}$, $\mathrm{MGD}({\gamma},{\beta})$ closely follows
$\mathrm{MGF}(\lambda)$.

Figure 4 also clearly shows that the asymptotic balancedness decreases as the
key quantity ${\lambda}$ increases over an interval $[0,\lambda^{\star}]$
where $\lambda^{\star}$ denotes the parameter inducing the best generalisation
performances. Then, for ${\lambda}$ above ${\lambda}^{\star}$, the magnitude
of ${\Delta}_{\infty}$ starts to grow and the sparsity of the solutions
deteriorates. We expect proving this phenomenon to be very challenging. Such a
proof would require a fine-grained analysis of the sums $S_{\pm}$, which
becomes already quite involved when $\beta=0$ as performed by Even et al.
(2023).

Now, similar to the continuous-time result, the following corollary shows that
if the iterates do not change sign, then the asymptotic balancedness becomes
smaller than the initial balancedness.

######  Corollary 3.

For $\gamma,\beta>0$, if the iterates $w_{\pm,k}=(u_{k}\pm v_{k})$ do not
change sign during training, then $|\tilde{\theta}_{0}|<\alpha^{2}$ and
${\Delta}_{\infty}<{\Delta}_{0}$.

The above corollary implies that the recovered solution
${\theta}^{\texttt{SMGD}}$ must perform at least as well as the gradient flow
interpolator ${\theta}^{\texttt{GF}}$. However, in contrast to the continuous
case and even though we believe it to be true, we were unable to prove that
the SMGD iterates do not change sign for small values of $\lambda$.

##  5 Conclusion

Considering an appropriate second-order differential equation which
discretises into MGD, we highlight the existence of a single key quantity
$\lambda=\gamma/(1-\beta)^{2}$ which fully determines the trajectory of MGF.
This continuous-time perspective also provides a simple acceleration rule and
insight into several relevant optimisation regimes. Then, focusing on
$2$-layer diagonal linear networks, we prove that the asymptotic balancedness
$\Delta_{\infty}$ solely governs the generalisation performances of MGF and
SMGD. We additionally prove that small values of $\lambda$ aid the recovery of
sparse MGF solutions. Future work should consider MGF/MGD optimisation on more
complex architectures and understand precisely the non-trivial effect of
$\lambda$ on the asymptotic balancedness $\Delta_{\infty}$.

Acknowledgements. The authors would like to thank Jana Vuckovic for
investigating various discretisation methods for MGF, Yoana Nakeva and Mathieu
Even for providing feedback on the original version of the paper, and Aditya
Varre for the helpful discussions which led to the proofs of Theorem 1 and
Theorem 2. This project was supported by the Swiss National Science Foundation
(grant number 212111).

## References

  * Alvarez (2000) Felipe Alvarez.  On the minimizing property of a second order dissipative system in hilbert spaces.  _SIAM Journal on Control and Optimization_ , 38(4):1102–1119, 2000.
  * Apidopoulos et al. (2022) Vassilis Apidopoulos, Nicolò Ginatta, and Silvia Villa.  Convergence rates for the heavy-ball continuous dynamics for non-convex optimization, under polyak–lojasiewicz condition.  _Journal of Global Optimization_ , 84(3):563–589, 2022.  ISSN 1573-2916.
  * Arora et al. (2019) Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo.  _Implicit Regularization in Deep Matrix Factorization_.  Curran Associates Inc., 2019.
  * Attouch et al. (2000) H. Attouch, X. Goudou, and P. Redont.  The heavy ball with friction method, i. the continuous dynamical system: Global exploration of the local minima of a real-valued function by asymptotic analysis of a dissipitive dynamical system.  _Communications in Contemporary Mathematics_ , 2(1):1–34, 2000.
  * Bauschke and Borwein (1997) Heinz G. Bauschke and Jonathan M. Borwein.  Legendre functions and the method of random bregman projections.  _Journal of Convex Analysis_ , 4:27–67, 1997.
  * Bauschke et al. (2017) Heinz H. Bauschke, Jérôme Bolte, and Marc Teboulle.  A descent lemma beyond lipschitz gradient continuity: First-order methods revisited and applications.  _Math. Oper. Res._ , 42:330–348, 2017.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Can et al. (2019) Bugra Can, Mert Gurbuzbalaban, and Lingjiong Zhu.  Accelerated linear convergence of stochastic momentum methods in wasserstein distances.  In _International Conference on Machine Learning_ , pages 891–901. PMLR, 2019.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_ , 2022.
  * Cutkosky and Mehta (2020) Ashok Cutkosky and Harsh Mehta.  Momentum improves normalized sgd.  In _International conference on machine learning_ , pages 2260–2268. PMLR, 2020.
  * Defazio (2020) Aaron Defazio.  Understanding the role of momentum in non-convex optimization: Practical insights from a lyapunov analysis.  _ArXiv_ , 2020.
  * Du et al. (2018) Simon S Du, Wei Hu, and Jason D Lee.  Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced.  In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 31. Curran Associates, Inc., 2018.
  * Even et al. (2023) Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion.  (s)gd over diagonal linear networks: Implicit regularisation, large stepsizes and edge of stability.  _NeurIPS 2023_ , 2023.
  * Flammarion and Bach (2015) Nicolas Flammarion and Francis Bach.  From averaging to acceleration, there is only a step-size.  In _Conference on Learning Theory_ , pages 658–695. PMLR, 2015.
  * Ghadimi et al. (2015) Euhanna Ghadimi, Hamid Reza Feyzmahdavian, and Mikael Johansson.  Global convergence of the heavy-ball method for convex optimization.  In _2015 European control conference (ECC)_ , pages 310–315. IEEE, 2015.
  * Ghai et al. (2020) 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, 2020.
  * Ghosh et al. (2023) Avrajit Ghosh, He Lyu, Xitong Zhang, and Rongrong Wang.  Implicit regularization in heavy-ball momentum accelerated stochastic gradient descent.  In _The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023_. OpenReview.net, 2023.
  * Goudou and Munier (2009) X. Goudou and J. Munier.  The gradient and heavy ball with friction dynamical systems: the quasiconvex case.  _Mathematical Programming_ , 116(1):173–191, 2009.
  * Gunasekar et al. (2018) Suriya Gunasekar, Jason D. Lee, Daniel Soudry, and Nathan Srebro.  Characterizing implicit bias in terms of optimization geometry.  In Jennifer G. Dy and Andreas Krause, editors, _Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018_ , volume 80 of _Proceedings of Machine Learning Research_ , pages 1827–1836. PMLR, 2018.
  * HaoChen et al. (2021) 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.
  * Haraux and Jendoubi (1998) A. Haraux and M.A. Jendoubi.  Convergence of solutions of second-order gradient-like systems with analytic nonlinearities.  _Journal of Differential Equations_ , 144(2):313–320, 1998.
  * Hoffmann et al. (2022) Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al.  Training compute-optimal large language models.  _arXiv preprint arXiv:2203.15556_ , 2022.
  * Jelassi and Li (2022) Samy Jelassi and Yuanzhi Li.  Towards understanding how momentum improves generalization in deep learning.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 9965–10040. PMLR, 2022.
  * Jin et al. (2018) Chi Jin, Praneeth Netrapalli, and Michael I Jordan.  Accelerated gradient descent escapes saddle points faster than gradient descent.  In _Conference On Learning Theory_ , pages 1042–1085. PMLR, 2018.
  * Kidambi et al. (2018) Rahul Kidambi, Praneeth Netrapalli, Prateek Jain, and Sham Kakade.  On the insufficiency of existing momentum schemes for stochastic optimization.  In _2018 Information Theory and Applications Workshop (ITA)_ , pages 1–9. IEEE, 2018.
  * Kovachki and Stuart (2021) Nikola B. Kovachki and Andrew M. Stuart.  Continuous time analysis of momentum methods.  _J. Mach. Learn. Res._ , 22(1), 2021.
  * Leclerc and Madry (2020) Guillaume Leclerc and Aleksander Madry.  The two regimes of deep network training.  _arXiv preprint arXiv:2002.10376_ , 2020.
  * Lemaire (1996) B Lemaire.  An asymptotical variational principle associated with the steepest descent method for a convex function.  _Journal of Convex Analysis_ , 3(1):63–70, 1996.
  * Li et al. (2022) Zhiyuan Li, Tianhao Wang, Jason D. Lee, and Sanjeev Arora.  Implicit bias of gradient descent on reparametrized models: On equivalence to mirror descent.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, _Advances in Neural Information Processing Systems 35 - 36th Conference on Neural Information Processing Systems, NeurIPS 2022_ , Advances in Neural Information Processing Systems. Neural information processing systems foundation, 2022.
  * Liu et al. (2020) Yanli Liu, Yuan Gao, and Wotao Yin.  An improved analysis of stochastic gradient descent with momentum.  _Advances in Neural Information Processing Systems_ , 33:18261–18271, 2020.
  * Mai and Johansson (2020) Vien Mai and Mikael Johansson.  Convergence of a stochastic gradient method with momentum for non-smooth non-convex optimization.  In _International conference on machine learning_ , pages 6630–6639. PMLR, 2020.
  * Moucer et al. (2023) Céline Moucer, Adrien Taylor, and Francis Bach.  A systematic approach to lyapunov analyses of continuous-time models in convex optimization.  _SIAM Journal on Optimization_ , 33(3):1558–1586, 2023.
  * Nacson et al. (2022) Mor Shpigel Nacson, Kavya Ravichandran, Nathan Srebro, and Daniel Soudry.  Implicit bias of the step size in linear diagonal neural networks.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 16270–16295. PMLR, 2022.
  * Nemirovsky and Yudin (1979) A. S. Nemirovsky and D. B. Yudin.  _Problem Complexity and Method Efficiency in Optimization_.  John Wiley & Sons, New York, 1979.
  * Orvieto et al. (2020) Antonio Orvieto, Jonas Kohler, and Aurelien Lucchi.  The role of memory in stochastic optimization.  In _Uncertainty in Artificial Intelligence_ , pages 356–366. PMLR, 2020.
  * Pesme et al. (2021) Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion.  Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity.  _Advances in Neural Information Processing Systems_ , 34:29218–29230, 2021.
  * Pillaud-Vivien et al. (2022) 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.
  * Polyak and Shcherbakov (2017) Boris Polyak and Pavel Shcherbakov.  Lyapunov functions: An optimization theory perspective.  _IFAC-PapersOnLine_ , 50(1):7456–7461, 2017.
  * Polyak (1964) B.T. Polyak.  Some methods of speeding up the convergence of iteration methods.  _USSR Computational Mathematics and Mathematical Physics_ , 4(5):1–17, 1964.
  * Radford et al. (2021) Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pages 8748–8763. PMLR, 2021.
  * Sanz Serna and Zygalakis (2021) Jesús María Sanz Serna and Konstantinos C Zygalakis.  The connections between lyapunov functions for some optimization algorithms and differential equations.  _SIAM Journal on Numerical Analysis_ , 59(3):1542–1565, 2021.
  * Sebbouh et al. (2021) Othmane Sebbouh, Robert M Gower, and Aaron Defazio.  Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball.  In _Conference on Learning Theory_ , pages 3935–3971. PMLR, 2021.
  * Shi et al. (2021) Bin Shi, Simon S Du, Michael I Jordan, and Weijie J Su.  Understanding the acceleration phenomenon via high-resolution differential equations.  _Mathematical Programming_ , pages 1–70, 2021.
  * Su et al. (2014) Weijie Su, Stephen Boyd, and Emmanuel Candes.  A differential equation for modeling nesterov’s accelerated gradient method: theory and insights.  _Advances in neural information processing systems_ , 27, 2014.
  * Sun et al. (2019) Tao Sun, Dongsheng Li, Zhe Quan, Hao Jiang, Shengguo Li, and Yong Dou.  Heavy-ball algorithms always escape saddle points.  In _Proceedings of the 28th International Joint Conference on Artificial Intelligence_ , IJCAI’19, page 3520–3526. AAAI Press, 2019.
  * Sutskever et al. (2013) Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton.  On the importance of initialization and momentum in deep learning.  In Sanjoy Dasgupta and David McAllester, editors, _Proceedings of the 30th International Conference on Machine Learning_ , volume 28 of _Proceedings of Machine Learning Research_ , pages 1139–1147, Atlanta, Georgia, USA, 2013. PMLR.
  * Vaškevičius et al. (2019) Tomas Vaškevičius, Varun Kanade, and Patrick Rebeschini.  Implicit regularization for optimal sparse recovery.  In _Proceedings of the 33rd International Conference on Neural Information Processing Systems_ , 2019.
  * Wang et al. (2023) Li Wang, Zhiguo Fu, Yingcong Zhou, and Zili Yan.  The implicit regularization of momentum gradient descent in overparametrized models.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 37(8):10149–10156, 2023.
  * Wibisono et al. (2016) Andre Wibisono, Ashia C Wilson, and Michael I Jordan.  A variational perspective on accelerated methods in optimization.  _proceedings of the National Academy of Sciences_ , 113(47):E7351–E7358, 2016.
  * Wilson et al. (2021) Ashia C Wilson, Ben Recht, and Michael I Jordan.  A lyapunov analysis of accelerated methods in optimization.  _The Journal of Machine Learning Research_ , 22(1):5040–5073, 2021.
  * Wind et al. (2023) Johan S Wind, Vegard Antun, and Anders C Hansen.  Implicit regularization in ai meets generalized hardness of approximation in optimization–sharp results for diagonal linear networks.  _arXiv preprint arXiv:2307.07410_ , 2023.
  * Woodworth et al. (2020) 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 Jacob Abernethy and Shivani Agarwal, editors, _Proceedings of Thirty Third Conference on Learning Theory_ , volume 125 of _Proceedings of Machine Learning Research_ , pages 3635–3673. PMLR, 2020.
  * Zhang et al. (2017) Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning requires rethinking generalization.  In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net, 2017.

## Organisation of the Appendix.

The appendix is organised as follows:

  * •

In Appendix A, we introduce additional notation and provide further comments
on the discretisation methods.

  * •

In Appendix B, we present a useful reparametrisation of our problem.

  * •

In Appendix C, we offer proofs for our continuous-time results, specifically
Theorem 1, Corollary 2, and Proposition 4.

  * •

In Appendix D, we detail the proofs for our discrete-time results, namely
Lemma 2, Theorem 2 and Corollary 3.

  * •

In Appendix E, we introduce technical lemmas necessary for our main results
and prove Proposition 3.

  * •

In Appendix F, we provide more details on the main-paper experiments and
showcase further experimental results.

##  Appendix A Additional Notations and Comments on Discretisation Methods

Vector Operations. Moving forward, all arithmetic operations and real-valued
functions will be considered as being applied coordinate-wise. In other words,
if $a$ and $b$ are vectors in ${\mathbb{R}}^{d}$ and $p,q\in{\mathbb{Q}}$,
then $a^{p}b^{q}\in{\mathbb{R}}^{d}$ will be used as a shorthand for the
vector with entries $\\{a_{i}^{p}b_{i}^{q}\\}_{i=1}^{d}$. And for any
$f:{\mathbb{R}}\to{\mathbb{R}}$, $f(a)$ will represent the vector with entries
$\\{f(a_{i})\\}_{i=1}^{d}$. Inequalities between vectors will also be
interpreted as holding coordinate-wise.

Mirror Maps. Various definitions of a mirror map
$\Phi:{\mathbb{R}}^{d}\to(-\infty,+\infty]$ exist in the optimization
literature (see Nemirovsky and Yudin, 1979; Li et al., 2022), and a common one
coincides with the concept of a Legendre function (see Bauschke et al., 2017;
Bauschke and Borwein, 1997). In our proofs, we do not deal with extended real-
valued functions, and the term mirror map is applied to $C^{\infty}$-smooth
strictly convex functions with coercive gradients. In particular, our mirror
maps are of Legendre type.

For such a mirror map $\Phi:{\mathbb{R}}^{d}\to{\mathbb{R}}$, we define the
Bregman divergence $D_{\Phi}(\theta_{1},\theta_{2})$ for
$\theta_{1},\theta_{2}\in{\mathbb{R}}^{d}$ as

$D_{\Phi}(\theta_{1},\theta_{2})=\Phi(\theta_{1})-\Phi(\theta_{2})-\langle\nabla\Phi(\theta_{2}),\theta_{1}-\theta_{2}\rangle.$

Notice that due to the strict convexity of $\Phi$,
$D_{\Phi}(\theta_{1},\theta_{2})>0$ whenever $\theta_{1}\neq\theta_{2}$.

Modified Cauchy Principal Value. Let $f:{\mathbb{R}}_{\geq
0}\to[-\infty,+\infty]$ be an extended real-valued function with a finite set
of poles $\mathcal{T}=\\{T_{1},T_{2},\dots,T_{N}\\}$ (i.e. points
$t\in{\mathbb{R}}_{\geq 0}$ at which $f(t)=\pm\infty$) such that $f$ is
continuous on ${\mathbb{R}}_{\geq 0}\setminus\mathcal{T}$. Let
$0<T_{1}<\dots<T_{N}$. Let $T\in\mathcal{T}$ and let ${\varepsilon}>0$ be
small enough such that
$(T-{\varepsilon},T+{\varepsilon})\cap\mathcal{T}=\\{T\\}$. Recall that,
provided the limit below exists, the Cauchy principal value
$\mathrm{p.v.}\int_{T-{\varepsilon}}^{T+{\varepsilon}}f(t)\mathrm{d}t$ is
defined as

$\mathrm{p.v.}\int_{T-{\varepsilon}}^{T+{\varepsilon}}f(t)\mathrm{d}t\coloneqq\lim_{{\delta}\to
0}\left[\int_{T-{\varepsilon}}^{T-{\delta}}f(t)\mathrm{d}t+\int_{T+{\delta}}^{T+{\varepsilon}}f(t)\mathrm{d}t\right].$

Now, let ${\varepsilon}_{m}>0$ be such that
$(T_{m}-{\varepsilon}_{m},T_{m}+{\varepsilon}_{m})\cap\mathcal{T}=\\{T_{m}\\}$
for $m\in[N]$. Moreover, let $T_{0}={\varepsilon}_{0}=0$ and
$T_{N+1}=+\infty$. Suppose $f$ has finite Cauchy principal values at all
poles. Then, for any $\tau\geq 0$ such that $\tau\notin\mathcal{T}$, we could
define $\mathrm{p.v.}\int_{0}^{\tau}f(t)\mathrm{d}t$ as

$\mathrm{p.v.}\int_{0}^{\tau}f(t)\mathrm{d}t\coloneqq\sum_{m:T_{m+1}<\tau}\left[\mathrm{p.v.}\int_{T_{m}-{\varepsilon}_{m}}^{T_{m}+{\varepsilon}_{m}}f(t)\mathrm{d}t+\int_{T_{m}+{\varepsilon}_{m}}^{T_{m+1}-{\varepsilon}_{m+1}}f(t)\mathrm{d}t\right]+\mathrm{p.v.}\int_{T_{k}-{\varepsilon}_{k}}^{T_{k}+{\varepsilon}_{k}}f(t)\mathrm{d}t+\int_{T_{k}+{\varepsilon}_{k}}^{\tau}f(t)\mathrm{d}t,$

where $T_{k}<\tau<T_{k+1}$.

For our proofs of Lemma 1 and Theorem 1, we require a modification to the
Cauchy principal value. For the aforementioned function $f$ with the described
properties and for $T\in\mathcal{T},\ {\varepsilon}>0$ such that
$(T-{\varepsilon},T+{\varepsilon})\cap\mathcal{T}=\\{T\\}$, we define the
modified principal value
$\mathrm{m.p.v.}\int_{T-{\varepsilon}}^{T+{\varepsilon}}f(t)\mathrm{d}t$ as

$\mathrm{m.p.v.}\int_{T-{\varepsilon}}^{T+{\varepsilon}}f(t)\mathrm{d}t\coloneqq\lim_{{\delta}\to
0}\left[\int_{T-{\varepsilon}}^{T-{\delta}}f(t)\mathrm{d}t\cdot
e^{\frac{{\delta}}{{\lambda}}}+\int_{T+{\delta}}^{T+{\varepsilon}}f(t)\mathrm{d}t\cdot
e^{-\frac{{\delta}}{{\lambda}}}\right],$ (11)

where ${\lambda}$ denotes our familiar MGF parameter. We also extend the
$\mathrm{m.p.v.}$ definition to integrals $\int_{0}^{\tau}f(t)\mathrm{d}t$ for
arbitrary $\tau\geq 0$ by mimicking the Cauchy-principal-value construction:

$\mathrm{m.p.v.}\int_{0}^{\tau}f(t)\mathrm{d}t\coloneqq\sum_{m:T_{m+1}<\tau}\left[\mathrm{m.p.v.}\int_{T_{m}-{\varepsilon}_{m}}^{T_{m}+{\varepsilon}_{m}}f(t)\mathrm{d}t+\int_{T_{m}+{\varepsilon}_{m}}^{T_{m+1}-{\varepsilon}_{m+1}}f(t)\mathrm{d}t\right]+\mathrm{m.p.v.}\int_{T_{k}-{\varepsilon}_{k}}^{T_{k}+{\varepsilon}_{k}}f(t)\mathrm{d}t+\int_{T_{k}+{\varepsilon}_{k}}^{\tau}f(t)\mathrm{d}t,$

where $T_{k}<\tau<T_{k+1}$. Note that the above definition implies that
whenever $f$ has no poles on an interval $(a,b)\subset{\mathbb{R}}_{\geq 0}$,
then

$\mathrm{m.p.v.}\int_{a}^{b}f(t)\mathrm{d}t=\int_{a}^{b}f(t)\mathrm{d}t.$

Additional Comments on the Discretisation of $\mathrm{MGF}(\lambda)$.
Following our discussion from Section 2, we want to point out that that there
are other ways of discretising

$\displaystyle\lambda\ddot{w}_{t}+\dot{w}_{t}+\nabla F(w_{t})=0.$

Indeed, instead of discretising as (2) in the main paper

$\displaystyle\lambda\
\frac{w_{k+1}-2w_{k}+w_{k-1}}{\varepsilon^{2}}+\frac{w_{k}-w_{k-1}}{\varepsilon}+\nabla
F(w_{k})=0,$

one could also consider a central first-order difference:

$\displaystyle\lambda\
\frac{w_{k+1}-2w_{k}+w_{k-1}}{\varepsilon^{2}}+\frac{w_{k+1}-w_{k-1}}{2\varepsilon}+\nabla
F(w_{k})=0.$

Rearranging, this leads to

$\displaystyle
w_{k+1}=w_{k}-\frac{\varepsilon^{2}}{{\lambda}(1+\frac{\varepsilon}{2\lambda})}\nabla
F(w_{k})+\frac{1-\frac{\varepsilon}{2\lambda}}{1+\frac{\varepsilon}{2\lambda}}(w_{k}-w_{k-1}),$

which corresponds to momentum with
$\gamma=\frac{\varepsilon^{2}}{{\lambda}(1+\frac{\varepsilon}{2\lambda})}$ and
$\beta=\frac{1-\frac{\varepsilon}{2\lambda}}{1+\frac{\varepsilon}{2\lambda}}$.
Solving for $\varepsilon$ and $\lambda$, we get

$\displaystyle\lambda=\frac{(1+\beta)\gamma}{2(1-\beta)^{2}}\qquad\text{and}\qquad\varepsilon=\frac{\gamma}{1-\beta}.$

Hence, we obtain the same discretisation step $\varepsilon$ as in Proposition
1 and a slightly different expression for $\lambda$. However, note that the
two versions of $\lambda$ become indistinguishable for large values of $\beta$
since $\frac{1+\beta}{2}\to_{\beta\to 1}1$. Experimentally, running
$\mathrm{MGF}(\lambda)$ with the two different values for $\lambda$ leads to
similar results. Thus, the discretisation scheme from the main paper was
chosen due to the more concise definition of $\lambda$ in this case.

##  Appendix B $(w_{+},w_{-})$-Reparametrisation

MGF Reparametrisation. We recall that we consider momentum gradient flow
$\mathrm{MGF}(\lambda)$ with parameter $\lambda>0$ over the diagonal-linear-
network loss $F((u,v)))={\mathcal{L}}(u\odot v)$:

$\displaystyle\begin{split}&\lambda\ddot{u}_{t}+\dot{u}_{t}+\nabla
L(\theta_{t})\odot v_{t}=0;\\\ &\lambda\ddot{v}_{t}+\dot{v}_{t}+\nabla
L(\theta_{t})\odot u_{t}=0.\end{split}$

For proof-writing convenience, we consider the simple reparametrisation
outlined below.

In order to eliminate the cross-dependencies in $(u,v)$ in the above
equations, it is natural to consider the quantities $(w_{+,t},w_{-,t})$ where
$w_{\pm,t}=u_{t}\pm v_{t}$ for $t\geq 0$. Hence, we get the following
reparametrised ODE:

$\left\\{\begin{aligned}
&{\lambda}\ddot{w}_{\pm,t}+\dot{w}_{\pm,t}\pm\nabla{\mathcal{L}}({\theta}_{t})\odot
w_{\pm,t}=0;\\\ &w_{\pm,0}=u_{0}\pm
v_{0},\enspace\dot{w}_{\pm,0}=0.\end{aligned}\right.$ (12)

Notice that with these new quantities, we have

$\displaystyle{\theta}_{t}=\frac{w_{+,t}^{2}-w_{-,t}^{2}}{4}\quad\text{and}\quad{\Delta}_{t}=|w_{+,t}w_{-,t}|.$

MGD Reparametrisation. For the discrete-time setting, we follow the same
reparametrisation from the MGD recursion:

$\displaystyle\begin{split}&u_{k+1}=u_{k}-{\gamma}\nabla L({\theta}_{k})\odot
v_{k}+{\beta}(u_{k}-u_{k-1});\\\ &v_{k+1}=v_{k}-{\gamma}\nabla
L({\theta}_{k})\odot u_{k}+{\beta}(v_{k}-v_{k-1}).\end{split}$

We let $w_{\pm,k}=u_{k}\pm v_{k}$ for $k\geq 0$. Then, for $k\geq 1$, the
equations above transform into

$\left\\{\begin{aligned}
&w_{\pm,k+1}=w_{\pm,k}\mp{\gamma}\nabla{\mathcal{L}}({\theta}_{k})\odot
w_{\pm,k}+{\beta}(w_{\pm,k}-w_{\pm,k-1});\\\ &w_{\pm,1}=w_{\pm,0}=u_{0}\pm
v_{0}.\end{aligned}\right.$ (13)

Again, with the newly defined quantities, we have

$\displaystyle{\theta}_{k}=\frac{w_{+,k}^{2}-w_{-,k}^{2}}{4}\quad\text{and}\quad{\Delta}_{k}=|w_{+,k}w_{-,k}|.$

##  Appendix C Continuous-Time Theorems

###  C.1 Convergence of Momentum Gradient Flow

Momentum gradient flow (with ${\lambda}>0$),

${\lambda}\ddot{w_{t}}+\dot{w_{t}}+\nabla F(w_{t})=0,$

also known in the optimisation literature as the heavy-ball with friction ODE
or the heavy-ball dynamical system with constant damping coefficient, has been
the object of extensive mathematical study over the years (Haraux and
Jendoubi, 1998; Attouch et al., 2000; Alvarez, 2000; Goudou and Munier, 2009;
Polyak and Shcherbakov, 2017; Apidopoulos et al., 2022). If we abstract away
from the diagonal linear network setting and consider an unspecified loss
$F\in C^{1}({\mathbb{R}}^{D},{\mathbb{R}}_{\geq 0})$ with locally Lipschitz
gradient, we can still identify a useful Lyapunov function, which perhaps
motivated the study of the ODE in the first place. The function in question
happens to be the energy of the system

$E_{t}=F(w_{t})+\frac{{\lambda}}{2}\lVert\dot{w}_{t}\rVert_{2}^{2},$ (14)

whose nonpositive time-derivative
$\dot{E}_{t}=-\lVert\dot{w}_{t}\rVert_{2}^{2}$ allows us to prove the global
existence and uniqueness of a solution to MGF [Attouch et al. (2000), Theorem
3.1] in this more general setting. We note that by an easy inductive argument,
when the function $F$ is $C^{k}$-smooth, the MGF solution $w_{t}$ is
$C^{k+1}$-smooth. Hence, in our setting where the diagonal-neural-network loss
$F$ is $C^{\infty}$-smooth, the learning trajectory $w_{t}$ is also
$C^{\infty}$-smooth.

Convergence under 1.

Under the assumption of a bounded trajectory – $w_{t}\in
L^{\infty}(0,\infty)$, one can prove the following convergences (Attouch et
al., 2000):

$\lim_{t\to\infty}\dot{w}_{t}=\lim_{t\to\infty}\nabla F(w_{t})=0.$

However, even when bounded, the iterates $w_{t}$ need not converge as
demonstrated by the coercive function from Section 4.3 in (Attouch et al.,
2000). Nevertheless, when the loss $F$ is also analytic, as in the case of
diagonal linear networks, assuming boundedness, one can further prove iterate
convergence $\lim_{t\to\infty}w_{t}=w_{\infty}$ [Haraux and Jendoubi (1998)].

Unfortunately, without assuming boundedness, iterate convergence has been
established only in the cases of convex loss [Alvarez (2000)], quasiconvex
loss [Goudou and Munier (2009)], and loss satisfying the Polyak-Lojasiewicz
inequality [Apidopoulos et al. (2022)]. Thus, the square loss for a diagonal
linear network (and neural networks in general) falls out of the scope of
these few favorable cases due to non-convexity and an abundance of local and
global minima. For that reason, we posit Assumption 1, which holds true
empirically in all our experiments on diagonal linear networks.

Convergence to 0 Loss under 2.

Let us now go back to the specific case of diagonal linear networks where the
loss is given by $F(w)={\mathcal{L}}(u\odot v)$ for $w=(u,v)$. Notice that
from the discussion above, if we assume boundedness of the trajectory, we have

$\lim_{t\to\infty}\nabla F(w_{t})=(\nabla{\mathcal{L}}(\theta_{\infty})\odot
v_{\infty},\nabla{\mathcal{L}}({\theta}_{\infty})\odot u_{\infty})=0.$

Therefore, since
$\nabla{\mathcal{L}}({\theta}_{\infty})\odot{\Delta}_{\infty}=0$, if the
balancedness at infinity ${\Delta}_{\infty}$ has nonzero coordinates, we can
conclude that $\nabla{\mathcal{L}}({\theta}_{\infty})=0$. Recalling that
${\mathcal{L}}$ is convex, we get that ${\mathcal{L}}({\theta}_{\infty})=0$.
Hence, ${\theta}_{\infty}$ interpolates the dataset.

###  C.2 Proof of Time-Varying Momentum Mirror Flow

In our discussion in Section C.1, we saw that assuming

  1. 1)

iterate boundedness: $u_{t},v_{t}\in L^{\infty}(0,\infty)$, and

  2. 2)

nonzero balancedness at infinity: ${\Delta}_{\infty,i}\neq 0,\enspace\forall
i\in[d]$,

we can prove that MGF over a diagonal linear network (4) converges to an
interpolator ${\theta}_{\infty}$.444Note that we also refer to
${\theta}_{\infty}$ as ${\theta}^{\texttt{MGF}}$. Before we jump into the
proof of Proposition 4, we need to establish the following lemma.

######  Lemma 3.

Assuming that $u_{t},v_{t}\in L^{\infty}(0,\infty)$ and
${\Delta}_{\infty,i}\neq 0,\enspace\forall i\in[d]$, the following integral
limit exists:

$\lim_{t\to\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})ds=\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t})\mathrm{d}t.$

Consequently,

$\lim_{t\to\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{{\lambda}}}ds=0.$

###### Proof.

Let us consider the $(w_{+},w_{-})$-reparametrisation of MGF (4) given by
Equation 12:

${\lambda}\ddot{w}_{\pm,t}+\dot{w}_{\pm,t}\pm\nabla{\mathcal{L}}({\theta}_{t})\odot
w_{\pm,t}=0.$

Since we assumed that ${\Delta}_{\infty}$ has nonzero coordinates, there
exists $T\geq 0$ such that for all $t\geq T$, $w_{\pm,t}$ have nonzero
coordinates. Hence, for $t\geq T$, we can safely divide by $w_{\pm,t}$ to
obtain

${\lambda}\frac{\mathrm{d}^{2}\ln|w_{\pm,t}|}{\mathrm{d}t^{2}}+\frac{\mathrm{d}\ln|w_{\pm,t}|}{\mathrm{d}t}\pm\nabla{\mathcal{L}}({\theta}_{t})+{\lambda}\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}=0.$

Let us notice a couple of things. First, as we discussed in Section C.1, the
boundedness of the iterates forces $w_{\pm,t}$ to converge to some vectors
$w_{\pm,\infty}$ with nonzero coordinates since we assumed the coordinates of
${\Delta}_{\infty}=w_{+,\infty}w_{-,\infty}$ are nonzero. Hence,

$\lVert\frac{\dot{w}_{\pm,t}^{2}}{w_{\pm,t}^{2}}\rVert_{\infty}\leq\texttt{const}\cdot(\lVert\dot{u}_{t}^{2}\rVert_{\infty}+\lVert\dot{v}_{t}^{2}\rVert_{\infty}),$

where the RHS is integrable as we saw in the proof of Proposition 3. Second,
from the discussion in Section C.1, we know that
$\lim_{t\to\infty}\dot{w}_{\pm,t}=0$, so
$\lim_{t\to\infty}\frac{\mathrm{d}\ln|w_{\pm,t}|}{\mathrm{d}t}=0$.

Now, for $t\geq T$,

$\displaystyle\int_{T}^{t}\nabla{\mathcal{L}}({\theta}_{s})ds$
$\displaystyle=\mp\left({\lambda}\int_{T}^{t}\frac{\mathrm{d}^{2}\ln|w_{\pm,s}|}{\mathrm{d}t^{2}}ds+\int_{T}^{t}\frac{\mathrm{d}\ln|w_{\pm,s}|}{\mathrm{d}t}ds+\int_{T}^{t}\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}ds\right)$
$\displaystyle=\mp\left({\lambda}\frac{\mathrm{d}\ln|w_{\pm,s}|}{\mathrm{d}t}\Big{|}_{T}^{t}+\ln|w_{\pm,s}|\Big{|}_{T}^{t}+\int_{T}^{t}\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}ds\right).$

So, using the above observations and letting $t\to\infty$ yields

$\lim_{t\to\infty}\int_{T}^{t}\nabla{\mathcal{L}}({\theta}_{s})ds=\mp\left(-{\lambda}\frac{\mathrm{d}\ln|w_{\pm,T}|}{\mathrm{d}t}-\ln|w_{\pm,T}|+\ln|w_{\pm,\infty}|-\int_{T}^{\infty}\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}\mathrm{d}t\right).$

Thus, we conclude that
$\lim_{t\to\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})ds$ exists, and
therefore,
$\lim_{t\to\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{{\lambda}}}ds=0$.
∎

We are now well-equipped to prove Proposition 4. We note that we phrased
Proposition 4 rather succinctly in the main part of the paper due to space
considerations. In what follows, we restate Proposition 4 by precisely
specifying the underlying assumptions.

######  Proposition.

Assume the solution $(u_{t},v_{t})$ of MGF (4) is bounded. If we also assume
that the balancedness at infinity ${\Delta}_{\infty}$ has nonzero coordinates,
then there exists a time $T\geq 0$, after which the predictors
$\theta_{t}=u_{t}\odot v_{t}$ follow a momentum mirror flow with time-varying
potentials $\Phi_{t}$:

$\displaystyle\lambda\frac{\mathrm{d}^{2}\nabla\Phi_{t}(\theta_{t})}{\mathrm{d}t^{2}}+\frac{\mathrm{d}\nabla\Phi_{t}(\theta_{t})}{\mathrm{d}t}+\nabla{\mathcal{L}}(\theta_{t})=0.$

Furthermore, if we assume that the balancedness ${\Delta}_{t}$ remains nonzero
for $t\in[0,+\infty]$, then the momentum mirror flow holds for every $t\geq
0$.

###### Proof.

We will consider the $(w_{+},w_{-})$-reparametrisation of momentum gradient
flow (4) introduced in Appendix B. For convenience of the reader, we recall
this reparametrisation here:

${\lambda}\ddot{w}_{\pm,t}+\dot{w}_{\pm,t}\pm\nabla{\mathcal{L}}({\theta}_{t})\odot
w_{\pm,t}=0.$

Now, let $\xi:{\mathbb{R}}_{\geq 0}\to{\mathbb{R}}^{d}$ be the
$C^{\infty}(0,\infty)$ solution of the following ODE:

${\lambda}\ddot{\xi}_{t}+\dot{\xi}_{t}+\nabla{\mathcal{L}}({\theta}_{t})=0,$

with the constraint $\xi_{0}=\dot{\xi}_{0}=0$. Hence, by Lemma 5,

$\xi_{t}=-\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})(1-e^{-\frac{t-s}{{\lambda}}})\mathrm{d}s,$

and by Lemma 3,

$\xi_{\infty}=-\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t})\mathrm{d}t.$

Thus, $\xi_{t}\in\texttt{span}(x_{1},\dots,x_{n})$, $\forall t\in[0,+\infty]$.

Having fixed $\xi_{t}$, we define the quantities ${\alpha}_{\pm,t}$ for every
$t\in[0,+\infty]$ through the following relation:

${\alpha}_{\pm,t}=w_{\pm,t}\exp(\mp\xi_{t}).$

So, ${\Delta}_{t}=|w_{+,t}w_{-,t}|=|{\alpha}_{+,t}{\alpha}_{-,t}|$.
Furthermore,

$\displaystyle{\theta}_{t}$
$\displaystyle=\frac{1}{4}(w_{+,t}^{2}-w_{-,t}^{2})$
$\displaystyle=\frac{1}{4}({\alpha}_{+,t}^{2}\exp(2\xi_{t})-{\alpha}_{-,t}^{2}\exp(-2\xi_{t}))$
$\displaystyle=\frac{1}{2}{\Delta}_{t}\sinh\left(2\xi_{t}+\ln\frac{|{\alpha}_{+,t}|}{|{\alpha}_{-,t}|}\right)$

Since we assumed that ${\Delta}_{\infty}$ has nonzero coordinates, there
exists $T\geq 0$ such that for all $t\geq T$, $w_{\pm,t}$ have nonzero
coordinates. Hence, for $t\geq T$, the logarithm
$\ln\frac{|{\alpha}_{+,t}|}{|{\alpha}_{-,t}|}$ is well-defined. If we assume
positive balancedness for $t\in[0,+\infty]$, then we can choose $T=0$. From
now until the end of the proof, whenever a time-dependent quantity features
division by ${\Delta}_{t}$, we will tacitly assume that $t\geq T$.

Let us now introduce the helper quantity $\phi_{t}$ through the following
identity:

$\phi_{t}=\frac{1}{2}\ln\frac{|{\alpha}_{+,t}|}{|{\alpha}_{-,t}|}=\frac{1}{2}\operatorname{arcsinh}\left(\frac{{\alpha}_{+,t}^{2}-{\alpha}_{-,t}^{2}}{2{\Delta}_{t}}\right).$

Then,

$\frac{1}{2}\operatorname{arcsinh}\left(\frac{2{\theta}_{t}}{{\Delta}_{t}}\right)-\phi_{t}=\xi_{t}\in\texttt{span}(x_{1},\dots,x_{n}).$

So, if we consider the time-varying potential

$\begin{split}\Phi_{t}({\theta})&=\frac{1}{4}\sum_{i=1}^{d}\left(2{\theta}_{i}\operatorname{arcsinh}\left(\frac{2{\theta}_{i}}{{\Delta}_{t,i}}\right)-\sqrt{4{\theta}_{i}^{2}+{\Delta}_{t,i}^{2}}+{\Delta}_{t,i}\right)-\langle\phi_{t},{\theta}\rangle\\\
&=\psi_{{\Delta}_{t}}({\theta})-\langle\phi_{t},{\theta}\rangle,\end{split}$
(15)

where $\psi_{{\Delta}_{t}}$ is the hyperbolic entropy defined in Equation 6,
then

$\nabla\Phi_{t}({\theta})=\frac{1}{2}\operatorname{arcsinh}\left(\frac{2{\theta}}{{\Delta}_{t}}\right)-\phi_{t}.$

Notice that
$\nabla^{2}\Phi_{t}=\texttt{diag}\left(1/\sqrt{4{\theta}^{2}+{\Delta}_{t}^{2}}\right)\succ
0$. Hence, $\Phi_{t}$ is a mirror map. Furthermore,
$\nabla\Phi_{t}({\theta}_{t})=\xi_{t}$ for $t\geq T$, so

$\lambda\frac{\mathrm{d}^{2}\nabla\Phi_{t}(\theta_{t})}{\mathrm{d}t^{2}}+\frac{\mathrm{d}\nabla\Phi_{t}(\theta_{t})}{\mathrm{d}t}+\nabla{\mathcal{L}}(\theta_{t})=0.$

∎

###  C.3 Proof of Theorem 1

We are now ready to prove our main result for the implicit bias of momentum
gradient flow on diagonal linear networks.

See 1

We split the proof into two parts for conceptual clarity. In the first part,
we utilise the time-varying mirror flow from Proposition 4 to derive the
implicit regularisation
${\theta}^{\texttt{MGF}}=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\
D_{\psi_{{\Delta}_{\infty}}}({\theta}^{\star},\tilde{{\theta}}_{0})$. Then, in
the second part, we prove that the integral quantities $I_{\pm}$ from Lemma 1
are well-defined, and we give the trajectory-dependent characterisations of
the asymptotic balancedness ${\Delta}_{\infty}$ and the perturbed
initialisation $\tilde{{\theta}}_{0}$.

####  C.3.1 Proof of Implicit Regularisation.

In Proposition 4, we proved that whenever the MGF trajectory is bounded and
the coordinates of ${\Delta}_{\infty}$ are nonzero, there exists a time $T\geq
0$, after which the predictors ${\theta}_{t}$ follow a momentum mirror flow
with potentials given by Equation 15. Recall that for $t\geq T$,

$\nabla\Phi_{t}({\theta}_{t})=\frac{1}{2}\operatorname{arcsinh}\left(\frac{2{\theta}}{{\Delta}_{t}}\right)-\phi_{t}=-\xi_{t}\in\texttt{span}(x_{1},\dots,x_{n}).$

where
$\xi_{t}=-\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})(1-e^{-\frac{t-s}{{\lambda}}})ds$,
${\alpha}_{\pm,t}=w_{\pm,t}\exp(\mp\xi_{t})$, and
$\phi_{t}=\frac{1}{2}\operatorname{arcsinh}\left(\frac{{\alpha}_{+,t}^{2}-{\alpha}_{-,t}^{2}}{2{\Delta}_{t}}\right)$.

Now, as $t\to\infty$, $\xi_{t}$ and the MGF iterates converge, so we know that
$\nabla\Phi_{\infty}({\theta}_{\infty})\in\texttt{span}(x_{1},\dots,x_{n})$,
where
$\Phi_{\infty}({\theta})=\psi_{{\Delta}_{\infty}}({\theta})-\langle\phi_{\infty},{\theta}\rangle$.
Thus, we can use the familiar Bregman-Cosine-Theorem trick to characterise the
interpolator ${\theta}_{\infty}$. We proceed with this characterisation.

Let $\tilde{{\theta}}_{0}$ be a perturbation term such that
$\nabla\Phi_{\infty}(\tilde{{\theta}}_{0})=0$. Equivalently,

$\displaystyle\frac{1}{2}\operatorname{arcsinh}\left(\frac{2\tilde{{\theta}}_{0}}{{\Delta}_{\infty}^{2}}\right)-\phi_{\infty}=0\iff$
$\displaystyle\operatorname{arcsinh}\left(\frac{2\tilde{{\theta}}_{0}}{{\Delta}_{\infty}^{2}}\right)-\operatorname{arcsinh}\left(\frac{{\alpha}_{+,\infty}^{2}-{\alpha}_{-,\infty}^{2}}{2{\Delta}_{\infty}^{2}}\right)=0\iff$
$\displaystyle\tilde{{\theta}}_{0}=\frac{{\alpha}_{+,\infty}^{2}-{\alpha}_{-,\infty}^{2}}{4}.$

Note that
${\alpha}_{\pm,\infty}=w_{\pm,\infty}\exp(\pm\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t})\mathrm{d}t)$
by Lemma 3 and ${\Delta}_{\infty}=|{\alpha}_{+,\infty}{\alpha}_{-,\infty}|$.

Now, let ${\theta}^{\star}\in\mathcal{S}$ be an arbitrary interpolator of the
dataset. Then,
${\theta}^{\star}-{\theta}_{\infty}\in\mathrm{ker}(X)=\texttt{span}(x_{1},\dots,x_{n})^{\perp}$.
Hence, the Bregman Cosine Theorem yields

$\displaystyle D_{\Phi_{\infty}}({\theta}^{\star},\tilde{{\theta}}_{0})$
$\displaystyle=D_{\Phi_{\infty}}({\theta}^{\star},{\theta}_{\infty})+D_{\Phi_{\infty}}({\theta}_{\infty},\tilde{{\theta}}_{0})+\langle{\theta}^{\star}-{\theta}_{\infty},\nabla\Phi({\theta}_{\infty})-\nabla\Phi(\tilde{{\theta}}_{0})\rangle$
$\displaystyle=D_{\Phi_{\infty}}({\theta}^{\star},{\theta}_{\infty})+D_{\Phi_{\infty}}({\theta}_{\infty},\tilde{{\theta}}_{0}),$

where we used that
$\nabla\Phi({\theta}_{\infty})-\nabla\Phi(\tilde{{\theta}}_{0})\in\texttt{span}(x_{1},\dots,x_{n})$.
Thus,

$\displaystyle{\theta}_{\infty}$
$\displaystyle=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}D_{\Phi_{\infty}}({\theta}^{\star},\tilde{{\theta}}_{0})$
$\displaystyle=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\Phi_{\infty}({\theta}^{\star}).$

Finally, notice that
$\nabla\psi_{{\Delta}_{\infty}}(\tilde{{\theta}}_{0})=\frac{1}{2}\operatorname{arcsinh}\left(\frac{2\tilde{{\theta}}_{0}}{{\Delta}_{\infty}}\right)=\phi_{\infty}$
as we showed above. Hence,

$D_{\psi_{{\Delta}_{\infty}}}({\theta},\tilde{{\theta}}_{0})=\Phi_{\infty}({\theta})-\psi_{{\Delta}_{\infty}}(\tilde{{\theta}}_{0})+\langle\nabla\psi_{{\Delta}_{\infty}}(\tilde{{\theta}}_{0}),\tilde{{\theta}}_{0}\rangle.$

Thus, we conclude that

${\theta}_{\infty}=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\Phi_{\infty}({\theta}^{\star})=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}D_{\psi_{{\Delta}_{\infty}}}({\theta}^{\star},\tilde{{\theta}}_{0}).$

∎

####  C.3.2 Proof of Trajectory-Dependent Characterisation.

We just showed that the recovered interpolator by MGF solves the constrained
minimisation problem
${\theta}_{\infty}=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}D_{\psi_{{\Delta}_{\infty}}}({\theta}^{\star},\tilde{{\theta}}_{0})$,
where ${\Delta}_{\infty}=|{\alpha}_{+,\infty}{\alpha}_{-,\infty}|$,
$\tilde{{\theta}}_{0}=({\alpha}_{+,\infty}^{2}-{\alpha}_{-,\infty}^{2})/4$,
and
${\alpha}_{\pm,\infty}=w_{\pm,\infty}\exp(\pm\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t})\mathrm{d}t)$.
Clearly, these opaque characterisations of ${\Delta}_{\infty}$ and
$\tilde{{\theta}}_{0}$ prevent us from describing how the magnitude of these
quantities compares to the magnitude of the initial balancedness
${\Delta}_{0}$ and the initialisation scale ${\alpha}=\max(|u_{0}|,|v_{0}|)$.
Ideally, we would like to find formulas for ${\Delta}_{\infty}$ and
$\tilde{{\theta}}_{0}$ which show that
$\tilde{{\theta}}_{0}\ll{\theta}^{\star}$,
$\forall{\theta}^{\star}\in\mathcal{S}$ and ${\Delta}_{\infty}<{\Delta}_{0}$
so that we can conclude that
${\theta}^{\texttt{MGF}}\approx\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\psi_{{\Delta}_{\infty}}({\theta}^{\star})$
enjoys better sparsity guarantees than
${\theta}^{\texttt{GF}}\approx\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\psi_{{\Delta}_{0}}({\theta}^{\star})$.
In what follows, we derive such formulas.

In our subsequent arguments, for a vector $z\in{\mathbb{R}}^{d}$ and a
coordinate $i\in[d]$, we will denote with $z(i)$ the $i^{\text{th}}$
coordinate of $z$ in order to reduce the index bloat. Let us consider again
the $(w_{+},w_{-})$-reparametrisation of MGF discussed in Appendix B:

${\lambda}\ddot{w}_{\pm,t}+\dot{w}_{\pm,t}\pm\nabla{\mathcal{L}}({\theta}_{t})\odot
w_{\pm,t}=0.$ (16)

Notice that if for some $T>0$ and $i\in[d]$, $w_{+,T}(i)=0$, then
$\dot{w}_{+,T}(i)$ must be nonzero. Indeed, as we argued in Section C.1, MGF
(4) admits a unique global solution. And if $w_{+,T}(i)=\dot{w}_{+,T}(i)=0$,
then we could construct another solution $(w_{+,t}^{\prime},w_{-,t}^{\prime})$
of MGF such that
$w_{+,t}^{\prime}(i)=\dot{w}_{+,t}^{\prime}(i)=0,\enspace\forall t\geq 0$, and
$w_{\pm,T}=w_{\pm,T}^{\prime}$. By uniqueness, we get that
$w_{\pm,t}=w_{\pm,t}^{\prime},\enspace\forall t\geq 0$ However, the newly
constructed solution will not be consistent with the imposed initialisation
${\Delta}_{0}\neq 0$,. Hence, ${\Delta}_{0}\neq 0$ prevents $w_{+,t}(i)$ and
$\dot{w}_{+}(i)$ from hitting 0 simultaneously. Similarly, this situation
cannot occur for $w_{-,t}$.

Until further notice, we fix a coordinate $i\in[d]$ and consider eq. 16 only
in the $i^{\text{th}}$ coordinate without explicit mention. If $w_{\pm,T}=0$
for some $T>0$, then $\dot{w}_{\pm,T}\neq 0$. Hence, for some small
${\delta}>0$, $\dot{w}_{\pm,t}$ does not change sign on
$[T-{\delta},T+{\delta}]$, so $w_{\pm,t}$ either strictly increases or
decreases on $[T-{\delta},T+{\delta}]$. Therefore, $w_{\pm,t}\neq 0$ on
$[T-{\delta},T+{\delta}]\setminus\\{T\\}$ implying that $w_{\pm,t}$ equals 0
at most a countable number of times. Recall that by 2, there exists a time
$T_{\infty}$ after which $w_{\pm}$ does not change sign. Therefore, if we
assume that $w_{\pm}$ vanishes on infinitely many points
$T_{1}<T_{2}<\dots<T_{\infty}$, then by compactness, the limit
$\tau=\lim_{m\to\infty}T_{m}$ exists. Since $w_{\pm}$ is continuous, we infer
that $w_{\pm,\tau}=0$. Moreover, by the Mean Value Theorem, for every $m\geq
1$, there exists $T_{m}^{\prime}\in(T_{m},T_{m+1})$ such that
$\dot{w}_{T_{m}^{\prime}}=0$. Notice that
$\lim_{m\to\infty}T_{m}^{\prime}=\tau$ as well. Hence, by continuity,
$w_{\pm,\tau}=\dot{w}_{\pm,\tau}=0$ – a contradiction.

Hence, $w_{\pm}$ vanishes on a finite set of points. Let us order these
vanishing times as $0<T_{1}<\dots<T_{N}$ and let $T_{0}=0$ and
$T_{N+1}=+\infty$. Observe that for $t\notin\mathcal{T}=\\{T_{i}:i\in[N]\\}$,
we can safely divide both sides of eq. 16 by $w_{\pm,t}$ to obtain

${\lambda}\frac{\ddot{w}_{\pm,t}}{w_{\pm,t}}+\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\pm\nabla{\mathcal{L}}({\theta}_{t})=0.$

The last expression is equivalent to

${\lambda}\left(\frac{\ddot{w}_{\pm,t}}{w_{\pm,t}}-\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}\right)+\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\pm\nabla{\mathcal{L}}({\theta}_{t})+{\lambda}\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}=0,$

which can be rewritten as

${\lambda}\frac{\mathrm{d}^{2}\ln(\operatorname{sgn}(w_{\pm,t})w_{\pm,t})}{\mathrm{d}t^{2}}+\frac{\mathrm{d}\ln(\operatorname{sgn}(w_{\pm,t})w_{\pm,t})}{\mathrm{d}t}\pm\nabla{\mathcal{L}}({\theta}_{t})+{\lambda}\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}=0.$

Let us define a new function $g_{\pm}:{\mathbb{R}}_{\geq
0}\setminus\mathcal{T}\to{\mathbb{R}}^{d}$ through the relation
$g_{\pm,t}=\ln(\operatorname{sgn}(w_{\pm,t})w_{\pm,t})$. Then, $g_{\pm}$ is
$C^{\infty}$-smooth on ${\mathbb{R}}_{\geq 0}\setminus\mathcal{T}$ and
satisfies the following ODE:

${\lambda}\ddot{g}_{\pm,t}+\dot{g}_{\pm,t}\pm\nabla{\mathcal{L}}({\theta}_{t})+{\lambda}\left(\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\right)^{2}=0.$
(17)

##### Induction on Vanishing Times.

Now, we proceed to prove by induction on $N-1\geq m\geq 0$ that for
$\tau\in(T_{m},T_{m+1})$ the following 3 things hold:

  * •

The following integral quantities555See Equation 11 for the definition of
$\mathrm{m.p.v.}$ exist and are finite:

$\mathrm{m.p.v.}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}e^{-\frac{\tau-s}{\lambda}}\operatorname{sgn}(w_{\pm,\tau}w_{\pm,t})\mathrm{d}t\quad\text{and}\quad\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t.$

  * •

The following identity holds:

$\dot{g}_{\pm,\tau}=-\mathrm{m.p.v.}\int_{0}^{\tau}\left[\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{t})\right]e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\pm,\tau}w_{\pm,t})\mathrm{d}t-\frac{1}{{\lambda}}\sum_{k=1}^{m}(-1)^{m-k}e^{-\frac{\tau-
T_{k}}{{\lambda}}}.$

  * •

The following identity holds:

$\displaystyle g_{\pm,\tau}$
$\displaystyle=g_{\pm,0}-\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t$ $\displaystyle-\sum_{k=1}^{m}(-1)^{m-k}\left(1-e^{-\frac{\tau-
T_{k}}{{\lambda}}}\right)+2\sum_{1\leq i<j\leq
m}(-1)^{j-i}\left(1-e^{-\frac{T_{j}-T_{i}}{{\lambda}}}\right).$

Recall that $\nabla{\mathcal{L}}({\theta}_{t})$ is a bounded function, so if
the modified principal value from the first bullet point exists, then the
modified principal values in the above identities are also well-defined.

Base case: $m=0$. Recall from the proof of Proposition 4 in Section C.2 that
$\dot{w}_{\pm}\in L^{2}(0,\infty)$. Now, since $w_{\pm,t}$ does not change
signs on the interval $(T_{0},\tau)$, we know that $1/w_{\pm,t}=\Omega(1)$.
Hence,

$\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\in
L^{1}(0,\infty).$

Similarly, $\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}$ is integrable
on all intervals $[T_{i}+{\varepsilon},T_{i+1}-{\varepsilon}]$ for any small
${\varepsilon}>0$. Consequently, the integral quantities

$\mathrm{m.p.v.}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,\tau}w_{\pm,t})\mathrm{d}t=\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\mathrm{d}t$

and

$\displaystyle\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t$
$\displaystyle=\int_{0}^{\tau}\int_{0}^{t}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\mathrm{d}s\
\mathrm{d}t$
$\displaystyle=\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}(1-e^{-\frac{\tau-t}{{\lambda}}})\mathrm{d}t$

are well-defined. Moreover, after applying Lemma 5 to eq. 17, we get

$\displaystyle\dot{g}_{\pm,\tau}=-\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}e^{-\frac{\tau-t}{{\lambda}}}\mathrm{d}t\
\mp\frac{1}{{\lambda}}\int_{0}^{\tau}\nabla{\mathcal{L}}({\theta}_{t})e^{-\frac{\tau-t}{{\lambda}}}\mathrm{d}t$
$\displaystyle
g_{\pm,\tau}=g_{\pm,0}-{\lambda}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}(1-e^{-\frac{\tau-t}{{\lambda}}})\mathrm{d}t\
\mp\int_{0}^{\tau}\nabla{\mathcal{L}}({\theta}_{t})(1-e^{-\frac{\tau-t}{{\lambda}}})\mathrm{d}t,$

which concludes the proof of the base case.

Induction step: $m\to m+1$. For $m\geq 0$, assume that for every
$\tau\in[0,T_{m+1})\setminus\mathcal{T}$ the expressions

$\displaystyle\dot{g}_{\pm,\tau}=-\mathrm{m.p.v.}\int_{0}^{\tau}\left[\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{t})\right]e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\pm,\tau}w_{\pm,t})\mathrm{d}t-\frac{1}{{\lambda}}\sum_{k=1}^{m}(-1)^{m-k}e^{-\frac{\tau-
T_{k}}{{\lambda}}}$

and

$\displaystyle g_{\pm,\tau}$
$\displaystyle=g_{\pm,0}-\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t$ $\displaystyle-\sum_{k=1}^{m}(-1)^{m-k}\left(1-e^{-\frac{\tau-
T_{k}}{{\lambda}}}\right)+2\sum_{1\leq i<j\leq
m}(-1)^{j-i}\left(1-e^{-\frac{T_{j}-T_{i}}{{\lambda}}}\right).$

are true and well-defined. We now want to extend the validity of these
identities to $\tau\in(T_{m+1},T_{m+2})$. For ease of notation during the
induction step, let $T_{m+1}=T$, $w_{\pm}=w$, and $g_{\pm}=g$. Let
${\varepsilon}>0$ and let $T_{\pm}=T\pm{\varepsilon}$.

Now, applying Lemma 5 to eq. 17 yields

$\displaystyle\dot{g}_{\tau}=\dot{g}_{T_{+}}e^{-\frac{\tau-
T_{+}}{{\lambda}}}-\int_{T_{+}}^{\tau}\left[\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{t})\right]e^{-\frac{\tau-t}{\lambda}}\mathrm{d}t$
$\displaystyle
g_{\tau}=g_{T_{+}}+\dot{g}_{T_{+}}\int_{T_{+}}^{\tau}e^{-\frac{t-T_{+}}{{\lambda}}}\mathrm{d}t-\int_{T_{+}}^{\tau}\int_{T_{+}}^{t}\left[\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\mathrm{d}s\
\mathrm{d}t.$

For further ease of notation and with some abuse of notation, let
$f_{t}=\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{t})$
on ${\mathbb{R}}_{\geq 0}\setminus\mathcal{T}$. We will shortly prove that
$g_{T_{+}}-g_{T_{-}}=O({\varepsilon})$ and
$\dot{g}_{T_{+}}+\dot{g}_{T_{-}}+\frac{1}{{\lambda}}=O({\varepsilon})$.666Whenever
we write an equation of the form $A=B+O({\varepsilon}^{r})$ for some $r>0$, we
mean that $A=B+C$, where $|C|=O({\varepsilon}^{r})$. Hence, the following
limits will hold:

$\displaystyle\dot{g}_{\tau}=\lim_{{\varepsilon}\to
0}\left[-\frac{1}{{\lambda}}e^{-\frac{\tau-
T_{+}}{{\lambda}}}-\dot{g}_{T_{-}}e^{-\frac{\tau-
T_{+}}{{\lambda}}}-\int_{T_{+}}^{\tau}f_{t}e^{-\frac{\tau-t}{\lambda}}\mathrm{d}t\right]$
$\displaystyle g_{\tau}=\lim_{{\varepsilon}\to
0}\left[g_{T_{-}}-\frac{1}{{\lambda}}\int_{T_{+}}^{\tau}e^{-\frac{t-T_{+}}{{\lambda}}}\mathrm{d}t-\dot{g}_{T_{-}}\int_{T_{+}}^{\tau}e^{-\frac{t-T_{+}}{{\lambda}}}\mathrm{d}t-\int_{T_{+}}^{\tau}\int_{T_{+}}^{t}f_{s}e^{-\frac{t-s}{\lambda}}\mathrm{d}s\
\mathrm{d}t\right].$

Induction step for $\dot{g}_{\tau}$. Let us begin to untangle the first limit
by substituting $\dot{g}_{T_{-}}$ with its integral formula given by the
induction hypothesis. Notice that

$\displaystyle\dot{g}_{T_{-}}e^{-\frac{\tau-T_{+}}{{\lambda}}}$
$\displaystyle=-\mathrm{m.p.v.}\int_{0}^{T_{-}}f_{t}e^{-\frac{T_{-}-t}{\lambda}}\operatorname{sgn}(w_{T_{-}}w_{t})\mathrm{d}t\cdot
e^{-\frac{\tau-T_{+}}{{\lambda}}}-\frac{e^{-\frac{\tau-
T_{+}}{{\lambda}}}}{{\lambda}}\sum_{k=1}^{m}(-1)^{m-k}e^{-\frac{T_{-}-T_{k}}{{\lambda}}}$
$\displaystyle=\mathrm{m.p.v.}\int_{0}^{T_{-}}f_{t}e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t\cdot
e^{\frac{2{\varepsilon}}{{\lambda}}}+\frac{e^{\frac{2{\varepsilon}}{{\lambda}}}}{{\lambda}}\sum_{k=1}^{m}(-1)^{(m+1)-k}e^{-\frac{\tau-
T_{k}}{{\lambda}}},$

where we used that $\operatorname{sgn}(\tau)=-\operatorname{sgn}(T_{-})$ since
$w$ changes signs at $T$. Hence, we have that

$\displaystyle\dot{g}_{\tau}$ $\displaystyle=-\lim_{{\varepsilon}\to
0}\Big{[}\frac{1}{{\lambda}}e^{-\frac{\tau-
T_{+}}{{\lambda}}}+\frac{e^{\frac{2{\varepsilon}}{{\lambda}}}}{{\lambda}}\sum_{k=1}^{m}(-1)^{(m+1)-k}e^{-\frac{\tau-
T_{k}}{{\lambda}}}$
$\displaystyle+\mathrm{m.p.v.}\int_{0}^{T_{-}}f_{t}e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t\cdot
e^{\frac{2{\varepsilon}}{{\lambda}}}+\int_{T_{+}}^{\tau}f_{t}e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t\Big{]}$
$\displaystyle=\mp\int_{0}^{\tau}\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{t})\operatorname{sgn}(w_{\tau}w_{t})-\frac{1}{{\lambda}}\sum_{k=1}^{m+1}(-1)^{(m+1)-k}e^{-\frac{\tau-
T_{k}}{{\lambda}}}$ $\displaystyle-\lim_{{\varepsilon}\to
0}\left[\mathrm{m.p.v.}\int_{0}^{T_{-}}\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t\cdot
e^{\frac{2{\varepsilon}}{{\lambda}}}+\int_{T_{+}}^{\tau}\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t\right],$

where the limit on the last line formally equals the modified principal value
$\mathrm{m.p.v.}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}e^{-\frac{\tau-s}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t$
whose existence we want to prove as part of the induction step. In fact,
notice that we just proved the existence of
$\mathrm{m.p.v.}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}e^{-\frac{\tau-s}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t$
since both $\dot{g}_{\tau}$ and
$\mp\int_{0}^{\tau}\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{t})\operatorname{sgn}(w_{\tau}w_{t})$
are finite quantities. Hence, for $\tau\in(T_{m+1},T_{m+2})$,

$\dot{g}_{\tau}=-\mathrm{m.p.v.}\int_{0}^{\tau}\left[\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{t})\right]e^{-\frac{\tau-t}{\lambda}}\operatorname{sgn}(w_{\tau}w_{t})\mathrm{d}t-\frac{1}{{\lambda}}\sum_{k=1}^{m+1}(-1)^{(m+1)-k}e^{-\frac{\tau-
T_{k}}{{\lambda}}}.$

Induction step for $g_{\tau}$. We move on to untangle the limit which equals
$g_{\tau}$. By the induction hypothesis,

$\displaystyle\dot{g}_{T_{-}}=-\mathrm{m.p.v.}\int_{0}^{T_{-}}f_{t}e^{-\frac{T_{-}-t}{\lambda}}\operatorname{sgn}(w_{T_{-}}w_{t})\mathrm{d}t-\frac{1}{{\lambda}}\sum_{k=1}^{m}(-1)^{m-k}e^{-\frac{T_{-}-T_{k}}{{\lambda}}}$
$\displaystyle
g_{T_{-}}=g_{0}-\int_{0}^{T_{-}}\mathrm{m.p.v.}\int_{0}^{t}f_{s}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t$
$\displaystyle-\sum_{k=1}^{m}(-1)^{m-k}\left(1-e^{-\frac{T_{-}-T_{k}}{{\lambda}}}\right)+2\sum_{1\leq
i<j\leq m}(-1)^{j-i}\left(1-e^{-\frac{T_{j}-T_{i}}{{\lambda}}}\right).$

Again, we can substitute $\operatorname{sgn}(w_{T_{-}})$ with
$-\operatorname{sgn}(w_{\tau})$, and after performing the familiar integral
and limit manipulations, we obtain

$\displaystyle g_{\tau}$
$\displaystyle=g_{0}\mp\int_{0}^{\tau}\int_{0}^{t}\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t-\lim_{{\varepsilon}\to
0}\left[A_{\varepsilon}+B_{\varepsilon}+C_{\varepsilon}\right]$
$\displaystyle-\sum_{k=1}^{m+1}(-1)^{(m+1)-k}\left(1-e^{-\frac{\tau-
T_{k}}{{\lambda}}}\right)+2\sum_{1\leq i<j\leq
m+1}(-1)^{j-i}\left(1-e^{-\frac{T_{j}-T_{i}}{{\lambda}}}\right),$

where

$\displaystyle
A_{\varepsilon}=\int_{0}^{T_{-}}\mathrm{m.p.v.}\int_{0}^{t}\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t$ $\displaystyle
B_{\varepsilon}=\int_{T_{+}}^{\tau}\mathrm{m.p.v.}\int_{0}^{T_{-}}\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t\cdot e^{\frac{2{\varepsilon}}{{\lambda}}}$ $\displaystyle
C_{\varepsilon}=\int_{T_{+}}^{\tau}\int_{T_{+}}^{t}\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t.$

Figure 5: A visualisation of the areas over which we integrate
$\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})$
in the above limit.

Notice that formally the limit $\lim_{{\varepsilon}\to
0}[A_{\varepsilon}+B_{\varepsilon}+C_{\varepsilon}]$ equals the integral
quantity

$\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t,$

whose existence we just proved as a consequence of the fact that

$g_{0}\mp\int_{0}^{\tau}\int_{0}^{t}\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t-\sum_{k=1}^{m+1}(-1)^{(m+1)-k}\left(1-e^{-\frac{\tau-
T_{k}}{{\lambda}}}\right)+2\sum_{1\leq i<j\leq
m+1}(-1)^{j-i}\left(1-e^{-\frac{T_{j}-T_{i}}{{\lambda}}}\right)-g_{\tau}$

is well-defined and finite. Thus, for $\tau\in(T_{m+1},T_{m+2})$,

$\displaystyle g_{\tau}$
$\displaystyle=g_{0}-\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t$
$\displaystyle-\sum_{k=1}^{m+1}(-1)^{(m+1)-k}\left(1-e^{-\frac{\tau-
T_{k}}{{\lambda}}}\right)+2\sum_{1\leq i<j\leq
m+1}(-1)^{j-i}\left(1-e^{-\frac{T_{j}-T_{i}}{{\lambda}}}\right).$

Proof of bounds. In order to conclude the induction step, we still have to
prove the following bounds:

$g_{T_{+}}-g_{T_{-}}=O({\varepsilon})\quad\text{and}\quad\dot{g}_{T_{+}}+\dot{g}_{T_{-}}+\frac{1}{{\lambda}}=O({\varepsilon}).$

Recall that $g_{T\pm{\varepsilon}}=\log|w_{T\pm{\varepsilon}}|$ and that
$w_{T}=0,\ \dot{w_{T}}\neq 0$. From the Taylor expansion of $w_{t}$, we know
that

$w_{T\pm{\varepsilon}}=\pm{\varepsilon}\dot{w}_{T}+O({\varepsilon}^{2}).$

Hence, $|w_{T+{\varepsilon}}/w_{T-{\varepsilon}}|=1+O({\varepsilon})$.
Therefore, using the Taylor expansion of the logarithm around 1, we get that

$|g_{T_{+}}-g_{T_{-}}|=|\log(1+O({\varepsilon}))|=O({\varepsilon}).$

Now, recall that
$\dot{g}_{T\pm{\varepsilon}}=\dot{w}_{T\pm{\varepsilon}}/w_{T\pm{\varepsilon}}$
and observe that

$\displaystyle
w_{T\pm{\varepsilon}}=\pm{\varepsilon}\dot{w}_{T}+\frac{1}{2}{\varepsilon}^{2}\ddot{w}_{T}+O({\varepsilon}^{3})$
$\displaystyle\dot{w}_{T\pm{\varepsilon}}=\dot{w}_{T}\pm{\varepsilon}\ddot{w}_{T}+O({\varepsilon}^{2}).$

Hence,

$\frac{w_{T+{\varepsilon}}\dot{w}_{T-{\varepsilon}}+w_{T-{\varepsilon}}\dot{w}_{T+{\varepsilon}}}{w_{T+{\varepsilon}}w_{T-{\varepsilon}}}=\frac{-{\varepsilon}^{2}\dot{w}_{T}\ddot{w}_{T}+O({\varepsilon}^{3})}{-{\varepsilon}^{2}\dot{w}_{T}^{2}+O({\varepsilon}^{3})}=\frac{\ddot{w}_{T}}{\dot{w}_{T}}+O({\varepsilon}).$

Since,
${\lambda}\ddot{w}_{T}+\dot{w}_{T}\pm\nabla{\mathcal{L}}({\theta}_{T})\odot
w_{T}=0$ and $w_{T}=0$, we get that
$\frac{\ddot{w}_{T}}{\dot{w}_{T}}=-\frac{1}{{\lambda}}$, which concludes the
induction step.

##### Proof of Lemma 1.

Thus, we proved that for $\tau\in(T_{m},T_{m+1})$, $m\in\\{0,1,\dots,N\\}$,

$\displaystyle\ln|w_{\tau}|$
$\displaystyle=\ln|w_{0}|-\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t$ $\displaystyle-\sum_{k=1}^{m}(-1)^{m-k}\left(1-e^{-\frac{\tau-
T_{k}}{{\lambda}}}\right)+2\sum_{1\leq i<j\leq
m}(-1)^{j-i}\left(1-e^{-\frac{T_{j}-T_{i}}{{\lambda}}}\right).$

Recall that throughout our inductive proof we worked with a fixed coordinate
$i\in[d]$ of $w_{\pm}$. Different coordinates of $w_{\pm}$ vanish at different
points in time, so writing the sum the last line in a coordinate-agnostic way
becomes impossible. Thus, deriving a simple expression for the full
$d$-dimensional vector $w_{\pm,\tau}$ for any $\tau\in{\mathbb{R}}_{\geq 0}$
also becomes impossible. However, remembering that the finite nonzero limits
$\lim_{\tau\to\infty}|w_{\pm,\tau}|=|w_{\pm,\infty}|$ exist and letting
$\tau\to\infty$ yields an interesting result for the weights at infinity.
Indeed, notice that for every vanishing time $T$,
$\lim_{\tau\to\infty}\left(1-e^{-\frac{\tau-T_{k}}{{\lambda}}}\right)=1$.
Hence,

$\frac{1}{{\lambda}}\sum_{k=1}^{N}(-1)^{N-k}\left(1-e^{-\frac{\tau-
T_{k}}{{\lambda}}}\right)=\mathbf{1}_{\\{N-\text{ odd}\\}}.$

For every $i\in[d]$, let $N_{\pm}(i)$ denote the number of vanishing points
for the coordinate $w_{\pm}(i)$. Let us define the $d$-dimensional parity
vectors $P_{\pm}\in\\{0,1\\}^{d}$ such that $P_{\pm}(i)\equiv N_{\pm}(i)\mod
2$. Let us also define the $d$-dimensional vectors
$Q_{\pm}\in{\mathbb{R}}^{d}$ such that for each coordinate $k\in[d]$,

$Q_{\pm}(k)\coloneqq-2\sum_{1\leq i<j\leq
N_{\pm}(k)}(-1)^{j-i}\left(1-e^{-\frac{T_{\pm,k}(j)-T_{\pm,k}(i)}{{\lambda}}}\right),$

where $0<T_{\pm,k}(1)<\dots<T_{\pm,k}(N_{\pm}(k))$ denote the vanishing times
of the weight $w_{\pm}(k)$. Hence, we obtain the formula

$|w_{\pm,\infty}|=|w_{\pm,0}|e^{-(P_{\pm}+Q_{\pm})}\exp\left(-\int_{0}^{\infty}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t\right).$ (18)

Recall that in Lemma 3, we proved that the limit

$\lim_{t\to\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})ds=\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t})\mathrm{d}t=\frac{1}{{\lambda}}\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\mathrm{d}s\
\mathrm{d}t$

exists and is finite. Therefore, we can decouple
$\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}$ and
$\nabla{\mathcal{L}}({\theta}_{s})$ from the above integral and show that the
following integral limits exist and are finite:

$\displaystyle\int_{0}^{\infty}\mathrm{m.p.v.}\int_{0}^{t}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t\quad\text{and}\quad\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t.$

Hence, the integral quantities ${\Omega}_{\pm}$ from Lemma 1 are well-defined
and finite. Thus, we finally proved Lemma 1.

##### Trajectory-Dependent Characterisation.

We started this section with a promise for more insightful representations of
${\Delta}_{\infty}=|{\alpha}_{+,\infty}{\alpha}_{-,\infty}|$ and
$\tilde{{\theta}}_{0}=({\alpha}_{+,\infty}^{2}-{\alpha}_{-,\infty}^{2})/4$. We
now deliver on that promise.

Recall that
${\alpha}_{\pm,\infty}=w_{\pm,\infty}\exp\left(\pm\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t})\mathrm{d}t\right)$
and notice that

$\displaystyle\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\mathrm{d}s\
\mathrm{d}t-\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t=2\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\mathbf{1}_{\\{w_{\pm,t}w_{\pm,s}<0\\}}\mathrm{d}s\
\mathrm{d}t.$

Hence, using the formula for $w_{\pm}$ from Equation 18, we derive the
following:

$\displaystyle|{\alpha}_{\pm,\infty}|=|w_{\pm,0}|e^{-(P_{\pm}+Q_{\pm})}$
$\displaystyle\odot\exp\left(-\int_{0}^{\infty}\mathrm{m.p.v.}\int_{0}^{t}\Big{(}\frac{\dot{w}_{\pm,s}}{w_{\pm,s}}\Big{)}^{2}e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{\pm,t}w_{\pm,s})\mathrm{d}s\
\mathrm{d}t\right)$
$\displaystyle\odot\exp\left(\pm\frac{2}{{\lambda}}\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\mathbf{1}_{\\{w_{\pm,t}w_{\pm,s}<0\\}}\mathrm{d}s\
\mathrm{d}t\right).$

Now, let

$\Lambda_{\pm}\coloneqq\mp\frac{2}{{\lambda}}\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\mathbf{1}_{\\{w_{\pm,t}w_{\pm,s}<0\\}}\mathrm{d}s\
\mathrm{d}t+P_{\pm}+Q_{\pm},$ (19)

where the quantities $P_{\pm}$ and $Q_{\pm}$ were defined in the previous
paragraph. Notice that as we promised underneath Lemma 1, $\Lambda_{\pm}$
vanish whenever the balancedness ${\Delta}_{t}$ remains strictly positive.
Using the abbreviation $I_{\pm}={\Omega}_{\pm}+\Lambda_{\pm}$, we get that

$\displaystyle|{\alpha}_{\pm,\infty}|=|w_{\pm,0}|\odot\exp\left(-I_{\pm}\right).$

Multiplying $|{\alpha}_{+,\infty}|$ by $|{\alpha}_{-,\infty}|$, we derive a
formula for the asymptotic balancedness:

$\displaystyle\begin{split}{\Delta}_{\infty}={\Delta}_{0}e^{-(P_{+}+P_{-}+Q_{+}+Q_{-})}&\odot\exp\left(-\int_{0}^{\infty}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{+,s}}{w_{+,s}}\Big{)}^{2}+\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{+,t}w_{+,s})\mathrm{d}s\
\mathrm{d}t\right)\\\
&\odot\exp\left(-\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{-,s}}{w_{-,s}}\Big{)}^{2}-\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{-,t}w_{-,s})\mathrm{d}s\
\mathrm{d}t\right)\\\
&\odot\exp\left(\frac{2}{{\lambda}}\int_{0}^{\infty}\int_{0}^{t}\nabla{\mathcal{L}}({\theta}_{s})e^{-\frac{t-s}{\lambda}}\left[\mathbf{1}_{\\{w_{+,t}w_{+,s}<0\\}}-\mathbf{1}_{\\{w_{-,t}w_{-,s}<0\\}}\right]\mathrm{d}s\
\mathrm{d}t\right).\end{split}$ (20)

Now, we can write ${\Delta}_{\infty}$ and $\tilde{{\theta}}_{0}$ more
succinctly as

${\Delta}_{\infty}={\Delta}_{0}\odot\exp\Big{(}-(I_{+}+I_{-})\Big{)}$

and

$\tilde{{\theta}}_{0}=\frac{1}{4}\Big{(}w_{+,0}^{2}\odot\exp\left(-2I_{+}\right)-w_{-,0}^{2}\odot\exp\left(-2I_{-}\right)\Big{)},$

which concludes the proof of Theorem 1. ∎

####  C.3.3 Consequences for Generalisation.

We just proved that whenever MGF on a diagonal linear network converges and
the balancedness at infinity is nonzero, we can characterize the recovered
interpolator through the implicit regularization problem

$\displaystyle{\theta}^{\texttt{MGF}}$
$\displaystyle=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\
D_{\psi_{{\Delta}_{\infty}}}({\theta}^{\star},\tilde{{\theta}}_{0})$
$\displaystyle=\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\left[\psi_{{\Delta}_{\infty}}({\theta}^{\star})-\langle\nabla\psi_{{\Delta}_{\infty}}(\tilde{{\theta}}_{0}),{\theta}^{\star}\rangle\right].$

Since

$\psi_{{\Delta}_{\infty}}({\theta})=\frac{1}{4}\sum_{i=1}^{d}\left(2{\theta}_{i}\operatorname{arcsinh}\left(\frac{2{\theta}_{i}}{{\Delta}_{\infty,i}}\right)-\sqrt{4{\theta}_{i}^{2}+{\Delta}_{\infty,i}^{2}}+{\Delta}_{\infty,i}\right)$

and

$\nabla\psi_{{\Delta}_{\infty}}({\theta})=\frac{1}{2}\operatorname{arcsinh}\left(\frac{2{\theta}}{{\Delta}_{\infty}}\right),$

for a small asymptotic balancedness
${\Delta}_{\infty}=O({\Delta}_{0})=O({\alpha}^{2})$ and small perturbed
initialisation $|\tilde{{\theta}}_{0}|=O({\alpha}^{2})\ll|{\theta}^{\star}|$,
we would expect $\psi_{{\Delta}_{\infty}}({\theta}^{\star})$ to dominate
$\langle\nabla\psi_{{\Delta}_{\infty}}(\tilde{{\theta}}_{0}),{\theta}^{\star}\rangle$.
More formally, for a fixed ${\theta}^{\star}\in\mathcal{S}$ and small
${\Delta}_{\infty}$ and $\tilde{{\theta}}_{0}$, we have the following
asymptotic equivalence:

$\psi_{{\Delta}_{\infty}}({\theta}^{\star})\underset{\scriptscriptstyle{\alpha}\to
0}{\sim}D_{\psi_{{\Delta}_{\infty}}}({\theta}^{\star},\tilde{{\theta}}_{0}).$

Hence,
${\theta}^{\texttt{MGF}}\approx\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\psi_{{\Delta}_{\infty}}({\theta}^{\star})={\theta}_{{\Delta}_{\infty}}^{\texttt{GF}}$
as we discussed in Section 3.1. So, if ${\Delta}_{\infty}<{\Delta}_{0}$, Lemma
4 implies that the MGF predictor will benefit from better sparsity guarantees
than the GF solution.

Therefore, to recap, for a small initialisation scale ${\alpha}$ and provided
that the bounds ${\Delta}_{\infty}=O({\alpha}^{2})$ and
$\tilde{{\theta}}_{0}=O({\alpha}^{2})$ hold, we conclude that the asymptotic
balancedness at infinity ${\Delta}_{\infty}$ roughly controls the sparsity of
the recovered interpolator. And when ${\Delta}_{\infty}<{\Delta}_{0}$,
${\theta}^{\texttt{MGF}}$ will be sparser than
${\theta}_{{\alpha}}^{\texttt{GF}}$. Unfortunately, without the assumption
that the balancedness ${\Delta}_{t}$ remains strictly positive for all
$t\in[0,+\infty]$, we cannot formally compare ${\Delta}_{\infty}$ and
$\tilde{{\theta}}_{0}$ with ${\alpha}$.

Note that even without the bounds ${\Delta}_{\infty}=O({\alpha}^{2})$ and
$\tilde{{\theta}}_{0}=O({\alpha}^{2})$, if
$|\tilde{{\theta}}_{0}|\ll|{\theta}^{\star}|$, then
$\psi_{{\Delta}_{\infty}}({\theta}^{\star})$ still dominates
$\langle\nabla\psi_{{\Delta}_{\infty}}(\tilde{{\theta}}_{0}),{\theta}^{\star}\rangle$.
Indeed, our experiments clearly show that the perturbation term
$\tilde{{\theta}}_{0}$ can safely be ignored since
${\theta}^{\texttt{MGF}}\approx\operatorname*{argmin}_{{\theta}^{\star}\in\mathcal{S}}\psi_{{\Delta}_{\infty}}({\theta}^{\star})$
(see the discussion around Figure 8.)

###  C.4 Non-Vanishing Balancedness

If we work under the assumption that the balancedness
${\Delta}_{t}=|w_{+,t}w_{-,t}|$ never vanishes, then much of the analysis from
Section C.3.2 greatly simplifies. First, the integral quantities $P_{\pm}$ and
$Q_{\pm}$ from the previous subsection become 0. Second, the multipliers
$\operatorname{sgn}(w_{\pm,t}w_{\pm,s})$ become equal to 1 for all
$t,s\in{\mathbb{R}}_{\geq 0}$. Hence, using Fubini’s Theorem as in the proof
of Lemma 5, we get that

$\displaystyle\int_{0}^{\tau}\mathrm{m.p.v.}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]$
$\displaystyle
e^{-\frac{t-s}{\lambda}}\operatorname{sgn}(w_{t}w_{s})\mathrm{d}s\
\mathrm{d}t$
$\displaystyle=\int_{0}^{\tau}\int_{0}^{t}\left[\Big{(}\frac{\dot{w}_{s}}{w_{s}}\Big{)}^{2}\pm\frac{1}{{\lambda}}\nabla{\mathcal{L}}({\theta}_{s})\right]e^{-\frac{t-s}{\lambda}}\mathrm{d}s\
\mathrm{d}t$
$\displaystyle={\lambda}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{t}}{w_{t}}\Big{)}^{2}\left(1-e^{-\frac{\tau-t}{\lambda}}\right)\mathrm{d}t\pm\int_{0}^{\tau}\nabla{\mathcal{L}}({\theta}_{t})\left(1-e^{-\frac{\tau-t}{\lambda}}\right)\mathrm{d}t.$

Therefore, referencing eq. 18, we can express the evolution of the iterates as
follows:

$w_{\pm,\tau}=w_{\pm,0}\exp\left(-{\lambda}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}\left(1-e^{-\frac{\tau-t}{\lambda}}\right)\mathrm{d}t\right)\exp\left(\mp\int_{0}^{\tau}\nabla{\mathcal{L}}({\theta}_{s})\left(1-e^{-\frac{\tau-t}{\lambda}}\right)\mathrm{d}t\right).$
(21)

Thus, the balancedness evolves as

${\Delta}_{t}={\Delta}_{0}\exp\left(-{\lambda}\int_{0}^{\tau}\left[\Big{(}\frac{\dot{w}_{+,t}}{w_{+,t}}\Big{)}^{2}+\Big{(}\frac{\dot{w}_{-,t}}{w_{-,t}}\Big{)}^{2}\right]\left(1-e^{-\frac{\tau-t}{\lambda}}\right)\mathrm{d}t\right).$
(22)

Now, from Lemma 3, we know that
$\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}$ is integrable and that

$\lim_{\tau\to\infty}\int_{0}^{\tau}\nabla{\mathcal{L}}({\theta}_{s})\left(1-e^{-\frac{\tau-t}{\lambda}}\right)\mathrm{d}t=\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{s})\mathrm{d}t$

exists. Furthermore, from Lemma 6, we know that that

$\lim_{\tau\to\infty}\int_{0}^{\tau}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}\left(1-e^{-\frac{\tau-t}{\lambda}}\right)\mathrm{d}t=\int_{0}^{\infty}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2})\mathrm{d}t.$

Therefore, letting $\tau\to\infty$, we obtain the formulas

$\displaystyle
w_{\pm,\tau}=w_{\pm,0}\exp\left(-{\lambda}\int_{0}^{\infty}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}\mathrm{d}t\right)\exp\left(\mp\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{s})\mathrm{d}t\right)$
(23)
$\displaystyle{\Delta}_{\infty}={\Delta}_{0}\exp\left(-{\lambda}\int_{0}^{\infty}\left[\Big{(}\frac{\dot{w}_{+,t}}{w_{+,t}}\Big{)}^{2}+\Big{(}\frac{\dot{w}_{-,t}}{w_{-,t}}\Big{)}^{2}\right]\mathrm{d}t\right).$
(24)

Hence, clearly, ${\Delta}_{\infty}<{\Delta}_{0}$.

Finally, let us consider how the perturbed initialisation
$\tilde{{\theta}}_{0}$ looks like when ${\Delta}_{t}$ remains nonzero. Recall
that $\tilde{{\theta}}_{0}=({\alpha}_{+}^{2}-{\alpha}_{-}^{2})/4$, where
${\alpha}_{\pm,\infty}=w_{\pm,\infty}\exp\left(\pm\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t})\mathrm{d}t\right)$.
Thus,

${\alpha}_{\pm,\infty}=w_{\pm,0}\exp\left(-{\lambda}\int_{0}^{\infty}\Big{(}\frac{\dot{w}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}\mathrm{d}t\right)$

and

$\tilde{{\theta}}_{0}=\frac{1}{4}\left[w_{+,0}^{2}\exp\left(-2{\lambda}\int_{0}^{\infty}\Big{(}\frac{\dot{w}_{+,t}}{w_{+,t}}\Big{)}^{2}\mathrm{d}t\right)-w_{\\-,0}^{2}\exp\left(-2{\lambda}\int_{0}^{\infty}\Big{(}\frac{\dot{w}_{-,t}}{w_{-,t}}\Big{)}^{2}\mathrm{d}t\right)\right].$
(25)

Now, ${\alpha}_{\pm,\infty}<w_{\pm,0}\leq 2{\alpha}$, where
${\alpha}=\max(\lVert u_{0}\rVert_{\infty},\lVert v_{0}\rVert_{\infty})$ stood
for the initialisation scale. Hence, $|\tilde{{\theta}}_{0}|<{\alpha}^{2}$.

Therefore, we just proved

See 2

###  C.5 Behaviour of ${\Delta}_{\infty}$ for Small Values of $\lambda$

Since a precise asymptotic result for small $\lambda$ is technically
difficult, in this section we focus on giving some qualitative results. For
$\lambda>0$, recall that our iterates follow

${\lambda}\ddot{w}_{\pm,t}^{(\lambda)}+\dot{w}_{\pm,t}^{(\lambda)}\pm\nabla{\mathcal{L}}({\theta}_{t}^{(\lambda)})\odot
w_{\pm,t}^{(\lambda)}=0,$

where we explicitly highlight the dependency on $\lambda$. Therefore, we have

$\frac{\dot{w}_{\pm,t}^{(\lambda)}}{w_{\pm,t}^{(\lambda)}}=\mp\nabla{\mathcal{L}}({\theta}_{t}^{(\lambda)})-{\lambda}\frac{\ddot{w}^{(\lambda)}_{\pm,t}}{w_{\pm,t}}$

and

$\Big{(}\frac{\dot{w}_{\pm,t}^{(\lambda)}}{w_{\pm,t}^{(\lambda)}}\Big{)}^{2}=\nabla{\mathcal{L}}({\theta}_{t}^{(\lambda)})^{2}+{\lambda}^{2}\Big{(}\frac{\ddot{w}^{(\lambda)}_{\pm,t}}{w_{\pm,t}}\Big{)}^{2}\pm
2\lambda{\mathcal{L}}({\theta}_{t}^{(\lambda)})\Big{(}\frac{\ddot{w}^{(\lambda)}_{\pm,t}}{w_{\pm,t}}\Big{)}.$

Informally, we expect $(t\mapsto\nabla
L({\theta}_{t}^{(\lambda)}))_{0<\lambda\leq 1}\in L^{2}(0,+\infty)$ and
$(t\mapsto\lambda\frac{\ddot{w}^{(\lambda)}_{\pm,t}}{w_{\pm,t}})_{\lambda}\underset{\lambda\to
0}{\longrightarrow}0$ in $L^{2}$-norm (see Theorem 5.1 in Attouch et al.
(2000)). Hence, we get

$\int_{0}^{\infty}\Big{(}\frac{\dot{w}_{\pm,t}^{(\lambda)}}{w_{\pm,t}^{(\lambda)}}\Big{)}^{2}\underset{\lambda\to
0}{\sim}\int_{0}^{\infty}\nabla{\mathcal{L}}({\theta}_{t}^{(\lambda)})^{2}\mathrm{d}t$

and

$\displaystyle{\Delta}_{\infty}\underset{\lambda\to
0}{\approx}\Delta_{0}\exp\big{(}-2\lambda\int_{0}^{\infty}\nabla{\mathcal{L}}(\theta_{s}^{(\lambda)})^{2}\mathrm{d}s\big{)}.$

##  Appendix D Discrete-Time Results

In this section, we cover the proofs of our discrete-time results from Section
4. We first recall the SMGD recursion (13) with the $w_{\pm}$-parametrisation
from Appendix B. Initialised at $w_{\pm,0}=w_{\pm,1}\in{\mathbb{R}}^{d}$, for
$k\geq 1$, the iterates follow

$\displaystyle
w_{\pm,k+1}=w_{\pm,k}\mp{\gamma}\nabla{\mathcal{L}}_{\mathcal{B}_{k}}({\theta}_{k})\odot
w_{\pm,k}+{\beta}(w_{\pm,k}-w_{\pm,k-1}).$ (26)

In what follows, we will adapt our continuous-time proof technique to the
discrete case and identify a quantity which follows a momentum mirror descent
with time-varying potentials. Our proofs closely follow the proof techniques
from (Even et al., 2023) which considers SGD without momentum.

###  D.1 Proof of Lemma 2, Theorem 2 and Corollary 3

We start by recalling the main-paper results. The first lemma introduces two
convergent series which will appear in our main result.

See 2

The proof of the lemma can be found in the proof of the following main
theorem.

See 2

##### Proving Convergence towards an Interpolator.

By 3, we have that the iterates $w_{\pm,k}$ converge towards limiting weights
$w_{\pm,\infty}$ and that the predictors converge towards a vector
$\theta^{\texttt{MGF}}$. Taking the limit in Equation 26, we get that
$\lim_{k\to\infty}\nabla{\mathcal{L}}_{\mathcal{B}_{k}}(\theta_{k})\odot
w_{\pm,k}=0$. By 4, $w_{\pm,\infty}$ have non-zero coordinates. Therefore,
$\lim_{k\to\infty}\nabla{\mathcal{L}}_{\mathcal{B}_{k}}(\theta_{k})=0$. For
any fixed batch $\mathcal{B}\subset\\{1,\cdots,n\\}$, the sampling with or
without replacement is such that (almost surely) the set
$M_{k}\coloneqq\\{k\geq 0,\mathcal{B}_{k}=\mathcal{B}\\}$ is infinite. Hence,
by continuity of $\nabla L_{\mathcal{B}}$, $\lim_{k\to\infty,k\in M_{k}}\nabla
L_{\mathcal{B}}(\theta_{k})=\nabla L_{\mathcal{B}}(\theta^{\texttt{SMGD}})$.
Therefore, for all fixed batches $\mathcal{B}$, $\nabla
L_{\mathcal{B}}(\theta^{\texttt{SMGD}})=0$ and hence $\theta^{\texttt{SMGD}}$
interpolates the dataset.

From here on now, for ease of notation, we do the proof for deterministic MGD.
The proof for stochastic MGD is exactly the same after replacing
$\nabla{\mathcal{L}}(\theta_{k})$ with
$\nabla{\mathcal{L}}_{\mathcal{B}_{k}}(\theta_{k})$.

##### Deriving the Momentum Mirror Descent.

Recall that the set of pairs $(\gamma,\beta)$ such that there exists $k$ where
$w_{\pm,k}=0$ is negligible in ${\mathbb{R}}^{2}$. We can hence assume that
the iterates are never exactly zero, and we consider the logarithmic
reparametrisation of the iterates $w_{\pm,k}$ as

$g_{\pm,k}=\begin{cases}\ln(w_{\pm,k}),&\text{if }w_{\pm,k}>0,\\\
\ln(|w_{\pm,k}|)+i\pi,&\text{if }w_{\pm,k}<0.\end{cases}$

This way we have that that $w_{\pm,k}=\exp(g_{\pm,k})$ for all $k$. Equation
26 then becomes

$\displaystyle\exp(g_{\pm,k+1})=\exp(g_{\pm,k})\mp{\gamma}\nabla{\mathcal{L}}({\theta}_{k})\odot\exp(g_{\pm,k})+{\beta}(\exp(g_{\pm,k})-\exp(g_{\pm,k-1})).$

Dividing by $\exp(g_{\pm,k})$ yields

$\displaystyle\exp(g_{\pm,k+1}-g_{\pm,k})=1\mp{\gamma}\nabla{\mathcal{L}}({\theta}_{k})+{\beta}(1-\exp(-(g_{\pm,k}-g_{\pm,k-1})).$

Now, for $k\geq 1$, let $\delta_{\pm,k}=g_{\pm,k}-g_{\pm,k-1}$ so that we can
more compactly write the above recurrence as

$\displaystyle\exp(\delta_{\pm,k+1})=1\mp{\gamma}\nabla{\mathcal{L}}({\theta}_{k})+{\beta}(1-\exp(-\delta_{\pm,k})).$

The trick, inspired by Even et al. (2023), is to consider the function
$q(z)=\exp(z)-(1+z)$ for $z\in{\mathbb{C}}$. Importantly, note that $q(z)\geq
0$ for $z\in{\mathbb{R}}$. Using this function, we can now rewrite the
recurrence as

$\displaystyle\delta_{\pm,k+1}+q(\delta_{\pm,k+1})=\mp{\gamma}\nabla{\mathcal{L}}({\theta}_{k})+{\beta}(\delta_{\pm,k}-q(-\delta_{\pm,k})).$

Setting the residues $Q_{\pm,k}\coloneqq
q(\delta_{\pm,k+1})+{\beta}q(-\delta_{\pm,k})$ leads to

$\displaystyle\delta_{\pm,k+1}={\beta}\delta_{\pm,k}\mp{\gamma}\nabla{\mathcal{L}}({\theta}_{k})-Q_{\pm,k}.$

This can be seen as a first-order recurrence relation with variable
coefficients. For ${\beta}=0$ we exactly recover the analysis from Even et al.
(2023). For ${\beta}>0$, since $\delta_{\pm,1}=0$, for $m\geq 1$, we can
expand the relation as

$\displaystyle\delta_{\pm,m+1}=-\sum_{k=1}^{m}{\beta}^{m-k}\left[\pm{\gamma}\nabla{\mathcal{L}}({\theta}_{k})+Q_{\pm,k}\right].$

Summing over $m$, we now get for $N\geq 1$ the following expression:

$\displaystyle g_{\pm,N+1}-g_{\pm,1}$
$\displaystyle=\sum_{m=1}^{N}\delta_{\pm,m+1}$