Training Dynamics

When all $W_l$ are diagonal, the product $W=W_L\cdots W_1$ is also diagonal with entries $(W{)}_{\alpha \alpha }={\prod }_{l=1}^L{w}_l^{(\alpha )}$. The teacher $M$ is diagonal with entries $s_{\alpha }$. Then:

Since the different modes $\alpha$ share no parameters, the loss decomposes into $d$ independent scalar problems. $\square$

For $L=2$, the gradient equations are:

Case 1: $s-w_1{w}_2=0$, i.e. $w_1{w}_2=s$. This is the global minimum manifold (a hyperbola $w_2=s/w_1$) with $\ell =0$.

Case 2: If $s-w_1{w}_2\ne 0$, then the first equation requires $w_2=0$ and the second requires $w_1=0$. (For instance, if $w_2\ne 0$ and $w_1=0$, the first equation gives $-s{w}_2=0$, which contradicts $s>0$ and $w_2\ne 0$.) So the only other critical point is $w_1=w_2=0$.

The Hessian at $(0,0)$. We need the second derivatives of $\ell =\frac{1}{2}(s-w_1{w}_2{)}^2$:

At $(0,0)$:

The eigenvalues are $\pm s$. Since $s>0$, $H$ has one positive and one negative eigenvalue. The origin is a strict saddle point. $\square$

At the critical point $W=P_{S}M$ with $S=\{1,\dots ,r\}$, the student retains only the first $r$ modes: $W=\text{diag}(s_1,\dots ,s_r,0,\dots ,0)$. The loss is:

This is strictly positive whenever $r<d$ (since all $s_{\alpha }>0$), so these are not global minima. By the Hessian analysis (extending Exercise 2.2), the direction corresponding to "switching on" a missing mode $\alpha \notin S$ is a descent direction, making these critical points saddle points. All critical points with $|S|=d$ have $W=M$ and $\mathcal{L}=0$: these are the global minima. Therefore there are no spurious local minima — every local minimum is global. $\square$

Under the group action:

All $g_l$ and $g_l^{-1}$ factors cancel telescopically. This means every parameter point in the orbit ${\mathcal{O}}_{\theta }=G{L}_h\cdot \theta$ maps to the same student $W$. $\square$

For the dimension of the set of global minima, this invariance means the following. The global minima form the fiber ${\mu }^{-1}(M)$, which contains the entire orbit $G{L}_h\cdot {\theta }^{\ast }$ for any global minimizer ${\theta }^{\ast }$. The orbit has dimension ${\sum }_{l=1}^{L-1}{d}_l^2$ (the dimension of $G{L}_h$), so the set of global minima is a continuous manifold of very high dimension — far from being isolated points. $\square$

We have $\mathcal{L}=\frac{1}{2}\Vert M-W_2{W}_1{\Vert }_F^2=\frac{1}{2}\text{Tr}\left[(M-W_2{W}_1{)}^{\top }(M-W_2{W}_1)\right]$.

General matrix derivation. Expand:

Differentiating with respect to $W_1$, using $\frac{\partial }{\partial A}\text{Tr}(B^{\top }A)=B$ and $\frac{\partial }{\partial A}\text{Tr}(A^{\top }CA)=(C+C^{\top })A$:

So ${\dot{W}}_1=-{\nabla }_{W_1}\mathcal{L}=W_2^{\top }(M-W_2{W}_1)$.

Similarly, ${\nabla }_{W_2}\mathcal{L}=-(M-W_2{W}_1)W_1^{\top }$, giving ${\dot{W}}_2=(M-W_2{W}_1)W_1^{\top }$.

Diagonal shortcut. For diagonal matrices $W_1=\text{diag}(a_{\alpha })$, $W_2=\text{diag}(b_{\alpha })$, $M=\text{diag}(s_{\alpha })$: the loss decouples as $\mathcal{L}=\frac{1}{2}{\sum }_{\alpha }(s_{\alpha }-b_{\alpha }{a}_{\alpha }{)}^2$. Then ${\dot{a}}_{\alpha }=-\partial \mathcal{L}/\partial a_{\alpha }=b_{\alpha }(s_{\alpha }-b_{\alpha }{a}_{\alpha })$ and ${\dot{b}}_{\alpha }=a_{\alpha }(s_{\alpha }-b_{\alpha }{a}_{\alpha })$, which is the diagonal version of the matrix equations above. $\square$

Let $E:=M-W_2{W}_1$ denote the residual. Compute:

Substitute the gradient flow equations. Using ${\dot{W}}_2=E{W}_1^{\top }$ and ${\dot{W}}_1=W_2^{\top }E$:

Therefore:

The terms cancel pairwise. $G$ is conserved. $\square$

Apply the product rule:

Writing $W=W_2{W}_1$:

Balancedness gives $W_2^{\top }{W}_2=W_1{W}_1^{\top }$, so

Both $W^{\top }W$ and $W_1^{\top }{W}_1$ are positive semidefinite. A positive semidefinite matrix has a unique positive semidefinite square root, so it follows that

Analogous to Exercise 3.4.

Substituting into Exercise 3.3:

With $F=M-W$, the equation from Exercise 3.6 reads:

The general depth-$L$ result $\dot{W}={\sum }_{k=1}^L(W{W}^{\top }{)}^{(L-k)/L}(M-W)(W^{\top }W{)}^{(k-1)/L}$ follows from the same approach applied to the $L$-fold balanced conditions $W_{l+1}^{\top }{W}_{l+1}=W_l{W}_l^{\top }$ for all $l$.

Under the alignment assumption, $W=U\text{diag}(w_{\alpha })V^{\top }$. Then:

Since $w_{\alpha }\ge 0$, the positive matrix square root is:

Similarly:

The residual is immediate:

The NTK operator acts on any matrix $F=U\text{diag}(f_{\alpha })V^{\top }$ as:

So $K[F]$ stays in the SVD basis of $M$ — the NTK is aligned to the task. $\square$

Substitute into $\dot{W}=(W{W}^{\top }{)}^{1/2}(M-W)+(M-W)(W^{\top }W{)}^{1/2}$:

First term:

Second term:

Summing:

Since $W=U\text{diag}(w_{\alpha })V^{\top }$ and the singular vectors are constant by assumption, we read off:

The NTK perspective makes the structure transparent: the factor $2w_{\alpha }$ is the NTK eigenvalue for mode $\alpha$, which amplifies learning in directions that are already strong, while $(s_{\alpha }-w_{\alpha })$ is the residual. $\square$

From the logistic solution $w(t)=\frac{s}{1+(s/w_0-1)e^{-2st}}$, we invert to find $t$ as a function of $w$. The derivation is equivalent to integrating $\dot{w}=2w(s-w)$ by separation of variables:

Using partial fractions $\frac{1}{w(s-w)}=\frac{1}{s}\left(\frac{1}{w}+\frac{1}{s-w}\right)$:

Therefore:

For a fixed ratio $w_f/s_{\alpha }$ and fixed $w_0$, the learning time scales as $t_{\alpha }\propto 1/s_{\alpha }$: modes with larger singular values are learned faster. The network learns features in decreasing order of their strength — a strong separation of timescales.

Contrast with the lazy regime (anticipating Section 5): in the lazy regime the NTK is frozen at $K_0=2w_0\gg 1$, so ${\dot{w}}_{\alpha }\approx 2w_0(s_{\alpha }-w_{\alpha })$. All modes converge at the same exponential rate $2w_0$, independent of $s_{\alpha }$. The rich regime has state-dependent NTK ($2w_{\alpha }$), creating a positive feedback loop — modes that are already large learn even faster — which amplifies the differences between $s_{\alpha }$ into a hierarchy of timescales $t_1\ll t_2\ll \cdots \ll t_r$. $\square$

With uniform initialization $w_{\alpha }(0)=w_0$ for all $\alpha$, the sigmoid solution shows that mode $\alpha$ remains near $w_0$ until time $t\approx \frac{1}{2s_{\alpha }}\ln (s_{\alpha }/w_0)$, then transitions rapidly to $w_{\alpha }\approx s_{\alpha }$.

Early times ($t\sim \frac{1}{2s_1}\ln (s_1/w_0)$): Only mode 1 has grown significantly, so the spectrum is one singular value near $s_1$, with all the others still near $w_0\ll s_r$. The student is approximately $W(t)\approx w_1(t)u_1{v}_1^{\top }$, effectively rank 1.

Intermediate times: As $t$ increases, $w_2$ switches on, then $w_3$, etc. The singular values switch on one at a time, in decreasing order of $s_{\alpha }$, each staying near $w_0$ and then rising to $s_{\alpha }$ over a narrow window. The student is approximately $W(t)\approx {\sum }_{\alpha =1}^{k(t)}{w}_{\alpha }(t)u_{\alpha }{v}_{\alpha }^{\top }$ where $k(t)$ increases stepwise, so the effective rank increases incrementally.

Long times ($t\to \infty$): All modes converge, $w_{\alpha }\to s_{\alpha }$, and $W\to M$.

This is an implicit bias toward low-rank (simple) solutions: at any finite time, the network represents the best rank-$k$ approximation to the teacher. The gradient flow effectively performs a greedy SVD.

Connection to the saddle structure (Exercise 2.3): The saddle points $P_{S}M$ with $|S|=k$ are exactly the best rank-$k$ approximations to $M$, with loss $\frac{1}{2}{\sum }_{\alpha >k}{s}_{\alpha }^2$. The gradient flow trajectory passes near these saddles as it incrementally recruits modes. The plateaus in the loss curve correspond to time spent near saddle points (only $k$ modes active), and the sharp transitions correspond to escaping along the unstable direction that activates the next mode. $\square$

Starting from the same ODE $\dot{w}=2w(s-w)$, with $w_0\gg s>0$, the weight $w$ needs to decrease from $w_0$ to $s$. As long as $w$ has not changed much from $w_0$, we can write $w=w_0+\delta w$ with $|\delta w|\ll w_0$, so:

The ODE becomes:

This is a linear ODE — the NTK factor $2w$ has been frozen at its initial value $2w_0$. Note that only the NTK prefactor is linearized; the residual $(s-w)$ is kept exact since it drives learning.

The linearized ODE $\dot{w}=2w_0(s-w)$ is first-order linear. Set $\tilde{w}:=w-s$, then $\dot{\tilde{w}}=-2w_0\tilde{w}$, with solution $\tilde{w}(t)=(w_0-s)e^{-2w_{0}t}$. Therefore:

This is exponential convergence to $s$ (contrast with the sigmoid of Section 4). The learning timescale is:

It depends on the initialization scale $w_0$ but not on the target $s$. This is the defining feature of the lazy regime: the learning speed is set by the NTK at initialization, not by the structure of the task.

Self-consistency: During the learning phase ($t\lesssim 1/(2w_0)$), the displacement is $|w(t)-w_0|\le |s-w_0|\approx w_0$ (since $w_0\gg s$). The relative change in the NTK is $\Delta (2w)/(2w_0)=(w_0-s)/w_0=1-s/w_0\approx 1$, so the NTK does change significantly in absolute terms. However, the key point is that the dynamics remain well approximated by the linear ODE because the convergence rate is dominated by $2w_0$, which is large and approximately constant throughout training. In the rich regime, by contrast, the NTK changes from $2w_0\approx 0$ to $2s$ — a change of order $s/w_0\to \infty$ relative to the initial value — making the linearization completely invalid. $\square$

Restoring the mode index, each mode satisfies ${\dot{w}}_{\alpha }\approx 2w_0(s_{\alpha }-w_{\alpha })$ with solution:

All modes converge at the same exponential rate $2w_0$, regardless of $s_{\alpha }$. The time for mode $\alpha$ to go from $w_0$ to within $(1-ϵ)$ of $s_{\alpha }$ is:

which is independent of $s_{\alpha }$. All modes reach their targets simultaneously.

Contrast with the rich regime:

In the lazy regime, all components of $M$ are learned simultaneously: signal and noise alike. The network converges directly to the full OLS solution without passing through low-rank intermediates. There is no mechanism to separate signal from noise based on singular value structure. With early stopping, the rich regime recovers a low-rank signal while ignoring noise; the lazy regime cannot. $\square$

The interpolating ODE is ${\dot{w}}_{\alpha }=2\sqrt{w_{\alpha }^2+{\tau }^2}(s_{\alpha }-w_{\alpha })$.

Rich limit ($\tau \to 0$): The square root reduces to $\sqrt{w_{\alpha }^2}=|w_{\alpha }|$, giving:

For $w_{\alpha }>0$ this is $2w_{\alpha }(s_{\alpha }-w_{\alpha })$, exactly the logistic ODE from Exercise 4.2. $\square$

Lazy limit ($\tau \to \infty$, $w_{\alpha }$ bounded): The square root is dominated by $\tau$: $\sqrt{w_{\alpha }^2+{\tau }^2}\approx \tau$, giving:

This is a linear ODE with rate $2\tau$, independent of $s_{\alpha }$ — the frozen-NTK regime of Section 5 (with $\tau$ playing the role of $w_0$). $\square$

At initialization, $w_{\alpha }(0)\sim {\sigma }^2\ll \tau ={\sigma }^2\sqrt{n}$ (since $\sqrt{n}\gg 1$). So initially all modes satisfy $|w_{\alpha }|\ll \tau$.

Early phase (lazy): $\sqrt{w_{\alpha }^2+{\tau }^2}\approx \tau$ for all $\alpha$. Every mode evolves at the same linear rate $2\tau$:

(assuming $w_{\alpha }(0)\approx 0$). The dynamics are approximately linear and the NTK is approximately constant. During this phase, the network aligns with the task: each $w_{\alpha }$ grows toward $s_{\alpha }$ at the same rate. There is no timescale separation.

Late phase (rich): Once some $w_{\alpha }$ grow past $\tau$, the square root transitions to $\sqrt{w_{\alpha }^2+{\tau }^2}\approx |w_{\alpha }|$, and those modes enter the rich regime with the sigmoidal, self-accelerating dynamics ${\dot{w}}_{\alpha }\approx 2w_{\alpha }(s_{\alpha }-w_{\alpha })$. The state-dependent NTK creates timescale separation, and the modes that have crossed the threshold converge rapidly to their targets.

In summary: the network starts in a lazy phase where all modes grow uniformly (alignment), then transitions to a rich phase where modes accelerate individually (incremental learning). $\square$

In the lazy phase, each mode grows as $w_{\alpha }(t)\approx s_{\alpha }(1-e^{-2\tau t})$. At any given time, $w_{\alpha }(t)\propto s_{\alpha }$: modes with larger $s_{\alpha }$ are larger. Since the lazy-to-rich transition occurs when $|w_{\alpha }|\sim \tau$, the time for mode $\alpha$ to cross the threshold satisfies:

For $s_{\alpha }>\tau$, this crossing time exists and is smaller for larger $s_{\alpha }$. For $s_{\alpha }<\tau$, the mode never reaches the threshold and remains permanently lazy.

This means that different modes can be in different regimes simultaneously: large-$s_{\alpha }$ modes have crossed into the rich regime and are converging rapidly, while small-$s_{\alpha }$ modes are still in the lazy phase (or may never leave it). The network interpolates between the two regimes mode by mode.

Connection to grokking: This lazy-to-rich transition provides a mechanism for grokking. During the lazy phase, the network fits the training data via kernel regression with the initial (misaligned) NTK — it memorizes. The test loss plateaus because the frozen kernel cannot capture the task structure. Later, as modes cross the threshold into the rich regime, feature learning kicks in: the NTK rotates to align with the task, and the test loss drops — generalization is achieved. The grokking delay is controlled by the time it takes the relevant modes to cross $\tau$. Increasing the initialization scale $\sigma$ or the width $n$ increases $\tau$, which extends the lazy phase and widens the grokking gap. This is the same mechanism identified by Kumar et al., Grokking as the Transition from Lazy to Rich Training Dynamics (arXiv:2310.06110). $\square$

Setting $\mathbf{j}=0$ (thermal equilibrium) with isotropic noise $\Sigma ={\sigma }^{2}I$:

Rearranging:

The left-hand side is $\nabla \ln p^{\ast }$, so:

Integrating both sides:

This is the Boltzmann distribution with inverse temperature $\beta =2/(\eta {\sigma }^2)$. $\square$

The effective temperature is $T=\eta {\sigma }^2/2$.

Small $\eta$ (low temperature, $\beta \to \infty$): The distribution concentrates sharply on the global minima of $\mathcal{L}$. SGD converges to the lowest-loss minimizer without exploring.

Large $\eta$ (high temperature, $\beta \to 0$): The distribution becomes nearly uniform over parameter space. SGD explores broadly and does not settle into any particular minimum.

Flat vs. sharp minima: At moderate temperature, the Boltzmann distribution assigns more probability mass to broad basins (flat minima) than to narrow ones. This is because a flat minimum occupies a larger volume of parameter space at any given loss level — the width of the basin acts as an entropic contribution. Flat minima tend to generalize better because small perturbations to the parameters (or slight distribution shift in the data) do not dramatically change the loss. Thus the implicit bias of SGD noise toward flat minima is beneficial for generalization. $\square$