SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Results in Section 3.2 indicate that diagonal linear networks exhibit different implicit bias in comparison with rank-1 and standard linear networks, e.g., both rank-1 and standard linear networks prefer  $\ell_2$  minimum norm solution for GD when the initialization is nearly-zero while, on the contrary, diagonal linear networks prefer such solution when the initialization is sufficiently large. In this section, we empirically compare the implicit bias for rank-1 linear networks and diagonal linear networks to show this phenomenon.
In particular, we use the same settings as in the experiments for rank-1 linear networks in A.1 while only change the model to diagonal linear networks. As in previous works [3, 22], the reparameterization of diagonal linear network is
$$
\theta = w _ {+} \odot w _ {+} - w _ {-} \odot w _ {-},
$$
where  $w_{+} \in \mathbb{R}^{100}$ ,  $w_{-} \in \mathbb{R}^{100}$  and  $\odot$  is the elementwise product. Let  $\mathbf{e} = (1, \dots, 1)^T \in \mathbb{R}^{100}$ , we set the initialization as
$$
w _ {+} (0) = C \mathbf {e}, \quad w _ {-} (0) = C \mathbf {e},
$$
where  $C$  is a positive constant measuring the initialization scale. For each diagonal linear network with different  $C$ , we run GD for 20000 iterations and calculate  $D(\theta(\infty), \theta_{\ell_2}^*)$  and  $D(\theta(\infty), \theta^*)$ . The results are plotted in Fig. 8, where, for convenience of comparing the implicit bias of GD for rank-1 linear networks with that of diagonal linear networks, we also plot the results of rank-1 linear networks of Fig. 5(a) and Fig. 5(b) in Fig. 8(a) and Fig. 8(b), respectively.
As shown in Fig. 8(a), as the initialization scale  $(\|\theta(0)\|$  for rank-1 linear networks,  $C$  for diagonal linear networks) increases,  $D(\theta(\infty), \theta_{l_2}^*)$  decreases for rank-1 linear linear networks trained with
(b)
Figure 8: For different initialization scale  $(\|\theta(0)\|$  for rank-1 linear networks and  $C$  for diagonal linear networks): (a)  $D(\theta(\infty), \theta_{\ell_2}^*)$  for rank-1 linear nets and diagonal linear networks (the green solid line). (b)  $D(\theta(\infty), \theta^*)$  for rank-1 linear nets and diagonal linear nets (the green solid line).
both GD and SGD, while it increases for diagonal linear networks trained with GD. This indicates the drastic difference between the implicit bias of GD exhibited by diagonal linear networks and rank-1 linear networks (also standard linear networks).
A.3 Additional Experiments for the "Alleviating" Effect in Theorem 3
Recall the form of  $V^{\mathcal{S}}(\theta ,t)$  in Theorem 3
$$
V ^ {\mathcal {S}} (\theta , t) = \frac {1}{\Omega_ {L}} \| \theta \| _ {\Omega_ {L}} ^ {\Omega_ {L}} + \frac {\theta^ {T} \theta (0)}{\| \theta (0) \| ^ {\lambda_ {L}}} - \frac {2 \lambda_ {L} \eta \theta^ {T}}{n b} \int_ {0} ^ {t} \frac {\mathcal {L} (\theta (s)) \operatorname {t r} (P _ {\perp} (\theta (s)) X ^ {T} X)}{\| \theta (s) \| ^ {2 - \lambda_ {L}}} \theta (s) d s,
$$
we let
$$
p _ {\theta} (t) = \frac {\theta^ {T} (t) \theta (0)}{\| \theta (0) \| ^ {\lambda_ {L}}}, \tag {13}
$$
$$
q _ {\theta} (t) = \frac {2 \lambda_ {L} \eta \theta^ {T} (t)}{n b} \int_ {0} ^ {t} \frac {\mathcal {L} (\theta (s)) \operatorname {t r} (P _ {\perp} (\theta (s)) X ^ {T} X)}{\| \theta (s) \| ^ {2 - \lambda_ {L}}} \theta (s) d s. \tag {14}
$$
To quantitatively measure the "alleviating" effect of the SGD sampling noise, we train 3 rank-1 linear networks with different initialization scales using SGD. The batch size is 4 and the learning rate
Figure 9: SGD sampling noise alleviates the dependence on the initialization. Numbers after the comma denote the initialization scales. The solid lines are for  $p_{\theta}$  (Eq. (13)) and the dotted lines are for  $p_{\theta} - q_{\theta}$  (Eq. (14)).
is  $5 \times 10^{-5}$ . For the rank-1 linear network  $f(x; u, v) = w_{L}^{T} W_{L-1} \cdots W_{1} x$  where  $W_{k} = u_{k} v_{k}^{T} \in \mathbb{R}^{d_{k} \times d_{k+1}}$ , we let  $L = 3$  and  $\forall k \in \{1, \dots, L\}: d_{k} = 100$ . We calculate both  $p_{\theta}(t)$  and  $q_{\theta}(t)$  along training, where  $q_{\theta}(t)$  measures the alleviating effect of the SGD sampling noise and their difference  $p_{\theta}(\infty) - q_{\theta}(\infty)$  is the alleviated initialization dependence of the SGD solution compared to GD solution.
As shown in Fig. 9, the effect coming from the SGD sampling noise,  $q_{\theta}$ , is equivalent to make the dependence of  $V^{S}$  on the initialization closer to 0 (after about 1000 iterations, every dotted line is closer to the x-axis compared to the corresponding solid line with the same color), thus it controls the dependence of the SGD solution on the initialization. This phenomenon further verifies our claims.
A.3.1 Training loss for Fig. 4(b)
Fig. 4(b) indicates that the final direction of the integral term in Eq. (12) highly depends on the initialization  $\theta(0)$  since the loss decays along training. To further support this argument, here we present the training loss when we perform the experiments of Fig. 4(b) in Fig. 10. It can be seen that, for a random initialization, the loss, the magnitude of the speed of  $\alpha$ , has a high value at the start of the training, and decays very quickly, which explains why the direction of  $\theta(0)$  is crucial to that of  $\alpha(\infty)$ .
A.4 Additional Experiments for Biased Initialization
In this section, we provide additional experiments to show that our conclusion still holds when removing the balanced initialization condition (Definition 1).
To make the initialization unbalanced, we add a small perturbation to the balanced initialization. Specifically, we define
$$
\Delta = \frac {1}{2 L - 1} \sum_ {k = 1} ^ {L - 1} \frac {\left| \left| v _ {k + 1} \right| ^ {2} - \left| \left| u _ {k} \right| ^ {2} \right| \right.}{\left| \left| u _ {k} \right| ^ {2} \right.}
$$
as the scale of the perturbation to the balanced initialization (larger  $\Delta$  implies that the initialization is more unbalanced). All the other experiment details are kept unchanged as in Section 5. As shown in Fig. 11 and Fig. 12, we still observe similar phenomenons as in the case of the balanced initialization, e.g., SGD solutions are closer to the  $\ell_2$ -norm minimization solution compared to GD, when a small perturbation is added to the balanced initialization. Thus the implicit bias is not unique to the balanced initialization. In particular:
Figure 10: The empirical loss  $\mathcal{L}(\theta)$  along training for Fig. 4(b).
- We report  $D(\theta(\infty), \theta_{\ell_2}^*)$  for both GD and SGD for different levels of perturbation  $\Delta$  (denoted in the title of each figure) in Fig. 11. In the last figure, we fix the initialization scale and report  $D((\theta(\infty), \theta_{\ell_2}^*)$  of both GD and SGD for different  $\Delta$ . It can be seen that, without the balanced initialization, GD and SGD still prefer  $\ell_2$ -norm minimization solution  $\theta_{\ell_2}^*$  for small initialization, while the SGD solution is closer to  $\theta_{\ell_2}^*$  due to its initialization reduction effect.
- We report  $D(\theta(t), \theta_{\ell_2}^*)$  during optimization for both GD and SGD for different levels of perturbation  $\Delta$  and the same scale of initialization  $(\| \theta(0) \| = 0.8696)$  in Fig. 12, which further clearly reveals that there are still similar phenomenons when  $\Delta \neq 0$  as in the case when the initialization is balanced.
B More Discussions
Our work proposes the rank-1 linear network which is a plausible proxy of standard linear networks with some neurons fully connected with neurons in its last and next layers. By showing that the proposed rank-1 linear networks are standard linear networks with special initialization, our conclusions may be generalized to standard linear networks. In comparison, the diagonal linear network, a special kind of linear networks that receives a lot of attention recently, does not have fully connected neurons. Furthermore, we find that the implicit bias of both GD and SGD for diagonal linear networks are not consistent with ours. The diagonal linear networks also exhibit drastically different implicit bias of GD when compared to standard linear networks, while the conclusions for rank-1 linear networks are consistent with those of standard linear networks. We also reveal the key role of the over-parameterization in characterizing the implicit bias of SGD, namely that it will only be different with that of GD for over-parameterization model.
The inconsistency between the implicit bias of GD and SGD for diagonal linear networks and rank-1 linear networks leads us to suggest intriguing questions for future work such as what about other architectures and is there any unified analytical approach for studying implicit bias of GD and SGD for different architectures? And it is interesting to reveal whether the "alleviating" initialization effect of the SGD sampling noise is general across different architectures.
We precisely characterize the implicit bias of both GD and SGD for rank-1 linear networks, where the dependence on the initialization and depth is explicit and clear. In this sense, we take a step forward in the direction of characterizing the implicit bias of optimization algorithms.
Finally, our analysis characterizes the implicit bias of SGD through analyzing the overall parametrization  $\theta$ . This is different with another line of recent work [4, 10, 19, 27, 9] which focused on the
Figure 11:  $D(\theta(\infty), \theta_{\ell_2^*})$  for different  $\| \theta(0) \|$  when the initialization is unbalanced ( $\Delta \neq 0$ , larger  $\Delta$  means the initialization is more unbalanced). We use solid lines for the results of GD and dashed lines for SGD. For results under the balanced initialization, we use blue lines; for the results when a small perturbation is added to the balanced initialization, i.e.,  $\Delta \neq 0$ , we use red lines.
Figure 12:  $D(\theta(t), \theta_{\ell_2}^*)$  along training for rank-1 linear networks when the initialization is unbalanced ( $\Delta \neq 0$ ).
flatness of the loss landscape by directly analyzing the independent model parameters, i.e.,  $u$  and  $v$  for the rank-1 linear networks, which is also crucial to fully understand the learning dynamics.
It is worth to mention that removing the balanced initialization (Definition 1) is also a promising direction. As verified by the numerical experiments in Appendix A.4, similar phenomena exist for unbalanced initialization. From the theoretical aspect, balanced initialization enables us to derive the exact dynamics of the overall parameter  $\theta$ , which is necessary to precisely characterize the implicit bias of GD/SGD. And it is difficult to discuss arbitrary initialization without the balanced initialization assumption. The effect of removing the balanced initialization is that the induced mirror flow potential should be composed of two parts: the original potential presented in Section 3 and a perturbation due to the imbalance of the initialization to it. This implies that the  $\ell_2$ -norm solution is still returned for small initialization. On the other hand, the case for SGD is much more complicated: the Brownian motion term of the corresponding SDE will also be affected by the imbalance of the initialization, which in turn induces a much more complex time varying mirror flow potential. We believe that the exact theoretical characterization of the implicit bias of SGD without the balanced initialization is a valuable future direction.
There are also some limitations in the current work. For example, although the conclusions of numerical experiments conducted on non-linear networks resemble that of rank-1 and standard linear networks, we can not directly generalize the current theoretical analysis to non-linear neural networks, which normally do not have overall parametrization vectors as  $\theta$  that is necessary for our analysis. Moreover, the exact characterization of the stochastic integral is also absent in the current work, while we expect that the integral term  $\mathcal{L}(\theta)\mathrm{tr}\left(P_{\perp}(\theta)X^{T}X\right)\frac{\theta}{\|\theta\|}$  in Eq. (12) might have close relation with the property of the training dynamics.
In this section, we present the technical details of Section 3. In particular, Section C.1 discusses the balanced initialization, Section C.2 proves Theorem 1 for rank-1 linear networks and Section C.3 proves Theorem 2 for standard linear networks.
For a rank-1 linear network  $f(x;u,v) = w_{L}^{T}W_{L - 1}\dots W_{1}x$  where  $W_{k} = u_{k}v_{k}^{T}$  for any  $k\in \{1,\ldots ,L - 1\}$ , recalling the definition Eq. (3) and that the network  $f(x;u,v)$  can be written as  $f(x;u,v) = \rho_k v_{k + 1}^T u_k\rho_{-k}v_1^T x$ . For convenience, we let  $\xi = w_{L}^{T}W_{L - 1}\dots W_{2}u_{1}$ .
C.1 Balanced initialization
For a rank-1 linear network Eq. (2), dynamics of gradient flow is given by
$$
\frac {d u _ {k}}{d t} = - \frac {2}{n} \rho_ {k} \rho_ {- k} v _ {1} ^ {T} X ^ {T} r v _ {k + 1}, \tag {15}
$$
$$
\frac {d v _ {k + 1}}{d t} = - \frac {2}{n} \rho_ {k} \rho_ {- k} v _ {1} ^ {T} X ^ {T} r u _ {k}. \tag {16}
$$
Based on this set of dynamics, we first discuss the following useful lemma that characterizes the dynamics of norms of model parameters:
Lemma 1. For  $f(x; u, v)$  trained with gradient flow, we have
$$
\forall k \in \{1, \dots , L - 1 \}: \frac {d \| u _ {k} \| ^ {2}}{d t} = \frac {d \| v _ {k} \| ^ {2}}{d t} = \frac {d \| v _ {k + 1} \| ^ {2}}{d t}, \tag {17}
$$
i.e., layer norms grow at the same rate. Furthermore, if  $\forall k\in \{1,\dots ,L - 1\} : \| u_k(0)\| = \| v_{k + 1}(0)\| = \| v_k(0)\|$ , we have
$$
\forall k \in \{2, \dots L - 1 \}: \frac {d \left\langle v _ {k + 1} , u _ {k} \right\rangle^ {2}}{d t} = \frac {d \left\langle v _ {k} , u _ {k - 1} \right\rangle^ {2}}{d t}. \tag {18}
$$
Proof. Using Eq.(16), we have
$$
\frac {1}{2} \frac {d \| u _ {k} \| ^ {2}}{d t} = \left(\frac {d u _ {k}}{d t}\right) ^ {T} u _ {k} = - \frac {2}{n} \rho_ {k} \rho_ {- k} v _ {1} ^ {T} X ^ {T} r v _ {k + 1} ^ {T} u _ {k} = - \frac {2}{n} \xi v _ {1} ^ {T} X ^ {T} r, \tag {19}
$$
$$
\frac {1}{2} \frac {d \| v _ {k + 1} \| ^ {2}}{d t} = \left(\frac {d u _ {k}}{d t}\right) ^ {T} u _ {k} = - \frac {2}{n} \rho_ {k} \rho_ {- k} v _ {1} ^ {T} X ^ {T} r u _ {k} ^ {T} v _ {k + 1} = - \frac {2}{n} \xi v _ {1} ^ {T} X ^ {T} r. \tag {20}
$$
Therefore, both  $\frac{d\|u_k\|^2}{dt}$  and  $\frac{d\|v_{k + 1}\|^2}{dt}$  do not depend on  $k$  and are same then Eq. (17) follows.
We now discuss Eq. (18). Since we assume  $\forall k\in \{1,\dots ,L - 1\} :||u_k(0)|| = ||v_{k + 1}(0)|| = ||v_k(0)||$  and Eq. (17) implies that for any  $t > 0$
$$
\left\| u _ {k} (t) \right\| ^ {2} - \left\| u _ {k} (0) \right\| ^ {2} = \left\| v _ {k + 1} (t) \right\| ^ {2} - \left\| v _ {k + 1} (0) \right\| ^ {2} = \left\| v _ {k} (t) \right\| ^ {2} - \left\| v _ {k} (0) \right\| ^ {2}, \tag {21}
$$
$$
\left\| u _ {k} (t) \right\| ^ {2} = \left\| v _ {k + 1} (t) \right\| ^ {2} = \left\| v _ {k} (t) \right\| ^ {2} = \left\| u _ {1} (t) \right\| ^ {2} \tag {22}
$$
To show Eq. (18), we note that
$$
\begin{array}{l} \frac {d \left\langle v _ {k + 1} , u _ {k} \right\rangle}{d t} = \left(\frac {d v _ {k + 1}}{d t}\right) ^ {T} u _ {k} + v _ {k + 1} ^ {T} \frac {d u _ {k}}{d t} \\ = - \frac {2}{n} \rho_ {k} \rho_ {- k} v _ {1} ^ {T} X ^ {T} r (\| u _ {k} \| ^ {2} + \| v _ {k + 1} \| ^ {2}) \\ = - \frac {2}{n} \xi v _ {1} ^ {T} X ^ {T} r \frac {\left\| u _ {k} \right\| ^ {2} + \left\| v _ {k + 1} \right\| ^ {2}}{\left\langle v _ {k + 1} , u _ {k} \right\rangle} = - \frac {4}{n} \xi v _ {1} ^ {T} X ^ {T} r \frac {\left\| u _ {1} \right\| ^ {2}}{\left\langle v _ {k + 1} , u _ {k} \right\rangle}, \tag {23} \\ \end{array}
$$
where we use Eq.(16) in the second equality and the third equality is because  $\xi = \rho_k\rho_{-k}\langle v_{k + 1},u_k\rangle$  As a result, the above equation implies that
$$
\frac {1}{2} \frac {d \left(\left\langle v _ {k + 1} , u _ {k} \right\rangle\right) ^ {2}}{d t} = - \frac {4}{n} \xi v _ {1} ^ {T} X ^ {T} r \| u _ {1} \| ^ {2}, \tag {24}
$$
which does not depend on  $k$ , and Eq. (18) follows.
To simplify the analysis, in Theorem 1, we have required the balanced initialization across layers (Definition 1). Recall that the balanced initialization is defined as
$$
\frac {\left\langle v _ {k + 1} (0) , u _ {k} (0) \right\rangle^ {2}}{\| v _ {k + 1} (0) \| ^ {2} \| u _ {k} (0) \| ^ {2}} = 1, \tag {25}
$$
$$
\| v _ {k + 1} (0) \| = \| u _ {k} (0) \| = \| v _ {1} (0) \|. \tag {26}
$$
Eq. (25) states that  $u_{k}$  of the  $k$ -th layer is aligned with  $v_{k+1}$  of the  $(k+1)$ -th layer in direction while Eq. (26) means they have the same magnitudes as  $v_{1}(0)$ . The balanced initialization has been suggested by several previous works [3, 2, 6, 28] for standard linear networks defined as follows.
Definition 2 (Balanced initialization for standard linear networks). Given an  $L$ -layer standard linear network  $f(x;W) = w_{L}^{T}W_{L - 1}\dots W_{1}x$ , for any  $k\in \{1,\ldots ,L - 1\}$ , the balanced initialization means that
$$
W _ {k + 1} ^ {T} (0) W _ {k + 1} (0) = W _ {k} (0) W _ {k} ^ {T} (0)
$$
for any  $k\in \{1,\ldots ,L\}$
This directly means that  $W_{k + 1}(0)$  and  $W_{k}(0)$  share the same singular values and  $W_{k + 1}$ 's right singular vector aligns with the left singular vector of  $W_{k}(0)$ . In our case, such reasoning gives us
$$
\frac {v _ {k + 1} (0)}{\| v _ {k + 1} (0) \|} = \frac {u _ {k} (0)}{\| u _ {k} (0) \|} \tag {27}
$$
$$
\| v _ {k + 1} (0) \| \| u _ {k + 1} (0) \| = \| v _ {k} (0) \| \| u _ {k} (0) \|, \tag {28}
$$
where Eq. (27) is similar to Eq. (25), which shows that  $u_{k}(0)$  aligns with  $v_{k + 1}(0)$ , and we adapt the condition (28) to Eq. (26) for rank-1 linear net since  $v_{k}$  and  $u_{k}$  are the independent model parameters.
A nice property of GD is that the balanced property across layers will be maintained during training [6, 11, 2], i.e.,  $W_{k + 1}^{T}(t)W_{k + 1}(t) = W_{k}(t)W_{k}^{T}(t)$  for  $t > 0$  and  $k \in \{1, \dots, L\}$ , which can be showed by taking derivative with respect to time on  $W_{k + 1}^{T}(t)W_{k + 1}(t)$  and  $W_{k}(t)W_{k}^{T}(t)$ . For the rank-1 linear network case, according to Lemma 1 and Eq. (25),  $\langle v_{k + 1}(t), u_k(t) \rangle$  are the same for all  $k \in \{1, \dots, L - 1\}$ , thus GD also maintains the balanced property for rank-1 deep linear networks. Although the balanced initialization conditions are slightly strict, it can be approximately accurate if the initialization scale is not large, which is rather common in practice. Under such initialization conditions, we are able to precisely characterize the implicit bias of GD and focus more on the effects coming from the overall initialization of model parameters, rather than the difference between layers.
In this section, we prove Theorem 1. Basically the idea is to show the existence of a potential function  $V(\theta)$  such that the model parameter  $\theta$  follows a mirror descent with respect to  $V(\theta)$ :
$$
\theta (\infty) = \arg \min  _ {\theta} V (\theta) \quad s. t. X \theta = y. \tag {29}
$$
This method is called infinitesimal mirror descent (IMD) approach and can be found in, e.g., [3]. Note that, to apply this method, the dynamics of  $\theta$  should satisfy certain conditions, which might be strict. For example, given the linear model  $f(x;\theta) = \theta^T x$  for  $\theta \in \mathbb{R}^d$ , both parameterization of  $\theta$  with standard linear networks, i.e.  $\theta^T = w_L^T W_{L-1}\dots W_1$ , and a much simpler one  $\theta = cv$  for  $c \in \mathbb{R}$  and  $v \in \mathbb{R}^d$  do not satisfy the condition of applying the IMD approach, while the rank-1 linear networks satisfy the condition, which implies that the parameterization of rank-1 linear networks is different with a scalar times a vector. Furthermore, the above example also confirms our motivation of studying rank-1 linear networks as a proxy of standard linear networks, especially considering that the implicit bias of SGD for rank-1 linear networks is more amenable.
We first present a useful Lemma in [20]:
Lemma 2. If  $H$  has rank 1 and  $G$  is invertible, then
$$
(G + H) ^ {- 1} = G ^ {- 1} - \frac {1}{1 + g} G ^ {- 1} H G ^ {- 1}. \tag {30}
$$
where  $g = tr\left(HG^{-1}\right)$
We now prove Theorem 1.
Proof. Recall that  $r \in \mathbb{R}^n$  with  $r_i = (f_i - y_i)$ , we let
$$
\Phi = \sum_ {k = 1} ^ {L - 1} \phi_ {k}, \tag {31}
$$
$$
\phi_ {k} = \rho_ {k} ^ {2} \rho_ {- k} ^ {2} \left(\left\| u _ {k} \right\| ^ {2} + \left\| v _ {k + 1} \right\| ^ {2}\right). \tag {32}
$$
The key step is to derive the dynamics of the overall parameter  $\theta$  which can be done by noting that
$$
\frac {d \theta}{d t} = v _ {1} \frac {d \xi}{d t} + \xi \frac {d v _ {1}}{d t},
$$
where, according to Lemma 1,
$$
\begin{array}{l} \frac {d \xi}{d t} = \sum_ {k = 1} ^ {L - 1} \rho_ {k} \rho_ {- k} \frac {d \left\langle v _ {k + 1} , u _ {k} \right\rangle}{d t} \\ = - \frac {2}{n} r ^ {T} X v _ {1} \sum_ {k = 1} ^ {L - 1} \rho_ {k} ^ {2} \rho_ {- k} ^ {2} \left(\left\| u _ {k} \right\| ^ {2} + \left\| v _ {k + 1} \right\| ^ {2}\right) \\ = - \frac {2}{n} \Phi v _ {1} ^ {T} X ^ {T} r \tag {33} \\ \end{array}
$$
$$
\begin{array}{l} \frac {d v _ {1}}{d t} = - \frac {2}{n} \xi X ^ {T} r (34) \\ \Longrightarrow \frac {d \theta}{d t} = - \frac {2}{n} \left(\xi^ {2} I + \Phi v _ {1} v _ {1} ^ {T}\right) X ^ {T} r. (35) \\ \end{array}
$$
Eq. (35) can be rewritten as
$$
\left(\xi^ {2} I + \Phi v _ {1} v _ {1} ^ {T}\right) ^ {- 1} \frac {d \theta}{d t} = - \frac {2}{n} X ^ {T} r. \tag {36}
$$
According to Lemma 2, note that
$$
\operatorname {t r} \left(\Phi \frac {v _ {1} v _ {1} ^ {T}}{\xi^ {2}}\right) = \frac {\Phi \| v _ {1} \| ^ {2}}{\xi^ {2}},
$$
the inverse appeared in Eq. (36) is
$$
\left(\xi^ {2} I + \Phi v _ {1} v _ {1} ^ {T}\right) ^ {- 1} = \frac {1}{\xi^ {2}} I - \frac {\xi^ {- 2} v _ {1} v _ {1} ^ {T} \Phi \xi^ {- 2}}{1 + \frac {\Phi \| v _ {1} \| ^ {2}}{\xi^ {2}}} = \frac {1}{\xi^ {2}} I - \frac {\theta \theta^ {T}}{\frac {\xi^ {6}}{\Phi} + \xi^ {2} \| \theta \| ^ {2}}. \tag {37}
$$
It is now left for us to express  $\xi^2$  and  $\Phi$  in terms of  $\theta$ . In the following, we assume the balanced initialization in Theorem 1 and apply Lemma 1.
1.  $\xi^2$ . This can be done by noting that  $\| \theta \| ^2 = \xi^2 \| v_1\| ^2$ , where  $\xi^2$  is given by
$$
\xi^ {2} = \prod_ {k = 1} ^ {L - 1} \left\langle v _ {k + 1}, u _ {k} \right\rangle^ {2}. \tag {38}
$$
Note that  $\langle v_{k + 1},u_k\rangle^2$  grow at the same rate for different  $k$  according to Lemma 1 and  $\langle v_{k + 1},u_k\rangle^2$  are the same at initialization for different  $k$  due to our assumption, we have  $\langle v_{k + 1},u_k\rangle^2 = \langle v_2,u_1\rangle^2$  and  $\xi^2 = \langle v_2,u_1\rangle^{2(L - 1)}$ . Note that
$$
\frac {1}{2} \frac {d \left\langle v _ {2} , u _ {1} \right\rangle^ {2}}{d t} = - \frac {4}{n} \xi v _ {1} ^ {T} X ^ {T} r \| u _ {1} \| ^ {2} = - \frac {4}{n} \xi v _ {1} ^ {T} X ^ {T} r \| v _ {1} \| ^ {2} = \frac {1}{2} \frac {d \| v _ {1} \| ^ {4}}{d t}, \tag {39}
$$
we have  $\langle v_2, u_1 \rangle^2 - \langle v_2(0), u_1(0) \rangle^2 = \|v_1\|^4 - \|v_1(0)\|^4$ . Since we have  $\langle v_2, u_1 \rangle^2 = \|u_1\|^4 = \|v_1\|^4$  at initialization according to our assumption,  $\xi^2$  can be finally written as
$$
\xi^ {2} = \left\| v _ {1} \right\| ^ {4 (L - 1)}. \tag {40}
$$
$$
\left\| \theta \right\| = \left\| v _ {1} \right\| ^ {2 L - 1} \Longrightarrow \left\| v _ {1} \right\| = \left\| \theta \right\| ^ {\frac {1}{2 L - 1}}, \quad \xi^ {2} = \left\| \theta \right\| ^ {\frac {4 (L - 1)}{2 L - 1}}. \tag {41}
$$
2.  $\xi^6 /\Phi$  . By taking some simple algebra, we have
$$
\Phi = \frac {2 (L - 1) \xi^ {2}}{\| v _ {1} \| ^ {2}} \Longrightarrow \frac {\xi^ {6}}{\Phi} = \frac {\xi^ {4} \| v _ {1} \| ^ {2}}{2 (L - 1)}. \tag {42}
$$
Now Eq. (36) becomes
$$
\left\| \theta \right\| ^ {- \frac {2 (L - 1)}{2 L - 1}} \left(I - \frac {\theta \theta^ {T}}{\frac {\| \theta \| ^ {2}}{2 (L - 1)} + \| \theta \| ^ {2}}\right) \frac {d \theta}{d t} = - \frac {2 | \xi |}{n} X ^ {T} r. \tag {43}
$$