SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Enric Boix-Adsera\*1,2 Etai Littwin\*1
Emmanuel Abbe<sup>1,3</sup> Samy Bengio<sup>1</sup> Joshua Susskind<sup>1</sup>
1Apple 2MIT 3EPFL
eboix@mit.edu, emmanuel.abbe@epfl.ch
{elittwin,bengio,jsusskind}@apple.com
We identify incremental learning dynamics in transformers, where the difference between trained and initial weights progressively increases in rank. We rigorously prove this occurs under the simplifying assumptions of diagonal weight matrices and small initialization. Our experiments support the theory and also show that phenomenon can occur in practice without the simplifying assumptions.
The transformer architecture achieves state of the art performance in various domains, yet we still lack a solid theoretical understanding of its training dynamics [VSP+17, DCLT19, LOG+19, DBK+20]. Nevertheless, the theoretical toolbox has matured over the last years and there are promising new approaches. One important line of work examines the role that initialization scale plays on the trajectory taken by gradient descent [JGH18, COB18, GSJW19, MGW+20, JGS+21, SS21, KC22]. When the weights are initialized small, it has been shown for simple networks that an incremental learning behaviour occurs, where functions of increasing complexity are learned in stages. This regime is known to be richer than the large-initialization regime<sup>1</sup>, but the incremental learning dynamics are difficult to analyze, and are so far understood only for extremely simple architectures. Can we apply this analysis to transformers? Namely:
Are there incremental learning dynamics when training a transformer architecture?
An obstacle is that past work on incremental learning has mainly studied linear networks [Ber22, ACHL19, MKAA21, LLL20, WGL $^{+}$ 19, JGS $^{+}$ 21, GSSD19, SKZ $^{+}$ 23, PF23], with one paper studying nonlinear 2-layer fully-connected networks [BPVF22]. In contrast, transformers have nonlinear attention heads that do not fall under previous analyses: given  $X \in \mathbb{R}^{n \times d}$ , an attention head computes
$$
\operatorname {a t t e n t i o n} (\boldsymbol {X}; \boldsymbol {W} _ {K}, \boldsymbol {W} _ {Q}, \boldsymbol {W} _ {V}, \boldsymbol {W} _ {O}) = \operatorname {s m a x} (\boldsymbol {X} \boldsymbol {W} _ {K} \boldsymbol {W} _ {Q} ^ {\top} \boldsymbol {X} ^ {\top}) \boldsymbol {X} \boldsymbol {W} _ {V} \boldsymbol {W} _ {O} ^ {\top} \tag {1}
$$
where  $\mathbf{W}_K, \mathbf{W}_Q, \mathbf{W}_V, \mathbf{W}_O \in \mathbb{R}^{d \times d'}$  are trainable matrices, and the softmax is applied row-wise. A transformer is even more complex, since it is formed by stacking alternating layers of attention heads and feedforward networks, along with residual connections.
Main finding Our main finding is that transformers exhibit incremental learning dynamics, where the difference between the trained and initial weights incrementally increases in rank. Our results have a theoretical component and an experimental component.
<sup>1</sup>In the large-initialization regime, deep learning behaves as a kernel method [JGH18, COB18]. Various separations with kernels are known for smaller initialization: e.g., [GMMM19, ABM22, MKAS21].
Figure 1: For an attention head in ViT trained on (a) CIFAR-10, and (b) ImageNet, we plot the normalized spectra of  $\mathbf{W}_K\mathbf{W}_Q^\top$  at initialization (in red), and of the learned perturbations to  $\mathbf{W}_K\mathbf{W}_Q^\top$  at different iterations (in green).
Theoretical contributions For our theory, we study a simplification of the transformer architecture, where the attention head weights are diagonal matrices: i.e., in each attention head we have  $\mathbf{W}_K = \mathrm{diag}(\mathbf{w}_K)$ , where  $\mathbf{w}_K \in \mathbb{R}^d$  are trainable weights, and similarly for  $\mathbf{W}_Q$ ,  $\mathbf{W}_V$  and  $\mathbf{W}_O$ . We rigorously establish the training dynamics of this architecture under gradient flow when the initialization is small. We prove that dynamics occur in discrete stages: (1) during most of each stage, the loss plateaus because the weights remain close to a saddle point, and (2) at the end, the saddle point is quickly escaped and the rank of the weights increases by at most one.
This theoretical result on transformers follows from a general theorem characterizing the learning dynamics of networks  $f_{\mathbb{N}\mathbb{N}}$  that depend on the product of parameters  $\pmb{u}, \pmb{v} \in \mathbb{R}^p$  as
$$
f _ {\mathrm {N N}} (\boldsymbol {x}; \boldsymbol {u}, \boldsymbol {v}) = h (\boldsymbol {x}; \boldsymbol {u} \odot \boldsymbol {v}), \tag {2}
$$
where  $\pmb{x}$  is the input,  $\odot$  denotes the elementwise product, and  $h$  is a smooth function.
Theorem 1.1 (Informal statement of incremental learning dynamics). Let  $f_{\mathbb{N}\mathbb{N}}$  be a network of the form (2), and suppose that the weights are initialized very small: i.e., the entries of  $\pmb{u}, \pmb{v}$  are initialized on the order  $\Theta(\alpha)$  for some small  $\alpha > 0$ . Then the dynamics of gradient flow training effectively proceeds in discrete stages, each one lasting time  $\Theta(\log(1 / \alpha))$ . In each stage, the number of nonnegligible entries of  $\pmb{u} \odot \pmb{v}$  increases by at most one.
A transformer with diagonal weight matrices falls under this result when we only train the attention head weights. For example, if the transformer has one attention head, then we can take  $\pmb{u} = [\pmb{w}_K, \pmb{w}_V] \in \mathbb{R}^{2d}$  and  $\pmb{v} = [\pmb{w}_Q, \pmb{w}_O] \in \mathbb{R}^{2d}$  to be concatenations of the diagonal entries of the weights of the head; see Example 3.2 for more details and the extension to transformers with many heads. Then, using Theorem 1.1, we see that in each stage either  $\pmb{W}_K\pmb{W}_Q^\top = \mathrm{diag}(\pmb{w}_K)\mathrm{diag}(\pmb{w}_Q)$  or  $\pmb{W}_V\pmb{W}_O^\top = \mathrm{diag}(\pmb{w}_V)\mathrm{diag}(\pmb{w}_O)$  increases in effective rank by at most one.
Experimental contributions In our experiments, we first validate our theoretical results, which require the simplifying assumptions of small initialization and diagonal weight matrices.
Then, we conduct experiments on vision and language transformers in settings closer to practice, without any of the assumptions required by our theoretical analysis. Perhaps surprisingly, we again observe incremental learning dynamics, even though the assumptions of the theory are not met. The difference between trained and initial weights has low rank, and the rank of this difference grows gradually during training; see Figure 1. The incremental nature of the dynamics is easier to see for ImageNet, since for CIFAR-10 the rank of the weight difference does not grow as much.
Relation to LoRA We note an intriguing connection to the LoRA algorithm, where a pretrained base model is cheaply fine-tuned by training a low-rank perturbation of the weights [LFLY18, AZG20,  $\mathrm{HSW}^{+}21$ ]. The method is surprisingly powerful, and recently LoRA has been fundamental to allowing the open-source community to inexpensively fine-tune language models [PA23,  $\mathrm{TGZ}^{+}23$ ]. On the other hand, in our work we observe that the trained weights are a low-rank perturbation of the initial weights due to the training dynamics, without having to apply an explicit rank constraint as in LoRA. This raises an exciting open question for future work: can we explain and improve algorithms like LoRA by better understanding and quantifying the incremental dynamics of large transformers?
2We also remark that Theorem 1.1 is interesting in its own right and may have other applications beyond transformers. It qualitatively recovers the incremental dynamics result of [Ber22, PF23] when specialized to linear diagonal networks, i.e., when  $f_{\mathbb{N}\mathbb{N}}(\pmb {x};\pmb {u},\pmb {v}) = \sum_{i = 1}^{p}u_{i}v_{i}x_{i}$
Low-rank bias in nonlinear networks For 2-layer networks, it is known that low-rank bias in the weights emerges if the target function depends on a low-dimensional subspace of the input [ABM22, ABM23, DLS22, BBSS22, MHPG+22]. The results of [ABM22, ABM23] are especially relevant, since they show that the rank of the weights increases in a sequential manner, determined by the "leap complexity" of the target function, which is reminiscent of our empirical observations on transformers. See also  $\mathrm{[FVB^{+}22}$ , TVS23] for more investigations of low-rank bias in 2-layer networks under different assumptions. For transformers, [YW23] report that empirically the trained weights (using default initialization) are not low-rank. This is consistent with our claim that the difference between initial and trained weights is low-rank, since the initial weights might not be low-rank.
Incremental learning dynamics Several works prove incremental learning behaviour in deep linear networks when the initialization is small. [GBLJ19] has shown that gradient descent dynamics on a 2-layer linear network with  $L_{2}$  loss effectively solve a reduced-rank regression problem with gradually increasing rank. [GSSD19] prove a dynamical depth separation result, allowing for milder assumptions on initialization scale. [ACHL19, MKAA21] show implicit bias towards low rank in deep matrix and tensor factorization. [LLL20] show deep matrix factorization dynamics with small initialization are equivalent to a greedy low-rank learning (GLRL) algorithm. And [JGS+21] independently provides a similar description of the dynamics, but without requiring balanced initialization. Finally, [Ber22, JLL+23, PF23] overcome a technical hurdle from previous analyses by proving incremental learning for the entire training trajectory, rather than just the first stage. In contrast to our result, these prior works apply only to linear networks with certain convex losses, whereas our result applies to nonlinear networks. In order to make our extension to nonlinear networks possible, we must make stronger assumptions on the training trajectory, which we verify hold empirically. As far as we are aware, one other work on incremental learning handles nonlinear networks: [BPVF22] proves that a 2-layer network learns with a two-stage incremental dynamic; but that result needs the stylized assumption that all data points are orthogonal.
Sections 2, 3, and 4 contain theoretical preliminaries, definitions of the models to which our theory applies, and our main theoretical result on incremental dynamics. Section 5 provides experiments which verify and extend the theory. Section 6 discusses limitations and future directions.
We consider training a network  $f_{\mathrm{NN}}(\cdot ;\theta)$  parametrized by a vector of weights  $\theta$ , to minimize a loss
$$
\mathcal {L} (\boldsymbol {\theta}) = \mathbb {E} _ {\boldsymbol {x}, \boldsymbol {y}} [ \ell (\boldsymbol {y}, f _ {\mathrm {N N}} (\boldsymbol {x}; \boldsymbol {\theta})) ],
$$
where the expectation is over samples  $(\pmb{x},\pmb{y})\in \mathbb{R}^{d_x}\times \mathbb{R}^{d_y}$  from a training data distribution, and  $\ell :\mathbb{R}^{d_y}\times \mathbb{R}^{d_{out}}\to \mathbb{R}$ . Consider a solution  $\theta (t)$  to the gradient flow
$$
\boldsymbol {\theta} (0) = \alpha \boldsymbol {\theta} _ {0}, \quad \frac {d \boldsymbol {\theta}}{d t} = - \nabla_ {\boldsymbol {\theta}} \mathcal {L} (\boldsymbol {\theta}) \tag {3}
$$
where  $\alpha > 0$  is a parameter governing the initialization scale, that we will take small. For our theory, we henceforth require the following mild regularity assumption on the loss and data.
Assumption 2.1 (Regularity of data distribution and loss). The function  $\ell(\boldsymbol{y},\boldsymbol{\zeta})$  is continuously twice-differentiable in the arguments  $[\boldsymbol{y},\boldsymbol{\zeta}] \in \mathbb{R}^{d_y + d_{out}}$ . There exists  $C > 0$  such that almost surely the data is bounded by  $\| \boldsymbol{x}\|$ ,  $\| \boldsymbol{y}\| \leq C$ .
The assumption on  $\ell$  is satisfied in typical cases such as the square and the cross-entropy losses. The data boundedness is often satisfied in practice (e.g., if the data is normalized).
We also use the notation  $\operatorname{supp}(\pmb{a}) \coloneqq \{i : a_i \neq 0\}$  to denote the support of a vector  $\pmb{a}$ .
Gradient flow training can be obtained as a limit of SGD or GD training as the learning rate tends to 0 (see, e.g., [Bac20]). It is a popular testbed for studying learning dynamics (see e.g., [SMG13, ACH18, RC20]), since it is generally simpler to analyze than SGD.
3 Neural networks with diagonal weights
Our theory analyzes the training dynamics of networks that depend on products of diagonal weight matrices. We use  $\odot$  to denote elementwise vector product.
Definition 3.1. A network  $f_{\mathbb{N}\mathbb{N}}$  is smooth with diagonal weights  $\pmb{\theta} = (\pmb{u},\pmb{v}) \in \mathbb{R}^{2p}$  if it is of the form
$$
f _ {\mathrm {N N}} (\boldsymbol {x}; \boldsymbol {\theta}) = h (\boldsymbol {x}; \boldsymbol {u} \odot \boldsymbol {v})
$$
where  $h: \mathbb{R}^{d_x} \times \mathbb{R}^p \to \mathbb{R}^{d_{out}}$  is continuously twice-differentiable in its arguments in  $\mathbb{R}^{d_x + p}$ .
The assumption on  $h$  precludes the use of the ReLU function since it is not continuously-differentiable. Otherwise the assumption is fairly mild since any  $h$  can be used to express an architecture of any depth as long as the nonlinearities are twice-differentiable, which includes for example GeLUs (as used in ViT). We describe how to express a transformer with diagonal weights.
Example 3.2 (Transformer with diagonal weights). A transformer with  $L$  layers and  $H$  attention heads on each layer is defined inductively by  $Z_0 = X \in \mathbb{R}^{n \times d}$  and
- (Attention layer)  $\tilde{\pmb{Z}}_{\ell} = \pmb{Z}_{\ell -1} + \sum_{i = 1}^{H}\mathrm{attention}(\pmb{Z}_{\ell -1};\pmb{W}_{K}^{\ell ,i},\pmb{W}_{Q}^{\ell ,i},\pmb{W}_{V}^{\ell ,i},\pmb{W}_{O}^{\ell ,i})$
- (Feedforward layer)  $\mathbf{Z}_{\ell} = \tilde{\mathbf{Z}}_{\ell} + \sigma (\tilde{\mathbf{Z}}_{\ell}\mathbf{W}_{A}^{\ell})(\mathbf{W}_{B}^{\ell})^{\top}$ ,
where  $W_{K}^{\ell,i}, W_{Q}^{\ell,i}, W_{V}^{\ell,i}, W_{O}^{\ell,i} \in \mathbb{R}^{d \times d'}$  are attention parameters, and  $W_{A}^{\ell}, W_{B}^{\ell} \in \mathbb{R}^{d \times d'}$  are the feedforward parameters, and  $\sigma$  is a continuously twice-differentiable activation. Suppose that the attention parameters are diagonal matrices: i.e.,  $W_{K}^{\ell,i} = \mathrm{diag}(w_{K}^{\ell,i}) \in \mathbb{R}^{d \times d}$ , and similarly for the  $W_{Q}^{\ell,i}, W_{V}^{\ell,i}, W_{O}^{\ell,i}$  matrices. Then by the definition of the attention layer (1), the final output of the transformer  $Z_{L}$  only depends on the attention parameters through the elementwise products  $\boldsymbol{w}_{K}^{\ell,i} \odot \boldsymbol{w}_{Q}^{\ell,i}$  and  $\boldsymbol{w}_{V}^{\ell,i} \odot \boldsymbol{w}_{O}^{\ell,i}$ . In other words, we can write
$$
\boldsymbol {Z} _ {L} = h (\boldsymbol {X}; \boldsymbol {u} \odot \boldsymbol {v}),
$$
for vectors  $\pmb{u} = [\pmb{w}_K^{\ell,i},\pmb{w}_V^{\ell,i}]_{(\ell,i)\in [L]\times[H]} \in \mathbb{R}^{2dHL}$  and  $\pmb{v} = [\pmb{w}_Q^{\ell,i},\pmb{w}_O^{\ell,i}]_{(\ell,i)\in [L]\times[H]} \in \mathbb{R}^{2dHL}$ , and some smooth model  $h$ . Thus, if only the attention layers are trained, the diagonal transformer fits under Definition 3.1.
4 Incremental learning in networks with diagonal weights
We prove that if the initialization scale  $\alpha$  is small, then learning proceeds in incremental stages, where in each stage the effective sparsity of the weights increases by at most one. These stages are implicitly defined by Algorithm 1 below, which constructs a sequence of times  $0 = T_0 < T_1 < \dots < T_k < \dots$  and weight vectors  $\pmb{\theta}^0, \pmb{\theta}^1, \dots, \pmb{\theta}^k, \dots \in \mathbb{R}^{2p}$  that define the stages. We prove the following:
Theorem 4.1 (Incremental dynamics at small initialization). Let  $f_{\mathbb{N}\mathbb{N}}$  be a model with diagonal weights as in Definition 3.1. For any stage  $k$  and time  $t \in (T_k, T_{k+1})$  the following holds under Assumptions 2.1, 4.3, 4.4 and 4.5. There is  $\alpha_0(t) > 0$  such that for all  $\alpha < \alpha_0$ , there exists a unique solution  $\theta : [0, t\log(1/\alpha)] \to \mathbb{R}^p$  to the gradient flow (3) and
$$
\lim  _ {\alpha \rightarrow 0} \boldsymbol {\theta} (t \cdot \log (1 / \alpha)) \rightarrow \boldsymbol {\theta} ^ {k},
$$
and at each stage the sparsity increases by at most one:  $\mathrm{supp}(\pmb{\theta}^{k + 1})\setminus \mathrm{supp}(\pmb{\theta}^k)\subseteq \{i_k\}$ .<sup>4</sup>
Application: transformer with diagonal weights Before giving the intuition for this theorem and stating the assumptions formally, let us discuss its application to the diagonal transformer model from Example 3.2. As a corollary of Theorem 4.1, the gradient flow on a diagonal transformer with small initialization will learn in stages, where in each stage there will be at most one head  $i \in [H]$  in one layer  $\ell \in [L]$  such that either the rank of  $\mathbf{W}_K^{\ell,i}(\mathbf{W}_Q^{\ell,i})^\top = \mathrm{diag}(\mathbf{w}_K^{\ell,i})\mathrm{diag}(\mathbf{w}_Q^{\ell,i})$  or the rank of  $\mathbf{W}_V^{\ell,i}(\mathbf{W}_O^{\ell,i})^\top = \mathrm{diag}(\mathbf{w}_V^{\ell,i})\mathrm{diag}(\mathbf{w}_O^{\ell,i})$  increases by at most one. In Figure 2, we illustrate these dynamics in the toy case of a single attention head trained in a student-teacher setup.
4Abusing notation, for  $\pmb {\theta} = (\pmb {u},\pmb {v})\in \mathbb{R}^p\times \mathbb{R}^p$  , we write  $\mathrm{supp}(\pmb {\theta}) = \mathrm{supp}(\pmb {u})\cup \mathrm{supp}(\pmb {v})$
Algorithm 1 Incremental learning in networks with diagonal weights
1:  $\pmb{b}^0, \pmb{\theta}^0 \gets \mathbf{0} \in \mathbb{R}^p, T_0 \gets 0$
2: for stage number  $k = 0,1,2,\ldots$  do
3: # (A) Pick new coordinate  $i_k \in [p]$  to activate.
4: For each  $i$ , define time  $\Delta_k(i)$  until active using (10).
5: Pick winning coordinate  $i_k$  using (11)
6: Calculate time  $T_{k + 1}$  using (11) and break if  $\infty$
7: Update logarithmic weight approximation  $\pmb{b}^{k + 1}$  using (12)
8: (#B) Train activated coordinates to stationarity.
9:  $\pmb{\theta}^{k + 1}\gets$  limiting dynamics point from (13)
(b)
Figure 2: (a) Loss versus rescaled time in the toy task of learning an attention head with diagonal weights, for various initialization scales  $\alpha$ . The loss curves converge as  $\alpha \to 0$  to a curve with stagewise loss plateaus and sharp decreases, as predicted by the theory; some stagewise learning behavior is already clear with  $\alpha = 0.01$ . (b) Each line shows the evolution of one of the entries of  $\mathrm{diag}(\boldsymbol{w}_Q)\mathrm{diag}(\boldsymbol{w}_K)$  and  $\mathrm{diag}(\boldsymbol{w}_V)\mathrm{diag}(\boldsymbol{w}_O)$  over rescaled time, demonstrating that the rank of these matrices increases incrementally; see Appendix A for experimental details and further results.
4.1 Intuition for incremental learning dynamics
We develop an informal intuition for Theorem 4.1 and fill out the definition of Algorithm 1. A model  $f_{\mathbb{N}\mathbb{N}}$  with diagonal weights  $\pmb{\theta} = (\pmb{u},\pmb{v})$  as in Definition 3.1 evolves under the gradient flow (3) as
$$
\frac {d \boldsymbol {u}}{d t} = \boldsymbol {v} \odot \boldsymbol {g} (\boldsymbol {\theta}), \quad \frac {d \boldsymbol {v}}{d t} = \boldsymbol {u} \odot \boldsymbol {g} (\boldsymbol {\theta}) \quad \text {w h e r e} \tag {4}
$$
$$
\boldsymbol {g} (\boldsymbol {\theta}) = - \mathbb {E} _ {\boldsymbol {x}, \boldsymbol {y}} \left[ D \ell (\boldsymbol {y}, h (\boldsymbol {x}; \boldsymbol {u} \odot \boldsymbol {v})) ^ {\top} D h (\boldsymbol {x}; \boldsymbol {u} \odot \boldsymbol {v}) ^ {\top} \right].
$$
Here  $D\ell(\boldsymbol{y}, \cdot) \in \mathbb{R}^{1 \times d_{out}}$  is the derivative of  $\ell$  in the second argument and  $Dh(\boldsymbol{x}, \cdot) \in \mathbb{R}^{d_{out} \times p}$  is the derivative of  $h$  in the second argument. The first key observation is a conservation law that simplifies the dynamics. It can be viewed as the balancedness property for networks with linear activations [ACH18, DHL18], specialized to the case of diagonal layers.
Lemma 4.2 (Conservation law). For any  $i \in [p]$  and any time  $t$ , we have
$$
u _ {i} ^ {2} (t) - v _ {i} ^ {2} (t) = u _ {i} ^ {2} (0) - v _ {i} ^ {2} (0). \tag {5}
$$
Proof. This follows from  $\frac{d}{dt} (u_i^2 - v_i^2) = u_i v_i g_i(\pmb{\theta}) - u_i v_i g_i(\pmb{\theta}) = 0$ .
The conservation law reduces the degrees of freedom and means that we need only keep track of  $p$  parameters in total. Specifically, if we define  $w_{i}(t) \coloneqq u_{i}(t) + v_{i}(t)$ , then the vector  $\boldsymbol{w}(t) = \boldsymbol{u}(t) + \boldsymbol{v}(t)$  evolves by
$$
\frac {d \boldsymbol {w}}{d t} = \boldsymbol {w} \odot \boldsymbol {g} (\boldsymbol {\theta}). \tag {6}
$$
Using the conservation law (5), we can keep track of the weights in terms of the initialization and  $\boldsymbol{w}(t)$ :
$$
\boldsymbol {\theta} (t) = \left(\frac {1}{2} (\boldsymbol {w} (t) + \frac {\boldsymbol {u} ^ {\odot 2} (0) - \boldsymbol {v} ^ {\odot 2} (0)}{\boldsymbol {w} (t)}, \frac {1}{2} (\boldsymbol {w} (t) - \frac {\boldsymbol {u} ^ {\odot 2} (0) - \boldsymbol {v} ^ {\odot 2} (0)}{\boldsymbol {w} (t)})\right) \tag {7}
$$
Therefore it suffices to analyze the dynamics of  $\pmb{w}(t)$ .
4.1.1 Stage 1 of dynamics
Stage 1A of dynamics: loss plateau for time  $\Theta (\log (1 / \alpha))$  At initialization,  $\theta (0)\approx 0$  because the weights are initialized on the order of  $\alpha$  which is small. This motivates the approximation  $g(\theta (t))\approx g(0)$ , under which the dynamics solve to:
$$
\boldsymbol {w} (t) \approx \boldsymbol {w} (0) \odot e ^ {\boldsymbol {g} (\mathbf {0}) t}. \tag {8}
$$
Of course, this approximation is valid only while the weights are still close to the small initialization. The approximation breaks once one of the entries of  $\theta(t)$  reaches constant size. By combining (7) and (8), this happens at time  $t \approx T_1 \cdot \log(1 / \alpha)$  for
$$
T _ {1} = \min  _ {i \in [ p ]} 1 / | g _ {i} (\mathbf {0}) |.
$$
Until this time, the network remains close to its initialization, and so we observe a loss plateau.
Stage 1B of dynamics: nonlinear dynamics for time  $O(1)$ . Subsequently, the loss decreases nonlinearly during a  $O(1)$  time-scale, which is vanishingly short relative to the time-scale of the loss plateau. To prove this, we make the non-degeneracy assumption that there is a unique coordinate  $i_0$  such that  $1 / |g_{i_0}(\mathbf{0})| = T_1$ . Under this assumption, in stage 1A all weights except for those at coordinate  $i_0$  remain vanishingly small, on the order of  $o_\alpha(1)$ . Concretely, for any small  $\epsilon > 0$ , there is a time  $t_1(\epsilon) \approx T_1 \cdot \log(1 / \alpha)$  and sign  $s \in \{+1, -1\}$  such that
$$
u _ {i _ {0}} (\underline {{t}} _ {1}) \approx \epsilon , v _ {i _ {0}} (\underline {{t}} _ {1}) \approx s \epsilon \quad \text {a n d} \quad | u _ {i} (\underline {{t}} _ {1}) |, | v _ {i} (\underline {{t}} _ {1}) | = o _ {\alpha} (1) \text {f o r a l l} i \neq i _ {0}.
$$
Because all coordinates except for  $i_0$  have vanishingly small  $o_{\alpha}(1)$  weights after stage 1A, we may perform the following approximation of the dynamics. Zero out the weights at coordinates except for  $i_0$ , and consider the training dynamics starting at  $\tilde{\pmb{\theta}} = (\epsilon \pmb{e}_{i_0}, s \epsilon \pmb{e}_{i_0})$ . After  $O(1)$  time, we should expect these dynamics to approach a stationary point. Although the evolution is nonlinear, all entries remain zero except for the  $i_0$  entries, so the stationary point is also sparse. Mathematically, there is a time  $\bar{t}_1 = \underline{t}_1 + O(1) \approx T_1 \cdot \log (1 / \alpha)$  such that
$$
\boldsymbol {\theta} (\bar {t} _ {1}) \approx (a \boldsymbol {e} _ {i _ {0}}, s a \boldsymbol {e} _ {i _ {0}}) := \boldsymbol {\theta} ^ {1},
$$