SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The practice of training transformer models often deviates substantially from the assumptions made in our theoretical analysis, and it is a priori unclear to what extent gradual rank increase behaviour and a low rank bias are manifested in setups more common in practical applications. To gauge the relevancy of our findings we conduct experiments on popular vision and language benchmarks, using algorithms and hyperparameters common in the literature. We use the stable rank of a matrix  $\mathbf{W}$  given by  $\frac{\| \mathbf{W} \|_F^2}{\| \mathbf{W} \|_2^2}$  as a smooth approximation of rank. We track the value of the stable rank for the different attention matrices throughout training. Although we do not expect our theoretical results to hold precisely in practice, we find evidence of gradual increase in stable rank, leading to a low rank bias in Figures 12, 13, 15, 17 and 19. In these experiments we use off-the-shelf vision transformers (ViT) [DBK+20] trained on popular vision benchmarks, as well as off-the-shelf GPT-2 trained on a popular language benchmark. We use no weight decay or dropout in our experiments. All models were initialized using the default initialization scale.
B.1 SGD-trained transformers
CIFAR-10/100 We trained a 6-layer ViT with 8 heads per layer, embedding dimension 512, head dimension 128, and MLP dimension 512 and patch-size 4 for 500 epochs on CIFAR10/CIFAR100 with SGD and learning rate 3e-1 and warmup. See Figures 12 and 13. Each run took 2 hours on one A100 GPU.
Figure 12: CIFAR-10, ViT trained with SGD: Stable rank of  $\Delta W_{K}W_{Q}^{\top}$  (blue) and  $\Delta W_{V}W_{O}^{\top}$  (orange) throughout training. Mean and standard deviation (shaded area) are computed across 8 heads per attention layer.
Figure 13: CIFAR-100, ViT trained with SGD: Stable rank of  $\Delta W_{K}W_{Q}^{\top}$  (blue) and  $\Delta W_{V}W_{O}^{\top}$  (orange) throughout training. Mean and standard deviation (shaded area) are computed across 8 heads per attention layer.
B.2 Adam-trained transformers
CIFAR-10/100 For the CIFAR-10/100 datasets we use a VIT with 6 layers, patchsize of 4, 8 heads per self attention layer, an embedding and MLP dimension of 512, and a head dimension of 128. We train the model using the Adam optimizer for 500 epochs with a base learning rate of 1e-4, a cyclic learning rate decay with a linear warmup schedule for 15 epochs and a batchsize of 512. Our results are summarized in Figures 14 and 15 for CIFAR-10, and Figures 16 and 17 for CIFAR-100.
Figure 14: CIFAR-10, ViT trained with Adam: normalized spectrum at different stages of training. (a) - (c) Normalized spectrum of  $\mathbf{W}_K\mathbf{W}_Q^\top$  at initialization and  $\Delta \mathbf{W}_K\mathbf{W}_Q^\top$  during training for different attention heads at different layers. (d) - (e) equivalent figures for  $\mathbf{W}_V\mathbf{W}_O^\top$ .
ImageNet For ImageNet, we use the VIT-Base/16 from  $\mathrm{[DBK^{+}20]}$  trained with Adam for 360 epochs with a base learning rate of 3e-3, a cyclic learning rate decay with a linear warmup schedule for 15 epochs and a batchsize of 4096. Our results are summarized in Figures 18 and 19 for ImageNet.
Figure 15: CIFAR-10, ViT trained with Adam: Stable rank of  $\Delta W_{K}W_{Q}^{\top}$  (blue) and  $\Delta W_{V}W_{O}^{\top}$  (red) throughout training. Mean and standard deviation (shaded area) are computed across 8 heads per attention layer.
Figure 16: CIFAR-100, ViT trained with Adam: normalized spectrum at different stages of training. (a) - (c) Normalized spectrum of  $\mathbf{W}_K\mathbf{W}_Q^\top$  at initialization and  $\Delta \mathbf{W}_K\mathbf{W}_Q^\top$  during training for different attention heads at different layers. (d) - (e) equivalent figures for  $\mathbf{W}_V\mathbf{W}_O^\top$ .
Figure 17: CIFAR-100, ViT trained with Adam: Stable rank of  $\Delta W_{K}W_{Q}^{\top}$  (blue) and  $\Delta W_{V}W_{O}^{\top}$  (red) throughout training. Mean and standard deviation (shaded area) are computed across 8 heads per attention layer.
Figure 18: ImageNet, ViT trained with Adam: normalized spectrum at different stages of training. (a) - (c) Normalized spectrum of  $\mathbf{W}_K\mathbf{W}_Q^\top$  at initialization and  $\Delta \mathbf{W}_K\mathbf{W}_Q^\top$  during training for different attention heads at different layers. (d) - (e) equivalent figures for  $\mathbf{W}_V\mathbf{W}_O^\top$ .
Figure 19: ImageNet, ViT trained with Adam: Stable rank of  $\Delta W_{K}W_{Q}^{\top}$  (blue) and  $\Delta W_{V}W_{O}^{\top}$  (red) throughout training. Mean and standard deviation (shaded area) are computed across 12 heads per attention layer.
Wikitext-103 The gradual rank increase phenomenon also occurs in the NLP setting with language transformers. We trained GPT-2 on Wikitext-103 using the HuggingFace training script with Adam learning rate 3e-4, per-GPU batch-size 8, and block-length 256. We trained for 3 epochs on 2 A100 GPUs, which took 12 hours. See Figure 20.
Figure 20: Wikitext-103, GPT-2 trained with Adam: Stable rank of  $\Delta W_V W_O^\top$  and  $\Delta W_Q W_K^\top$ , versus training iteration. Stable rank of the perturbation increases gradually, but remains small throughout training.
C Proof for dynamics of networks with diagonal parametrization (Theorem 4.1)
C.1 Assumptions
Recall we have defined  $\pmb{\theta}^0, \dots, \pmb{\theta}^k, \dots \in \mathbb{R}^{2p}$  as the sequence of weights such that  $\pmb{\theta}^0 = \mathbf{0}$  and  $\pmb{\theta}^{k+1}$  is defined inductively as follows. Consider the dynamics of  $\psi^k(t,\epsilon) \in \mathbb{R}^{2p}$  initialized at  $\psi^k(0,\epsilon) = \pmb{\theta}^k + (\epsilon \pmb{e}_{i_k}, \operatorname{sgn}(g_i(\pmb{\theta}^k))\epsilon \pmb{e}_{i_k})$  and evolving according to the gradient flow  $\frac{d\psi^k(t,\epsilon)}{dt} = -\nabla_\pmb{\theta} \mathcal{L}(\psi^k)$ . We assume that there is a limiting point  $\pmb{\theta}^{k+1}$  of these dynamics as  $\epsilon$  is taken small and the time is taken large:
$$
\lim _ {\epsilon \to 0} \lim _ {t \to \infty} \psi^ {k} (t, \epsilon) = \theta^ {k + 1}.
$$
Under the above assumption that this sequence  $\pmb{\theta}^0, \dots, \pmb{\theta}^k, \dots$  is well-defined, we can derive a useful property of it for free. Namely, the conservation law (5) implies that  $\pmb{u} \odot \pmb{u} - \pmb{v} \odot \pmb{v}$  is preserved. It follows that for each  $k$  we have that  $\pmb{\theta}^k = (\pmb{u}^k, \pmb{v}^k)$  satisfies  $|\pmb{u}^k| = |\pmb{v}^k|$  entrywise. In other words, there is  $\pmb{s}^k \in \{+1, -1\}^p$  satisfying
$$
\boldsymbol {\theta} ^ {k} = \left(\boldsymbol {u} ^ {k}, \boldsymbol {s} ^ {k} \odot \boldsymbol {u} ^ {k}\right) \in \mathbb {R} ^ {2 p}.
$$
We also abuse notation and write  $\mathrm{supp}(\pmb{\theta}^k) \coloneqq \mathrm{supp}(\pmb{u}^k) \subseteq [p]$ , since the support of  $\pmb{\theta}^k$  on the first  $p$  coordinates matches its support on the last  $p$  coordinates.
Having fixed this notation, we now recall the main assumptions of the theorem.
Assumption C.1 (Nondegeneracy of dynamics in part (A); Assumption 4.3). 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$ .
Assumption C.2 (Stationary points are strict local minima; Assumption 4.4). For stage  $k$ , there exist  $\delta_k > 0$  and  $c_k > 0$  such that for  $\tilde{\boldsymbol{u}} \in B(\boldsymbol{u}^k, \delta)$  supported on  $\mathrm{supp}(\boldsymbol{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}.
$$
Assumption C.3 (Noise-robustness of dynamics in part (B); Assumption 4.5). For stage  $k$ , either  $T_{k+1} = \infty$  or the following holds. For 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{\pmb{u}}, s^{k+1} \odot \tilde{\pmb{u}})$ , and at times  $t \geq \tau(\tilde{u}_{i_k})$ ,
$$
\left\| \boldsymbol {\psi} (t) - \boldsymbol {\theta} ^ {k + 1} \right\| <   \epsilon .
$$
C.2 Rescaling time for notational convenience
For ease of notation, we rescale time
$$
\boldsymbol {u} _ {\alpha} (0) = \alpha \boldsymbol {u} (0), \quad \boldsymbol {v} _ {\alpha} (0) = \alpha \boldsymbol {v} (0)
$$
$$
\frac {d \boldsymbol {u} _ {\alpha}}{d t} = \log (1 / \alpha) \boldsymbol {v} _ {\alpha} \odot \boldsymbol {g} (\boldsymbol {u} _ {\alpha}, \boldsymbol {v} _ {\alpha}), \quad \frac {d \boldsymbol {v} _ {\alpha}}{d t} = \log (1 / \alpha) \boldsymbol {u} _ {\alpha} \odot \boldsymbol {g} (\boldsymbol {u} _ {\alpha}, \boldsymbol {v} _ {\alpha}). \tag {14}
$$
$$
\boldsymbol {\theta} _ {\alpha} (t) = \left(\boldsymbol {u} _ {\alpha} (t), \boldsymbol {v} _ {\alpha} (t)\right) \in \mathbb {R} ^ {2 p}.
$$
Because of this time-rescaling, we equivalently state Theorem 4.1 as:
Theorem C.4 (Restatement of Theorem 4.1). Let  $K \in \mathbb{Z}_{\geq 0}$  be such that Assumptions 4.3 and 4.4 hold for all  $k \leq K$  and Assumption 4.5 holds for all  $k < K$ . Then for any  $k \leq K$  and time  $t \in (T_k, T_{k+1})$  the following holds. There is  $\alpha_0(t) > 0$  such that for all  $\alpha < \alpha_0$ , there exists a unique solution  $\theta_\alpha : [0, t] \to \mathbb{R}^p$  to the gradient flow (14) and
$$
\lim _ {\alpha \to 0} \boldsymbol {\theta} _ {\alpha} (t) \to \boldsymbol {\theta} ^ {k},
$$
where at each stage  $|\operatorname{supp}(\boldsymbol{u}^k) \setminus \operatorname{supp}(\boldsymbol{u}^{k-1})| \leq 1$ .
For shorthand, we also write
$$
S _ {k} = \operatorname {s u p p} \left(\boldsymbol {u} ^ {k}\right) \text {a n d} S _ {k} ^ {c} = [ p ] \backslash \operatorname {s u p p} \left(\boldsymbol {u} ^ {k}\right).
$$
C.3 Simplifying problem without loss of generality
For each coordinate  $i \in [p]$  we have  $|u_{\alpha, i}(0)| \neq |v_{\alpha, i}(0)|$  by the non-degeneracy Assumption 4.3. So we can assume  $|u_{\alpha, i}(0)| > |v_{\alpha, i}(0)|$  without loss of generality. Furthermore, we can assume the entrywise inequality
$$
\boldsymbol {u} _ {\alpha} (0) > 0
$$
by otherwise training weights  $\tilde{\pmb{u}}_{\alpha}(t)$ ,  $\tilde{\pmb{v}}_{\alpha}(t)$  initialized at  $\tilde{\pmb{u}}_{\alpha}(0) = \mathrm{sgn}(\pmb{u}_{\alpha}(0))\pmb{u}_{\alpha}(0)$  and  $\tilde{\pmb{v}}_{\alpha}(0) = \mathrm{sgn}(\pmb{v}_{\alpha}(0))\pmb{v}_{\alpha}(0)$ , as  $\tilde{\pmb{u}}_{\alpha}(t) \odot \tilde{\pmb{v}}_{\alpha}(t) = \pmb{u}_{\alpha}(t) \odot \pmb{v}_{\alpha}(t)$  at all times.
Since  $u_{\alpha,i}^2(t) - v_{\alpha,i}^2(t) = u_{\alpha,i}^2(0) - v_{\alpha,i}^2(0)$  by the conservation law (5), it holds that  $|u_{\alpha,i}(t)| > |v_{\alpha,i}(t)|$  throughout. So by continuity
$$
\boldsymbol {u} _ {\alpha} (t) > 0
$$
throughout training.
C.4 Tracking the sum of the weights
$$
\boldsymbol {w} _ {\alpha} (t) = \boldsymbol {u} _ {\alpha} (t) + \boldsymbol {v} _ {\alpha} (t).
$$
The reason for this definition is that during training we have
$$
\frac {d \boldsymbol {w} _ {\alpha}}{d t} = \log (1 / \alpha) \boldsymbol {w} _ {\alpha} \odot \boldsymbol {g} \left(\boldsymbol {\theta} _ {\alpha}\right), \tag {15}
$$
Notice that since that we have assumed  $u_{\alpha, i}(0) > |v_{\alpha, i}(0)|$  for each  $i \in [p]$  we have  $\boldsymbol{w}_{\alpha}(0) > 0$  entrywise. So, by (15) for all  $t > 0$ ,
$$
\boldsymbol {w} _ {\alpha} (t) > 0.
$$
It suffices to track  $\pmb{w}_{\alpha}(t)$  because we can relate the log-scale magnitude of  $\pmb{w}_{\alpha}(t)$  to the magnitudes of the corresponding coordinates in  $\pmb{u}_{\alpha}(t)$  and  $\pmb{v}_{\alpha}(t)$  - see technical Lemmas D.1 D.2 and D.3.
C.5 Claimed invariants in proof of Theorem C.4
In order to prove Theorem C.4, we consider any gradient flow  $\theta_{\alpha}:[0,T^{*}]\to \mathbb{R}^{p}$  solving (14) where  $T^{*}\in (T_{K},T_{K + 1})$ . For now, we focus only on proving properties of this gradient flow, and defer its existence and uniqueness to Section C.8.
We show the following invariants inductively on the stage  $k$ . For any  $\epsilon > 0$ , any stage  $k \leq K$ , there is  $\alpha_k := \alpha_k(\epsilon) > 0$  such that for all  $\alpha < \alpha_k$  the following holds. There are times  $\bar{t}_k := \bar{t}_k(\alpha, \epsilon)$  and  $t_{k+1} := t_{k+1}(\alpha, \epsilon)$ , such that
$$
\bar {t} _ {k} \in [ T _ {k} - \epsilon , T _ {k} + \epsilon ], \tag {16}
$$
$$
\underline {{t}} _ {k + 1} \in \left\{ \begin{array}{l l} \left[ T _ {k + 1} - \epsilon , T _ {k + 1} + \epsilon \right], & \text {i f} T _ {k + 1} <   \infty \\ \left\{T ^ {*} \right\}, & \text {i f} T _ {k + 1} = \infty \end{array} . \right. \tag {17}
$$
and the weights approximate the greedy limit for all times  $t \in [\bar{t}_k, \underline{t}_{k+1}]$
$$
\left\| \boldsymbol {\theta} _ {\alpha} (t) - \boldsymbol {\theta} ^ {k} \right\| <   \epsilon , \tag {18}
$$
and the weights at times  $\bar{t}_k$  and  $\underline{t}_{k + 1}$  are correctly estimated by the incremental learning dynamics on the logarithmic-scale
$$
\left\| \log_ {\alpha} \left(\boldsymbol {w} _ {\alpha} \left(\bar {t} _ {k}\right)\right) - \boldsymbol {b} ^ {k} \right\| <   \epsilon \tag {19}
$$
and if  $T_{k + 1} < \infty$  then
$$
\left\| \log_ {\alpha} \left(\boldsymbol {w} _ {\alpha} \left(\underline {{t}} _ {k + 1}\right)\right) - \boldsymbol {b} ^ {k + 1} \right\| <   \epsilon . \tag {20}
$$
Base case  $k = 0$ : Take  $\bar{t}_0(\alpha, \epsilon) = 0$ . Then statement (16) holds since  $T_0 = 0$ . Notice that as  $\alpha \to 0$  we have that  $\pmb{u}_{\alpha}(0), \pmb{v}_{\alpha}(0) \to \pmb{0} = \pmb{u}^0$ , and also  $\log_{\alpha} \pmb{w}_{\alpha}(0) \to \pmb{1} = \pmb{b}^0$ . So statement (19) follows if we take  $\alpha_0$  small enough. In Section C.6 we show how to construct time  $\underline{t}_1$  such that (18) and (20) hold.
Inductive step: Suppose that (16), (18), (19) and (20) hold for some iteration  $k < K$ . We prove them for iteration  $k + 1$ . In Section C.7 we construct time  $\bar{t}_k$ . In Section C.6 we construct time  $\underline{t}_{k+1}$ .
C.6 Dynamics from time  $\bar{t}_k$  to time  $\underline{t}_{k + 1}$  (Linear dynamics for  $O(\log (1 / \alpha))$  unrescaled time)
Let  $k \leq K$ , and suppose that we know that for any  $\bar{\epsilon}_k > 0$ , there is  $\bar{\alpha}_k(\bar{\epsilon}_k) > 0$  such that for all  $0 < \alpha < \bar{\alpha}_k$ , there is a time  $\bar{t}_k = \bar{t}_k(\alpha, \bar{\epsilon}_k)$  satisfying
$$
\left| T _ {k} - \bar {t} _ {k} \right| <   \bar {\epsilon} _ {k}
$$
$$
\left\| \boldsymbol {\theta} _ {\alpha} (\bar {t} _ {k}) - \boldsymbol {\theta} ^ {k} \right\| <   \bar {\epsilon} _ {k}
$$
$$
\left\| \log_ {\alpha} \left(\boldsymbol {w} _ {\alpha} (\bar {t} _ {k})\right) - \boldsymbol {b} ^ {k} \right\| <   \bar {\epsilon} _ {k}.
$$
C.6.1 Analysis in case where  $T_{k + 1} < \infty$
Consider first the case where  $T_{k + 1} < \infty$ . We show that, for any  $\underline{\epsilon}_{k + 1} > 0$ , there is  $\rho_{k + 1}(\underline{\epsilon}_{k + 1}) > 0$  such that for all  $0 < \rho < \rho_{k + 1}(\bar{\epsilon}_{k + 1})$  there is  $\underline{\alpha}_{k + 1}(\rho, \underline{\epsilon}_{k + 1}) > 0$  such that for all  $\alpha < \underline{\alpha}_{k + 1}$ , there is a time  $t_{k + 1} = t_{k + 1}(\alpha, \rho, \underline{\epsilon}_{k + 1})$  satisfying
$$
\left| T _ {k + 1} - \underline {{t}} _ {k + 1} \right| <   \underline {{\epsilon}} _ {k + 1} \tag {21}
$$
$$
\left\| \boldsymbol {\theta} _ {\alpha} (t) - \boldsymbol {\theta} ^ {k} \right\| <   \underline {{\epsilon}} _ {k + 1} \text {f o r a l l} t \in [ \bar {t} _ {k}, \underline {{t}} _ {k + 1} ] \tag {22}
$$
$$
\left\| \log_ {\alpha} \left(\boldsymbol {w} _ {\alpha} \left(\underline {{t}} _ {k + 1}\right)\right) - \boldsymbol {b} ^ {k + 1} \right\| <   \underline {{\epsilon}} _ {k + 1} \tag {23}
$$
$$
u _ {\alpha , i _ {k}} \left(\underline {{t}} _ {k + 1}\right) \in [ \rho , 3 \rho ], \tag {24}
$$
$$
\operatorname {s g n} \left(v _ {\alpha , i _ {k}} \left(\underline {{t}} _ {k + 1}\right)\right) = s _ {i _ {k}} ^ {k + 1}. \tag {25}
$$
For any  $\rho, \alpha$ , let  $\bar{\epsilon}_k = \rho_{\underline{\epsilon}_{k+1}} / (4p)$  and choose  $\bar{t}_k = \bar{t}_k(\alpha, \bar{\epsilon}_k)$ . Then define
$$
\begin{array}{l} \underline {{t}} _ {k + 1} = \underline {{t}} _ {k + 1} (\alpha , \rho , \underline {{\epsilon}} _ {k + 1}) \tag {26} \\ = \inf  \left\{t \in [ \bar {t} _ {k}, \infty): \| \boldsymbol {u} _ {\alpha , S _ {k} ^ {c}} (t) - \boldsymbol {u} _ {\alpha , S _ {k} ^ {c}} (\bar {t} _ {k}) \| + \| \boldsymbol {v} _ {\alpha , S _ {k} ^ {c}} (t) - \boldsymbol {v} _ {\alpha , S _ {k} ^ {c}} (\bar {t} _ {k}) \| > 4 \rho \right\}. \\ \end{array}
$$
Now we show that the weights  $\theta_{\alpha}(t)$  cannot move much from time  $\bar{t}_k$  to  $\underline{t}_{k + 1}$ . The argument uses the local Lipschitzness of the loss  $\mathcal{L}$  (from technical Lemma D.7), and the strictness of  $\theta^k$  as a stationary point (from Assumption 4.4).
Lemma C.5 (Stability of active variables during part (A) of dynamics). There is  $\rho_{k + 1}$  small enough and  $\underline{\alpha}_{k + 1}(\rho)$  small enough depending on  $\rho$ , such that for all  $\rho < \rho_{k + 1}$  and  $\alpha < \underline{\alpha}_{k + 1}$  and all  $t \in [\overline{t}_k, \underline{t}_{k + 1})$ ,
$$
\left\| \boldsymbol {\theta} _ {\alpha} (t) - \boldsymbol {\theta} ^ {k} \right\| <   \rho^ {\prime} := \max  \left(2 4 \rho , 1 8 \sqrt {\rho K _ {R _ {k}} / c _ {k}}\right). \tag {27}
$$
where  $c_k$  is the strict-minimum constant from Assumption 4.4 and  $K_{R_k}$  is the Lipschitzness constant from Lemma D.7 for the ball of radius  $R_k = \| \pmb{\theta}^k\| +1$ .
Proof. Assume by contradiction that (27) is violated at some time  $t < \underline{t}_{k+1}$ . Let us choose the first such time
$$
t ^ {*} = \inf  \left\{t \in [ \bar {t} _ {k}, \underline {{t}} _ {k + 1}): \| \boldsymbol {u} _ {\alpha} (t ^ {*}) - \boldsymbol {u} ^ {k} \| + \| \boldsymbol {v} _ {\alpha} (t ^ {*}) - \boldsymbol {s} ^ {k} \odot \boldsymbol {u} ^ {k} \| \geq \rho^ {\prime} \right\}.
$$
Define  $\tilde{\pmb{\theta}} = (\tilde{u},\tilde{v})$  by
$$
\tilde {u} _ {i} = \left\{ \begin{array}{l l} u _ {\alpha , i} (t ^ {*}), & i \in S _ {k} \\ 0, & i \notin S _ {k} \end{array} \right. \quad \text {a n d} \quad \tilde {v} _ {i} = \left\{ \begin{array}{l l} v _ {\alpha , i} (t ^ {*}), & i \in S _ {k} \\ 0, & i \notin S _ {k} \end{array} \right..
$$
By the definition of  $\underline{t}_{k + 1}$ , this satisfies
$$
\begin{array}{l} \left\| \tilde {\boldsymbol {u}} - \boldsymbol {u} _ {\alpha} \left(t ^ {*}\right) \right\| = \left\| \boldsymbol {u} _ {\alpha , S _ {k} ^ {c}} \left(t ^ {*}\right) \right\| \leq 4 \rho + \left\| \boldsymbol {u} _ {\alpha , S _ {k} ^ {c}} \left(\bar {t} _ {k}\right) \right\| \leq 4 \rho + \underline {{\epsilon}} _ {k} <   5 \rho , \\ \left\| \tilde {\boldsymbol {v}} - \boldsymbol {v} _ {\alpha} \left(t ^ {*}\right) \right\| = \left\| \boldsymbol {v} _ {\alpha , S _ {k} ^ {c}} \left(t ^ {*}\right) \right\| \leq 4 \rho + \left\| \boldsymbol {v} _ {\alpha , S _ {k} ^ {c}} \left(\bar {t} _ {k}\right) \right\| \leq 4 \rho + \underline {{\epsilon}} _ {k} <   5 \rho . \\ \end{array}
$$
$$
\| \tilde {\boldsymbol {u}} - \boldsymbol {u} ^ {k} \| + \| \tilde {\boldsymbol {v}} - \boldsymbol {s} ^ {k} \odot \boldsymbol {u} ^ {k} \| = \| \boldsymbol {u} _ {\alpha , S _ {k}} (t ^ {*}) - \boldsymbol {z} _ {S _ {k}} ^ {k} \| + \| \boldsymbol {v} _ {\alpha , S _ {k}} (t ^ {*}) - \boldsymbol {s} _ {S _ {k}} ^ {k} \odot \boldsymbol {z} _ {S _ {k}} ^ {k} \| \geq \rho^ {\prime} - 1 0 \rho \geq \rho^ {\prime} / 2.
$$
Using (a) the strict minimum Assumption 4.4 with constant  $c_k$ , since  $\| \tilde{\pmb{\theta}} - \pmb{\theta}^k \| \leq \rho'$  and we take  $\rho'$  small enough,
$$
\begin{array}{l} \mathcal {L} \left(\boldsymbol {\theta} _ {\alpha} \left(t ^ {*}\right)\right) \geq \mathcal {L} (\tilde {\boldsymbol {\theta}}) - 4 \rho K _ {R _ {k}} \overset {(a)} {\geq} \mathcal {L} \left(\boldsymbol {\theta} ^ {k}\right) - 4 \rho K _ {R _ {k}} + \frac {c _ {k} \left(\rho^ {\prime}\right) ^ {2}}{1 6} \\ \geq \mathcal {L} \left(\boldsymbol {\theta} _ {\alpha} (\bar {t} _ {k})\right) - (4 \rho + \bar {\epsilon} _ {k}) K _ {R _ {k}} + \frac {c _ {k} \left(\rho^ {\prime}\right) ^ {2}}{1 6} > \mathcal {L} \left(\boldsymbol {\theta} _ {\alpha} (\bar {t} _ {k})\right). \\ \end{array}
$$
This is a contradiction because  $\mathcal{L}$  is nondecreasing along the gradient flow.
Lemma C.6 (Log-scale approximation is correct during part (A)). There are functions  $\rho_{k+1}(\underline{\epsilon}_{k+1}) > 0$  and  $\underline{\alpha}_{k+1}(\rho, \underline{\epsilon}_{k+1}) > 0$  such that for all  $\rho < \rho_{k+1}$  and  $\alpha < \underline{\alpha}_{k+1}$ , and for all  $t \in (\bar{t}_k, \underline{t}_{k+1})$  we have for a constant  $C$  depending on  $k$ ,
$$
\left\| \log_ {\alpha} \left(\boldsymbol {w} _ {\alpha} (t)\right) - \boldsymbol {b} ^ {k} + \left(t - \bar {t} _ {k}\right) \boldsymbol {g} \left(\boldsymbol {\theta} ^ {k}\right) \right\| <   \rho \epsilon_ {k + 1} + C \rho^ {\prime} \left(t - \bar {t} _ {k}\right). \tag {28}
$$
Furthermore, for all  $i \in S_k^c$  and  $t \in (\bar{t}_k, \underline{t}_{k+1})$  we have
$$
\operatorname {s g n} \left(g _ {i} \left(\boldsymbol {\theta} _ {\alpha} (t)\right)\right) = \operatorname {s g n} \left(g _ {i} \left(\boldsymbol {\theta} ^ {k}\right)\right). \tag {29}
$$
Proof. By Lemma C.5 and Lemma D.7, there is a constant  $C$  depending on  $\theta^k$  such that for all  $t \in (\bar{t}_k, \underline{t}_{k+1})$ ,
$$
\left\| \boldsymbol {g} \left(\boldsymbol {\theta} _ {\alpha} (t)\right) - \boldsymbol {g} \left(\boldsymbol {\theta} ^ {k}\right) \right\| \leq C \rho^ {\prime}.
$$
For shorthand, write  $\bar{\pmb{g}} (\pmb{\theta}^k) = \pmb {g}(\pmb{\theta}^k) + C\rho '\mathbf{1}$  and  $\underline{\pmb{g}} (\pmb{\theta}^{k}) = \underline{\pmb{g}} (\pmb{\theta}^{k}) - C\rho^{\prime}\mathbf{1}$ . Since  $\boldsymbol {w}_{\alpha}(t) > 0$  entrywise as we have assumed without loss of generality (see Section C.3), we have the following entrywise inequalities
$$
\underline {{\boldsymbol {g}}} \left(\boldsymbol {\theta} ^ {k}\right) \odot \boldsymbol {w} _ {\alpha} (t) <   \boldsymbol {g} \left(\boldsymbol {\theta} _ {\alpha} (t)\right) \odot \boldsymbol {w} _ {\alpha} (t) <   \bar {\boldsymbol {g}} \left(\boldsymbol {\theta} ^ {k}\right) \odot \boldsymbol {w} _ {\alpha} (t). \tag {30}
$$
Since the dynamics are given by  $\frac{dw_{\alpha}}{dt} = \log (1 / \alpha)\pmb {g}(\pmb {w}_{\alpha})\odot \pmb{w}_{\alpha}$
$$
\boldsymbol {w} _ {\alpha} (\bar {t} _ {k}) e ^ {(t - \bar {t} _ {k}) \log (1 / \alpha)} \underline {{\boldsymbol {g}}} (\boldsymbol {\theta} ^ {k}) \leq \boldsymbol {w} _ {\alpha} (t) \leq \boldsymbol {w} _ {\alpha} (\bar {t} _ {k}) e ^ {(t - \bar {t} _ {k}) \log (1 / \alpha)} \bar {\boldsymbol {g}} (\boldsymbol {\theta} ^ {k}).
$$
Taking the logarithms with base  $\alpha \in (0,1)$
$$
\left(t - \bar {t} _ {k}\right) \underline {{\boldsymbol {g}}} \left(\boldsymbol {u} ^ {k}\right) \leq \log_ {\alpha} \left(\boldsymbol {w} _ {\alpha} \left(\bar {t} _ {k}\right)\right) - \log_ {\alpha} \left(\boldsymbol {w} _ {\alpha} (t)\right) \leq \left(t - \bar {t} _ {k}\right) \bar {\boldsymbol {g}} \left(\boldsymbol {u} ^ {k}\right).
$$
The bound (28) follows since  $\| \log_{\alpha}(\pmb{w}_{\alpha}(\bar{t}_k)) - \pmb{b}^k\| < \bar{\epsilon}_k < \rho \underline{\epsilon}_{k + 1}$ .