# Mixed Dynamics In Linear Networks:
Unifying the Lazy and Active Regimes

Zhenfeng Tu
Courant Institute
New York University
New York, NY 10012
<EMAIL_ADDRESS>
&Santiago Aranguri
Courant Institute
New York University
New York, NY 10012
<EMAIL_ADDRESS>
&Arthur Jacot
Courant Institute
New York University
New York, NY 10012
<EMAIL_ADDRESS>
###### Abstract

The training dynamics of linear networks are well studied in two distinct
setups: the lazy regime and balanced/active regime, depending on the
initialization and width of the network. We provide a surprisingly simple
unyfing formula for the evolution of the learned matrix that contains as
special cases both lazy and balanced regimes but also a mixed regime in
between the two. In the mixed regime, a part of the network is lazy while the
other is balanced. More precisely the network is lazy along singular values
that are below a certain threshold and balanced along those that are above the
same threshold. At initialization, all singular values are lazy, allowing for
the network to align itself with the task, so that later in time, when some of
the singular value cross the threshold and become active they will converge
rapidly (convergence in the balanced regime is notoriously difficult in the
absence of alignment). The mixed regime is the ‘best of both worlds’: it
converges from any random initialization (in contrast to balanced dynamics
which require special initialization), and has a low rank bias (absent in the
lazy dynamics). This allows us to prove an almost complete phase diagram of
training behavior as a function of the variance at initialization and the
width, for a MSE training task.

##  1 Introduction

Whether in linear networks or nonlinear ones, there has been a lot of interest
in the distinction between the lazy regime jacot2018neural  and the active
regime Chizat2018 ; Rotskoff2018 ; chizat2018note ; yang2020feature ;
bordelon2022deep_MF  as the number of neurons grows towards infinity. In the
lazy regime the training dynamics become linear, so that they can be easily
described in terms of the Neural Tangent Kernel (NTK) jacot2018neural ;
exact_arora2019 ; yang2020NTK ; liu2020_NTK-PL-inequ , while the active regime
exhibits complex nonlinear dynamics. While our understanding of the active
regime remains much more limited, it appears to be characterized by the
emergence of feature learninggeiger2019disentangling ; yang2020feature , and
of a form of sparsity advani_2017_independent_diag ; bach2017_F1_norm ;
abbe_2021_staircase ; abbe_2022_staircase  (the type of sparsity observed
depends on the network type bach2017_F1_norm ; dai_2021_repres_cost_DLN ;
jacot_2022_BN_rank_short ; jacot_2023_bottleneck2 , but we will focus on
fully-connected linear networks which exhibit a rank sparsity in the learned
linear map arora_2018_depth_speed_Lp ; li2020towards ; jacot-2021-DLN-Saddle ;
wang_2023_bias_SGD_L2 ) which are both absent in the lazy regime .

Note that even though it is common to talk of the ‘the’ active regime, we do
not know yet whether there is only one or multiple active regimes. Indeed the
term active regime is usually used to describe any regime that differs from
the lazy regime and exhibit some form of feature learning. Though we do not
have an complete understanding of where the lazy regimes ends and the active
regime(s) start, we know that the lazy regime requires extreme
overparametrization (a large number of neurons in comparison to the number of
datapoints) Allen-Zhu2018 ; Du2019 , a ‘large’ initialization of the weights
chizat2018note , a small learning rate, and early stopping when using a cross-
entropy loss or weight decay . Indeed, active regimes have been observed by
breaking either of these requirements: taking limits with mild or no
overparametrization arous2022summary_stats , taking smaller or even
vanishingly small initializations li2020towards ; jacot-2021-DLN-Saddle ,
using large learning rates lewkowycz_2020_large_lr  or SGD
Pesme_2021_SGD_bias_diag_nets ; wang_2023_bias_SGD_L2 , or studying the late
training dynamics with the cross-entropy loss Ji_2018_directional ;
chizat_2020_implicit_bias  or weight decay lewkowycz2020_L2_training ;
Ongie_2020_repres_bounded_norm_shallow_ReLU_net ; jacot_2022_L2_reformulation
; jacot_2022_BN_rank . Though each of these can lead to active regimes with
significantly different dynamics, they often lead to similar types of feature
learning and sparsity.

In this paper, we study this transition in the context of linear networks and
focus mainly on the effects of the width $w$ and the variance of the weights
at initialization $\sigma^{2}$, and give a precise and almost complete phase
diagram, showing the transitions between lazy and active regimes. In this
setting, we will show that there typically is only ‘one’ active regime, which
is the same (up to approximation) as the already well-studied balanced regime
saxe_2014_exact ; arora_2018_depth_speed_Lp ; arora_2019_matrix_factorization
.

But our result also paint a more subtle picture than the lazy/active
dichotomy. We propose a more granular approach, where at a certain time some
part of the network can be in the lazy regime, while others are in the active
or balanced regime. More precisely the network is lazy along the singular
values of the matrix represented by the network that are smaller than
$\sigma^{2}w$, and in the active regime along the singular values larger than
$\sigma^{2}w$.

###  1.1 Contributions

We consider the training dynamics of shallow linear networks
$A_{\theta}=W_{2}W_{1}$ and show that for large enough width $w$ (the inner
dimension), and a iid $\mathcal{N}(0,\sigma^{2})$ initialization of all
weights, the dynamics of $A_{\theta(t)}$ as a result of training the
parameters $\theta=(W_{1},W_{2})$ with GD/GF on the loss
$\mathcal{L}(\theta)=C(A_{\theta})$ for a general matrix cost $C$ is
approximately given by the self-consistent dynamics

$\partial_{t}A_{\theta(t)}\approx-\sqrt{A_{\theta}A_{\theta}^{T}+\sigma^{4}w^{2}I}\nabla
C(A_{\theta})-\nabla
C(A_{\theta})\sqrt{A_{\theta}^{T}A_{\theta}+\sigma^{4}w^{2}I}.$ (1)

These dynamics contain as special cases both the lazy dynamics

$\partial_{t}A_{\theta(t)}\approx-2\sigma^{2}w\nabla C(A_{\theta})$

when $\sigma^{2}w\gg\lambda_{max}(A_{\theta})$ and the balanced dynamics

$\partial_{t}A_{\theta(t)}=\sqrt{A_{\theta}A_{\theta}^{T}}\nabla
C(A_{\theta})-\nabla C(A_{\theta})\sqrt{A_{\theta}^{T}A_{\theta}}$

when $\sigma^{2}w\ll\lambda_{min}(A_{\theta})$. But it also reveals the whole
spectrum of mixed dynamics in between, where some singular values of
$A_{\theta}$ are below the $\sigma^{2}w$ threshold and some are above it.

This suggests that the lazy/active transition is best understood at a more
granular level, where at each time $t$ every singular value of $A_{\theta}$
can either be lazy or active/balanced. The mixed regime is the best of both
worlds: on one hand, since $\sqrt{A_{\theta}A_{\theta}^{T}+\sigma^{4}w^{2}I}$
is always positive definite, the network can never get stuck at a saddle/local
minimum as can happen in the balanced regime, on the other hand there is a
momentum effect where the dynamics along large singular values is much faster
than along the small ones, leading to an incremental learning behavior and a
low-rank bias, which is absent in lazy learning. By choosing the threshold
$\sigma^{2}w$ adequately, one can best take advantage of these two phenomenon.

Finally, we focus on the task of recovering a low-rank $d\times d$ matrix
$A^{*}$ from noisy observations $A^{*}+E$, training on the MSE error
$\frac{1}{d^{2}}\left\|A_{\theta}-(A^{*}+E)\right\|^{2}_{F}$ in the limit as
the dimension $d$, width $w$ and variance $\sigma^{2}$ scale together with
scaling laws $w=d^{\gamma_{w}}$ and $\sigma^{2}=d^{\gamma_{\sigma^{2}}}$. We
describe the training dynamics for almost all reasonable scalings
$\gamma_{w},\gamma_{\sigma^{2}}$ leading to a phase diagram with two main
regimes:

  * •

Lazy ($1<\gamma_{\sigma^{2}}+\gamma_{w}$) where all singular values remain
below the threshold $\sigma^{2}w$ throughout training, and where the network
fails to recover $A^{*}$ due to the absence of low-rank bias.

  * •

_Active (__$1 >\gamma_{\sigma^{2}}+\gamma_{w}$__)__where
$K=\mathrm{Rank}A^{*}$ pass the threshold and fit $A^{*}$ before the other
singular values have time to fit the noise $E$, leading to the recovery of
$A^{*}$._

There are two other degenerate regimes that we avoid: the underparametrized
regime when $w<d$ (or $\gamma_{w}\ll 1$) where the rank is constrained by the
network architecture rather than the training dynamics, and the noisy regime
$2\gamma_{\sigma^{2}}+\gamma_{w}+1>0$ where the variance of the entries of
$A_{\theta(0)}$ at initialization is infinite.

Figure 1: For both plots, we train either using gradient descent or the self-
consistent dynamics from equation (1), with the scaling
$\gamma_{\sigma^{2}}=-1.85,$ $\gamma_{w}=2.25$ which lies in the active
regime. (Left panel): We plot train and test error for both dynamics. We
observe that the train/test error for gradient descent is very close to the
train/test error for the self-consistent dynamics. (Right panel): We plot with
a solid line the singular values of $A_{\theta(t)}$ when running the self-
consistent dynamics, and use a dashed line for the singular values from
running gradient descent. In this experiment, $\text{Rank}A^{\star}=5.$ We use
different colors for the $5$ largest singular values and the same color for
the remaining singular values. We can see how the $5$ largest singular values
‘speed up’ as they cross the $\sigma^{2}w$ threshold, allowing them to
converge earlier than the rest. The minimal test error is achieved in the
short period where the large singular values have converged but not the rest.

###  1.2 Previous Works

Linear networks have been used as a testing ground, a stepping stone on the
way to understand nonlinear networks. Linear networks and their training
dynamics are many ways much simpler than nonlinear ones, but in spite of a
long research history, our understanding remains limited.

The setting that is best understood is that of diagonal linear networks where
the dynamics decouple along the diagonal entries leading to an incremental
learning behavior and a sparsity bias saxe_2014_exact ;
advani_2017_independent_diag ; Saxe_2019_independent_diag ;
gidel_2019_independent_diag , some of this analysis has been extended to
include effects of initialization scale woodworth_2020_diagnet_bias  and SGD
Pesme_2021_SGD_bias_diag_nets . While the same decoupling happens in general
linear with diagonal initializations and diagonal task, it remains an
extremely strong assumption.

Some work has been done to prove similar incremental learning dynamics outside
the diagonal case li2020towards ; jacot-2021-DLN-Saddle ;
jin_2023_incremental_learning . For shallow networks, the training dynamics
with MSE can be explicitly solved Fukumizu1998_solution_shallow_DLN  but
remain very complex so that one needs to assume some form of alignment to
guarantee convergence Braun_2022_exact_DLN . For deeper networks there exists
explicit formulas in the mean-field limit where the number of neurons grows to
infinity chizat_2022_mean_field_DLN .

A very powerful tool in the analysis of linear network is its training
invariants, and the balancedness condition which greatly simplifies the
dynamics arora_2018_DLN_convergence ; arora_2018_depth_speed_Lp . Balanced
networks exhibit a momentum effect, where the training dynamics along a
singular value $s_{i}$ have ‘speed’ proportional to $s_{i}$ itself (or $s_{i}$
to some power), while this momentum effect seems to be key to understand the
low-rank bias of linear networks arora_2019_matrix_factorization , it also
means that one needs to guarantee that the dynamics never approach zero, which
is one the main hurdle towards proving convergence in balanced networks. To
solve this issue, recent work has focused on initialization that slightly
imbalanced xiong_2023_overparam_balanced_hurt_conv ;
min_2021_imbalance_shallow_DLN ; tarmoun_2021_GF_matrix_fact ;
min_2023_imbalance_deep ; xu_2023_imbalance_GD . This suggests that it is key
to find the right balance between balancedness and imbalancedness to obtain
both fast convergence and low-rank bias.

###  1.3 Setup

We will study shallow linear networks (or matrix factorization) where a
$d_{out}\times d_{in}$ matrix $A_{\theta}$ is represented as the product of
two matrices $A_{\theta}=W_{2}W_{1}$, where the weight matrices $W_{1}$ and
$W_{2}$ are respectively $w\times d_{in}$ and $d_{out}\times w$ dimensional,
for some width $w$. The parameters $\theta$ of the network are the
concatenation of the entries of both submatrices $\theta=(W_{1},W_{2})$.

The parameters $\theta$ are learned in the following manner: they are
initialized as i.i.d. Gaussian $\mathcal{N}(0,\sigma^{2})$, and then optimized
with gradient descent to minimize a loss $\mathcal{L}(\theta)=C(A_{\theta})$.
Though most of our analysis works for general convex costs
$C:\mathbb{R}^{d_{out}\times d_{in}}\to\mathbb{R}$ on matrices, we will in the
second part focus on the task of recovering a low-rank matrix $A^{*}$ from
noisy observations $A^{*}+E$, by training a linear network $A_{\theta}$ on the
MSE loss

$\mathcal{L}(\theta)=\frac{1}{d^{2}}\left\|A_{\theta}-(A^{*}+E)\right\|_{F}^{2}.$

The width $w$ allows us to control the over parametrzation, indeed the set of
matrices that can be represented by a network of width $w$ is the set
$\mathcal{M}_{\leq w}$ of matrices of rank $w$ or less. The overparametrized
regime is when $w\geq\text{min}\\{d_{in},d_{out}\\}$ because all matrices can
be represented in this case.

###  1.4 Lazy Dynamics

The evolution of the weight matrices during gradient descent with learning
rate $\eta$ is given by

$\displaystyle W_{1}(t+1)$ $\displaystyle=W_{1}(t)-\eta W_{2}^{T}(t)\nabla
C(A_{\theta(t)})$ $\displaystyle W_{2}(t+1)$
$\displaystyle=W_{2}(t)-\eta\nabla C(A_{\theta(t)})W_{1}^{T}(t)$

where we view the gradient $\nabla C(A_{\theta(t)})$ of the cost $C$ as a
$d_{out}\times d_{in}$ matrix, which for the MSE cost equals $\nabla
C(A_{\theta(t)})=2d^{-2}(A_{\theta(t)}-(A^{*}+E))$.

But we care more about the evolution of the complete matrix
$A_{\theta(t)}=W_{2}(t)W_{1}(t)$ induced by the evolution of
$W_{1}(t),W_{2}(t)$, which can be approximated by

$A_{\theta(t+1)}=A_{\theta(t)}-\eta W_{2}(t)W_{2}^{T}(t)\nabla
C(A_{\theta(t)})-\eta\nabla C(A_{\theta(t)})W_{1}^{T}(t)W_{1}(t)+O(\eta^{2}).$
(2)

Thus we see that if we can describe the matrices $C_{1}=W_{1}^{T}W_{1}$ and
$C_{2}=W_{2}W_{2}^{T}$ throughout training, then we can describe the evolution
of $A_{\theta(t)}$.

When $w$ is very large, we end up in the lazy regime where the parameters move
enough up to a time $t$ to change $A_{\theta(t)}$, but not enough to change
$C_{1},C_{2}$111 To be more precise the direction in parameter space that
change $C_{1},C_{2}$ are approximately orthogonal to those that change
$A_{\theta}$, and GD/GF only moves along the later direction., allowing us to
make the approximation $C_{i}(t)\approx C_{i}(0)$. Furthermore at
initialization these matrices concentrate as $w\to\infty$ around their
expectations $\mathbb{E}\left[C_{1}\right]=\sigma^{2}wI_{d_{in}}$,
$\mathbb{E}\left[C_{2}\right]=\sigma^{2}wI_{d_{out}}$. The GD dynamics can
then be approximated by the much simpler dynamics:

$A_{\theta(t+1)}=A_{\theta(t)}-2\eta\sigma^{2}w\nabla C(A_{\theta(t)}),$

which are equivalent to doing GD on the cost $C$ directly with a learning rate
of $2\eta\sigma^{2}w$.

One can then easily prove exponential convergence for any convex cost $C$
following the convergence analysis of traditional linear models. But we can
see the absence of feature learning from the fact that the covariance $C_{1}$
of the ‘feature map’ $W_{1}$ is (approximately) constant. More problematic in
the context of low-rank matrix recovery is the absence of low-rank bias,
indeed one can easily solve the dynamics to obtain

$A_{\theta(t)}=(A^{*}+E)+(1-4d^{-2}\eta\sigma^{2}w)^{t}(A_{\theta(0)}-(A^{*}+E)),$

and since $\mathbb{E}A_{\theta(0)}=0$ we obtain

$\mathbb{E}\left[A_{\theta(t)}\right]=\left(1-(1-4d^{-2}\eta\sigma^{2}w)^{t}\right)(A^{*}+E).$

The expected test error $\mathbb{E}\left\|A_{\theta(t)}-A^{*}\right\|^{2}$ is
therefore lower bounded by

$\left\|\mathbb{E}A_{\theta(t)}-A^{*}\right\|^{2}=\left\|(1-4d^{-2}\eta\sigma^{2}w)^{t}A^{*}+(1-(1-4d^{-2}\eta\sigma^{2}w)^{t})E\right\|^{2}$

which never approaches zero.

In linear networks, there is no advantage to being in the lazy regime, as we
simply recover a simple linear model at an additional cost of more parameters
and thus more compute. But we will see that a short period of lazy regime at
the beginning of training plays a crucial role in making sure that the
subsequent active regime starts from an ‘aligned’ state.

###  1.5 Balanced Dynamics

There has been much more focus on so-called balanced linear networks, which
are networks that satisfy the balanced condition
$W_{1}W_{1}^{T}=W_{2}^{T}W_{2}$. If the network is balanced at initialization,
it remains so throughout training, because, the difference
$W_{1}W_{1}^{T}-W_{2}^{T}W_{2}$ is an invariant of GF (and an approximate
invariant of GD with small enough learning rate).

First observe that the balanced condition implies the following shared
eigendecomposition $W_{1}W_{1}^{T}=W_{2}^{T}W_{2}=USU^{T}$. This implies the
following shared SVD decompositions $W_{1}=U\sqrt{S}U_{in}^{T}$,
$W_{2}=U_{out}\sqrt{S}U^{T}$ and $A_{\theta}=U_{out}SU_{in}^{T}$. Furthermore,
we have $C_{1}=U_{in}SU_{in}^{T}=\sqrt{A_{\theta}^{T}A_{\theta}}$ and
$C_{2}=U_{out}SU_{out}^{T}=\sqrt{A_{\theta}A_{\theta}^{T}}$, which leads to
self-consistent dynamics for $A_{\theta(t)}$:

$A_{\theta(t+1)}=A_{\theta(t)}-\eta\sqrt{A_{\theta(t)}A_{\theta(t)}^{T}}\nabla
C(A_{\theta(t)})-\eta\nabla
C(A_{\theta(t)})\sqrt{A_{\theta(t)}^{T}A_{\theta(t)}}+O(\eta^{2}).$

Now these dynamics are quite complex in general, and it remains difficult to
prove convergence. Indeed one can easily find initializations $A_{\theta(0)}$
that will not converge, for example if $A_{\theta(0)}=0$ then GD will remain
stuck there. A lot of work has been dedicated to finding conditions that
guarantee the convergence of the above dynamics arora_2018_DLN_convergence ;
Braun_2022_exact_DLN , but these assumptions are often quite strong.

The simplest initialization that guarantees convergence (and the one that will
be most relevant to our analysis) is the positively aligned initialization. If
at initialization $A_{\theta(0)}$ and $A^{*}+E$ are ‘aligned’, i.e. shares the
same singular vectors $A_{\theta(0)}=U_{out}SU_{in}^{T}$ and
$A^{*}+E=U_{out}S^{*}U_{in}^{T}$, then they will remain aligned throughout
training $A_{\theta(t)}=U_{out}S(t)U_{in}^{T}$ and the dynamics decouple along
each singular value

$s_{i}(t+1)=s_{i}(t)+2\eta\left|s_{i}(t)\right|\left(s_{i}^{*}-s_{i}(t)\right)+O(\eta^{2}).$

Since we always have $s_{i}^{*}\geq 0$, then for small enough learning rates
$\eta$, we see that if $s_{i}(0)\in(0,s_{i}^{*}]$ it will grow monotonically
and converge to $s_{i}^{*}$; if $s_{i}(0)>s_{i}^{*}$ it will decrease
monotonically to $s_{i}^{*}$, and if $s_{i}(0)\leq 0$ it will increase and
converge to $0$. Thus one can guarantee convergence if we further assume
positive alignment $s_{i}(0)>0$.

The advantage is that there is a momentum effect in the form of the prefactor
$\left|s_{i}(t)\right|$, which implies that the dynamics along large singular
values are faster than along small ones. As a result, if all singular values
are initialized with the same small value, then they will at first grow very
slowly until they reach a critical size where the momentum effect will make
them converge very fast. The singular values aligned with the top singular
values of $A^{*}+E$ will reach this threshold much faster, and they will
therefore converge to approximately their final value $s_{i}=s_{i}^{*}$ at a
time when the other singular values are still basically zero. If we stop
training at this time then the linear network will have essentially learned
only the top $K$ singular values of $A^{*}+E$, which is a good approximation
for $A^{*}$, leading to a small test error (see gidel_2019_independent_diag
for details).

But this analysis relies on the very strong assumption of positive alignment
at initialization. If we do not assume a positive alignment and assume that
the $s_{i}$ are random (i.i.d. w.r.t. a symmetric distribution), then each
$s_{i}$ has probability $\nicefrac{{1}}{{2}}$ of starting with a negative
alignment and getting stuck at zero, which means that with high probability
training will fail to recover $A^{*}$ and will recover only a random subset of
the singular values of $A^{*}$. The presence of these attractive saddles shows
the complexity of the balanced dynamics.

A limitation of this approach is that it requires a quadratic cost and a very
specific initialization, and in the case of positive alignment, an
initialization that requires knowledge of the (SVD of the) true function
$A^{*}$. Nevertheless, the positively aligned and balanced dynamics seem to
capture some qualitative phenomenon that has been observed empirically outside
of this restricted setting. This is the phenomenon of incremental learning,
where if the singular values are initialized as very small, they first grow
very slowly, but the multiplicative momentum will lead to come up one by one
in a very abrupt manner, and this leads to a low rank bias where the network
first only fits the largest singular value, then two largest, and so on. More
generally, this can be interpreted as the network performing a greedy low-rank
algorithm li2020towards .

Our analysis will confirm the fact that positive alignment happens naturally
as a result of a short period of lazy training, allowing us to prove similar
decoupling and incremental learning for a general random initialization.

######  Remark.

We can define the time dependent map $\Theta(G;t)=C_{2}(t)G+GC_{1}(t)$, so
that the GD dynamics can be rewritten as
$A_{\theta(t+1)}=A_{\theta(t)}-\eta\Theta(\nabla
C(A_{\theta(t)}),t)+O(\eta^{2})$. The map $\Theta$ is none other than the NTK
for shallow linear networks, but it has also been called the preconditioning
matrix in previous work arora_2018_depth_speed_Lp . The lazy regime is then
characterized by the NTK $\Theta$ being approximately equal to the time-
independent NTK $\Theta^{\text{lazy}}(G)=2\sigma^{2}wG$, whereas the balanced
regime is characterized by the time-dependent
$\Theta^{\text{bal}}(G;t)=\sqrt{A_{\theta(t)}A_{\theta(t)}^{T}}G+G\sqrt{A_{\theta(t)}^{T}A_{\theta(t)}}$,
with the distinction that the time dependence is only through $A_{\theta(t)}$.

##  2 Mixed Lazy/Balanced Dynamics

Both lazy and balanced dynamics have the surprising but very useful property
that the evolution of the network matrix $A_{\theta}$ is approximately self-
consistent: the evolution of $A_{\theta}$ can be expressed in terms of itself.
The lazy approximation becomes correct for a sufficiently large
initialization, while the balanced one is correct for a balanced
initialization. However, for most initializations, neither of these
approximations are correct.

We fill this gap by providing a self-consistent evolution of $A_{\theta}$ that
applies for any initialization scale:

$\partial_{t}A_{\theta(t+1)}\approx-\eta\sqrt{A_{\theta(t)}A_{\theta(t)}^{T}+\sigma^{4}w^{2}I}\nabla
C(A_{t})-\eta\nabla
C(A_{t})\sqrt{A_{\theta(t)}^{T}A_{\theta(t)}+\sigma^{4}w^{2}I}.$

This approximation is formalized in the following theorem, denoting
$\hat{C}_{1}(t)=\sqrt{A_{\theta(t)}^{T}A_{\theta(t)}+\sigma^{4}w^{2}I}$ and
$\hat{C}_{2}(t)=\sqrt{A_{\theta(t)}A_{\theta(t)}^{T}+\sigma^{4}w^{2}I}$

######  Theorem 1.

For a linear net $A_{\theta}=W_{2}W_{1}$ with width $w$, initialized with
i.i.d. $\mathcal{N}(0,\sigma^{2})$ weights and trained with Gradient Flow, we
have with prob. $1-\delta$ for all times $t$

$\displaystyle\left\|C_{1}(t)-\hat{C}_{1}(t)\right\|_{op},\left\|C_{2}(t)-\hat{C}_{2}(t)\right\|_{op}$
$\displaystyle\leq\text{min}\left\\{O(\sigma^{2}w),O\left(\sqrt{\frac{d}{w}}\left\|C_{1}\right\|_{op}\right)\right\\}.$

###### Proof.

(sketch) The quantity $W_{1}W_{1}^{T}-W_{2}^{T}W_{2}$ is invariant under GF
(and approximately so under GD) and it is approximately equal to
$\sigma^{2}w(P_{1}-P_{2})$ for two orthogonal projections $P_{1},P_{2}$ (at
initialization and for all subsequent times because of the invariance). We
therefore have

$W_{1}^{T}(W_{1}W_{1}^{T}-W_{2}^{T}W_{2})^{2}W_{1}\approx\sigma^{4}w^{2}W_{1}^{T}(P_{1}+P_{2})W_{1}\approx\sigma^{4}w^{2}C_{1}.$

Thus the pairs $C_{1},C_{2}$ approximately satisfy the following equations:

$\displaystyle 0$ $\displaystyle\approx
C_{1}^{3}-A_{\theta}^{T}A_{\theta}C_{1}-C_{1}A_{\theta}^{T}A_{\theta}-\sigma^{4}w^{2}C_{1}+A_{\theta}^{T}C_{2}A_{\theta}$
$\displaystyle 0$ $\displaystyle\approx
C_{2}^{3}-A_{\theta}A_{\theta}^{2}C_{2}-C_{2}A_{\theta}A_{\theta}^{T}-\sigma^{4}w^{2}C_{2}+A_{\theta}C_{1}A_{\theta}^{T}.$

The pair $\hat{C}_{1},\hat{C}_{2}$ is a solution of the above, and one can
show that $C_{1},C_{2}$ must approach them and not any of the other solutions.
∎

We see how both the lazy and balanced dynamics appear as special cases
depending on how large the variance at initialization $\sigma^{2}$ is in
comparison to the singular values of the matrix $A_{\theta(t)}$:

  * •

Lazy: When $\sigma^{2}w\gg s_{max}(A_{\theta(t)})$, then
$\hat{C}_{1}\approx\sigma^{2}wI_{d_{out}}$ and
$\hat{C}_{2}\approx\sigma^{2}wI_{d_{in}}$, recovering the lazy dynamics.

  * •

Balanced: When $\sigma^{2}w\ll s_{min}(A_{\theta(t)})$, then
$\hat{C}_{1}\approx\sqrt{A_{\theta(t)}^{T}A_{\theta(t)}}$ and
$\hat{C}_{2}\approx\sqrt{A_{\theta(t)}A_{\theta(t)}^{T}}$, recovering the
balanced dynamics.

But clearly there can be times when neither conditions are satisfied, when
some singular values of $A_{\theta(t)}$ are larger than the threshold
$\sigma^{2}w$ while others are smaller, in such cases we are in a mixed
regime, where the network is lazy along the small singular values of
$A_{\theta(t)}$ ($s_{i}\ll\sigma^{2}w$) and active/balanced along the large
ones ($s_{i}\gg\sigma^{2}w$).

At initialization, the singular values are of size $\sigma^{2}\sqrt{wd}$. This
implies that with overparametrization ($w\gg d$), all singular values start in
the lazy regime and follow the simple lazy dynamics, which may (or may not)
lead to some singular growing and crossing the $\sigma^{2}w$ threshold, at
which point they will switch to balanced dynamics (after a short transition
period when the singular value is around the threshold
$s_{i}\approx\sigma^{2}w$). Once a singular value is far past the threshold
$s_{i}\gg\sigma^{2}w$, training along this singular value will be much faster
than along the lazy singular values (this speed up can be seen in Figure 1).
This allows the newly active singular values to converge while the lazy
singular values remain almost constant. Once the active singular values have
converged, the slow training of the remaining lazy singular values continues
until some of these singular values reaches the threshold, or until GD
converges.

This type of behavior is illustrated by the following formula, which describes
the derivative in time of the $i$-th singular value $s_{i,t}$ of $A_{t}$, with
singular vectors $u_{i,t},v_{i,t}$:

$s_{i,t+1}-s_{i,t}\approx\eta_{t}u_{i,t}^{T}\partial_{t}A_{\theta(t)}v_{i,t}\approx-2\eta_{t}\sqrt{s_{i,t}^{2}+\sigma^{4}w^{2}}u_{i,t}^{T}\nabla
C(A_{\theta(t)})v_{i,t},$

where the prefactor $2\eta_{t}\sqrt{s_{i,t}^{2}+\sigma^{4}w^{2}}$ describes
the effective learning rate along the $i$-th singular value, which depends on
the $i$-th singular value $s_{i,t}$ itself.

This suggests that it is more natural to distinguish between the lazy and
active regime at a much more granular level: at every time $t$ a singular
value can be either active or lazy (or very close to the transition but this
typically only happens for a very short time). In contrast, the traditional
definition of the lazy regime was defined for a whole network and over the
whole training time. To avoid confusion, we call this the pure lazy regime,
where all singular values remain lazy throughout training. This begs the
question of whether a pure balanced regime also exists, but all singular
values will always be lazy for at least a short time period (assuming $w>d$),
and as we will see this short lazy period plays a crucial role in aligning the
network so that the subsequent balanced regime can learn successfully. A pure
balanced regime can only be obtained in the underparametrized regime, or by
taking a balanced initialization instead of the traditional i.i.d. random
initialization.

While this challenges the traditional lazy/active dichotomy, it also
reinforces it, as it shows that there is no fundamentally different third
regime, only lazy, active, and some mix of the two. Theorem 1 thus allows us
to revisit previous descriptions of lazy and balanced dynamics and ‘glue them
together’ to extend them to the general case. This simple strategy will allow
to almost fully ‘fill in the phase diagram’, i.e. describe the dynamics,
convergence and generalization properties of DLNs for almost all reasonable
initialization scales $\sigma^{2}$ and widths $w$.

######  Remark.

The transition of a singular value $s_{i}$ from lazy to active can be
understood as a form of alignment happening in the hidden layer: the two
vectors $W_{1}v_{i}$ and $W_{2}^{T}u_{i}$ for $u_{i},v_{i}$ the left and right
singular vectors of $s_{i}$ are orthogonal in the lazy regime and become
perpendicular in the balanced regime. Indeed the normalized scalar product of
these two vectors satisfies

$\frac{u_{i}^{T}W_{2}W_{1}v_{i}}{\left\|W_{2}^{T}u_{i}\right\|\left\|W_{1}v_{i}\right\|}=\frac{s_{i}}{\sqrt{u_{i}^{T}C_{2}u_{i}}\sqrt{u_{i}^{T}C_{1}u_{i}}}\approx\frac{s_{i}}{\sqrt{s_{i}^{2}+\sigma^{4}w^{2}}}$

which is close to zero for lazy singular values $s_{i}\ll\sigma^{2}w$ and
close to one for active ones $s_{i}\gg\sigma^{2}w$.

Figure 2: As a function of $\gamma_{\sigma^{2}},\gamma_{w},$ we run GD and
plot different quantities. Our theoretical results only apply to the top left
region for $\gamma_{w}>1$ and below the red line, although these plots suggest
that some results may extend to smaller $\gamma_{w}$s. (Top left panel): We
plot the smallest test error $\frac{1}{d^{2}}\|A_{\theta(t)}-A^{*}\|_{F}^{2}$
in the whole run. The active region (below the black line) has a small error
while the lazy region does not. (Top right panel): We plot the stable rank of
$A_{\theta(t)}$ (defined as
$\|A_{\theta(t)}\|_{F}^{2}/\|A_{\theta(t)}\|_{\text{op}}^{2}$) at the time of
minimal test error. In this experiment, we took $\text{Rank}A^{*}=5.$ We see
that the active region has approximately the correct rank while the lazy
region overestimates it. (Bottom left panel): We plot the number of iterations
until minimal test error, illustrating the trade-off between test error and
training time. (Bottom right panel): We compute
$\ln\left(\frac{1}{d^{2}}\|A_{\theta(t)}-\hat{A}_{\theta(t)}\|_{F}^{2}\right)$
where $A_{\theta(t)}$ comes from GD and $\hat{A}_{\theta(t)}$ from the self-
consistent dynamics. We observe that this distance is not only small for the
region where our theoretical results apply but also almost everywhere outside
this region.

###  2.1 Phase Diagram for MSE

To illustrate the power of Theorem 1 we provide a phase diagram of the
behavior of large shallow networks on a MSE task, for almost all (reasonable)
choices of width $w$ and variance $\sigma^{2}$ scalings.

We want to recover a rank $K$ and $d\times d$-dimensional matrix $A^{*}$ with
$s_{i}(A^{*})=da_{i}$ for some $a_{1}\geq a_{2}\geq\dots\geq a_{K}$
independent of the dimension $d$. We however only observe a noisy version
$A^{*}+E$ for some $E$ such that $\left\|E\right\|_{op}\leq c_{0}d^{\delta}$.
One could imagine $E$ to have iid random Gaussian entries $\mathcal{N}(0,1)$
in which case $\left\|E\right\|_{op}\leq c_{0}\sqrt{d}$ with high probability.

As the dimension $d$ grows, the size of the network needs to scale too, as
well as the initialization variance, but it is unclear what is the optimal way
to choose $w$ and $\sigma^{2}$. We will therefore consider general scalings
$w=d^{\gamma_{w}}$ and $\sigma^{2}=d^{\gamma_{\sigma^{2}}}$. We will now
describe the $(\gamma_{w},\gamma_{\sigma^{2}})$-phase diagram which features 4
regimes: underparametrized, infinite-noise, lazy and mixed/active.

We can identify a region of ‘reasonable’ pairs
$(\gamma_{\sigma^{2}},\gamma_{w})$ by ruling out degenerate behavior. First,
the width $w$ needs to be larger than the dimension $d$, since a network of
width $w$ can only represent matrices of rank $w$ or less, this means that we
need $\gamma_{w}\geq 1$. Another constraint comes from the variance of
$A_{\theta}$ at initialization: the entries $A_{\theta(0),ij}$ at
initialization have variance $\sigma^{4}w$. We want this variance to not blow
up as $d$ grows which implies that we need
$2\gamma_{\sigma^{2}}+\gamma_{w}\leq 0$.

Now within this reasonable region we observe two regimes, the pure lazy regime
for $1<\gamma_{\sigma^{2}}+\gamma_{w}$ where the network simply fits $A^{*}+E$
thus failing to learn $A^{*}$ and the mixed regime for
$1>\gamma_{\sigma^{2}}+\gamma_{w}$ where the dynamics are lazy for a short
amount of time until $K$ singular values grow large enough to switch to the
balanced dynamics and fit the true matrix $A^{*}$.

######  Theorem 2.

For pairs $\gamma_{w},\gamma_{\sigma^{2}}$ such that $\gamma_{w}>1$ and
$2\gamma_{\sigma^{2}}+\gamma_{w}<0$, we have two regimes:

  * •

Lazy ($1<\gamma_{\sigma^{2}}+\gamma_{w}$): with a learning rate
$\eta\ll\frac{d^{2}}{\sigma^{2}w}$ we have that for all time $t$,
$\frac{1}{d^{2}}\left\|A_{\theta(t)}-A^{*}\right\|_{F}^{2}\geq c$.

  * •

Active ($1>\gamma_{\sigma^{2}}+\gamma_{w}$): with a learning rate
$\eta\ll\frac{d^{2}}{s_{1}(A^{*})}\sim d$, and at time

$t=\frac{1}{\eta}\left(\frac{\Delta}{a_{K}}+\frac{2\text{max}(1,2\Delta)}{c(a_{1},\ldots,a_{K})}+\frac{\text{max}(1,2\Delta)}{2a_{K}}\right)d\log
d+\eta^{-1}O(d\log\log d),$

for $\Delta=1-\gamma_{\sigma^{2}}-\gamma_{w}>0$, we have that

$\frac{1}{d^{2}}\left\|A_{\theta(t)}-A^{*}\right\|_{F}^{2}\leq
O(\sigma^{4}w+\frac{\sigma^{4}w^{2}\log^{2}d}{d^{2}}+d^{-\frac{1}{2}}+\frac{\sigma^{2}w}{d}+\eta^{2}\frac{\log^{2}d}{d^{2}}),$

for $c(a_{1},\ldots,a_{K})=\frac{\text{min}_{k,j:a_{k}\neq
a_{j}}\absolutevalue{a_{k}-a_{j}}a_{K}^{2}}{\text{max}_{k,j:a_{k}\neq
a_{j}}\absolutevalue{a_{k}^{2}-a_{j}^{2}}}.$

Note that all the terms inside the final $O(\dots)$ term vanish:
$\sigma^{4}w\to 0$ because $\gamma_{\sigma^{2}}+\gamma_{w}<0$,
$\frac{\sigma^{4}w^{2}\log^{2}d}{d^{2}}+\frac{\sigma^{2}w}{d}\to 0$ since
$1>\gamma_{\sigma^{2}}+\gamma_{w}$, and $\eta^{2}\frac{\log^{2}d}{d^{2}}\to 0$
since we assumed $\eta\ll d$.

This shows that the lazy regime only appears for very large widths
$\gamma_{w}>2$ (or at least the lazy regime with finite variance at
initialization). Indeed the choice $\gamma_{w}=2,\gamma_{\sigma^{2}}=-1$ is at
the boundary of the lazy regime with the smallest $\gamma_{w}$. This could
explain why it is rare to observe the lazy regime in practice.

Our theoretical results applies to the overparametrized regime $w\gg d$, but
actually we only want to fit $A^{*}$ which has a much smaller rank $r$, and so
we might only need $w\gg r$. Figure 2, top left panel, confirms this, since we
see a good generalization even for small widths $w<d$, and in particular when
$w\approx\mathrm{Rank}A^{*}$. But to leverage this underparametrized regime,
one would need to know the rank of the true matrix $A^{*}$ in advance, which
is typically not the case in practice. Nevertheless, the interesting behavior
we observe in the (mildly) underparametrized regime warrants further analysis,
and the fact that our self-consistent dynamics remain a good approximation in
this regime (Figure 2, bottom right panel), suggests that the analysis we
present here could be extended to this regime too.

Finally, we observe a trade-off between generalization error and training
time: on one hand the test error has terms that scale negatively with
$1-\gamma_{\sigma^{2}}-\gamma_{w}$, which is the distance to the lazy/active
transition, on the other hand, the time it takes to reach the minimal loss
point scales positively with the same term. This can be seen from Figure 2,
bottom left panel, which plots the number of steps required to reach minimal
test error, which increases as one goes further into the active regime.

######  Remark.

In general when trying to fit a matrix $B$ (instead of the special case
$B=A^{*}+E$), the transition between lazy and mixed regime is when
$\sigma^{2}w\approx\|B\|_{op}$. Thus the exact location of the transition is
task-dependent, so that the same variance $\sigma^{2}$ and width $w$ can lead
to NTK or mixed regimes depending on the task. For example, let us assume that
$A^{*}$ is full-rank instead of finite rank, then we expect
$\|A^{*}\|_{op}\sim\sqrt{d}$ instead of $\|A^{*}\|_{op}\sim d$, thus the
transition would be at $\frac{1}{2}=\gamma_{\sigma^{2}}+\gamma_{w}$ instead of
$1=\gamma_{\sigma^{2}}+\gamma_{w}$. This suggests that linear networks are
able to adapt themselves to the task: leveraging active dynamics when the true
data is low-rank to get better generalization, or remaining in the lazy
dynamics in the absence of low-rank structure, to take advantage of the faster
convergence. Note also that in the absence of sparsity, the lazy regime can be
attained with a smaller width ($\gamma_{w}>1$ instead of $\gamma_{w}>2$),
since the choice $\gamma_{w}=1,\gamma_{\sigma^{2}}=-\frac{1}{2}$ is already on
the boundary of the lazy regime.

##  3 Conclusion

We prove a surprisingly simple self-consistent dynamic for the evolution of
the matrix represented by a shallow linear network under gradient descent.
This description not only unifies the already known lazy and balanced
dynamics, but reveals the existence of a spectrum of mixed dynamics where some
of the singular values are lazy while others are balanced.

Thanks to this description we are able to give an almost complete phase
diagram of training dynamics as a function of the scaling of the width and
variance at initialization w.r.t. the dimension.

A natural question that comes out of these results is whether nonlinear
network also feature similar mixed regimes, and whether they could be the key
to understand the convergence of general DNNs.

## References

  * [1] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz.  The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks.  In Conference on Learning Theory, pages 4782–4887. PMLR, 2022.
  * [2] Emmanuel Abbe, Enric Boix-Adserà, Matthew Stewart Brennan, Guy Bresler, and Dheeraj Mysore Nagaraj.  The staircase property: How hierarchical structure can guide deep learning.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
  * [3] Madhu S. Advani and Andrew M. Saxe.  High-dimensional dynamics of generalization error in neural networks, 2017\.
  * [4] Bloemendal Alex, László Erdös, Antti Knowles, Horng-Tzer Yau, and Jun Yin.  Isotropic local laws for sample covariance and generalized Wigner matrices.  Electronic Journal of Probability, 19(none):1 – 53, 2014.
  * [5] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  PMLR, pages 242–252, 2019.
  * [6] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu.  A convergence analysis of gradient descent for deep linear neural networks.  In International Conference on Learning Representations, 2019.
  * [7] Sanjeev Arora, Nadav Cohen, and Elad Hazan.  On the optimization of deep networks: Implicit acceleration by overparameterization.  In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 244–253. PMLR, 10–15 Jul 2018\.
  * [8] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo.  Implicit regularization in deep matrix factorization.  Advances in Neural Information Processing Systems, 32, 2019.
  * [9] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  Advances in Neural Information Processing Systems, 32, 2019.
  * [10] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath.  High-dimensional limit theorems for SGD: Effective dynamics and critical scaling.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  * [11] Francis Bach.  Breaking the curse of dimensionality with convex neural networks.  The Journal of Machine Learning Research, 18(1):629–681, 2017.
  * [12] Rajendra Bhatia.  Matrix analysis, volume 169.  Springer Science & Business Media, 2013.
  * [13] Blake Bordelon and Cengiz Pehlevan.  Self-consistent dynamical field theory of kernel evolution in wide neural networks.  Advances in Neural Information Processing Systems, 35:32240–32256, 2022.
  * [14] Lukas Braun, Clémentine Dominé, James Fitzgerald, and Andrew Saxe.  Exact learning dynamics of deep linear networks with prior knowledge.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 6615–6629. Curran Associates, Inc., 2022.
  * [15] Lenaic Chizat and Francis Bach.  A note on lazy training in supervised differentiable programming.  arXiv preprint arXiv:1812.07956, 2018.
  * [16] Lénaïc Chizat and Francis Bach.  On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport.  In Advances in Neural Information Processing Systems 31, pages 3040–3050. Curran Associates, Inc., 2018.
  * [17] Lénaïc Chizat and Francis Bach.  Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss.  In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pages 1305–1338. PMLR, 09–12 Jul 2020.
  * [18] Lénaïc Chizat, Maria Colombo, Xavier Fernández-Real, and Alessio Figalli.  Infinite-width limit of deep linear neural networks.  Communications on Pure and Applied Mathematics, 2022.
  * [19] Zhen Dai, Mina Karzand, and Nathan Srebro.  Representation costs of linear neural networks: Analysis and design.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
  * [20] Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  In International Conference on Learning Representations, 2019.
  * [21] Kenji Fukumizu.  Effect of batch learning in multilayer neural networks.  In International Conference on Neural Information Processing, 1998\.
  * [22] Mario Geiger, Stefano Spigler, Arthur Jacot, and Matthieu Wyart.  Disentangling feature and lazy training in deep neural networks.  Journal of Statistical Mechanics: Theory and Experiment, 2020(11):113301, 2020.
  * [23] Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien.  Implicit regularization of discrete gradient dynamics in linear neural networks.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [24] Arthur Jacot.  Bottleneck structure in learned features: Low-dimension vs regularity tradeoff.  In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 23607–23629. Curran Associates, Inc., 2023.
  * [25] Arthur Jacot.  Implicit bias of large depth networks: a notion of rank for nonlinear functions.  ICLR, 2023.
  * [26] Arthur Jacot.  Implicit bias of large depth networks: a notion of rank for nonlinear functions.  In The Eleventh International Conference on Learning Representations, 2023.
  * [27] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural Tangent Kernel: Convergence and Generalization in Neural Networks.  In Advances in Neural Information Processing Systems 31, pages 8580–8589. Curran Associates, Inc., 2018.
  * [28] Arthur Jacot, François Ged, Berfin Şimşek, Clément Hongler, and Franck Gabriel.  Saddle-to-saddle dynamics in deep linear networks: Small initialization training, symmetry, and sparsity, 2022.
  * [29] Arthur Jacot, Eugene Golikov, Clément Hongler, and Franck Gabriel.  Feature learning in $l_{2}$-regularized dnns: Attraction/repulsion and sparsity.  In Advances in Neural Information Processing Systems, volume 36, 2022.
  * [30] Ziwei Ji and Matus Telgarsky.  Gradient descent aligns the layers of deep linear networks.  CoRR, abs/1810.02032, 2018.
  * [31] Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon Shaolei Du, and Jason D. Lee.  Understanding incremental learning of gradient descent: A fine-grained analysis of matrix sensing.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 15200–15238. PMLR, 23–29 Jul 2023.
  * [32] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari.  The large learning rate phase of deep learning: the catapult mechanism.  arXiv preprint arXiv:2003.02218, 2020.
  * [33] Aitor Lewkowycz and Guy Gur-Ari.  On the training dynamics of deep networks with $l\\_2$ regularization.  Advances in Neural Information Processing Systems, 33:4790–4799, 2020.
  * [34] Zhiyuan Li, Yuping Luo, and Kaifeng Lyu.  Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning.  In International Conference on Learning Representations, 2020.
  * [35] Chaoyue Liu, Libin Zhu, and Mikhail Belkin.  Toward a theory of optimization for over-parameterized systems of non-linear equations: the lessons of deep learning.  arXiv preprint arXiv:2003.00307, 2020.
  * [36] Hancheng Min, Salma Tarmoun, René Vidal, and Enrique Mallada.  On the explicit role of initialization on the convergence and implicit bias of overparametrized linear networks.  In International Conference on Machine Learning, pages 7760–7768. PMLR, 2021.
  * [37] Hancheng Min, Rene Vidal, and Enrique Mallada.  On the convergence of gradient flow on multi-layer linear models.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 24850–24887. PMLR, 23–29 Jul 2023.
  * [38] Greg Ongie, Rebecca Willett, Daniel Soudry, and Nathan Srebro.  A function space view of bounded norm infinite width relu nets: The multivariate case.  In International Conference on Learning Representations, 2020.
  * [39] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  arXiv, 2019.
  * [40] Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion.  Implicit bias of sgd for diagonal linear networks: a provable benefit of stochasticity.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 29218–29230. Curran Associates, Inc., 2021.
  * [41] Grant Rotskoff and Eric Vanden-Eijnden.  Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks.  In Advances in Neural Information Processing Systems 31, pages 7146–7155. Curran Associates, Inc., 2018.
  * [42] Andrew M Saxe, James L McClelland, and Surya Ganguli.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, 2014.
  * [43] Andrew M. Saxe, James L. McClelland, and Surya Ganguli.  A mathematical theory of semantic development in deep neural networks.  Proceedings of the National Academy of Sciences, 116(23):11537–11546, 2019.
  * [44] Salma Tarmoun, Guilherme Franca, Benjamin D Haeffele, and Rene Vidal.  Understanding the dynamics of gradient flow in overparameterized linear models.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10153–10161. PMLR, 18–24 Jul 2021.
  * [45] Zihan Wang and Arthur Jacot.  Implicit bias of SGD in $l_{2}$-regularized linear DNNs: One-way jumps from high to low rank.  In The Twelfth International Conference on Learning Representations, 2024.
  * [46] Blake Woodworth, Suriya Gunasekar, Pedro Savarese, Edward Moroshko, Itay Golan, Jason Lee, Daniel Soudry, and Nathan Srebro.  Kernel and rich regimes in overparametrized models, 2020.
  * [47] Nuoya Xiong, Lijun Ding, and Simon Shaolei Du.  How over-parameterization slows down gradient descent in matrix sensing: The curses of symmetry and initialization.  In OPT 2023: Optimization for Machine Learning, 2023.
  * [48] Ziqing Xu, Hancheng Min, Salma Tarmoun, Enrique Mallada, and Rene Vidal.  Linear convergence of gradient descent for finite width over-parametrized linear networks with general initialization.  In Francisco Ruiz, Jennifer Dy, and Jan-Willem van de Meent, editors, Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pages 2262–2284. PMLR, 25–27 Apr 2023.
  * [49] Greg Yang.  Tensor programs ii: Neural tangent kernel for any architecture.  arXiv preprint arXiv:2006.14548, 2020.
  * [50] Greg Yang and Edward J. Hu.  Feature learning in infinite-width neural networks, 2020.
  * [51] Yi Yu, Tengyao Wang, and Richard J Samworth.  A useful variant of the davis–kahan theorem for statisticians.  Biometrika, 102(2):315–323, 2015.
  * [52] Difan Zou, Philip M Long, and Quanquan Gu.  On the global convergence of training deep linear resnets.  arXiv preprint arXiv:2003.01094, 2020.

The appendix is structured as follows.

  * •

In section A, we introduce the notation, and establish several results about
how perturbing a matrix would impact its singular vectors.

  * •

In section B, we study the gradient flow dynamics of $A_{t}$ in the active
regime and prove that $A_{t}$ will be approximately aligned with $A^{*}$
throughout the Saddle-to-Saddle regime.

  * •

In section C, we prove theorem 1 for general cost.

  * •

In section D, we study the gradient flow dynamics for $A_{\theta(t)}$ in the
lazy regime.

  * •

In section E, we show that $A_{\theta(t)}$ is also approximately aligned with
$A^{*}$ throughout the Saddle-to-Saddle regime, using results in section B and
section C. In subsection E.3, we summarize the approximate dynamics of
$A_{\theta(t)}$ throughout training. In section E.4, we bound the final test
error.

  * •

In section F, we bound the error from gradient descent and prove theorem 2.

  * •

In section G, we describe the experimental setup.

##  Appendix A Preliminaries

###  A.1 Convention and Notation

Constants. $d$ is the input dimension, $K$ is the rank of matrix $A^{*}$, and
$w$ is the input and output dimension. $c$ and $C$ will usually denote
constants that are independent of $d$, and depending on the context, the value
of $c$ and $C$ might be different. If $x$ is a scalar that depends on $d$ and
$y$ is a scalar, then $x=O(y)$ means there exists a constant $c$ independent
of $d$, such that for $d$ sufficiently large we have $x\leq cy$. If $A$ is a
matrix, then $A=O(y)$ means there exists a constant $c$ independent of $d$
such that for $d$ sufficiently large, $\|A\|_{op}\leq cy$. $O(y)$ can be
either a matrix or a scalar, and its meaning will be always clear from
context.

Matrix. We use $\|\cdot\|_{op},\|\cdot\|_{F},\|\cdot\|$ to denote operator
norm of a matrix, Frobenius norm of a matrix, and the norm of vectors. For
every matrix $A$, we define $\text{max}A$ as the $L^{\infty}$ norm,
$\text{max}_{i,j}\absolutevalue{A_{ij}}$. We use $I$ to denote identity
matrix, the dimension of which is determined by context. We shall assume that
the signal singular values of $A^{*}$ are $s_{1}^{*},\ldots,s_{K}^{*}$.
$\forall i=1,2,\ldots,K$, $s_{i}=a_{i}d$ where $a_{1}\geq a_{2}\geq\ldots\geq
a_{K}$ are constants independent of $d$. By selecting proper basis in the
input and output space, we assume that $A^{*}=S^{*}$, where $S^{*}$ is the
diagonal matrix consisting of singular values of $A^{*}$, ordered from largest
to smallest. The $p,q$-th element of a matrix $A$ is denoted $A_{pq}$. We
reserve the notation $A(i,j)$ for the $i,j$-th block matrix of $A$, which we
shall define below.

Submatrix. Assume that $n_{0}=0$, and
$a_{1}=\ldots=a_{n_{1}}>a_{n_{1}+1}=\ldots a_{n_{2}}>\ldots=a_{n_{m}}=a_{K}$,
and let $n_{m+1}=d$. For a matrix $U$, we define the $k,j$-th sub-block
$U(k,j)$ of a matrix $U$ as $U_{n_{k-1}+1:n_{k+1},n_{j-1}+1:n_{j}}$, with both
sides included. Notice that $U^{T}(k,j)=U(j,k)^{T}$. In this notation, we can
write the singular value decomposition of a matrix $A$ as

$A(i,j)=\sum_{k:\text{signal}}U(i,k)S(k,k)V(j,k)+U(i,m+1)S(m+1,m+1)V^{T}(m+1,j).$

We call an index $k$ (of sub-block) "signal", if $k\leq m$. Index $m+1$ is
called "noise". Let $S(k,k)^{*}$ be block matrix
$A^{*}(n_{k}:n_{k+1},n_{k}:n_{k+1})$. Then
$A^{*}=\mathrm{diag}(S(1,1)^{*},\ldots,S(m,m)^{*})$ and
$S(k,k)^{*}=s_{n_{k}}I$ is the $k$-th sub-block of $A^{*}$. Each matrix has
only finitely many sub-blocks.

Indexing Conventions. Entries of matrices will usually be indexed by $p,q,r$
and sub-blocks of matrices will usually be indexed by $i,j,k,\ell$. Usually
$k$ ranges from 1 to $m$, and $j$ usually ranges from 1 to $m+1$.

###  A.2 Matrix Norms

######  Lemma A.1.

Assume that $A$, $B$ and $R$ are square matrices. Let
$R_{max}=\text{max}_{i,j}\absolutevalue{R_{ij}}$. Then

$tr[A(R\odot B)]=tr[BR^{T}\odot A],$

and

$\absolutevalue{tr[A(R\odot B)]}\leq
R_{\text{max}}\sqrt{tr(A^{T}A)}\sqrt{tr(B^{T}B)}$

In particular, if $\forall p,q,R_{pq}\geq R_{\text{min}}>0$, then

$tr[A^{T}(R\odot A)]\geq R_{\text{min}}tr(A^{T}A)$

###### Proof.

All are simple computations.

$\displaystyle tr[AR\odot B]=$ $\displaystyle\sum_{p,q}A_{pq}R_{qp}B_{qp}$
$\displaystyle tr[BR^{T}\odot A]=$ $\displaystyle\sum_{p,q}B_{pq}R_{pq}A_{qp}$

The two equations above proves the first claim.

$\displaystyle\absolutevalue{tr[A(R\odot B)]}\leq$
$\displaystyle\sqrt{tr(A^{T}A)}\sqrt{tr((R\odot B)^{T}R\odot B)}$
$\displaystyle\leq$ $\displaystyle
R_{\text{max}}\sqrt{tr(A^{T}A)}\sqrt{tr(B^{T}B)}$

This completes the second claim.

$\displaystyle tr[A^{T}(R\odot A)]$
$\displaystyle=\sum_{i,j}A_{ji}R_{ji}A_{ji}$ $\displaystyle\geq
R_{\text{min}}tr[A^{T}A]$

This completes the third claim. ∎

######  Lemma A.2.

$\sigma_{min}(A)\|B\|_{F}\leq\|AB\|_{F}\leq\sigma_{max}(A)\|B\|_{F}$.

###### Proof.

This is lemma B.3 of [52]. ∎

###  A.3 Perturbation of Singular Values and Singular Vectors

We will often use the following variant of Davis-Kahan $\sin\theta$ theorem.

######  Theorem A.3 (DK-$\sin\theta$ Theorem).

. Let $\Sigma,\hat{\Sigma}\in\mathbb{R}^{p\times p}$ be symmetric, with
eigenvalues $\lambda_{1}\geq\ldots\lambda_{p}$ and
$\hat{\lambda}_{1}\geq\ldots\hat{\lambda}_{p}$. Fix $1\leq r\leq s\leq p$, let
$d=r-s+1$ and let $V=(v_{r},\ldots,v_{s})$ and
$\hat{V}=(\hat{v}_{r},\ldots,\hat{v}_{s})$ have orthonormal columns satisfying
$\Sigma v_{j}=\lambda_{j}v_{j}$ and
$\Sigma\hat{v}_{j}=\hat{\lambda}_{j}v_{j}$. Let $\sigma_{1},\ldots,\sigma_{d}$
be the singular values of $\hat{V}^{T}V$. Let $\Theta(V,\hat{V})$ be the
diagonal matrix with $\cos\Theta(V,\hat{V})_{jj}=\sigma_{j}$ and
$\sin\Theta(V,\hat{V})$ be defined entry-wise. Then

$\|\sin\Theta(V,\hat{V})\|_{F}\leq\frac{2\text{min}(\sqrt{d}\|\hat{\Sigma}-\Sigma\|_{op},\|\hat{\Sigma}-\Sigma\|_{F})}{\text{min}(\lambda_{r-1}-\lambda_{r},\lambda_{s}-\lambda_{s+1})}.$

###### Proof.

This is theorem 2 in [51]. ∎

The term $\|\sin\Theta(V,\hat{V})\|_{F}$ describes the extend to which $V$ is
aligned with $\hat{V}$, but it is a bit complicated to take derivative. In
this paper we will use a different characterization of alignment.

######  Lemma A.4.

Let $\hat{\Sigma}$ be a $d\times d$ diagonal matrix, and let
$s_{1},\ldots,s_{K},\ldots,s_{d}$ be its diagonal entries. Assume that
$s_{1}=\ldots=s_{n_{1}}>s_{n_{1}+1}=\ldots=s_{n_{2}}\geq\ldots
s_{n_{m}}=s_{K}>s_{K+1}=\ldots=s_{d}$. Then $\hat{\Sigma}$ has $m+1$ blocks in
total. Assume that $\|X-\hat{\Sigma}\|_{op}=d^{\alpha}$, which is very small.
Let $X=USV^{T}$, and define

$x=4K-\sum_{k:\text{signal}}\mathrm{tr}(U^{T}(k,k)U(k,k)+V^{T}(k,k)V(k,k)+2U(k,k)V(k,k)^{T}).$

Then

$x\leq
K\frac{\|\hat{\Sigma}-X\|_{op}}{s_{K}}+\frac{ms_{1}}{s_{K}}\left(\frac{2\text{min}(\sqrt{K}(\|X\|_{op}+\|\hat{\Sigma}\|_{op}),\|X\|_{F}+\|\hat{\Sigma}\|_{F})}{\text{min}_{k}(s_{n_{k}}^{2}-s_{n_{k+1}}^{2})}\right)^{2}\|X-\hat{\Sigma}\|_{op}^{2}$

###### Proof.

Let $X=USV^{T}$. Then $X^{T}X=VS^{2}V^{T}$ is a symmetric matrix. Let
$V(:,k)=(v_{n_{k-1}+1},\ldots,v_{n_{k}})$ be the singular vectors of $X$
corresponding to $s_{n_{k-1}+1},\ldots,s_{n_{k}}$. There are some freedom to
choose $\hat{V}$, but for simplicity we pick
$\hat{V}(:,k)=(e_{n_{k-1}+1},\ldots,e_{n_{k}})$ where $e_{i}$ is the $i$-th
coordinate vector. As a result, $\hat{V}(:,k)^{T}V(:,k)=V(k,k)$. Apply
DK-$\sin\theta$ theorem to $V$ and $\hat{V}$, we see that

$\|\sin\Theta(V(:,k),\hat{V}(:,k))\|_{F}\leq\frac{2\text{min}(\sqrt{n_{k}-n_{k-1}}\|X^{T}X-\hat{\Sigma}^{2}\|_{op},\|X^{T}X-\hat{\Sigma}^{2}\|_{F})}{\text{min}(s_{n_{k-1}}^{2}-s_{n_{k-1}+1}^{2},s_{n_{k}}^{2}-s_{n_{k}+1}^{2})}$

Let $\sigma_{p}$ be the singular values of $V(k,k)$, for $1\leq p\leq
n_{k}-n_{k-1}$. Here we are using the notation for sub-block of a big matrix.
Then $\sin\Theta(V(:,k),\hat{V}(:,k))$ is the diagonal matrix, whose diagonal
entries are given by $\sqrt{1-\sigma_{p}^{2}}$. So
$\|\sin\Theta(V,\hat{V})\|_{F}^{2}=\sum_{p}(1-\sigma^{2}_{p})=tr[I-V(k,k)^{T}V(k,k)]$.
Now observe that

$\displaystyle\|X^{T}X-\hat{\Sigma}^{2}\|_{op}\leq$
$\displaystyle(\|X\|_{op}+\|\hat{\Sigma}\|_{op})\|X-\hat{\Sigma}\|_{op}$
$\displaystyle\|X^{T}X-\hat{\Sigma}^{2}\|_{F}\leq$
$\displaystyle(\|X\|_{F}+\|\hat{\Sigma}\|_{F})\|X-\hat{\Sigma}\|_{op}$

We conclude that

$tr[I-V(k,k)^{T}V(k,k)]\leq\left(\frac{2\text{min}(\sqrt{n_{k}-n_{k-1}}(\|X\|_{op}+\|\hat{\Sigma}\|_{op}),\|X\|_{F}+\|\hat{\Sigma}\|_{F})}{\text{min}(s_{n_{k-1}}^{2}-s_{n_{k-1}+1}^{2},s_{n_{k}}^{2}-s_{n_{k}+1}^{2})}\|X-\hat{\Sigma}\|_{op}\right)^{2}.$

Similar conclusion is true for $U(k,k)$. Next we bound
$tr(I-U(k,k)V(k,k)^{T})$.

$\|\hat{\Sigma}(k,k)-X(k,k)\|_{F}^{2}=\sum_{q_{1},q_{2}=n_{k-1}+1}^{n_{k}}(\hat{\Sigma}_{q_{1}q_{2}}-X_{q_{1}q_{2}})^{2}\leq(n_{k}-n_{k-1})\|\hat{\Sigma}-X\|_{op}^{2}.$
$\displaystyle\|\hat{\Sigma}(k,k)-X(k,k)\|_{F}^{2}=$
$\displaystyle\|\hat{\Sigma}(k,k)-\sum_{j}U(k,j)S(j,j)V^{T}(j,k)\|_{F}^{2}$
$\displaystyle=$ $\displaystyle\|\hat{\Sigma}(k,k)-\sum_{j\neq
k}U(k,j)S(j,j)V(k,j)^{T}-U(k,k)S(k,k)V(k,k)^{T}\|_{F}^{2}$ $\displaystyle\geq$
$\displaystyle\|\hat{\Sigma}(k,k)-U(k,k)\hat{\Sigma}(k,k)V(k,k)^{T}\|_{F}^{2}$
$\displaystyle-\|\sum_{j\neq k}U(k,j)S(j,j)V(k,j)^{T}\|_{F}^{2}$
$\displaystyle-\|U(k,k)(S(k,k)-\hat{\Sigma}(k,k))V(k,k)^{T}\|_{F}^{2}$
$\displaystyle\geq$ $\displaystyle
s_{n_{k}}^{2}\|I-U(k,k)V(k,k)^{T}\|_{F}^{2}$ $\displaystyle-\sum_{j\neq
k}\|U(k,j)\|_{F}^{2}\|S(j,j)V(k,j)\|_{op}^{2}$
$\displaystyle-\|X-\hat{\Sigma}\|_{op}^{2}$ $\displaystyle\geq$ $\displaystyle
s_{n_{k}}^{2}\|I-U(k,k)V(k,k)^{T}\|_{F}^{2}$ $\displaystyle-
s_{1}^{2}tr(I-U(k,k)^{T}U(k,k))tr(I-V(k,k)^{T}V(k,k))$
$\displaystyle-\|X-\hat{\Sigma}\|_{op}^{2}$

For every $n_{k}-n_{k-1}\times n_{k}-n_{k-1}$ matrix $M$, we have
$tr(A)\leq\sum_{p}\absolutevalue{\lambda_{p}}\leq\sqrt{n_{k}-n_{k-1}}\left(\sum\absolutevalue{\lambda_{p}}^{2}\right)^{\frac{1}{2}}=\sqrt{n_{k}-n_{k-1}}\|M\|_{F}$.
Using this inequality, we conclude that

$\displaystyle tr(I-U(k,k)V(k,k)^{T})\leq$
$\displaystyle(n_{k}-n_{k-1})\frac{\|\hat{\Sigma}-X\|_{op}}{s_{n_{k}}}$
$\displaystyle+\frac{s_{1}}{s_{n_{k}}}\left(\frac{2\text{min}(\sqrt{n_{k}-n_{k-1}}(\|X\|_{op}+\|\hat{\Sigma}\|_{op}),\|X\|_{F}+\|\hat{\Sigma}\|_{F})}{\text{min}(s_{n_{k-1}}^{2}-s_{n_{k-1}+1}^{2},s_{n_{k}}^{2}-s_{n_{k}+1}^{2})}\|X-\hat{\Sigma}\|_{op}\right)^{2}$

Summing on $k$, we conclude that

$x\leq
K\frac{\|\hat{\Sigma}-X\|_{op}}{s_{K}}+\frac{ms_{1}}{s_{K}}\left(\frac{2\text{min}(\sqrt{K}(\|X\|_{op}+\|\hat{\Sigma}\|_{op}),\|X\|_{F}+\|\hat{\Sigma}\|_{F})}{\text{min}_{k}(s_{n_{k}}^{2}-s_{n_{k+1}}^{2})}\right)^{2}\|X-\hat{\Sigma}\|_{op}^{2}$

∎

In the case where $d$ increases to infinity, $m$ is a constant,
$s_{1}=O(s_{K})$, $s_{K}>cs_{K+1}$, $K$ is a constant,
$\|X-\hat{\Sigma}\|_{op}=o(\|\hat{\Sigma}\|_{op})$, and
$\text{min}_{k}(s_{n_{k}}^{2}-s_{n_{k+1}}^{2})\geq cs_{K}^{2}$ for some
constant $c$, we have $x\leq
O(\frac{\|X-\hat{\Sigma}\|_{op}}{\|\hat{\Sigma}\|_{op}})$.

######  Lemma A.5.

Let $X=USV^{T}$, and define

$x=4K-\sum_{k:\text{signal}}\mathrm{tr}(U^{T}(k,k)U(k,k)+V^{T}(k,k)V(k,k)+2U(k,k)V(k,k)^{T}).$

Let $\Sigma$ be a $d\times d$ diagonal matrix whose diagonal entries are given
by $b_{1},\ldots,b_{K},0,\ldots,0$,where $b_{1}=\ldots
b_{n_{1}}>b_{n_{1}+1}=\ldots=b_{n_{2}}\geq\ldots=b_{n_{m}}=b_{K}$. Then

$\|\Sigma-U^{T}\Sigma
V\|_{op}\leq\|\Sigma\|_{op}((m+1)^{3}\sqrt{x}+2(m+1)^{3}x)$

###### Proof.

First observe that $\|\Sigma-U^{T}\Sigma
V\|_{op}\leq\sum_{i,j=1}^{m+1}\|\Sigma(i,j)-U^{T}\Sigma V(i,j)\|_{op}$. If
$i\neq j$, then $\Sigma(i,j)=0$,

$\displaystyle\|U^{T}\Sigma V(i,j)\|_{op}=$
$\displaystyle\|\sum_{\ell=1}^{m+1}U(\ell,i)^{T}\Sigma(\ell,\ell)V(\ell,j)\|_{op}$
$\displaystyle\leq$ $\displaystyle
b_{1}\sum_{\ell=1}^{m+1}\|U(\ell,i)\|_{op}\|V(\ell,j)\|_{op}$
$\displaystyle\leq$ $\displaystyle b_{1}(m+1)\sqrt{x}.$

If $i=j$, then $\Sigma(i,i)=b_{n_{i}}I$,

$\displaystyle\|U^{T}\Sigma V(i,i)-b_{n_{i}}I\|_{op}\leq$
$\displaystyle\|b_{n_{i}}U(i,i)^{T}V(i,i)-b_{n_{i}}I\|_{op}+b_{1}\sum_{\ell\neq
i}\|U(\ell,i)\|_{op}\|V(\ell,i)\|_{op}$ $\displaystyle\leq$ $\displaystyle
b_{1}(m+1)x+b_{1}(m+1)x.$

Therefore

$\|\Sigma-U^{T}\Sigma
V\|_{op}\leq\|\Sigma\|_{op}((m+1)^{3}\sqrt{x}+2(m+1)^{3}x).$

∎

This bound is not optimal in $m$, but throughout the paper, $m$ is of constant
order and it is fine to miss a constant factor when estimating error.

##  Appendix B Gradient Flow Dynamics of $A_{t}$ in Active Regime

###  B.1 Saddle to Saddle Regime

Let $A_{t}$ have the following dynamics:

$d^{2}\frac{d}{dt}A_{t}=(A^{*}-A_{t})\sqrt{A_{t}^{T}A_{t}+\sigma^{4}w^{2}I}+\sqrt{A_{t}A_{t}^{T}+\sigma^{4}w^{2}I}(A^{*}-A_{t}).$

The goal of this section is to prove that the singular vectors of $A_{t}$ is
well-aligned with the singular vectors of $A^{*}$, throughout the Saddle-to-
Saddle regime. In the rest of the section, we will assume the dependence of
$A_{t}$ on $t$ and use $A$ to represent $A_{t}$.

######  Definition B.1.

Define $P_{1}$ be the family of $d\times d$ matrices $A$ that satisfies the
following conditions.

  * •

$s_{K}\geq C\sigma^{2}w$, $s_{K+1}\leq C^{\prime}\sigma^{2}w$ and
$\frac{s_{K+1}}{s_{K}}\leq c<\frac{1}{2}$ for some $d$-independent constants
$c$, $C$ and $C^{\prime}$.

  * •

If $a_{k}>a_{k+1}$, then $s_{k}-s_{k+1}\geq cs_{k}$ for some $d$-independent
constant $c$.

Define $P_{1}^{\prime}$ to be the family of $w\times w$ matrix $A$ such that
$s_{k}-s_{k+1}\geq\frac{c}{2}s_{k}$ if $a_{k+1}<a_{k}$,
$\frac{s_{K+1}}{s_{K}}\leq\frac{3}{4}$ and $s_{K}\geq c\sigma^{2}w$. Let
$\gamma>0$ be constant. Define $P_{2}(C,\gamma)$ be the family of $w\times w$
matrices $A$ satisfying the following conditions.

  * •

(alignment of signals). Let $A=USV^{T}$. Define

$x=4K-\sum_{k:\text{signal}}\mathrm{tr}(U^{T}(k,k)U(k,k)+V^{T}(k,k)V(k,k)+2U(k,k)V(k,k)^{T}).$

$P_{2}(C,\gamma)$ is the family of matrix $A$ such that $x\leq Cd^{-\gamma}$.

######  Theorem B.2.

Assume that

$d^{2}\frac{dA}{dt}=(A^{*}-A)\sqrt{A^{T}A+\sigma^{4}w^{2}I}+\sqrt{AA^{T}+\sigma^{4}w^{2}I}(A^{*}-A)$
$A(0)\in P_{1}\cap P_{2}(C_{4},\gamma)$

Let $T=O(d\log d)$. Then $\forall t\in[0,T]$, we have $A_{t}\in P_{1}\cap
P_{2}(C,\text{min}(1,\gamma))$.

###### Proof.

A simple result of the induction lemma B.3, lemma B.4, lemma B.6 and lemma
B.7. ∎

We will use the following induction lemma to show that the singular vectors of
$A_{t}$ are roughly aligned with $A^{*}$.

######  Lemma B.3.

Assume that $P_{1}$ and $P_{2}$ be families of increasing sets, and let
$P_{1}^{\prime}\supset P_{1}$. Assume that we have a family of matrices
$A_{t}$, $0\leq t\leq T$ for some fixed number $T$. $T$ does not depend on the
family of matrices. Assume that $A_{0}\in P_{1}\cap P_{2}$. Let
$A_{[t_{1},t_{2}]}=\\{A_{t}:t_{1}\leq t\leq t_{2}\\}$. Assume the following
are true.

  1. 1.

If $A_{[0,t]}\in P_{1}$ then there exists a constant $\varepsilon>0$
independent of $A_{t}$ such that $A_{[0,t+\varepsilon]}\subset
P_{1}^{\prime}$.

  2. 2.

Let $t_{1}<t_{2}$. If $A_{[0,t_{2}]}\subset P_{1}^{\prime}$ and
$A_{[0,t_{1}]}\in P_{2}$, then $A_{[0,t_{2}]}\subset P_{2}$.

  3. 3.

Let $t_{1}<t_{2}$. If $A_{[0,t_{2}]}\subset P_{2}$ and $A_{[0,t_{1}]}\in
P_{1}$ then $A_{[0,t_{2}]}\subset P_{1}$.

Then $A_{t}\in P_{1}$ and $A_{t}\in P_{2}$, $\forall 0\leq t\leq T$.

###### Proof.

Since $A_{0}$ satisfies $P_{1}$, use condition 1 we have
$A_{[0,\varepsilon]}\subset P_{1}^{\prime}$. Using condition 2, we know that
$A_{[0,\varepsilon]}\subset P_{2}$ . By condition 3 we know that
$A_{[0,\varepsilon}]\subset P_{1}$ and by condition 2,
$A_{[0,\varepsilon]}\subset P_{2}$. Iterate the argument. ∎

######  Lemma B.4.

If $A_{t}\in P_{1}$, then there exists $\tau=\sigma^{2}\sqrt{dw}d^{-1}$ such
that $A_{[t,t+\tau]}\subset P_{1}^{\prime}$.

###### Proof.

Let $0\leq s\leq\tau$. Using the approximation in time $[t,s+t]$ we have

$\displaystyle d^{2}\frac{dS}{dt}=$ $\displaystyle
I\odot\left(U^{T}A^{*}V\sqrt{S^{2}+\sigma^{4}w^{2}}+\sqrt{S^{2}+\sigma^{4}w^{2}}U^{T}A^{*}V-2S\sqrt{S^{2}+\sigma^{4}w^{2}}\right)$

For $j=1,2,\ldots,K+1$,

$\displaystyle d^{2}\absolutevalue{\frac{ds_{j}}{dt}}\leq
Cd(s_{j}+\sigma^{2}w)$

for some constant $C$ independent of $w$ and $A_{t}$. The conclusion follows
from Grönwall.

∎

######  Lemma B.5.

Assume that

$\displaystyle d^{2}\frac{dA}{dt}=$
$\displaystyle(A^{*}-A+E)\sqrt{A^{T}A+\sigma^{2}wI}+\sqrt{AA^{T}+\sigma^{4}w^{2}I}(A^{*}-A+E).$

Then for signal $k$,

$d^{2}\frac{d}{dt}tr[U^{T}(k,k)U(k,k)]=tr[U^{T}(k,k)\sum_{j\neq
k}U(k,j)D(j,k)]$

Here for $j\neq k$,

$D(j,k)=R(k,j)\odot C(j,k)-R(j,k)\odot C(k,j)^{T}$

with $C=U^{T}(A^{*}+E)V$,
$R_{pq}=\frac{s_{p}\tilde{s}_{p}+s_{p}\tilde{s}_{q}}{s_{p}^{2}-s_{q}^{2}}$ for
$1\leq p,q\leq w$, $p\neq q$, and $R(k,j)=R_{n_{k}:n_{k+1},n_{j}:n_{j+1}}$
being the $k,j$-th block of $R$.

###### Proof.

Use SVD dervative.

$\displaystyle\frac{dU}{dt}=$ $\displaystyle
U\left(F\odot\left[U^{T}\frac{dA}{dt}VS+SV^{T}\frac{dA^{T}}{dt}U\right]\right)$
(3) $\displaystyle=$ $\displaystyle
U\left(F\odot\left[(U^{T}(A^{*}+E)V-S)\sqrt{S^{2}+\sigma^{4}w^{2}}S+\sqrt{S^{2}+\sigma^{4}w^{2}}(U^{T}(A+E)^{*}V-S)S\right]\right)$
(4)
$\displaystyle+U\left(F\odot\left[\sqrt{S^{2}+\sigma^{4}w^{2}}S(V^{T}(A^{*}+E)^{T}U-S)+S(V^{T}(A^{*}+E)^{T}U-S)\sqrt{S^{2}+\sigma^{4}w^{2}}\right]\right)$
(5) $\displaystyle=$ $\displaystyle
U\left(F\odot\left[U^{T}(A^{*}+E)V\sqrt{S^{2}+\sigma^{4}w^{2}}S+\sqrt{S^{2}+\sigma^{4}w^{2}}U^{T}(A^{*}+E)VS\right]\right)$
(6)
$\displaystyle+U\left(F\odot\left[\sqrt{S^{2}+\sigma^{4}w^{2}}SV^{T}(A^{*}+E)^{T}U+SV^{T}(A^{*}+E)^{T}U\sqrt{S^{2}+\sigma^{4}w^{2}}\right]\right)$
(7)

Let $D=F\odot\left[U^{T}\frac{dA}{dt}VS+SV^{T}\frac{dA^{T}}{dt}U\right]$.
Since $F$ is anti-symmetric and the term in square bracket is symmetric, we
know $D$ is anti-symmetric. Then $\frac{dU}{dt}=UD$. Let
$\tilde{S}=\sqrt{S^{2}+\sigma^{4}w^{2}}$, $C=U^{T}(A^{*}+E)V$. As a result,
$\forall 1\leq p,q\leq w$,

$\displaystyle\frac{dU_{pq}}{dt}=$ $\displaystyle\sum_{r:r\neq
q}U_{pr}\frac{1}{s_{q}^{2}-s_{r}^{2}}\left[(s_{q}\tilde{s}_{q}+\tilde{s_{r}}s_{q})C_{rq}+(\tilde{s_{r}}s_{r}+s_{r}\tilde{s}_{k})C_{rq}\right]$

Let $R_{pq}=\frac{s_{p}\tilde{s}_{p}+s_{p}\tilde{s}_{q}}{s_{p}^{2}-s_{q}^{2}}$
if $p\neq q$ and $R_{pp}=0$. Then

$\displaystyle\frac{dU}{dt}(k,k)=\sum_{j:j\neq k}U(k,j)\left(R(k,j)\odot
C(j,k)-R(j,k)\odot C(k,j)^{T}\right)$

$D(j,k)=R(k,j)\odot C(j,k)-R(j,k)\odot C(k,j)^{T}.$

As a result, for $n_{k}\leq p,q<n_{k+1}$,

$\displaystyle\frac{1}{2}\frac{d}{dt}tr[U^{T}(k,k)U(k,k)]=$
$\displaystyle\sum_{p,q\in[n_{k},n_{k+1})}U_{pq}\sum_{r}U_{pr}D_{rq}$
$\displaystyle=$
$\displaystyle\sum_{p,q,r\in[n_{k},n_{k+1})}U_{pq}U_{pr}D_{rq}+\sum_{p,q\in[n_{k},n_{k+1})}U_{pq}\sum_{r\notin[n_{k},n_{k+1})}U_{pr}D_{rq}$
$\displaystyle=$
$\displaystyle\sum_{p,q\in[n_{k},n_{k+1})}U_{pq}\sum_{r\notin[n_{k},n_{k+1})}U_{pr}D_{rq}$
$\displaystyle=$ $\displaystyle tr[U^{T}(k,k)\sum_{j\neq k}U(k,j)D(j,k)]$

∎

######  Lemma B.6.

If $A_{[t_{1},t_{2}]}\subset P_{1}$ and $A_{t_{1}}\in P_{2}(C_{4},\gamma)$,
then

$d^{2}\frac{dx}{dt}\leq-cdx+O(\sqrt{dx}),$

where $\forall
c=\inf_{t\in[t_{1},t_{2}]}\text{min}(\frac{(s_{k}^{*}-s_{j}^{*})(s_{k}+s_{j})(\tilde{s}_{k}+\tilde{s}_{j})}{s_{k}^{2}-s_{j}^{2}},\frac{(s_{k}-s_{j})(s_{k}^{*}+s_{j}^{*})(\tilde{s}_{k}+\tilde{s}_{j})}{s_{k}^{2}-s_{j}^{2}}).$
In particular, $A_{[t_{1},t_{2}]}\subset P_{2}(C,\text{min}(\gamma,1))$ for
some constant $C$. If $\gamma<1$ then we can take $C=C_{4}$.

###### Proof.

Use previous lemma. Let

$x=4K-\sum_{k:\text{signal}}\mathrm{tr}(U^{T}(k,k)U(k,k)+V^{T}(k,k)V(k,k)+2U(k,k)V(k,k)^{T}).$

$x$ measures the alignment of singular vectors of $A$ with singular vectors of
$A^{*}$. From $UU^{T}=I$ we have

$\sum_{\ell}U(k,\ell)U(k,\ell)^{T}=I$

and for every $j\neq k$,

$U(k,j)U(k,j)^{T}\leq I-U(k,k)U(k,k)^{T}\leq O(x)$

In particular, $\|U(k,j)\|_{op}\leq O(\sqrt{x})$. We first estimate $C$. For
all $j=1,2,\ldots,w$, $a_{j}\neq a_{k}$,

$\displaystyle(U^{T}A^{*}V)(j,k)=$ $\displaystyle
U^{T}(j,j)S^{*}(j,j)V(j,k)+U^{T}(j,k)S^{*}(k,k)V(k,k)+\sum_{\ell\neq
j,k}U^{T}(j,\ell)S_{\ell}^{*}V(\ell,k)$ $\displaystyle=$ $\displaystyle
U^{T}(j,j)S^{*}(j,j)V(j,k)+U(k,j)^{T}S^{*}(k,k)V(k,k)$ $\displaystyle+O(dx)$
$\displaystyle(U^{T}EV)(j,k)=$ $\displaystyle\sum_{\ell_{2}\neq
k}U^{T}(j,j)E(j,\ell_{2})V(\ell_{2},k)+\sum_{\ell_{1}\neq
j}U^{T}(j,\ell_{1})E(\ell_{1},k)V(k,k)$
$\displaystyle+U^{T}(j,j)E(j,k)V(k,k)+\sum_{\ell_{1}\neq j,\ell_{2}\neq
k}U^{T}(j,\ell_{1})E(\ell_{1},\ell_{2})V(\ell_{2},k)$ $\displaystyle=$
$\displaystyle O(\sqrt{d}\sqrt{x})+O(\sqrt{d})+O(\sqrt{d}x)$ $\displaystyle=$
$\displaystyle O(\sqrt{d})$ $\displaystyle D(j,k)=$ $\displaystyle R(k,j)\odot
C(j,k)-R(j,k)\odot C(k,j)^{T}$ $\displaystyle=$ $\displaystyle
R(k,j)\odot{(U(j,j)^{T}S^{*}(j,j)V(j,k))}+R(k,j)\odot
U(k,j)^{T}S^{*}(k,k)V(k,k)$
$\displaystyle-R(j,k)\odot(V(k,j)^{T}S^{*}(k,k)U(k,k)+V(j,j)^{T}S^{*}(j,j)U(j,k))$
$\displaystyle+O(dx+\sqrt{d})$

By $UU^{T}=I$ we know that $\sum_{\ell}U(j,\ell)U^{T}(\ell,k)=0$. Therefore

$U(j,j)U(k,j)^{T}+U(j,k)U(k,k)^{T}+O(x)=0$
$U(j,k)=-O(x)-U(j,j)U(k,j)^{T}U(k,k).$

Similarly we have for $V$,

$V(k,j)=-O(x)-V(k,k)V(j,k)^{T}V(j,j).$

We can rewrite $D(j,k)$ as

$\displaystyle D(j,k)=$ $\displaystyle
R(k,j)\odot(U(j,j)^{T}S^{*}(j,j)V(j,k)+U(k,j)^{T}S^{*}(k,k)V(k,k))$
$\displaystyle+$ $\displaystyle
R(j,k)\odot(V(j,j)^{T}V(j,k)V(k,k)^{T}S^{*}(k,k)U(k,k)+V(j,j)^{T}S^{*}(j,j)U(j,j)U(k,j)^{T}U(k,k))$
$\displaystyle+$ $\displaystyle O(dx+\sqrt{d})$ $\displaystyle
d^{2}\frac{d}{dt}\mathrm{tr}\left[U(k,k)^{T}U(k,k)\right]$ $\displaystyle=$
$\displaystyle 2\mathrm{tr}\left[U(k,k)^{T}\sum_{j}U(k,j)D(j,k)\right]$
$\displaystyle=$ $\displaystyle
2\sum_{j}\mathrm{tr}\left[U(k,k)^{T}U(k,j)\left(R(k,j)\odot(U(j,j)^{T}S^{*}(j,j)V(j,k)+U(k,j)^{T}S^{*}(k,k)V(k,k))\right)\right]$
$\displaystyle+2\sum_{j}\mathrm{tr}\left[U(k,k)^{T}U(k,j)\left(R(j,k)\odot(V(j,j)^{T}V(j,k)V(k,k)^{T}S^{*}(k,k)U(k,k))\right)\right]$
$\displaystyle+2\sum_{j}\mathrm{tr}\left[U(k,k)^{T}U(k,j)\left(R(j,k)\odot(V(j,j)^{T}S^{*}(j,j)U(j,j)U(k,j)^{T}U(k,k))\right)\right]$
$\displaystyle+O(dx^{\frac{3}{2}}+\sqrt{d}\sqrt{x})$ $\displaystyle\geq$
$\displaystyle-2\sum_{j}\absolutevalue{S^{*}(j,j)R(k,j)+S^{*}(k,k)R(j,k)}\sqrt{tr(U(k,j)^{T}U(k,j)}\sqrt{tr(V(j,k)^{T}V(j,k)}$
$\displaystyle+2\sum_{j}S^{*}(k,k)\mathrm{tr}[U(k,k)^{T}U(k,j)R(k,j)\odot(U(k,j)^{T}U(k,k))]$
$\displaystyle+2\sum_{j}S^{*}(j,j)\mathrm{tr}[U(k,k)^{T}U(k,j)R(j,k)\odot(U(k,j)^{T}U(k,k))]$
$\displaystyle+O(dx^{\frac{3}{2}})+O(\sqrt{dx})$

We know that

$\displaystyle(S^{*}(k,k)R(k,j)+S^{*}(j,j)R(j,k))-\absolutevalue{S^{*}(j,j)R(k,j)+S^{*}(k,k)R(j,k)}$
$\displaystyle=$
$\displaystyle\text{min}(\frac{(s_{k}^{*}-s_{j}^{*})(s_{k}+s_{j})(\tilde{s}_{k}+\tilde{s}_{j})}{s_{k}^{2}-s_{j}^{2}},\frac{(s_{k}-s_{j})(s_{k}^{*}+s_{j}^{*})(\tilde{s}_{k}+\tilde{s}_{j})}{s_{k}^{2}-s_{j}^{2}})$
$\displaystyle\geq$ $\displaystyle cd$

for some positive constant $c$ independent of $w$. We conclude that

$d^{2}\frac{d}{dt}tr[U(k,k)^{T}U(k,k)]\geq
cdx+O(dx^{\frac{3}{2}})+O(\sqrt{dx}).$

The same trick applies to $V$. We similarly have

$d^{2}\frac{d}{dt}tr[V(k,k)^{T}V(k,k)]\geq
cdx+O(dx^{\frac{3}{2}})+O(\sqrt{dx}).$

It remains to show that $\frac{d}{dt}tr[V(k,k)^{T}U(k,k)]$ bounded below.

$\displaystyle
d^{2}\frac{d}{dt}\mathrm{tr}[U(k,k)^{T}V(k,k)]=tr[U(k,k)^{T}\frac{dV}{dt}(k,k)]+tr[V(k,k)^{T}\frac{dU}{dt}(k,k)]$

$\displaystyle tr[V(k,k)^{T}\frac{d}{dt}U(k,k)]=$ $\displaystyle
tr[U(k,k)^{T}\frac{d}{dt}U(k,k)]+tr[(V(k,k)^{T}-U(k,k)^{T})\frac{d}{dt}U(k,k)]$
$\displaystyle\geq$ $\displaystyle
tr[U(k,k)^{T}\frac{d}{dt}U(k,k)]-\|V(k,k)-U(k,k)\|_{F}\|\frac{d}{dt}U(k,k)\|_{F}$
$\displaystyle\geq$ $\displaystyle
O(dx^{\frac{3}{2}})+cdx+O(dx^{\frac{3}{2}})+O(\sqrt{dx})$ $\displaystyle\geq$
$\displaystyle cdx+O(dx^{\frac{3}{2}})+O(\sqrt{dx})$

Similarly,

$\displaystyle tr[U(k,k)^{T}\frac{d}{dt}V(k,k)]\geq
cdx+O(dx^{\frac{3}{2}})+O(\sqrt{dx}).$

This proves that

$\displaystyle d^{2}\frac{dx}{dt}\leq-cdx+O(dx^{\frac{3}{2}})+O(\sqrt{dx}).$
(8)

Observe that if $x\geq d^{-1+\varepsilon}$ for some $\varepsilon>0$, then
$\frac{dx}{dt}\leq-\frac{c}{2}dx$ and therefore $x$ decrease exponentially. On
the other hand, if $x\leq d^{-1-\varepsilon}$ for some $\varepsilon>0$, then
$\frac{dx}{dt}\geq O(\sqrt{dx})$ and therefore $x$ might increase. We
therefore conclude that $A_{[t_{1},t_{2}]}\subset
P_{2}(C,\text{min}(\gamma,1))$. ∎

######  Lemma B.7.

Assume that $A_{[0,t_{2}]}\subset P_{2}(C_{4},\gamma)$ and
$A_{[0,t_{1}]}\subset P_{1}$. Then $A_{[0,t_{2}]}\subset P_{1}$.

###### Proof.

By assumption we know that $t_{2}=O(d\log d)$.

$\displaystyle d^{2}\frac{dS}{dt}=$ $\displaystyle
I\odot\left(U^{T}(A^{*}+E)V\sqrt{S^{2}+\sigma^{4}w^{2}}+\sqrt{S^{2}+\sigma^{4}w^{2}}U^{T}(A^{*}+E)V-2S\sqrt{S^{2}+\sigma^{4}w^{2}}\right)$
$\displaystyle
d^{2}\frac{ds_{K+1}}{dt}=2\left(O(\sqrt{d})+O(d^{1-\gamma})-s_{K+1}\right)\sqrt{s_{K+1}^{2}+\sigma^{4}w^{2}}$
$\absolutevalue{s_{K+1}(t_{2})-s_{K+1}(t_{1})}\leq
O(\sigma^{2}w(d^{-\frac{1}{2}}+d^{-\gamma})\log d)<<\sigma^{2}w$

It remains to verify the gap between different families. Let $T$ be the first
time when $\|A\|_{op}=\frac{1}{2}s_{K}^{*}$ and assume that $t_{2}\leq T$.
Assume that $a_{k}>a_{j}$, then there exists some constant $\varepsilon$ such
that

$d^{2}\frac{ds_{k}}{dt}\geq(s_{k}^{*}(1-\varepsilon)-s_{k})s_{k}$
$d^{2}\frac{ds_{j}}{dt}\leq(s_{j}^{*}(1+\varepsilon)-s_{j})(s_{j}+\sigma^{2}w)$

An application of Grönwall inequality on $s_{k}$ and $s_{j}$ implies that
there exists some constant $\delta>0$ that depends only on $a_{j},a_{k}$ such
that $\frac{s_{j}(T)}{s_{k}(T)}\leq O(d^{-\delta})$. Next we deal with the
case when $t_{2}>T$. One important observation is that the
$\inf\\{t>T:s_{1}(t)\geq\frac{a_{k}d+a_{j}d}{2}\\}-T=O(d)$, as a simple result
of Grönwall. Moreover, $s_{2}(T+O(d))\leq\frac{1}{4}s_{K}^{*}$. This implies
that $s_{1}-s_{k}>cs_{1}$ for some constant $c$. Repeat this argument for all
remaining signal singular values completes the proof. ∎

###  B.2 First NTK Regime

At initialization, the singular values of $A$ are of order
$\sigma^{2}\sqrt{wd}$ by using Mechenko-Pastur law on $A^{T}A$, which is
infinitely smaller than $\sigma^{2}w$. As a result, there is a very short
period when the dynamics is very close to NTK.

######  Lemma B.8.

Let $t_{1}$ be the first when the first singular value of $A$ hits
$\sigma^{2}wd^{-\frac{\delta}{2}}$ for $\delta=\frac{\gamma_{w}-1}{2}$. Then
$A_{t_{1}}\in P_{2}(C,\text{min}(\frac{1}{2},\frac{\delta}{2}))$ for some
constant $C$.

###### Proof.

We approximate the dynamics with NTK dynamics. Assume that

$d^{2}\frac{dB}{dt}=2(A^{*}+E-B)\sigma^{2}w$

with $B(0)=A(0)$. Then

$\displaystyle d^{2}\frac{d(A-B)}{dt}=$ $\displaystyle
2(B-A)\sigma^{2}w(A^{*}+E-A)$
$\displaystyle+(\sqrt{A^{T}A+\sigma^{4}w^{2}}-\sigma^{2}w)+(\sqrt{AA^{T}+\sigma^{4}w^{2}}-\sigma^{2}w)(A^{*}+E-A).$

A simple application of Grönwall inequality implies that

$\displaystyle\|A-B\|_{op}(t)\leq$
$\displaystyle\frac{1}{d^{2}}e^{-2\sigma^{2}wt}\int_{0}^{t}e^{2\sigma^{2}w\tau}\frac{2s_{1}(\tau)^{2}}{\sigma^{2}w}\|A^{*}+E-A\|_{op}(\tau)d\tau$
$\displaystyle\leq$ $\displaystyle
3\frac{\|A_{t}\|^{2}_{op}}{\sigma^{2}w}s_{1}^{*}\frac{t}{d^{2}}$

We first check $P_{2}$. Clearly, $t_{1}$ is of order $d^{1-\frac{\delta}{2}}$.
Then we have

$\displaystyle\|A_{t_{1}}-\frac{t_{1}}{d^{2}}\sigma^{2}wA^{*}\|_{op}\leq$
$\displaystyle
2d^{-\delta}\|A_{t_{1}}\|_{op}+\|\frac{t_{1}}{d^{2}}\sigma^{2}wE\|_{op}+O(\sigma^{2}\sqrt{wd})$
$\displaystyle=$
$\displaystyle(O(d^{-\delta})+O(d^{-\frac{1}{2}})+O(d^{-\frac{\delta}{2}}))\|A_{t_{1}}\|_{op}$

By lemma B.4 we have $x\leq Cd^{-\text{min}(\frac{1}{2},\frac{\delta}{2})}$
for some constant $C$. ∎

######  Lemma B.9.

Let $t_{2}=cd$, where $c$ is a constant to be chosen. Then

$A_{t_{2}}\in P_{1}\cap
P_{2}(C,\text{min}(\frac{\gamma_{w}-1}{4},\frac{1}{2})).$

Moreover, there exists constants $b_{1}=\ldots
b_{n_{1}}>b_{n_{1}+1}=\ldots=b_{n_{2}}\geq\ldots=b_{n_{m}}=b_{K}$, such that
for a $d\times d$ diagonal matrix $\Sigma$ whose diagonal entries are given by
$b_{1},\ldots,b_{K},0,\ldots,0$, we have $\|A_{t_{2}}-\Sigma\|_{op}\leq
O(\sigma^{2}w)d^{-\text{min}(\frac{\gamma_{w}-1}{8},\frac{1}{4})}$.

###### Proof.

At time $t_{2}$, $\|B(t_{2})\|_{op}=2a_{1}c\sigma^{2}w(1+o(1))$. From previous
lemma we know that

$\|A-B\|_{op}(t)\leq
3\frac{\|A_{t}\|^{2}_{op}}{\sigma^{2}w}s_{1}^{*}\frac{t}{d^{2}}.$

Let $t=t_{2}$ we see that
$\absolutevalue{\|A(t_{2})\|_{op}-\|B(t_{2})\|_{op}}\leq\|A-B\|_{op}(t_{2})\leq
2\frac{\|A_{t_{2}}\|^{2}_{op}}{\sigma^{2}w}a_{1}c$. Therefore as long as $c$
is sufficiently small, we can guarantee that
$\|A-B\|_{op}(t)\leq\frac{\text{min}_{1\leq k\leq
K}(a_{k-1}-a_{k})}{100a_{1}}\|B\|_{op}(t)$. This guarantees that the singular
values of $A$ grows linearly in $[t_{1},t_{2}]$, that $s_{k}-s_{k+1}\geq
cs_{k}$ if $a_{k}>a_{k+1}$ for some constant $c$, that
$s_{K+1}<\frac{1}{2}s_{K}$ and that $A_{t_{2}}\in P_{1}$. It remains to check
that $A_{t_{2}}$ satisfies $P_{2}(C,\frac{\gamma_{w}-1}{2})$. We are ready to
use similar techniques as in lemma B.6. All computation are similar, and the
only difference is that $\absolutevalue{R(k,j)}=O(\frac{\sigma^{2}w}{s_{K}})$.
We have

$d^{2}\frac{d}{dt}\mathrm{tr}[U(k,k)^{T}U(k,k)]\geq
c\frac{\sigma^{2}w}{s_{K}}dx+O(dx^{\frac{3}{2}}\frac{\sigma^{2}w}{s_{K}})+O(\sqrt{dx}\frac{\sigma^{2}w}{s_{K}})$
$d^{2}\frac{dx}{dt}\leq
2C\frac{\sigma^{2}w}{a_{K}(\sigma^{2}wwt+\sigma^{2}wd^{-\frac{\delta}{2}})}(-cdx+dx^{\frac{3}{2}}+\sqrt{dx})$

Therefore $A_{[t_{1},t_{2}]}\subset
P_{2}(C,\text{min}(\frac{\gamma_{w}-1}{4},\frac{1}{2}))$. Let
$\gamma=\text{min}(\frac{\gamma_{w}-1}{4},\frac{1}{2}).$ From lemma B.8 we see
that if $n_{k_{1}}+1\leq p_{1},p_{2}\leq n_{k}$, then
$\absolutevalue{s_{p_{1}}-s_{p_{2}}}\leq O(d^{-\gamma})s_{K}$. Moreover, the
dynamics of $s_{p_{i}},i=1,2$ are both given by

$d^{2}\frac{ds_{p_{i}}}{dt}=2(O(\sqrt{d})+O(d^{1-\gamma})+a_{k}d-s_{p_{i}})\sqrt{s_{p_{i}}^{2}+\sigma^{4}w^{2}}.$

$\frac{d\absolutevalue{s_{p_{1}}-s_{p_{2}}}}{dt}\leq
O(d^{-1})\absolutevalue{s_{p_{1}}-s_{p_{2}}}$
$\absolutevalue{s_{p_{1}}-s_{p_{2}}}(t_{2})\leq
O(1)\absolutevalue{s_{p_{1}}-s_{p_{2}}}(t_{1})=\sigma^{2}wO(d^{-\gamma-\frac{\gamma_{w}-1}{4}}).$

Let $b_{n_{i}}=s_{i}$. By lemma A.5, there exists a $d\times d$ diagonal
matrix $\Sigma$, whose diagonal entries are given by
$b_{1}=\ldots=b_{n_{1}}>\ldots=b_{n_{2}}\geq\ldots=b_{K}$, such that
$\|A_{t_{2}}-\Sigma\|_{op}\leq O(\sigma^{2}w)d^{-\frac{\gamma}{2}}.$ ∎

##  Appendix C Proof of Theorem 1

###  C.1 Weak bound

We prove the following weak bound.

######  Proposition C.1.

For every $\varepsilon>0$ and every $t<T$, we have with high probability,

$\|W_{1}^{T}W_{1}-\sqrt{A^{T}A+\sigma^{4}w^{2}I}\|_{\mathrm{op}}\leq(1+\varepsilon)\sigma^{2}w$
(9)

Analogous results holds for $W_{2}$.

Our main tool is the following lemma.

######  Lemma C.2.

For every cost $C$ and every time $t$, we have

$W_{2}^{T}W_{2}(t)-W_{1}W_{1}^{T}(t)=W_{2}^{T}(0)W_{2}(0)-W_{1}(0)W_{1}^{T}(0)$
(10)

###### Proof.

Let $L=\|A^{*}-A\|^{2}$ be the loss function. Then

$\displaystyle\partial_{t}W_{1}$ $\displaystyle=2W_{2}^{T}\nabla C$
$\displaystyle\partial_{t}W_{2}$ $\displaystyle=2\nabla CW_{1}^{T}$

We see that

$\partial_{t}(W_{1}W_{1}^{T})=2W_{2}^{T}\nabla CW_{1}^{T}+W_{1}\nabla
C^{T}W_{2}=\partial_{t}(W_{2}^{T}W_{2})$

∎

Next, we show that at initialization, $W_{1}W_{1}^{T}$ and $W_{2}^{T}W_{2}$
are approximately orthogonal projections, up to a factor. Stating precisely,
we have the following lemma.

######  Lemma C.3.

There exists two projections $P_{1}$ and
$P_{2}:\mathbb{R}^{w}\to\mathbb{R}^{w}$ such that the following are true.

  1. 1.

The image of $P_{1}$ and $P_{2}$ are orthogonal to each other;

  2. 2.

With high probability,

$\|W_{1}W_{1}^{T}(0)-\sigma^{2}wP_{1}\|_{\mathrm{op}}=O(\sigma^{2}\sqrt{wd}\log
d)$ (11)
$\|W_{2}^{T}W_{2}(0)-\sigma^{2}wP_{2}\|_{\mathrm{op}}=O(\sigma^{2}\sqrt{wd}\log
d)$ (12)

###### Proof of Lemma C.3 .

: At initialization, the rank of $w$is $w$ with probability 1. Therefore
$W_{1}W_{1}^{T}$ has $w-d$ eigenvalues that are 0, and the $w$ non-zero
eigenvalues equals the eigenvalues of $d\times d$ matrix $W_{1}^{T}W_{1}$.

$W_{1}^{T}W_{1}(0)$ is a (scaled) Wishart ensemble, whose limiting
distribution is given by the Marchenko-Pastur law. The Marchenko-Pastur law,
as stated in [4], proves the following. Let $X$ be an $M\times N$ matrix with
complex-valued independent entries $X_{i\mu}$ such that

  1. 1.

$\mathbb{E}[X_{i\mu}]=0$;

  2. 2.

$\mathbb{E}[\absolutevalue{X_{i\mu}}^{2}]=\frac{1}{\sqrt{MN}}$;

  3. 3.

for every $p\in\mathbb{N}$, there exists a constant $C_{p}$ such that

$\mathbb{E}\left[\absolutevalue{(NM)^{\frac{1}{4}}X_{i\mu}}^{p}\right]\leq
C_{p}.$

Here, $M$ satisfies

$0<C^{-1}\leq\frac{\log M}{\log N}\leq C<\infty$

for some constant $C$ independent of $M$ and $N$. Let $\phi=\frac{M}{N}$,
which may or may not depend on $N$. Then the eigenvalues of $N\times N$ matrix
$X^{*}X$ has the same asymptotics as

$\rho_{\phi}(dx):=\frac{\sqrt{\phi}}{2\pi}\sqrt{\frac{[(x-\gamma_{-})(\gamma_{+}-x)]_{+}}{x^{2}}}dx+(1-\phi)_{+}\delta(dx)$

where

$\gamma_{\pm}:=\sqrt{\phi}+\frac{1}{\sqrt{\phi}}\pm 2$ (13)

In our situation, $W_{1}$ is a $w\times d$ matrix with independent and
identically distributed Gaussian entries whose variance is $\sigma^{2}.$ Let
$M=w$, $N=d$ and therefore $W_{1}^{T}W_{1}(0)$ has the same distribution as
$\sigma^{2}\sqrt{wd}X^{*}X$. Notice that for this choice of $M$ and $N$, the
asymptotic distribution eigenvalues of $X^{*}X$ is $\rho_{\frac{w}{d}}(dx)$.
Notice that $\rho_{\frac{w}{d}}$ is supported on interval
$\left[\sqrt{\frac{w}{d}}+\sqrt{\frac{d}{w}}-2,\sqrt{\frac{w}{d}}+\sqrt{\frac{d}{w}}+2\right]$,
from which we conclude that in the limit, all eigenvalues of
$W_{1}^{T}W_{1}(0)$ is approximately $\sqrt{\frac{d}{w}}$.

By Theorem 2.10 of [4], we have eigenvalue rigidity results for $X^{*}X$. Let
$\lambda_{k}^{\prime}$ be the $k$-th largest eigenvalue for $X^{*}X$. $\forall
k\in\\{1,2,\ldots,w\\}$, we have

$\absolutevalue{\lambda_{k}^{\prime}-\gamma_{k}}<d^{-\frac{2}{3}+\varepsilon}$
(14)

with high probability. Here $\gamma_{\alpha}$ is defined through

$\int_{\gamma_{k}}^{\infty}\rho_{\phi}(dx)=\frac{k}{d}$ (15)

Let $\lambda_{k}$ be the $k$-th largest eigenvalue of $W_{1}^{T}W_{1}(0)$. By
the relationship between $W_{1}^{T}W_{1}(0)$ and $X^{*}X$, we know
$\lambda_{k}$ has the same law as $\sigma^{2}\sqrt{wd}\lambda_{k}^{\prime}$.
Then for every $k$,

$\displaystyle\absolutevalue{\lambda_{k}-\sigma^{2}w}$
$\displaystyle\leq\sigma^{2}\sqrt{wd}\left(\absolutevalue{\lambda_{k}^{\prime}-\gamma_{k}}+\absolutevalue{\gamma_{k}-\sqrt{\frac{w}{d}}}\right)$
(16) $\displaystyle\leq O(\sigma^{2}\sqrt{wd})$ (17)

with high probability. We conclude that the first $w$ eigenvalues of
$W_{1}^{T}W_{1}(0)$ is at most $O(\sigma^{2}\sqrt{wd})$-away from 1 and all
other eigenvalues are 0. Therefore there exists a projection $P_{1}$ such that

$\|W_{1}W_{1}^{T}(0)-\sigma^{2}wP_{1}\|_{\mathrm{op}}=O(\sigma^{2}\sqrt{wd})$
(18)

Similarly, there exists a projection $\tilde{P}_{2}$ such that

$\|W_{2}^{T}W_{2}(0)-\sigma^{2}w\tilde{P}_{2}\|_{\mathrm{op}}=O(\sigma^{2}\sqrt{wd})$
(19)

Notice that $\tilde{P}_{2}$ is not exactly orthogonal to $P_{1}$, and it
remains to find a projection $P_{2}$ that is orthogonal to $P_{1}$ and is
close to $W_{2}^{T}W_{2}(0)$. Assume that the column vectors of $W_{1}(0)$ are
$u_{1},\ldots,u_{d}\in\mathbb{R}^{w}$ and column vectors of $W_{2}^{T}(0)$ are
$v_{1},\ldots,v_{d}\in\mathbb{R}^{w}$. For $k=1,2,\ldots,d$, we define vector
$v_{k}^{\prime}$ as

$v_{k}^{\prime}=v_{k}-P_{1}v_{k}.$ (20)

We claim that $P_{1}v_{k}$ is very small. By law of large numbers,
$\|v_{k}\|\leq\sigma\sqrt{w}\log d$ with high probability.

$\displaystyle\|P_{1}v_{k}\|\leq$
$\displaystyle\|\frac{1}{\sigma^{2}w}W_{1}W_{1}^{T}(0)v_{k}\|+O(\sqrt{\frac{d}{w}}\|v_{1}\|)$
(21) $\displaystyle=$ $\displaystyle\frac{1}{\sigma^{2}w}\|\langle
u_{1},v_{k}\rangle u_{1}+\ldots+\langle u_{w},v_{k}\rangle
u_{w}\|+O(\sigma\sqrt{d}\log d)$ (22)

Notice that $\langle u_{j},v_{k}\rangle u_{j},\forall j$ is a family of
independent and identically distributed random vectors. For each of these
random vectors, all entries have zero mean. The variance of any one of the
entries is given by

$\mathbb{E}\left[\langle u_{i},v_{k}\rangle^{2}\langle
u_{j},e_{\ell}\rangle^{2}\right]=O(\sigma^{6}w)$ (23)

We conclude that for each $\ell$ we have, by CLT,

$\frac{1}{\sigma^{3}w^{\frac{1}{2}}}\frac{\langle u_{1},v_{k}\rangle\langle
u_{1},e_{\ell}\rangle+\ldots+\langle u_{w},v_{k}\rangle\langle
u_{w},e_{\ell}\rangle}{\sqrt{d}}\xrightarrow{(d)}N(0,1)$ (24)

In particular,

$\displaystyle\mathbb{P}\\{\text{max}_{\ell}\sum_{i=1}^{w}\absolutevalue{\langle
u_{i},v_{k}\rangle\langle
u_{k},e_{\ell}\rangle}>100\sigma^{3}\sqrt{d}w^{\frac{1}{2}}\log d\\}\leq$
$\displaystyle w\mathbb{P}\\{\sum_{i=1}^{w}\absolutevalue{\langle
u_{i},v_{k}\rangle\langle
u_{k},e_{1}\rangle}>100\sigma^{3}\sqrt{d}w^{\frac{1}{2}}\log d\\}$
$\displaystyle\leq$ $\displaystyle 2w\mathbb{P}\\{N(0,1)>100\log d\\}$
$\displaystyle\leq$ $\displaystyle O(d^{-50})$

Therefore with high probability,
$\text{max}_{\ell}\sum_{i=1}^{w}\absolutevalue{\langle
u_{i},v_{k}\rangle\langle u_{k},e_{\ell}\rangle}\leq
100\sigma^{3}\sqrt{d}w^{\frac{1}{2}}\log d$. This implies that with high
probability,

$\|\langle u_{1},v_{k}\rangle u_{1}+\ldots+\langle u_{w},v_{k}\rangle
u_{w}\|\leq\sqrt{w}\sigma^{3}\sqrt{d}w^{\frac{1}{2}}\log d$ (25)
$\|P_{1}v_{k}\|\leq\sigma\sqrt{d}\log d$ (26)

Now let ${W_{2}^{\prime}}^{T}$ be the matrix with column vector
$v_{1}^{\prime},\ldots,v_{w}^{\prime}$ and let $P_{2}$ be the projection to
the column space of ${W_{2}^{\prime}}^{T}$ . By construction
$\|W_{2}-W_{2}^{\prime}\|_{\mathrm{op}}\leq O(\sigma\sqrt{d}\log d)$ , so
$\|W_{2}^{T}W_{2}-{W_{2}^{\prime}}^{T}W_{2}^{\prime}\|_{\mathrm{op}}\leq
O(\sigma^{2}\sqrt{wd}\log d)$. Since the nonzero eigenvalues of
$W_{2}^{T}W_{2}$ is at most $O(\sigma^{2}\sqrt{wd})$ from 1, we know that the
nonzero singular values of ${W_{2}^{\prime}}^{T}W_{2}^{\prime}$ is also at
most $O(\sigma^{2}\sqrt{wd})$ from 1. We conclude that
$\|P_{2}-W_{2}^{T}W_{2}\|_{\mathrm{op}}\leq O(\sigma^{2}\sqrt{wd}\log d)$. ∎

###### Proof of Proposition C.1.

: We have

$\displaystyle W_{1}^{T}W_{1}W_{1}^{T}W_{1}(t)$
$\displaystyle=W_{1}^{T}W_{2}^{T}W_{2}W_{1}+W_{1}^{T}(W_{1}W_{1}^{T}(t)-W_{2}^{T}W_{2}(t))W_{1}$
$\displaystyle=A^{T}A+W_{1}^{T}(W_{1}W_{1}^{T}(0)-W_{2}^{T}W_{2}(0))W_{1}$

From lemma C.3 we see that as positive semi-definite matrix, for every
constant $\varepsilon>0$,

$0\leq W_{1}W_{1}^{T}(0)\leq(1+\varepsilon)\sigma^{2}wI$ (27)

with high probability. Therefore

$-(1+\varepsilon)\sigma^{2}wW_{1}^{T}W_{1}\leq
W_{1}^{T}(P_{1}(0)-P_{2}(0))W_{1}\leq(1+\varepsilon)\sigma^{2}wW_{1}^{T}W_{1}.$
(28)

By moving terms around and corollary, we have

$\displaystyle(W_{1}^{T}W_{1})^{2}-(1+\varepsilon)\sigma^{2}wW_{1}^{T}W_{1}+(1+\varepsilon)^{2}\frac{\sigma^{4}w^{2}}{4}I$
(29) $\displaystyle\leq$ $\displaystyle
A^{T}A+\frac{(1+\varepsilon)^{2}}{4}\sigma^{4}w^{2}I$ (30) $\displaystyle\leq$
$\displaystyle(W_{1}^{T}W_{1})^{2}+(1+\varepsilon)W_{1}^{T}W_{1}+(1+\varepsilon)^{2}\frac{\sigma^{4}w^{2}}{4}I$
(31)

Theorem V.1.9 of [12] states that the square-root function is operator
monotone, which implies that if $A\geq B$ then $\sqrt{A}\geq\sqrt{B}$. Taking
square-root, we have

$W_{1}^{T}W_{1}-(1+\varepsilon)\frac{\sigma^{2}w}{2}Id\leq\sqrt{A^{T}A+(1+\varepsilon)^{2}\frac{\sigma^{4}w^{2}}{4}}\leq
W_{1}^{T}W_{1}+(1+\varepsilon)\frac{\sigma^{2}w}{2}$ (32)

∎

###  C.2 Strong Bound

Recall that the evolution of weight matrix in gradient descent is given by the
following.

$\frac{d}{dt}W_{1}(t)=\eta W_{2}^{T}\nabla C$ (33)
$\frac{d}{dt}W_{2}(t)=\eta\nabla CW_{1}^{T}$ (34)

The goal of this section is to prove that

$W_{1}^{T}W_{1}\approx\sqrt{A^{T}A+\sigma^{4}w^{2}I}$ (35)

For simplicity of notations, we shall assume that $C_{1}=W_{1}^{T}W_{1}$,
$C_{2}=W_{2}W_{2}^{T}$, $\hat{C}_{1}=\sqrt{A^{T}A+\sigma^{4}w^{2}I}$ and
$C_{2}=\sqrt{AA^{T}+\sigma^{4}w^{2}I}$. It is easy of see that $\hat{C}_{1}$
and $\hat{C}_{2}$ are invertible. Our main result for this section is the
following proposition.

######  Proposition C.4.

For every cost $C$ we have

$\|W_{1}^{T}W_{1}-\sqrt{A^{T}A+\sigma^{4}w^{2}I}\|_{\mathrm{op}}\leq\text{min}\\{O(\sigma^{2}w),O\left(\sqrt{\frac{d}{w}}\|W_{1}^{T}W_{1}\|_{op}\right)\\}$
(36)

###### Proof of Lemma C.4 .

: We start from the equations:

$\displaystyle
W_{1}^{T}\left[(W_{2}^{T}W_{2}-W_{1}W_{1}^{T})^{2}-\sigma^{4}w^{2}I\right]W_{1}$
$\displaystyle=C_{1}^{3}-A^{T}AC_{1}-C_{1}A^{T}A-\sigma^{4}w^{2}C_{1}+A^{T}C_{2}A$
$\displaystyle
W_{2}\left[(W_{2}^{T}W_{2}-W_{1}W_{1}^{T})^{2}-\sigma^{4}w^{2}I\right]W_{2}^{T}$
$\displaystyle=C_{2}^{3}-AA^{T}C_{2}-C_{2}AA^{T}-\sigma^{4}w^{2}C_{2}+AC_{1}A^{T}.$

Our goal is to show that $C_{1},C_{2}$ are close to the solution
$\hat{C}_{1}=\sqrt{A^{T}A+\sigma^{4}w^{2}I},\hat{C}_{2}=\sqrt{AA^{T}+\sigma^{4}w^{2}I}$
with

$\displaystyle 0$
$\displaystyle=\hat{C}_{1}^{3}-A^{T}A\hat{C}_{1}-\hat{C}_{1}A^{T}A-\sigma^{4}w^{2}\hat{C}_{1}+A^{T}\hat{C}_{2}A$
$\displaystyle 0$
$\displaystyle=\hat{C}_{2}^{3}-AA^{T}\hat{C}_{2}-\hat{C}_{2}AA^{T}-\sigma^{4}w^{2}\hat{C}_{2}+A\hat{C}_{1}A^{T}.$

For gradient flow, $W_{2}^{T}W_{2}-W_{1}W_{1}^{T}$ is preserved and there
exists projections $P_{1}$ and $P_{2}$ such that
$W_{2}^{T}W_{2}-W_{1}W_{1}^{T}=\sigma^{2}w(P_{2}-P_{1})+O(\sigma^{2}\sqrt{wd})$.
Therefore for every unit vector $v$ we have

$\displaystyle\|v^{T}W_{1}^{T}\left[(W_{2}^{T}W_{2}-W_{1}W_{1}^{T})^{2}-\sigma^{4}w^{2}I\right]W_{1}\|\leq\|C_{1}\|_{op}\sigma^{4}d^{\frac{1}{2}}w^{\frac{3}{2}}.$

Recall that $\frac{dw}{dt}=$ Substracting the second pair of equations from
the first pair and denoting $dC_{i}=C_{i}-\hat{C}$, we obtain:

$\displaystyle\|C_{1}\|_{op}O(\sigma^{4}w^{2}\sqrt{\frac{d}{w}})$
$\displaystyle=C_{1}^{3}-\hat{C}_{1}^{3}-A^{T}AdC_{1}-dC_{1}A^{T}A-\sigma^{4}w^{2}dC_{1}+A^{T}dC_{2}A$
$\displaystyle\|C_{2}\|_{op}O(\sigma^{4}w^{2}\sqrt{\frac{d}{w}})$
$\displaystyle=C_{2}^{3}-\hat{C}_{2}^{3}-AA^{T}dC_{2}-dC_{2}AA^{T}-\sigma^{4}w^{2}dC_{2}+AdC_{1}A^{T}.$

Now since

$\displaystyle C_{1}^{3}-\hat{C}_{1}^{3}$
$\displaystyle=\hat{C}_{1}^{2}dC_{1}+\hat{C}_{1}dC_{1}C_{1}+dC_{1}C_{1}^{2}$

we obtain

$\displaystyle\|C_{1}\|_{op}O(\sigma^{4}w^{2}\sqrt{\frac{d}{w}})$
$\displaystyle=\left(\hat{C}_{1}^{2}-A^{T}A\right)dC_{1}+\hat{C}_{1}dC_{1}C_{1}+dC_{1}\left(C_{1}^{2}-A^{T}A\right)-dC_{1}+A^{T}dC_{2}A$
$\displaystyle=\hat{C}_{1}dC_{1}C_{1}+dC_{1}\left(C_{1}^{2}-A^{T}A\right)+A^{T}dC_{2}A,$

and similarly for the other index. For any singular value $s_{i}$ of $A$, with
left and right singular vectors $u_{i},v_{i}$, we multiply the first equation
to the left by $v_{i}^{T}$ and the second by $u_{i}^{T}$ to obtain an equation
for $v_{i}^{T}dC_{1}$ and $u_{i}^{T}dC_{2}$:

$\displaystyle\|C_{1}\|_{op}O(\sigma^{2}\sqrt{wd})$
$\displaystyle=v_{i}^{T}dC_{1}\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}C_{1}+\frac{1}{\sigma^{2}w}(C_{1}^{2}-A^{T}A)\right)+\left(\frac{s_{i}}{\sigma^{2}w}\right)u_{i}^{T}dC_{2}A$
$\displaystyle\|C_{2}\|_{op}O(\sigma^{2}\sqrt{wd})$
$\displaystyle=u_{i}^{T}dC_{2,i}\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}C_{2}+\frac{1}{\sigma^{2}w}(C_{2}^{2}-AA^{T})\right)+\left(\frac{s_{i}}{\sigma^{2}w}\right)v_{i}^{T}dC_{1}A^{T}.$

Notice that
$C_{1}^{2}-A^{T}A=W_{1}^{T}(W_{1}W_{1}^{T}-W_{2}^{T}W_{2})W_{1}=\sigma^{2}wW_{1}^{T}(P_{1}-P_{2})W_{1}+\|C_{1}\|O(\sigma^{2}\sqrt{wd})$.
In the two equations above, by replacing $C_{1}^{2}-A^{T}A$ with
$\sigma^{2}wW_{1}^{T}(P_{1}-P_{2})W_{1}$, we are making an error of at most
$\|C_{1}\|_{op}\|dC_{1}\|_{op}O(\sqrt{\frac{d}{w}})$. From weak bound we know
that $\|dC_{1}\|_{op}\leq O(\sigma^{2}w)$. Therefore the error we made by
making the approximation on the right hand side can be absorbed into left hand
side.

To show that $\|v_{i}^{T}dC_{1}\|$ and $\|u_{i}^{T}dC_{2}\|$, we need to show
that the $(d_{in}+d_{out})\times(d_{in}+d_{out})$ block matrix

$\left(\begin{array}[]{cc}\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}C_{1}+W_{1}^{T}(P_{1}-P_{2})W_{1}&\left(\frac{s_{i}}{\sigma^{2}w}\right)A\\\
\left(\frac{s_{i}}{\sigma^{2}w}\right)A^{T}&\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}C_{2}+W_{2}(P_{2}-P_{1})W_{2}^{T}\end{array}\right)$

is full rank. This matrix can be further simplified to

$\left(\begin{array}[]{cc}W_{1}^{T}&0\\\
0&W_{2}\end{array}\right)\left(\begin{array}[]{cc}\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}I+P_{1}-P_{2}&\left(\frac{s_{i}}{\sigma^{2}w}\right)I\\\
\left(\frac{s_{i}}{\sigma^{2}w}\right)I&\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}I+P_{2}-P_{1}\end{array}\right)\left(\begin{array}[]{cc}W_{1}&0\\\
0&W_{2}^{T}\end{array}\right)$

The inner matrix can then be rewritten as $RR^{T}$ where $R$ is defined as

$R=\left(\begin{array}[]{c}\sqrt{\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}+1}P_{1}+\sqrt{\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}-1}P_{2}\\\
\sqrt{\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}-1}P_{1}+\sqrt{\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}+1}P_{2}\end{array}\right).$

Let $Q=\left(\begin{array}[]{cc}W_{1}^{T}&0\\\ 0&W_{2}\end{array}\right)$.
Then the smallest nonzero singular value of $QR(QR)^{T}$ is the same as the
smallest nonzero singular value of $(QR)^{T}QR$. We expand $(QR)^{T}QR$ as
follows.

$\displaystyle(QR)^{T}QR$ $\displaystyle=$
$\displaystyle\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}+1\right)P_{1}W_{1}W_{1}^{T}P_{1}+\left(\frac{s_{i}}{\sigma^{2}w}\right)P_{1}W_{1}W_{1}^{T}P_{2}$
$\displaystyle+\left(\frac{s_{i}}{\sigma^{2}w}\right)P_{2}W_{1}W_{1}^{T}P_{1}+\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}-1\right)P_{2}W_{1}W_{1}^{T}P_{2}$
$\displaystyle+\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}+1\right)P_{2}W_{2}^{T}W_{2}P_{2}+\left(\frac{s_{i}}{\sigma^{2}w}\right)P_{1}W_{2}^{T}W_{2}P_{2}$
$\displaystyle+\left(\frac{s_{i}}{\sigma^{2}w}\right)P_{2}W_{2}^{T}W_{2}P_{1}+\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}-1\right)P_{1}W_{2}^{T}W_{2}P_{1}$
$\displaystyle=$
$\displaystyle\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}+1\right)(P_{1}+P_{2})\sigma^{2}w$
$\displaystyle+2(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}-\left(\frac{s_{i}}{\sigma^{2}w}\right))(P_{1}W_{2}^{T}W_{2}P_{1}+P_{2}W_{1}W_{1}^{T}P_{2})$
$\displaystyle+\left(\frac{s_{i}}{\sigma^{2}w}\right)(W_{1}W_{1}^{T}+W_{2}^{T}W_{2}-\sigma^{2}wP_{1}-\sigma^{2}wP_{2})+O(\sigma^{2}\sqrt{wd})$
$\displaystyle=$
$\displaystyle\left(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}+1-\left(\frac{s_{i}}{\sigma^{2}w}\right)\right)(P_{1}+P_{2})\sigma^{2}w$
$\displaystyle+2(\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}-\left(\frac{s_{i}}{\sigma^{2}w}\right))(P_{1}W_{2}^{T}W_{2}P_{1}+P_{2}W_{1}W_{1}^{T}P_{2})$
$\displaystyle+\left(\frac{s_{i}}{\sigma^{2}w}\right)(W_{1}W_{1}^{T}+W_{2}^{T}W_{2})+O(\sigma^{2}\sqrt{wd}).$

where we used the fact that

$W_{1}W_{1}^{T}=P_{1}+P_{1}W_{2}^{T}W_{2}P_{1}+P_{1}W_{1}W_{1}^{T}P_{2}+P_{2}W_{1}W_{1}^{T}P_{1}+P_{2}W_{1}W_{1}^{T}P_{2}.$

Tthe $(d_{in}+d_{out})$-th eigenvalue of the above is lower bounded by
$\sigma^{2}w\sqrt{\left(\frac{s_{i}}{\sigma^{2}w}\right)^{2}+1}+\sigma^{2}w-\sigma^{2}w\left(\frac{s_{i}}{\sigma^{2}w}\right)\geq\sigma^{2}w$.
We conclude that

$\|dC_{1}\|_{op}+\|dC_{2}\|_{op}\leq O(\sqrt{\frac{d}{w}})\|C_{1}\|_{op},$

∎

Compared to lemma C.3, lemma C.4 gives a tighter bound on
$\|C_{1}-\hat{C}_{1}\|_{\mathrm{op}}$ when
$\|C_{1}\|_{\mathrm{op}}\leq\sigma^{2}w\sqrt{\frac{d}{w}}$.

##  Appendix D  The Gradient Flow Dynamics of $A_{\theta(t)}$ in Lazy Regime

######  Lemma D.1.

With high probability, for all time $t$,
$\|A_{\theta(t)}-A^{*}\|_{F}^{2}\geq\frac{1}{3}\text{min}(\|A^{*}\|_{F}^{2},\|E\|_{F}^{2})$.

###### Proof.

In gradient flow dynamics, $\|A_{\theta(t)}-A^{*}-E\|_{F}$ is decreasing with
time and therefore for all time $t$, $\|A_{\theta(t)}\|_{F}^{2}\leq
9(\|A^{*}\|_{F}^{2}+\|E\|_{F}^{2})$. In particular,
$\|A_{\theta(t)}\|_{op}\leq 9(\|A^{*}\|_{F}+\|E\|_{F})$. By assumption,
$\|\hat{C}_{1}-C_{1}\|_{op}\leq
O(\sqrt{\frac{d}{w}})\|C_{1}\|_{op}=O(d\sqrt{\frac{d}{w}})$, and the dynamics
of $A_{\theta(t)}$ can be written as

$d^{2}\frac{dA_{\theta(t)}}{dt}=2(A^{*}+E-A)\sigma^{2}w+(A^{*}+E-A)O(d\sqrt{\frac{d}{w}})+O(d\sqrt{\frac{d}{w}})(A^{*}+E-A)$

Assume that $B(0)=A_{\theta(0)}$ and

$d^{2}\frac{dB}{dt}=2(A^{*}+E-B)\sigma^{2}w.$

Then

$\displaystyle\frac{1}{2}d^{2}\frac{d}{dt}\|A_{\theta(t)}-B\|_{F}^{2}=$
$\displaystyle-2\sigma^{2}w\|A_{\theta(t)}-B\|_{F}^{2}+tr((A-B)^{T}(A^{*}+E-A)O(d\sqrt{\frac{w}{d}}))$
$\displaystyle+tr((A-B)^{T}O(d\sqrt{\frac{w}{d}})(A^{*}+E-A))$
$\displaystyle\leq$
$\displaystyle-2\sigma^{2}w\|A_{\theta(t)}-B\|_{F}^{2}+\|A_{\theta(t)}-B\|_{F}O(d^{2}\sqrt{\frac{d}{w}})$

This implies that for all time,

$\|A_{\theta(t)}-B\|_{F}\leq dO(\frac{d}{\sigma^{2}w}\sqrt{\frac{d}{w}}).$

The dynamics of $B$ is linear, and we have

$B_{t}=(A^{*}+E)(1-e^{-2\sigma^{2}wt/d^{2}})+B_{0}e^{-2\sigma^{2}wt/d^{2}}.$
$\|B_{t}-A^{*}\|_{F}^{2}=\|e^{-2\sigma^{2}t/d^{2}}A^{*}+(1-e^{-2\sigma^{2}wt/d^{2}})E+B_{0}e^{-2\sigma^{2}wt/d^{2}}\|_{F}^{2}\geq
0.99\|e^{-2\sigma^{2}t/d^{2}}A^{*}+(1-e^{-2\sigma^{2}wt/d^{2}})E\|_{F}^{2}$

Let $P$ be the projection to the image of $A^{*}$. Then with high probability,

$\displaystyle\|B_{t}-A^{*}\|_{F}^{2}=$
$\displaystyle\|e^{-2\sigma^{2}wt/d^{2}}A^{*}+(1-e^{-2\sigma^{2}wt/d^{2}})E+B_{0}e^{-2\sigma^{2}wt/d^{2}}\|_{F}^{2}$
$\displaystyle\geq$ $\displaystyle
0.99\|e^{-2\sigma^{2}t/d^{2}}A^{*}+(1-e^{-2\sigma^{2}wt/d^{2}})E\|_{F}^{2}$
$\displaystyle=$ $\displaystyle
0.99\|e^{-2\sigma^{2}wt/d^{2}}A^{*}+(1-e^{-2\sigma^{2}wt/d^{2}})PE\|_{F}^{2}+0.99\|(1-e^{-2\sigma^{2}wt/d^{2}})(I-P)E\|_{F}^{2}$
$\displaystyle\geq$
$\displaystyle\frac{1}{2}\text{min}(\|A^{*}\|_{F}^{2},\|E\|_{F}^{2})$

Since $\|A_{\theta(t)}-B\|_{F}^{2}=o(d^{2})$, we have

$\|A_{\theta(t)}-A^{*}\|_{F}^{2}\geq\frac{1}{3}\text{min}(\|A^{*}\|_{F}^{2},\|E\|_{F}^{2})$

∎

##  Appendix E The Gradient Flow Dynamics of $A_{\theta(t)}$ in Active Regime

###  E.1 Saddle-to-Saddle Regime

The goal of this section is to show that at the end of the first NTK regime,
the singular vectors of $A_{\theta(t)}$ are roughly aligned with $A^{*}$.
Moreover, the alignment remains to be good throughout the mixed regime. The
following generalization of lemma B.6 and lemma B.7 will be useful.

######  Lemma E.1.

Assume that $A^{\prime}_{[t_{1},t_{2}]}\subset P_{1}$ and $A_{t_{1}}\in
P_{2}(C_{4},\gamma)$. Let $s_{1}^{\prime},\ldots,s_{K}^{\prime}$ be the first
$K$ singular values of $A^{\prime}$. Assume that the dynamics of $A^{\prime}$
is the following.

$\displaystyle
d^{2}\frac{dA^{\prime}}{dt}=(A^{*}+E-A^{\prime})\sqrt{A^{{}^{\prime}T}A^{\prime}+\sigma^{4}w^{2}I}$
$\displaystyle+\sqrt{A^{\prime}A^{{}^{\prime}T}+\sigma^{4}w^{2}}(A^{*}+E-A^{\prime})$
$\displaystyle+O(d^{-\lambda}s_{K}^{\prime})(A^{*}+E-A^{\prime})+(A^{*}+E-A^{\prime})O(d^{-\lambda}s_{K}^{\prime})$

Here, $O(d^{-\lambda}s_{K}^{\prime})$ is a matrix whose operator norm is
bounded by $O(d^{-\lambda}s_{K}^{\prime})$. Then $A_{[t_{1},t_{2}]}\subset
P_{2}(C_{4},\text{min}(\gamma,1,2\lambda))$.

###### Proof.

Let $A^{\prime}=USV^{T}$. Computing the SVD derivative for $A^{\prime}$, we
have

$\displaystyle d^{2}\frac{dU}{dt}=$ $\displaystyle
U\left(F\odot\left[U^{T}A^{*}V\sqrt{S^{2}+\sigma^{4}w^{2}}S+\sqrt{S^{2}+\sigma^{4}w^{2}}U^{T}A^{*}VS\right]\right)$
$\displaystyle+U\left(F\odot\left[\sqrt{S^{2}+\sigma^{4}w^{2}}SV^{T}A^{*}U+V^{T}A^{*}US\sqrt{S^{2}+\sigma^{4}w^{2}}\right]\right)$
$\displaystyle+U\left(F\odot\left[O(s_{K}d^{-\lambda})(U^{T}A^{*}VS+SV^{T}A^{*}U-2S^{2})\right]\right)$

Assume that $a_{k}\neq a_{\ell}$. Notice that

$F_{k\ell}(O(s_{k}d^{-\lambda})U^{T}A^{*}VS)_{k\ell}=O(d^{1-\lambda})$

Using the same technique as B.6 we have

$\displaystyle\frac{1}{2}d^{2}\frac{d}{dt}tr[U(k,k)^{T}U(k,k)]$
$\displaystyle\geq$ $\displaystyle cdx+O(dx^{\frac{3}{2}})+O(\sqrt{dx})$
$\displaystyle+tr[U(k,k)^{T}\sum_{j\neq
k}U(k,j)(F\odot\left[O(s_{K}d^{-\lambda})(U^{T}A^{*}VS+SV^{T}A^{*}U-2S^{2})\right])]$
$\displaystyle\geq$ $\displaystyle
cdx+O(dx^{\frac{3}{2}})+O(d^{1-\lambda}\sqrt{x})+O(\sqrt{dx})$

Similar conclusion also holds for $tr[U(k,k)^{T}V(k,k)]$ and
$tr[V(k,k)^{T}V(k,k)]$. We conclude that

$d^{2}\frac{dx}{dt}\leq-
cdx+O(dx^{\frac{3}{2}})+O(d^{1-\lambda}\sqrt{x})+O(\sqrt{dx}),$

and $A^{\prime}_{[t_{1},t_{2}]}\subset
P_{2}(C_{4},\text{min}(\gamma,1,2\lambda))$. ∎

######  Lemma E.2.

Assume that $A^{\prime}_{[t_{1},t_{2}]}\subset P_{2}(C_{4},\gamma)$ and
$A_{t_{1}}\in P_{1}$. Let $s_{1}^{\prime},\ldots,s_{K}^{\prime}$ be the first
$K$ singular values of $A^{\prime}$. Assume that the dynamics of $A^{\prime}$
is the following.

$d^{2}\frac{dA^{\prime}}{dt}=(A^{*}-A^{\prime})\sqrt{A^{{}^{\prime}T}A^{\prime}+\sigma^{4}w^{2}I}+\sqrt{A^{\prime}A^{{}^{\prime}T}+\sigma^{4}w^{2}}(A^{*}-A^{\prime})+O(d^{-\lambda}s_{K}^{\prime})(A^{*}-A)+(A^{*}-A)O(d^{-\lambda}s_{K}^{\prime})$

Here, $O(d^{-\lambda}s_{K}^{\prime})$ is a matrix whose operator norm is
bounded by $O(d^{-\lambda}s_{K}^{\prime})$. Then $A_{[t_{1},t_{2}]}\subset
P_{1}$.

###### Proof.

Proceeding as in lemma B.7, we have

$d^{2}\frac{ds_{p}}{dt}=2(O(\sqrt{d})+O(d^{1-\gamma})+a_{p}d-s_{p})(\sqrt{s_{p}^{2}+\sigma^{4}w^{2}}+O(d^{-\lambda})s_{K})$

if $p$ is a signal, and

$d^{2}\frac{ds_{K+1}}{dt}=2(O(\sqrt{d})+O(d^{1-\gamma})-s_{K+1})(\sqrt{s_{K+1}^{2}+\sigma^{4}w^{2}}+O(d^{-\gamma})s_{K}).$

Now the conclusion follows from Grönwall. ∎

######  Theorem E.3.

Assume that $A_{\theta(t)}\subset P_{1}\cap P_{2}(C,\gamma)$. Then
$A_{[\theta([t,T])}\subset P_{1}\cap P_{2}(C,\gamma^{\prime})$ for some
constant $\gamma^{\prime}$.

###### Proof.

We use induction lemma B.3 to prove the theorem. For the first requirement,
notice that for every $i=1,2,\ldots,w$,

$\displaystyle d^{2}\absolutevalue{\frac{ds_{i}}{dt}}\leq
4\|A^{*}\|_{op}^{2}.$

If $A_{\theta(t)}\in P_{1}$, then $A_{\theta([t,t+\tau])}\subset
P_{1}^{\prime}$ for
$\tau=c\frac{\sigma^{2}\sqrt{wd}}{d^{2}\|A^{*}\|_{op}^{2}}$. Here $c$ is some
small constant that depends only on $P_{1}$. This proves the first
requirement. For the second requirement, we observe that $A_{\theta(t)}$
satisfies the dynamics described in lemma E.1 at each stage, as long as
$\|\hat{C}_{1}-C_{1}\|_{op}\leq O(s_{K}d^{-\lambda}$ for some $\lambda>0$. It
suffices to show that there exists some positive $\lambda$ independent of $w$
such that $\|\hat{C}_{1}-C_{1}\|_{op}\leq O(d^{-\lambda}s_{K})$ in $[0,T]$.
Let $T_{1}$ be the first time when $\|A\|_{op}=\sigma^{2}w$. In $[0,T_{1}]$,
all singular values are of the same order. Therefore

$\|\hat{C}_{1}-C_{1}\|_{op}\leq
O(\|C_{1}\|\sqrt{\frac{d}{w}})=O(s_{K}d^{-\frac{\gamma_{w}-1}{2}}).$

Now let $T_{2}$ be the first time when
$s_{1}=\sigma^{2}wd^{\frac{\gamma_{w}-1}{4}}$. In $[T_{1},T_{2}]$, we have

$\|\hat{C}_{1}-C_{1}\|_{op}\leq O(\|C_{1}\|\sqrt{\frac{d}{w}})\leq
O(s_{K}d^{-\frac{\gamma_{w}-1}{4}})$

Therefore we can pick $\lambda=\frac{\gamma_{w}}{4}$. During $[T_{1},T_{2}]$,
we have $s_{K}(0)\geq c\sigma^{2}w$ for some constant $c$, and

$d^{2}\frac{ds_{K}}{dt}\geq s_{K}(\frac{1}{2}a_{K}d-s_{K})$

Moreover, since $d^{2}\frac{ds_{1}}{dt}\leq(2a_{1}-s_{1})(s_{1}+\sigma^{2}w)$,
we have $T_{2}-T_{1}\geq c\frac{\log w}{w}$ for some $c$. Therefore we have
$s_{K}(T_{2})\geq\sigma^{2}wd^{-\nu}$ for some constant $\nu>0$. Now in
$[T_{2},T]$, we have weak bound $\|\hat{C}_{1}-\hat{C}\|_{op}\leq
O(\sigma^{2}w)$. We can therefore pick $\lambda=\nu$ for $[T_{2},T]$. The
third requirement is verified in lemma E.2. ∎

###  E.2 First NTK Regime

As in $A_{t}$, we need to show that at the end of the NTK regime,
$A_{\theta(t)}$ must be in $P_{1}$ and $P_{2}$. Based on what we already have
for $A_{t}$, this conclusion follows from Grönwall.

######  Lemma E.4.

Assume that $A_{0}=A_{\theta(0)}$. Let $T_{1}^{\prime}$ be the first time when
$\|A_{t}\|_{op}+\|A_{\theta(t)}\|_{op}$ reaches $\sigma^{2}w$. Then

$\|A_{T_{1}^{\prime}}-A_{\theta(T_{1}^{\prime})}\|_{op}\leq
O(\sigma^{2}\sqrt{wd}).$

In particular, $A_{\theta(T_{1}^{\prime})}\in P_{1}\cap
P_{2}(C,\text{min}(\frac{\gamma_{w}-1}{8},\frac{1}{4}))$.

###### Proof.

From dynamics of $A_{t}$ we see that $T_{1}^{\prime}\leq\frac{c}{w}$ for some
constant $c$. The dynamics of $A_{t}$ is the following.

$\displaystyle d^{2}\frac{dA_{t}}{dt}=$
$\displaystyle(A^{*}-A_{t})\sqrt{A^{T}_{t}A_{t}+\sigma^{2}wI}+\sqrt{A_{t}A_{t}^{T}+\sigma^{4}w^{2}I}(A^{*}-A_{t}).$

The dynamics of $A_{\theta(t)}$ can be written as follows.

$\displaystyle d^{2}\frac{dA_{\theta(t)}}{dt}=$
$\displaystyle(A^{*}-A_{\theta(t)})\sqrt{A_{\theta(t)}^{T}A_{\theta(t)}+\sigma^{4}w^{2}I}+\sqrt{A_{\theta(t)}A_{\theta(t)}^{T}+\sigma^{4}w^{2}I}(A^{*}-A_{\theta(t)})$
$\displaystyle+O(\sigma^{2}\sqrt{wd})(A^{*}-A_{\theta(t)})+(A^{*}-A_{\theta(t)})O(\sigma^{2}\sqrt{wd})$

Observe that

$\displaystyle\|A_{\theta(t)}^{T}A_{\theta(t)}-A_{t}^{T}A_{t}\|_{op}\leq$
$\displaystyle\|A_{\theta(t)}^{T}(A_{\theta(t)}-A_{t})\|_{op}+\|(A_{\theta(t)}-A_{t})^{T}A_{t}\|_{op}$
$\displaystyle\leq$ $\displaystyle\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}.$

This implies that we have the following inequality on positive definite
matrices.

$\displaystyle A_{t}^{T}A_{t}-\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}I\leq
A_{\theta(t)}^{T}A_{\theta(t)}\leq
A_{t}^{T}A_{t}+\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}I$

Assume that $A_{t}=USV^{T}$, then

$\displaystyle\|\sqrt{A_{\theta(t)}^{T}A_{\theta(t)}+\sigma^{4}w^{2}I}-\sqrt{A_{t}^{T}A+\sigma^{4}w^{2}I}\|_{op}$
$\displaystyle\leq$
$\displaystyle\|\sqrt{\sigma^{4}w^{2}I+A_{t}^{T}A_{t}+\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}I}-\sqrt{\sigma^{4}w^{2}I+A_{t}^{T}A_{t}-\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}I}\|_{op}$
$\displaystyle=$
$\displaystyle\|\sqrt{\sigma^{4}w^{2}I+\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}I+S^{2}}-\sqrt{\sigma^{4}w^{2}I-\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}I+S^{2}}\|_{op}$
$\displaystyle=$
$\displaystyle\text{max}_{i}\left(\sqrt{\sigma^{4}w^{2}+s_{i}^{2}+\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}}-\sqrt{\sigma^{4}w^{2}+s_{i}^{2}-\sigma^{2}w\|A_{\theta(t)}-A_{t}\|_{op}}\right)$
$\displaystyle\leq$ $\displaystyle 3\|A_{\theta(t)}-A_{t}\|_{op}$

We can control the difference between $A_{t}$ and $A_{\theta(t)}$.

$\displaystyle d^{2}\frac{d}{dt}\|A_{t}-A_{\theta(t)}\|_{op}\leq
O(d)\|A_{t}-A_{\theta(t)}\|+O(d)O(\sigma^{2}\sqrt{wd})$

Grönwall inequality implies that

$\displaystyle\|A_{T_{1}^{\prime}}-A_{\theta(T_{1}^{\prime})}\|_{op}\leq$
$\displaystyle O(\sigma^{2}\sqrt{wd})$

We check that $A_{\theta(T_{1}^{\prime})}$ satisfies the $P_{1}$. Since
$\|A_{T_{1}^{\prime}}-A_{\theta(T_{1}^{\prime})}\|_{op}\leq
O(\sigma^{2}\sqrt{wd})$,
$\|A_{\theta(T_{1}^{\prime})}\|_{op}\geq\frac{1}{3}\sigma^{2}w$. Let
$\sigma_{j}(A_{T_{1}^{\prime}})$ and $\sigma_{j}(A_{\theta(T_{1}^{\prime})})$
be the $j$-th singular value of $A_{T_{1}^{\prime}}$ and
$A_{\theta(T_{1}^{\prime})}$. If $a_{j}\neq a_{k}$, $j,k\leq K$, then
$\absolutevalue{\sigma_{j}(A_{T_{1}^{\prime}})-\sigma_{k}(A_{T_{1}^{\prime}})}\geq
c\sigma^{2}w$, and therefore

$\absolutevalue{\sigma_{j}(A_{\theta(T_{1}^{\prime})})-\sigma_{k}(A_{\theta(T_{1}^{\prime})})}\geq
c\sigma^{2}w+O(\sigma^{2}\sqrt{wd}).$

The noise is clearly of order $O(\sigma^{2}\sqrt{wd})$. This completes the
verification of $P_{1}$. By lemma B.9,
$\|A_{\theta(T_{1}^{\prime})}-\Sigma\|_{op}\leq\sigma^{2}wO(\sigma^{2}w)d^{-\text{min}(\frac{\gamma_{w}-1}{8},\frac{1}{4})}$.
By Lemma B.4 we are done.

∎

###  E.3 Summary of Approximate Dynamics of $A_{\theta(t)}$ at Each Stage

In previous sections we have proved that $A_{\theta(t)}$ satisfies $P_{1}$ and
$P_{2}(C,\gamma)$ for some $\gamma>0$. The $P_{1}$ and $P_{2}(C,\gamma)$
property actually implies that the dynamics of $A_{\theta(t)}$ is such that
each group of singular values evolve independently. We state the approximate
dynamics for each stage of the dynamics, and show that the alignment will be
improved, if the alignment is not already good enough.

  * •

Initialization. At initialization, $A_{\theta(0)}$ is a random matrix. The
mean of each entry is 0 and the variance of each entry is
$\sigma^{2}\sqrt{w}$. The gap between singular values is infinitely smaller
than the magnitude of singular values, and the singular vectors are not
aligned with $A^{*}$.

  * •

Initialization to $\|A_{\theta(t)}\|_{op}=\sigma^{2}w$. Let $T_{1}$ be the
first time when $\|A_{\theta(t)}\|_{op}$ reaches $\sigma^{2}w$. $[0,T_{1}]$
corresponds the very short NTK regime for the signals, and the dynamics of
$A_{\theta(t)}$ is approximately linear. The evolution of signal singular
values are roughly linear (i.e., bounded above and below by linear functions),
and at $T_{1}$, we have $A_{\theta(T_{1})}\in
P_{2}(C,\text{min}(\frac{\gamma_{w}-1}{8},\frac{1}{4}))$. Since the singular
values grows roughly linearly, $T_{1}=O(d)$.

  * •

$\|A_{\theta(t)}\|_{op}=\sigma^{2}w$ to
$\|A_{\theta(t)}\|_{op}=\sigma^{2}w\left(\frac{w}{d}\right)^{\frac{1}{4}}$.
Let $T_{2}$ be the first time when
$\|A_{\theta(t)}\|_{op}=\sigma^{2}w\left(\frac{w}{d}\right)^{\frac{1}{4}}$.
$[T_{1},T_{2}]$ is the early stage of saddle-to-saddle dynamics. The dynamics
of $s_{i}$ is given by

$d^{2}\frac{ds_{i}}{dt}=2(a_{i}d(1+o(1))-s_{i})\sqrt{s_{i}^{2}+\sigma^{4}w^{2}}(1+o(1)).$

By theorem 1, we have
$\|\hat{C}_{1}-C_{1}\|_{op}\leq\sigma^{2}w\left(\frac{w}{d}\right)^{-\frac{1}{4}}$.
Let
$d^{-\lambda}\sigma^{w}=\sigma^{2}w\left(\frac{w}{d}\right)^{-\frac{1}{4}}$,
we have $\lambda=\frac{\gamma_{w}-1}{4}$. By lemma 4,1, the dynamics of $x$
satisfies

$d^{2}\frac{dx}{dt}\leq-
cdx+O(dx^{\frac{3}{2}})+O(\sqrt{x}d^{1-\frac{\gamma_{w}-1}{4}})+O(\sqrt{dx})$

with $x(T_{1})=\frac{\gamma_{w}-1}{2}$. We conclude that

$x(T_{2})\leq O(d^{-\frac{\gamma_{w}-1}{2}}).$

  * •

$\|A_{\theta(t)}\|_{op}=\sigma^{2}w\left(\frac{w}{d}\right)^{\frac{1}{4}}$ to
$s_{K}(A_{\theta(t)})=\frac{1}{2}a_{K}d$. Let $T_{3}$ be the first time
$s_{K}(A_{\theta(t)})=\frac{1}{2}a_{K}d$. At time $T_{2}$, we have
$s_{K}\geq\sigma^{2}wd^{\delta}$ for some $\delta>0$ that depends only on
$a_{1},\ldots,a_{K}$. Then we have

$\|\hat{C}_{1}-C_{1}\|_{op}\leq O(s_{K}d^{-\delta}),$ $d^{2}\frac{dx}{dt}\leq-
cdx+O(dx^{\frac{3}{2}})+O(\sqrt{x}d^{1-\delta})+O(\sqrt{dx}).$

We conclude that $x(T_{3})=O(d^{-2\delta})$. The dynamics of $s_{i}$ is given
by

$d^{2}\frac{ds_{i}}{dt}=2(a_{i}d(1+O(d^{-2\delta}))-s_{i})\sqrt{s_{i}^{2}+\sigma^{4}w^{2}}(1+O(d^{-2\delta})).$
$s_{i}(a_{i}d-s_{i})\leq d^{2}\frac{ds_{i}}{dt}\leq
2(a_{i}d+Cd^{1-2\delta}-s_{i})(s_{i}+\sigma^{2}w).$

The bounds on $\frac{ds_{i}}{dt}$ implies that

$T_{3}-T_{1}=\frac{d\log\frac{a_{K}d}{\sigma^{2}w}}{a_{K}}+O(d)=\frac{1-\gamma_{\sigma^{2}}-\gamma_{w}}{a_{K}}d\log
d+O(d).$

  * •

Final Stage. Let $t\geq T_{3}$. Since $s_{K}=O(d)$, we have

$\|\hat{C}_{1}-C_{1}\|_{op}\leq
O(s_{K}\frac{\sigma^{2}w}{d})=O(s_{K}d^{\gamma_{\sigma^{2}}+\gamma_{w}-1});$
$d^{2}\frac{dx}{dt}\leq-
cdx+O(dx^{\frac{3}{2}})+O(\sqrt{x}d^{\gamma_{\sigma^{2}}+\gamma_{w}})+O(\sqrt{dx}).$

Recall that the constant $c$ in term $-cdx$ must satisfy

$c\leq\frac{1}{d}\text{min}_{k,j:a_{k}\neq
a_{j}}(\frac{(s_{k}^{*}-s_{j}^{*})(s_{k}+s_{j})(\tilde{s}_{k}+\tilde{s}_{j})}{s_{k}^{2}-s_{j}^{2}},\frac{(s_{k}-s_{j})(s_{k}^{*}+s_{j}^{*})(\tilde{s}_{k}+\tilde{s}_{j})}{s_{k}^{2}-s_{j}^{2}})$

As a result, we can take

$\displaystyle c=c(a_{1},\ldots,a_{K})=\frac{\text{min}_{k,j:a_{k}\neq
a_{j}}\absolutevalue{a_{k}-a_{j}}a_{K}^{2}}{\text{max}_{k,j:a_{k}\neq
a_{j}}\absolutevalue{a_{k}^{2}-a_{j}^{2}}}$

Let $c^{\prime}$ be a large constant such that if
$x>c^{\prime}(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})$, then
$\frac{dx}{dt}\leq-\frac{c(a_{1},\ldots,a_{K})d}{2}x$. Let $T_{4}$ be the
first time when $x\leq
c^{\prime}(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})$ after $T_{3}$.
Then
$T_{4}-T_{3}\leq-\frac{2d}{c(a_{1},\ldots,a_{K})}\log(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})+O(d)$.
Moreover, for every $t>T_{4}$, $x$ cannot be larger than
$c^{\prime}(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})$ because
$\frac{dx}{dt}<0$ if
$x=c^{\prime}(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})$. After $T_{4}$,
the dynamics of each singular value is given by

$d^{2}\frac{ds_{i}}{dt}=2((1+O(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)}))a_{i}d-s_{i})s_{i}(1+O(d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)}))$

Let $T_{5}$ be the first time after $T_{4}$ when
$\absolutevalue{s_{i}-a_{i}d}\leq
O(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})\log d$. Then

$T_{5}-T_{4}\leq-\frac{\text{max}(-1,2(\gamma_{\sigma^{2}}+\gamma_{w}-1))}{2a_{K}}d\log
d+O(d\log\log d).$

We conclude that at time

$\displaystyle T^{*}=$
$\displaystyle\left(\frac{1-\gamma_{\sigma^{2}}-\gamma_{w}}{a_{K}}+\frac{2\text{max}(1,2(-\gamma_{\sigma^{2}}-\gamma_{w}+1))}{c(a_{1},\ldots,a_{K})}+\frac{\text{max}(1,2(-\gamma_{\sigma^{2}}-\gamma_{w}+1))}{2a_{K}}\right)d\log
d$ $\displaystyle+O(d\log\log d)$

we have

$x(T^{*})\leq O(d^{\text{max}(-1,2(\gamma_{\sigma^{2}}+\gamma_{w}-1))}),$
$A_{\theta(T^{*})}\in
P_{2}(C,\text{max}\\{1,2(1-\gamma_{\sigma^{2}}-\gamma_{w})\\},$

and

$\absolutevalue{s_{i}(T^{*})-a_{i}d}\leq
O(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})\log d$

###  E.4 Analysis of Testing Error

In section E.3 we proved that $A_{\theta(t)}$ is almost aligned with $A^{*}$
throughout the training. With the alignment in hand, we are ready to give a
time to stop training and the testing error at the end of training.

######  Theorem E.5.

Assume that $A_{\theta(t)}$ follows the gradient flow dynamics. At time

$T^{*}=\left(\frac{1-\gamma_{\sigma^{2}}-\gamma_{w}}{a_{K}}+\frac{2\text{max}(1,2(-\gamma_{\sigma^{2}}-\gamma_{w}+1))}{c(a_{1},\ldots,a_{K})}+\frac{\text{max}(1,2(-\gamma_{\sigma^{2}}-\gamma_{w}+1))}{2a_{K}}\right)d\log
d+O(d\log\log d),$

the testing error

$\|A_{\theta(T^{*})}-A^{*}\|_{F}^{2}\leq
O(\sigma^{4}wd^{2})+O(\sigma^{4}w^{2}\log^{2}d)+O(d^{\frac{3}{2}})+O(\sigma^{2}wd)$

###### Proof.

Let $P_{K}$ be the projection to the largest $K$ singular values. Then

$\|A_{\theta(t)}-A^{*}\|_{F}^{2}\leq\|P_{K}A_{\theta(t)}-A^{*}\|_{F}^{2}+\|(I-P_{K})A_{\theta(t)}\|^{2}_{F}.$

Let $s_{1},\ldots,s_{d}$ be the singular values of $A_{\theta(t)}$. Then
$\|(I-P_{K})A_{\theta(t)}\|_{F}^{2}=\sum_{p\geq K+1}s_{p}^{2}$. The derivative
of $s_{1},\ldots,s_{d}$ reads

$\displaystyle d^{2}\frac{dS}{dt}=$ $\displaystyle
I\odot\left(U^{T}(A^{*}+E)V\sqrt{S^{2}+\sigma^{4}w^{2}}+\sqrt{S^{2}+\sigma^{4}w^{2}}U^{T}(A^{*}+E)V-2S\sqrt{S^{2}+\sigma^{4}w^{2}}\right)$
$\displaystyle+I\odot\left(U^{T}(A^{*}+E)VO(\sigma^{2}w_{1})+O(\sigma^{2}w_{1})U^{T}(A^{*}+E)V-2SO(\sigma^{2}w_{1})\right)$

If $p\geq K+1$, then $s_{p}\leq C^{\prime}\sigma^{2}w$ for some constant
$C^{\prime}$, and

$\displaystyle
d^{2}\absolutevalue{\frac{ds_{p}}{dt}}\leq\absolutevalue{(U^{T}A^{*}V)_{pp}+(U^{T}EV)_{pp}}O(\sigma^{2}w).$
$\displaystyle d^{2}\frac{d}{dt}\sum_{p\geq K+1}s_{p}^{2}\leq$
$\displaystyle\sum_{p}s_{p}\left(\sum_{q\leq
K}U_{qp}V_{qp}A^{*}(qq)+(U^{T}A^{*}V)_{pp}\right)O(\sigma^{2}w)$
$\displaystyle\leq$ $\displaystyle\sum_{p}s_{p}\left(\sum_{q\leq
K}\absolutevalue{U_{qp}}\absolutevalue{V_{qp}}O(d)+O(\sqrt{d})\right)O(\sigma^{2}w)$
$\displaystyle\leq$ $\displaystyle\sum_{q\leq
K}\left(\sum_{p}s_{p}^{2}\right)^{\frac{1}{2}}\left(\sum_{p}\absolutevalue{U_{qp}}^{2}\right)^{\frac{1}{2}}\text{max}_{p}\absolutevalue{V_{qp}}O(\sigma^{2}wd)+\left(\sum_{p}s_{p}^{2}\right)^{\frac{1}{2}}O(\sigma^{2}wd)$
$\displaystyle\leq$
$\displaystyle\left(\sum_{p}s_{p}^{2}\right)^{\frac{1}{2}}\left(O(d^{1-\gamma}\sigma^{2}w)+O(\sigma^{2}wd)\right)$

We conclude that

$d^{2}\frac{d}{dt}\|(I-P_{K})A_{\theta(t)}\|_{F}\leq O(\sigma^{2}wd),$

which implies that

$\|(I-P_{K})A_{\theta(t)}\|_{F}\leq\sqrt{d}O(\sigma^{2}\sqrt{wd})+O(\sigma^{2}wd\frac{\log
d}{d})=O(\sigma^{2}\sqrt{w}d+\sigma^{2}w\log d)$
$\|(I-P_{K})A_{\theta(t)}\|_{F}^{2}\leq
O(\sigma^{4}wd^{2}+\sigma^{4}w^{2}\log^{2}d)$

Next we estimate $\|P_{K}A_{\theta(t)}-A^{*}\|_{F}^{2}$. Notice that
$\|P_{K}A_{\theta(t)}-A^{*}\|_{F}^{2}=\sum_{i,j}\|P_{K}A_{\theta(t)}(i,j)-A^{*}(i,j)\|_{F}^{2}$.
If $i\neq j$, then $A^{*}(i,j)=0$, and

$\displaystyle\|P_{K}A_{\theta(t)}(i,j)\|_{F}^{2}\leq$
$\displaystyle\sum_{k:\text{signal},k\neq
i}\|U(i,k)\|_{F}^{2}\|S(k,k)V(j,k)^{T}\|_{F}^{2}$
$\displaystyle+\|U(i,j)S(j,j)\|_{F}^{2}+\|V(j,k)^{T}\|_{F}^{2}$
$\displaystyle\leq$ $\displaystyle O(xd^{2})$

If $i=j=m+1$, we also have $A^{*}(m+1,m+1)=0$, and

$\displaystyle\|P_{K}A_{\theta(t)}(m+1,m+1)\|_{F}^{2}\leq\sum_{k:\text{signal}}O(d^{2})\|U(m+1,k)\|_{F}^{2}+\|V(m+1,k)\|_{F}^{2}\leq
O(d^{2}x)$

Now assume that $i=j$ are both signals. Then

$\displaystyle\|P_{K}A_{\theta(t)}(i,i)-A^{*}(i,i)\|_{F}^{2}\leq$
$\displaystyle\|U(i,i)(S(i,i)-S^{*}(i,i))V^{T}(i,i)\|_{F}^{2}+O(d^{2})\|U(i,i)V(i,i)^{T}-I\|_{F}^{2}$
$\displaystyle\leq$ $\displaystyle
O(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})\log d+O(d^{2}\sqrt{x})$

Since there are only finitely many blocks in total, we conclude that at time
$T^{*}$,

$\displaystyle\|A_{\theta(T^{*})}-A^{*}\|_{F}^{2}\leq
O(\sigma^{4}wd^{2}+\sigma^{4}w^{2}\log^{2}d+d^{2}O(d^{\text{max}(-\frac{1}{2},(\gamma_{\sigma^{2}}+\gamma_{w}-1))}))+O(d^{-1}+d^{2(\gamma_{\sigma^{2}}+\gamma_{w}-1)})\log
d$ $\|A_{\theta(T^{*})}-A^{*}\|_{F}^{2}\leq
O(\sigma^{4}wd^{2})+O(\sigma^{4}w^{2}\log^{2}d)+O(d^{\frac{3}{2}})+O(\sigma^{2}wd)$
(37)

∎

##  Appendix F Gradient Descent Dynamics and Proof of Theorem 2

###  F.1 Gradient Descent vs Gradient Flow

To study the dynamics of $A_{t}$ under gradient flow, we show that if the
learning rate is small enough, then the gradient flow dynamics will be close
to the gradient descent dynamics.

######  Lemma F.1.

Assume that $A$ is a matrix (not necessarily square matrix). $F$ is a
function: $\mathbb{R}^{\dim A}\to\mathbb{R}^{\dim A}$. The norm $\|\cdot\|$
satisfies $\|AB\|\leq\|A\|\|B\|$ for all $A$ and $B$. In particular, operator
norm and Frobenius norm satisfies this property. Assume that
$\sup_{A}\|F(A)\|\leq C_{0}$ and $\|\nabla F\|\leq C_{1}$ for some constant.
Consider gradient flow dynamics and gradient descent dynamics.

$\text{Gradient Flow: }\frac{dA_{f}}{dt}=F(A_{f})$ $\text{Gradient Descent:
}\frac{A_{d}((k+1)\eta)-A_{d}(k\eta)}{\eta}=F(A_{d}(k\eta))$

Assume that $A_{f}(0)=A_{d}(0)$. Then

$\|A_{f}-A_{d}\|(k\eta)\leq((1+\eta C_{1})^{k-1}-1)\frac{1}{2}\eta C_{0}$

###### Proof.

Notice that we have

$A_{f}((k+1)\eta)-A_{f}(k\eta)=\int_{0}^{\eta}F(A_{f}(k\eta+t))dt$
$(A_{f}-A_{d})((k+1)\eta)-((A_{f}-A_{d})(k\eta))=\int_{0}^{\eta}F(A_{f}(k\eta+t))-F(A_{d}(k\eta))dt$
$\displaystyle\int_{0}^{\eta}F(A_{f}(k\eta+t))-F(A_{d}(k\eta))dt=$
$\displaystyle\int_{0}^{\eta}F(A_{f}(k\eta+t))-F(A_{f}(k\eta))dt$
$\displaystyle+\eta(F(A_{f}(k\eta))-F(A_{d}(k\eta)))$

Let $G(t)=\int_{0}^{t}F(A_{f}(k\eta+s))ds$. Then

$\displaystyle\int_{0}^{\eta}F(A_{f}(k\eta+t))-F(A_{f}(k\eta))dt=$
$\displaystyle G(\eta)-G(0)-\eta
G^{\prime}(0)-\frac{1}{2}\eta^{2}G^{\prime\prime}(\xi)$ $\displaystyle=$
$\displaystyle\frac{1}{2}\eta^{2}\frac{d}{dt}|_{t=\xi}F(A_{f}(k\eta+t))$
$\displaystyle=$ $\displaystyle\frac{1}{2}\eta^{2}\frac{\partial F}{\partial
A}(A_{f}(k\eta+\xi))\frac{d}{dt}|_{t=\xi}A_{f}(k\eta+t)$ $\displaystyle=$
$\displaystyle\frac{1}{2}\eta^{2}\frac{\partial F}{\partial
A}(A_{f}(k\eta+\xi))F(A_{f}(k\eta+\xi))$
$\|\int_{0}^{\eta}F(A_{f}(k\eta+t))-F(A_{f}(k\eta))dt\|\leq\frac{1}{2}\eta^{2}C_{0}C_{1}$
$\eta\|(F(A_{f}(k\eta))-F(A_{d}(k\eta)))\|\leq\eta\|\nabla
F\|\|A_{f}(k\eta)-A_{d}(k\eta)\|\leq\eta C_{1}\|A_{f}(k\eta)-A_{d}(k\eta)\|$

We conclude that

$\displaystyle\|A_{f}((k+1)\eta)-A_{d}((k+1)\eta)\|\leq$ $\displaystyle(1+\eta
C_{1})\|A_{f}(k\eta)-A_{d}(k\eta)\|+\frac{1}{2}\eta^{2}C_{0}C_{1}$
$\displaystyle\|A_{f}-A_{d}\|((k+1)\eta)+\frac{1}{2}\eta C_{0}\leq$
$\displaystyle(1+\eta C_{1})(\|A_{f}-A_{d}\|(k\eta)+\frac{1}{2}\eta C_{0})$
$\displaystyle\leq$ $\displaystyle(1+\eta C_{1})^{k}\frac{1}{2}\eta C_{0}$
$\|A_{f}-A_{d}\|(k\eta)\leq((1+\eta C_{1})^{k-1}-1)\frac{1}{2}\eta C_{0}$

∎

To apply the lemma above we need to prove that $\|A_{\theta(t)}\|_{F}^{2}\leq
O(d^{2})$ throughout the training.

######  Lemma F.2.

For both lazy and active regime, we always have $\|A_{\theta(t)}\|_{F}^{2}\leq
10(\|A^{*}\|_{F}^{2}+\|E\|_{F}^{2})$ throughout the training.

###### Proof.

The gradient descent dynamics of $A_{\theta(t)}$ is given by

$A_{\theta(t+1)}-A_{\theta(t)}=\eta
d^{-2}(A^{*}+E-A_{\theta(t)})C_{1}+C_{2}(A^{*}+E-A_{\theta(t)}).$
$\displaystyle
tr((A^{*}+E-A_{\theta(t+1)})^{T}(A^{*}+E-A_{\theta(t+1)})-(A^{*}+E-A_{\theta(t)})^{T}(A^{*}+E-A_{\theta(t)}))$
$\displaystyle=$
$\displaystyle-2tr((A^{*}+E-A_{\theta(t)}))^{T}(A_{\theta(t+1)}-A_{\theta(t)}))+\|A_{\theta(t+1)}-A_{\theta(t)}\|_{F}^{2}$
$\displaystyle\leq$ $\displaystyle-2\eta
d^{-2}tr((A^{*}+E-A_{\theta(t)})^{T}(C_{1}+C_{2})(A^{*}+E-A_{\theta(t)}))+\eta^{2}O(d^{-2}\|C_{1}\|_{op}^{2})$

From theorem 2, we know that $C_{1}+C_{2}\geq\frac{1}{3}\sigma^{2}wI$.
Therefore
$tr((A^{*}+E-A_{\theta(t)})^{T}(C_{1}+C_{2})(A^{*}+E-A_{\theta(t)}))\geq\|A^{*}+E-A_{\theta(t)}\|_{F}^{2}\frac{1}{3}\sigma^{2}w\geq
cd^{2}\sigma^{2}w$ for some constant $c$. Therefore for the lazy regime, we
always have
$\|A^{*}+E-A_{\theta(t)}\|_{F}^{2}(t+1)\leq\|A^{*}+E-A_{\theta(t)}\|_{F}^{2}(t)$
if $5(\|A^{*}\|_{F}^{2}+\|E\|_{F}^{2})\leq\|A_{\theta(t)}\|_{F}^{2}\leq
7(\|A^{*}\|_{F}^{2}+\|E\|_{F}^{2})$, which implies that
$\|A_{\theta(t)}\|_{F}^{2}\leq 10(\|A^{*}\|_{F}^{2}+\|E\|_{F}^{2})$ for all
time. For the active regime, we have

$\|A^{*}+E-A_{\theta(t)}\|_{F}^{2}(t+1)-\|A^{*}+E-A_{\theta(t)}\|_{F}^{2}(t)\leq\eta^{2}O(1).$

Since the training has at most $O(\frac{T^{*}}{\eta})$ steps, we see that
$\forall t$, $\|A^{*}+E-A_{\theta(t)}\|_{F}^{2}\leq
2(\|A^{*}\|_{F}^{2}+\|E\|_{F}^{2})+O(\eta T^{*})\leq
10(\|A^{*}\|_{F}^{2}+\|E\|_{F}^{2})$. ∎

###### Proof of main theorem.

We first consider the active regime. It suffices to consider the error from
considering gradient descent, rather than gradient flow. To apply lemma F.1,
it is more convenient to consider the dynamics for $W_{1}$ and $W_{2}$. By
lemma G.2, $\|W_{1}^{T}W_{1}\|_{F}^{2}+\|W_{2}W_{2}^{T}\|_{F}^{2}\leq
O(d^{2})$ and $\|W_{1}\|_{op}+\|W_{2}\|_{op}\leq O(\sqrt{d})$. The gradient
flow dynamics of $[W_{1},W_{2}^{T}]$ is given by

$\frac{d}{dt}[W_{1},W_{2}^{T}]=d^{-2}[W_{2}^{T}(A^{*}-W_{2}W_{1}),W_{1}(A^{*}-W_{2}W_{1})^{T}].$

Let
$F([W_{1},W_{2}^{T}])=d^{-2}[W_{2}^{T}(A^{*}-W_{2}W_{1}),W_{1}(A^{*}-W_{2}W_{1})^{T}].$
Then $\sup_{t}\|F(W_{1},W_{2}^{T})\|_{F}\leq O(d^{-\frac{1}{2}})$. Computing
the differential of $F$, we obtain that

$\displaystyle dF([W_{1},W_{2}^{T}])=d^{-2}[$ $\displaystyle-
dW_{2}^{T}(A^{*}-W_{2}W_{1})+W_{2}^{T}(-dW_{2}W_{1}-W_{2}dW_{1}),$
$\displaystyle-
dW_{1}(A^{*}-W_{2}W_{1})^{T}+W_{1}(-dW_{1}^{T}W_{2}-W_{1}^{T}dW_{2})]$

and therefore $\|\nabla F\|_{F}\leq O(d^{-1})$. In the active regime, the
total number of training steps is $\eta^{-1}O(d\log d)$ By lemma F.1,

$\|W_{1}^{flow}-W_{1}^{descent}\|_{F}(T^{*})\leq O(\eta
d^{-\frac{1}{2}})((1+\eta O(d^{-1}))^{\frac{O(d\log d)}{\eta}}-1)=O(\eta
d^{-\frac{1}{2}}\log d)$

and the same holds true for $W_{2}$. We conclude that

$\displaystyle\|W_{2}^{flow}W_{1}^{flow}-W_{2}^{descent}W_{1}^{descent}\|_{F}$
$\displaystyle\leq$
$\displaystyle(\|W_{2}^{descent}\|_{op}+\|W_{1}^{descent}\|_{op})O(\eta
d^{-\frac{1}{2}}\log d)$ $\displaystyle=$ $\displaystyle O(\eta\log d)$
$\|A_{\theta}^{flow}-A_{\theta}^{descent}\|_{F}^{2}\leq O(\eta^{2}\log^{2}d).$

In the lazy regime, from strong bound we have

$W_{1}^{T}W_{1}=(1+O(\sqrt{\frac{d}{w}}))\sqrt{\sigma^{4}w^{2}I+A^{T}A}.$

Therefore $\|W_{1}\|_{op}\leq O(\sigma\sqrt{w})$,
$\sup_{t}\|F([W_{1},W_{2}^{T}])\|_{F}\leq
d^{-2}(\|W_{1}\|_{op}+\|W_{2}\|_{op})O(d)=O(d^{-1}\sigma\sqrt{w})$. Similarly,
$\|\nabla F\|_{F}\leq
d^{-2}O(\|A^{*}-W_{2}W_{1}\|_{F}+\|W_{2}\|_{F}\|W_{1}\|_{F})=d^{-2}O(\sigma^{2}w)$.
By lemma F.1, at time $\frac{100d^{2}\log d}{\sigma^{2}w}$,

$\|W_{1}^{flow}-W_{1}^{descent}\|_{F}(\frac{100d^{2}\log d}{\sigma^{2}w})\leq
O(\eta d^{-1}\sigma\sqrt{w})((1+\eta O(d^{-2}\sigma^{2}w))^{\frac{100d^{2}\log
d}{\sigma^{2}w\eta}}-1)=O(\eta d^{-1}\sigma\sqrt{w}\log d).$

We conclude that

$\displaystyle\|W_{2}^{flow}W_{1}^{flow}-W_{2}^{descent}W_{1}^{descent}\|_{F}$
$\displaystyle\leq$
$\displaystyle(\|W_{2}^{descent}\|_{op}+\|W_{1}^{descent}\|_{op})O(\eta
d^{-1}\sigma\sqrt{w}\log d)$ $\displaystyle=$ $\displaystyle O(\eta
d^{-1}\sigma^{2}w\log d).$
$\|W_{2}^{flow}W_{1}^{flow}-W_{2}^{descent}W_{1}^{descent}\|_{F}\leq
O(\eta^{2}d^{-2}\sigma^{2}w\log^{2}d).$

Recall that in lemma D.1 we proved that

$\|A_{\theta(t)}^{flow}-B_{t}\|_{F}^{2}=o(d^{2}),$

and at time $100\frac{d^{2}\log d}{\sigma^{2}w}$,
$\|B_{t}-A^{*}-E\|_{F}^{2}\leq d^{-50}$, which implies that
$\|A_{\theta}^{descent}-A^{*}-E\|_{F}^{2}\leq o(d^{2})$. Before this time, we
have
$\|B_{t}-A^{*}\|_{F}^{2}\geq\frac{1}{3}\text{min}(\|A^{*}\|_{F}^{2},\|E\|_{F}^{2}).$