SlowGuess/NeurIPS_2023_OCR huggingface dataset:
for some  $a \in \mathbb{R}_{>0}$ , where  $\theta^1$  is a stationary point of the loss. Despite the nonlinearity of the dynamics, the approximation can be proved using Gronwall's inequality since  $\bar{t}_1 - \underline{t}_1 = O(1)$  is a constant time-scale.
To summarize, we have argued that the network approximately reaches stationary point that is 1-sparse, where only the weights at coordinate  $i_0$  are nonzero.
Without loss of generality, we can ensure that at initialization  $\pmb{u}(0)$  and  $\pmb{u}(0) + \pmb{v}(0)$  are nonnegative. This implies  $\pmb{u}(t)$  is nonnegative. The fact that  $u_{i_0}$  and  $v_{i_0}$  are roughly equal in magnitude but might differ in sign is due to the conservation law (5). See Appendix C.3 for details.
The entries of  $\pmb{u}$  and  $\pmb{v}$  are close in magnitude (but may differ in sign) because of the conservation law (5).
4.1.2 Later stages
We inductively extend the argument to any number of stages  $k$ , where each stage has a  $\Theta (\log (1 / \alpha))$ -time plateau, and then a  $O(1)$ -time nonlinear evolution, with the sparsity of the weights increasing by at most one. The argument to analyze multiple stages is analogous, but we must also keep track of the magnitude of the weights on the logarithmic scale, since these determine how much longer. Inductively on  $k$ , suppose that there is some  $T_{k}\in \mathbb{R},\pmb{b}^{k}\in \mathbb{R}^{p}$  and  $\pmb{\theta}^{k}\in \mathbb{R}^{2p}$  and a time  $\bar{t}_k\approx T_k\cdot \log (1 / \alpha)$  such that
$$
\log_ {\alpha} (\boldsymbol {w} (\bar {t} _ {k})) \approx \boldsymbol {b} ^ {k} \text {a n d} \boldsymbol {\theta} (\bar {t} _ {k}) \approx \boldsymbol {\theta} ^ {k},
$$
where  $\pmb{\theta}^k$  is a stationary point of the loss. Our inductive step shows that there is  $T_{k+1} \in \mathbb{R}$  such that during times  $t \in (\bar{t}_k, T_{k+1} \cdot \log(1 / \alpha) - \Omega(1))$  the weights remain close to the stationary point from the previous stage, i.e.,  $\pmb{\theta}(t) \approx \pmb{\theta}^k$ . And at a time  $\bar{t}_{k+1} \approx T_{k+1} \cdot \log(1 / \alpha)$  we have
$$
\log_ {\alpha} (\boldsymbol {w} (\bar {t} _ {k + 1})) \approx \boldsymbol {b} ^ {k + 1} \text {a n d} \boldsymbol {\theta} (\bar {t} _ {k + 1}) \approx \boldsymbol {\theta} ^ {k + 1},
$$
where  $\theta^{k + 1}$  and  $b^{k + 1}$  are defined below (summarized in Algorithm 1). Most notably,  $\theta^{k + 1}$  is a stationary point of the loss whose support grows by at most one compared to  $\theta^k$ .
Stage  $(k + 1)\mathbf{A}$ , loss plateau for time  $\Theta (\log (1 / \alpha))$  At the beginning of stage  $k + 1$ , the weights are close to the stationary point  $\theta^k$ , and so, similarly to stage 1A, linear dynamics are valid.
$$
\boldsymbol {w} (t) \approx \boldsymbol {w} \left(\bar {t} _ {k}\right) \odot e ^ {\boldsymbol {g} \left(\boldsymbol {\theta} ^ {k}\right) \left(t - \bar {t} _ {k}\right)}. \tag {9}
$$
Using the conservation law (7), we derive a "time until active" for each coordinate  $i \in [p]$ , which corresponds to the time for the weight at that coordinate to grow from  $o_{\alpha}(1)$  to  $\Theta(1)$  magnitude:
$$
\Delta_ {k} (i) = \left\{ \begin{array}{l l} \left(b _ {i} ^ {k} - 1 + \operatorname {s g n} \left(g _ {i} \left(\boldsymbol {\theta} ^ {k}\right)\right)\right) / g _ {i} \left(\boldsymbol {\theta} ^ {k}\right), & \text {i f} g _ {i} \left(\boldsymbol {\theta} ^ {k}\right) \neq 0 \\ \infty , & \text {i f} g _ {i} \left(\boldsymbol {\theta} ^ {k}\right) = 0 \end{array} \right. \tag {10}
$$
The linear dynamics approximation (9) breaks down at a time  $t \approx T_{k+1} \cdot \log(1 / \alpha)$ , where
$$
T _ {k + 1} = T _ {k} + \Delta_ {k} \left(i _ {k}\right), \quad i _ {k} = \underset {i \in [ p ]} {\arg \min } \Delta_ {k} (i), \tag {11}
$$
which corresponds to the first time at the weights at a coordinate grow from  $o_{\alpha}(1)$  to  $\Theta(1)$  magnitude. And at times  $t \approx T_{k+1} \cdot \log(1 / \alpha)$ , on the logarithmic scale  $w$  is given by
$$
\log_ {\alpha} (\boldsymbol {w} (t)) \approx \boldsymbol {b} ^ {k + 1} := \boldsymbol {b} ^ {k} - \boldsymbol {g} \left(\boldsymbol {\theta} ^ {k}\right) \Delta_ {k} \left(i _ {k}\right), \tag {12}
$$
Stage  $(k + 1)\mathbf{B}$  of dynamics: nonlinear dynamics for time  $O(1)$  Subsequently, the weights evolve nonlinearly during  $O(1)$  time. In a similar way to the analysis of Stage 1B, we show that at a time  $\bar{t}_{k + 1} = \underline{t}_{k + 1} + O(1)\approx T_{k + 1}\cdot \log (1 / \alpha)$ , we have
$$
\boldsymbol {\theta} \left(\bar {t} _ {k + 1}\right) \approx \boldsymbol {\theta} ^ {k + 1} := \lim  _ {\epsilon \rightarrow 0} \lim  _ {t \rightarrow \infty} \psi^ {k} (t, \epsilon), \tag {13}
$$
where the dynamics  $\psi^k (t,\epsilon)\in \mathbb{R}^{2p}$  are initialized at  $\psi^k (0,\epsilon) = \pmb {\theta}^k +(\epsilon \pmb{e}_{i_k},\mathrm{sgn}(g_i(\pmb{\theta}^k))\epsilon \pmb{e}_{i_k})$  and evolve according to the gradient flow  $\frac{d\psi^k(t,\epsilon)}{dt} = -\nabla_\theta \mathcal{L}(\psi^k)$ . This concludes the inductive step.
4.2 Assumptions for incremental dynamics
To make this intuition rigorous, we formalize below the assumptions required for Theorem 4.1. In Figure 3 and Appendix A, we provide experiments validating these assumptions on the toy model.
The first assumption is that the dynamics are non-degenerate, in the sense that two coordinates do not have weights that grow from  $o_{\alpha}(1)$  to  $\Theta(1)$  size at the same rescaled time. We also place a technical condition to handle the corner case when a coordinate leaves the support of the current stage's stationary point.
Assumption 4.3 (Nondegeneracy of dynamics in part (A)). The initialization satisfies  $|u_i(0)| \neq |v_i(0)|$  for all  $i$ . For stage  $k$ , either  $T_{k+1} = \infty$  or there is a unique minimizer  $i_k$  to  $\min_i \Delta_k(i_k)$  in (11). Finally, for all  $i \in \operatorname{supp}(\pmb{\theta}^{k-1}) \setminus \operatorname{supp}(\pmb{\theta}^k)$  we have  $g_i(\pmb{\theta}^k) \neq 0$ .
Figure 3: Validation of assumptions on the toy model of learning a single attention head. (a) Assumption 4.4: weights perturbed at a random time during training (solid lines) tend back to the near-stationary point (dashed lines). (b) Assumption 4.5: weights perturbed at the beginning of a stage (solid lines) have same nonlinear evolution as without perturbation (dashed lines). Details of these experiments and further validations are provided in Appendix A.
Next, we require that very small perturbations of the coordinates outside of  $\mathrm{supp}(\theta^k)$  do not change the dynamics. For this, it suffices that  $\theta^k$  be a strict local minimum.
Assumption 4.4 (Stationary points are strict local minima). For stage  $k$ , there exist  $\delta_k > 0$  and  $c_k > 0$  such that for  $\tilde{\pmb{u}} \in B(\pmb{u}^k, \delta)$  supported on  $\mathrm{supp}(\pmb{u}^k)$ , we have
$$
\mathcal {L} (\tilde {\boldsymbol {u}}, \boldsymbol {s} ^ {k} \odot \tilde {\boldsymbol {u}}) \geq c _ {k} \| \boldsymbol {u} ^ {k} - \tilde {\boldsymbol {u}} \| ^ {2}
$$
Finally, we require a robust version of the assumption that the limit (13) exists, asking for convergence to a neighborhood of  $\theta^{k + 1}$  even when the initialization is slightly noisy.
Assumption 4.5 (Noise-robustness of dynamics in part (B)). For any stage  $k$  with  $T_{k + 1} < \infty$  and any  $\epsilon >0$ , there are  $\delta >0$  and  $\tau :\mathbb{R}_{>0}\to \mathbb{R}$  such that the following holds. For any  $\tilde{\pmb{u}}\in B(\pmb {u}^k,\delta)\cap \mathbb{R}_{\geq 0}^p$  supported on  $\mathrm{supp}(\tilde{\pmb{u}})\subseteq \mathrm{supp}(\pmb {u}^k)\cup \{i_k\}$ , there exists a unique solution  $\psi :[0,\infty)\to \mathbb{R}^p$  of the gradient flow  $\frac{d\psi}{dt} = -\nabla_{\theta}\mathcal{L}(\psi)$  initialized at  $\psi (0) = (\tilde{u},s^{k + 1}\odot \tilde{u})$ , and at times  $t\geq \tau (\tilde{\psi}_{i_k})$
$$
\left\| \psi (t) - \boldsymbol {\theta} ^ {k + 1} \right\| <   \epsilon .
$$
We run experiments that go beyond the toy diagonal attention head model (see Figures 2 and 3) to test the extent to which low-rank incremental learning occurs in popular models used in practice. We conduct experiments with vision transformers (ViT)  $\left[\mathrm{DBK}^{+}20\right]$  trained on the CIFAR-10/100 and ImageNet datasets, and with the GPT-2 language transformer  $\left[\mathrm{BMR}^{+}20\right]$  trained on the Wikitext-103 dataset. Full experiments are deferred to Appendix B.
Gradual rank increase in vision and language models We train practical transformer architectures on vision and language tasks using Adam and the cross-entropy loss. We train all layers (including the feedforward layers). To capture the low-rank bias with a non-vanishing initialization scale, we study the spectrum of the difference  $\Delta W_{K}W_{Q}^{\top}$  and  $\Delta W_{V}W_{O}^{\top}$  between the weights post-training and their initial values. Specifically, in Figure 4, we plot the stable rank of the differences  $\Delta W_{K}W_{Q}^{\top}$  and  $\Delta W_{V}W_{O}^{\top}$ . The weight perturbation learned during the training process gradually increases in stable rank during training, and is ultimately low-rank when compared to the initial spectrum. Finally, for CIFAR-10, we plot the spectrum of  $\Delta W_{K}W_{Q}^{\top}$  against that of its initialized state in Figure 5 for different self-attention heads, illustrating that the weight perturbation learned during the training process is extremely low-rank when compared to the initial spectrum. In Appendix B, we also study optimization with SGD, which shows similar gradual rank increase behavior.
Effect of initialization scale We probe the effect of initialization scale on gradual rank increase dynamics for a ViT trained on CIFAR-10. We use a ViT of depth 6, with 8 self-attention heads per layer (with layer normalization). We use an embedding and MLP dimension of  $d_{\mathrm{emb}} = 512$ , and a head dimension of  $d_h = 128$  (i.e  $W_K, W_Q, W_V, W_O \in \mathbb{R}^{d_{\mathrm{emb}} \times d_h}$ ). We train the transformer
(a) ViT, CIFAR-10
(b) ViT, CIFAR-100
(c) ViT, ImageNet
(d) GPT-2, Wikitext-103
Figure 4: Stable rank of  $\Delta W_{K}W_{Q}^{\top}$  (blue) and  $\Delta W_{V}W_{O}^{\top}$  (orange) on an arbitrary chosen layer throughout training for four different pairs of networks and tasks. The stable rank of a matrix  $W$  is defined as  $\| W\| _F^2 /\| W\| _2^2$ , and gives a smooth approximation of the rank. Mean and standard deviation (shaded area) are computed across all heads in each attention layer. Full details and results are in Appendix B.
Figure 5: Spectrum of the weight perturbation  $\Delta W_{K}W_{Q}^{\top}$  vs. initialization in a vision transformer trained on CIFAR-10, using Adam and default initialization scale, in random self-attention heads in different layers. The learned perturbation exhibits extreme low-rank bias post-training even in default initialization scales. Analogous plots for CIFAR-100 and ImageNet are in Appendix B.
using Adam with the cross-entropy loss. We train all layers (including the feedforward layers) while varying the initialization scale of all layers by multiplying their initial values by a scale factor (we fix the scale of the initial token mapper). Figure 6 shows the evolution of the principal components of  $\Delta W_{K}W_{Q}^{\top}$  and  $\Delta W_{V}W_{O}^{\top}$  for a randomly-chosen self-attention head and layer throughout training, exhibiting incremental learning dynamics and a low-rank bias. Note that incremental learning and low-rank bias are increasingly evident with smaller initialization scales, as further demonstrated in Figure 7.
We have identified incremental learning dynamics in transformers, proved them rigorously in a simplified setting, and shown them experimentally in networks trained with practical hyperparameters.
Limitations There are clear limitations to our theory: the diagonal weights and small initialization assumptions. More subtly, the theory does not apply to losses with exponential-like tails because the weights may not converge to a finite value and so Assumption 4.4 is not met (this could possibly be addressed by adding regularization). Also, the architecture must be smooth, which precludes ReLUs – but allows for smoothed ReLUs such as the GeLUs used in ViT  $\left[\mathrm{DBK}^{+}20\right]$ . Finally, the theory is for training with gradient flow, while other optimizers such as Adam are used in practice instead [KB14]. Nevertheless, our experiments on ViTs indicate that the incremental learning dynamics occur even when training with Adam.
Figure 6: Training a vision transformer on CIFAR-10 using Adam, while varying the initialization scale (unit scale indicates default initialization). Plotted are the evolution of the eigenvalues of  $\Delta W_{K}W_{Q}^{\top}$  (a) - (c) and  $\Delta W_{V}W_{O}^{\top}$  (d) - (f) in a random self-attention head in the second layer throughout training. Incremental learning dynamics and a low-rank bias are evident for all scales, albeit more pronounced at smaller initialization scales.
Figure 7: Stable rank of  $\Delta W_{K}W_{Q}^{\top}$  per initialization scale (Unit scale refers to the default initialization) in different self-attention heads post-training, at layers 1, 3, 5. At each layer, the stable rank mean and standard deviation are computed across 8 heads per layer, for each initialization scale. All models were trained on CIFAR-10 using the Adam optimizer. Smaller initialization scales lead to lower-rank attention heads.
Future directions An interesting avenue of future research is to develop a theoretical understanding of the implicit bias in function space of transformers whose weights are a low-rank perturbation of randomly initialized weights. Another promising direction is to examine the connection between our results on incremental dynamics and the LoRA method  $\mathrm{[HSW^{+}21]}$ , with the goal of explaining and improving on this algorithm; see also the discussion in Section 1.1. Along this vein, a concurrent work  $\mathrm{[ZZC^{+}23]}$  independently observes gradual rank increase dynamics during transformer training and this inspires a low-rank training algorithm that obtains runtime and memory improvements over regular training. The results of  $\mathrm{[ZZC^{+}23]}$  are complementary to ours, since they study the feedforward layers of the transformer, and their theory applies to linear networks in the standard initialization scale; in contrast, we study the attention layers, and our theory applies to nonlinear networks with small initialization scale.
We would like to thank Vimal Thilak for his help in setting up the infrastructure for conducting experiments, and the anonymous reviewers for their helpful feedback.
[ABM22] 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, COLT, 2022.
[RC20] Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. Advances in neural information processing systems, 33:21174-21187, 2020.
[SKZ+23] James B Simon, Maksis Knutins, Liu Ziyin, Daniel Geisz, Abraham J Fetterman, and Joshua Albrecht. On the stepwise nature of self-supervised learning. arXiv preprint arXiv:2303.15438, 2023.
[SMG13] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
[SS21] Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. In Neural Information Processing Systems, 2021.
[TGZ+23] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Alpaca: A strong, replicable instruction-following model. Stanford Center for Research on Foundation Models. https://crfm.stanford.edu/2023/03/13/alpaca.html, 2023.
[TVS23] Nadav Timor, Gal Vardi, and Ohad Shamir. Implicit regularization towards rank minimization in relu networks. In International Conference on Algorithmic Learning Theory, pages 1429-1459. PMLR, 2023.
$\left[\mathrm{VSP}^{+}17\right]$  Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. ArXiv, abs/1706.03762, 2017.
[WGL+19] Blake E. Woodworth, Suriya Gunasekar, J. Lee, Edward Moroshko, Pedro H. P. Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. ArXiv, abs/2002.09277, 2019.
[YW23] Hao Yu and Jianxin Wu. Compressing transformers: Features are low-rank, but weights are not! 2023.
$\left[\mathrm{ZZC}^{+}23\right]$  Jiawei Zhao, Yifei Zhang, Beidi Chen, Florian Schäfer, and Anima Anandkumar. Inrank: Incremental low-rank learning. arXiv preprint arXiv:2306.11250, 2023.
1.1 Related work 2
1.2 Paper organization 3
3 Neural networks with diagonal weights 4
4 Incremental learning in networks with diagonal weights 4
4.1 Intuition for incremental learning dynamics 5
4.1.1 Stage 1 of dynamics 6
4.1.2 Later stages 7
4.2 Assumptions for incremental dynamics 7
5 Experimental results 8
6 Discussion 9
A Experimental validation of the assumptions in Theorem 4.1 15
B Further experiments on vision and language transformers 20
B.1 SGD-trained transformers 20
B.2 Adam-trained transformers 21
C Proof for dynamics of networks with diagonal parametrization (Theorem 4.1) 25
C.1 Assumptions 25
C.2 Rescaling time for notational convenience 25
C.3 Simplifying problem without loss of generality 26
C.4 Tracking the sum of the weights 26
C.5 Claimed invariants in proof of Theorem C.4 26
C.6 Dynamics from time  $\bar{t}_k$  to time  $\underline{t}_{k + 1}$  (Linear dynamics for  $O(\log (1 / \alpha))$  unrescaled time) 27
C.6.1 Analysis in case where  $T_{k + 1} <   \infty$  27
C.6.2 Analysis in case where  $T_{k + 1} = \infty$  30
C.7 Dynamics from time  $\underline{t}_k$  to time  $\overline{t}_k$  (Nonlinear evolution for  $O(1)$  unrescaled time). 30
C.8 Concluding the proof of Theorem C.4 32
D Technical lemmas 32
D.1 Relating the sum of the weights to the original weights using the conservation law. 32
D.2 Sign of gradients on coordinates that leave support 33
D.3 Local lipschitzness and smoothness 33
A Experimental validation of the assumptions in Theorem 4.1
In Figures 2, 8, and 9, we plot the evolution of the losses, of the entries of  $\mathbf{W}_K\mathbf{W}_Q^\top = \mathrm{diag}(\boldsymbol{w}_K)\mathrm{diag}(\boldsymbol{w}_Q)$ , and of the entries of  $\mathbf{W}_V\mathbf{W}_O^\top = \mathrm{diag}(\boldsymbol{w}_V)\mathrm{diag}(\boldsymbol{w}_O)$  in the toy task of training an attention head (1) with diagonal weights. The model is trained with SGD on the mean-squared error loss on 1000 random samples  $(\mathbf{X},\mathbf{y})$ . Each random sample has  $\mathbf{X}\in \mathbb{R}^{10\times 50}$ , which a sequence of 10 tokens, each of dimension 50, which is distributed as isotropic Gaussian. The label  $\mathbf{y}$  is given by a randomly-generated teacher model that is also an attention head (1) with diagonal weights. In Figures 2, 8, and 9, for  $\alpha \in \{0.1,0.01,0.0001,10^{-8},10^{-16},10^{-32}\}$  we plot the evolution of the loss and of the weights when initialized at  $\theta (0) = \alpha \theta_0$ , for some random Gaussian  $\theta_0$ . Qualitatively, as  $\alpha \to 0$  we observe that the loss curve and the trajectories of the weights appear to converge to a limiting stagewise dynamics, where there are plateaus followed by movement on short time-scales, as predicted by the theory.
Validation of Assumption 4.3 (non-degeneracy of dynamics) As  $\alpha \to 0$ , notice that the stages appear to separate and happen at distinct times. Furthermore, the extra technical condition on coordinates  $i \in \operatorname{supp}(\theta^k) \setminus \operatorname{supp}(\theta^{k-1})$  in Assumption 4.3 is satisfied since no coordinates ever leave the support of  $\theta^k$ .
Validation of Assumption 4.4 (stationary points are strict local minima) In Figure 10 we consider the  $\alpha = 10^{-32}$  trajectory, since this is closest to the dynamics in the  $\alpha \rightarrow 0$  limit. We randomly select several epochs. Since the transitions between stages are a vanishing fraction of the total training time, each of these randomly-selected epochs is likely during a plateau, as we see in the figure. For each epoch perform the following experiment. For each nonnegligible coordinate of the weights (those where the weight is of magnitude greater than the threshold  $\tau = 10^{-5}$ ), we perturb the weights by adding noise of standard deviation 0.05. We then run the training dynamics starting at this perturbed initialization for 1000 epochs. We observe that the training dynamics quickly converge to the original unperturbed initialization, indicating that the weights were close to a strict local minimum of the loss.
Validation of Assumption 4.5 (noise-robustness of dynamics) In Figure 11 we perform the same experiment as in Figure 10, except that the epochs we select to perturb the weights are those where there is a newly-nonnegligible coordinate (less than  $10^{-5}$  in magnitude in the previous epoch, and more than  $10^{-5}$  in magnitude in this epoch). We find that the nonlinear dynamics are robust and tend to the limiting endpoint even under a random Gaussian perturbation of standard deviation  $10^{-2}$  on each of the nonnegligible coordinates, supporting Assumption 4.5.
Figure 8: Evolution of  $\mathrm{diag}(\boldsymbol{w}_Q)\mathrm{diag}(\boldsymbol{w}_K)$  entries over rescaled time initializing at various scalings  $\alpha$ . Notice that as  $\alpha \to 0$ , the training trajectories tend to a limiting trajectory. Each line corresponds to a diagonal entry of  $\mathrm{diag}(\boldsymbol{w}_Q)\mathrm{diag}(\boldsymbol{w}_K)$ .
Figure 9: Evolution of  $\mathrm{diag}(\boldsymbol{w}_V)\mathrm{diag}(\boldsymbol{w}_O)$  entries in the toy task of learning an attention head with diagonal weights. Each line corresponds to the evolution of an entry of  $\mathrm{diag}(\boldsymbol{w}_V)\mathrm{diag}(\boldsymbol{w}_O)$  over rescaled time. Each plot corresponds to a different initialization magnitude  $\alpha$ . Notice that as  $\alpha \to 0$ , the training trajectories tend to a limiting trajectory.
Entries of  $W_{V}W_{O}^{T}$ , perturb time / log(1/alpha) = 373.4
Entries of  $W_{K}W_{O}^{T}$ , perturb time / log(1/alpha) = 373.4
Entries of  $W_{V}W_{O}^{T}$ , perturb time / log(1/alpha) = 541.2
Entries of  $W_{K}W_{O}^{T}$ , perturb time / log(1/alpha) = 541.2
Entries of  $W_{V}W_{O}^{T}$ , perturb time /  $\log(1/\text{alpha}) = 778.0$
Entries of  $W_{K}W_{O}^{T}$ , perturb time /  $\log(1/\text{alpha}) = 778.0$
Entries of  $W_{V}W_{O}^{T}$ , perturb time /  $\log(1/\text{alpha}) = 862.4$
Entries of  $W_{K}W_{O}^{T}$ , perturb time /  $\log(1/\text{alpha}) = 862.4$
Entries of  $W_{V}W_{O}^{T}$ , perturb time / log(1/alpha) = 942.2
Figure 10: Evolution of weights of toy attention model under perturbation, validating Assumption 4.4. At 5 different random times during training, we perturb the nonnegligible weight coordinates and continue to train with SGD. The evolution of each of the weights under the initial perturbation (solid line) is compared to the original evolution without perturbation (dashed line). Observe that the training dynamics quickly brings each weight back to the unperturbed weight trajectory, indicating that the weights are originally close to a strict local minimum.
Entries of  $W_{K}W_{Q}^{T}$ , perturb time / log(1/alpha) = 942.2
Figure 11: Validating Assumption 4.5 with the same experiment as in Figure 10, except that the epochs for the perturbation chosen are those where there is a newly nonnegligible coordinate. Perturbed dynamics (solid lines) are again robust to perturbation and track the original dynamics (dashed lines).
B Further experiments on vision and language transformers