SlowGuess/NeurIPS_2023_OCR huggingface dataset:
These conditions are now sufficient for us to find the form of the potential  $V(\theta)$ . Suppose that  $V(\theta)$  can be written as
$$
V (\theta) = \hat {V} (\| \theta \|) + h ^ {T} \theta \tag {44}
$$
for some vector  $h$  and satisfies the following relation:
$$
\begin{array}{l} \nabla_ {\theta} ^ {2} V (\theta) = \nabla_ {\theta} ^ {2} \hat {V} (\theta) \\ = \| \theta \| ^ {- \frac {2 (L - 1)}{2 L - 1}} \left(I - \frac {1}{1 + \frac {1}{2 (L - 1)}} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right), \tag {45} \\ \end{array}
$$
then Eq. (43) gives us
$$
\frac {d}{d t} (\nabla_ {\theta} V (\theta)) = - \frac {2}{n} X ^ {T} r \tag {46}
$$
and the integration relation
$$
\left. \nabla_ {\theta} V (\theta) - \nabla_ {\theta} V (\theta) \right| _ {\theta = \theta (0)} = \sum_ {i = 1} ^ {n} x _ {i} \int \tilde {r} _ {i} (\tau) d \tau \tag {47}
$$
where we let  $\tilde{r} = -2|\xi |r / n$ . Requiring  $\nabla_{\theta}V(\theta)|_{\theta = \theta (0)} = 0$  and denoting  $\lambda_{i} = \int_{0}^{\infty}\tilde{r}_{i}(\tau)d\tau$  gives us the condition at  $t = \infty$ :
$$
\left. \nabla_ {\theta} V (\theta) \right| _ {\theta = \theta (\infty)} = \sum_ {i = 1} ^ {n} x _ {i} \lambda_ {i}. \tag {48}
$$
Eq. (48) coincides with the KKT stationary condition of the optimization problem (29). Therefore, we can prove the theorem by deriving the explicit form of  $V(\theta)$ .
Solving  $V(\theta)$ . According to Eq. (44), we can derive the following relation:
$$
\partial_ {\theta} V (\theta) = \hat {V} ^ {\prime} \frac {\theta}{\| \theta \|} + h ^ {T} \tag {49}
$$
$$
\begin{array}{l} \partial_ {\theta} ^ {2} V (\theta) = \frac {1}{\| \theta \| ^ {2}} \left[ \left(\hat {V} ^ {\prime \prime} \frac {\theta \theta^ {T}}{\| \theta \|} + \hat {V} ^ {\prime} I\right) \| \theta \| - \hat {V} ^ {\prime} \frac {\theta \theta^ {T}}{\| \theta \|} \right] \\ = \frac {\hat {V} ^ {\prime}}{\| \theta \|} \left[ I - \left(1 - \| \theta \| \frac {\hat {V} ^ {\prime \prime}}{\hat {V} ^ {\prime}}\right) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right]. \tag {50} \\ \end{array}
$$
Comparing this with Eq. (45), we conclude that
$$
1 - \| \theta \| _ {\hat {V} ^ {\prime}} ^ {\hat {V} ^ {\prime \prime}} = \frac {1}{1 + \frac {1}{2 (L - 1)}} \Rightarrow \frac {\hat {V} ^ {\prime \prime}}{\hat {V} ^ {\prime}} = \frac {1}{\| \theta \|} \left(1 - \frac {1}{1 + \frac {1}{2 (L - 1)}}\right), \tag {51}
$$
which, by noting that  $\partial_{\| \theta \|}\ln \hat{V}^{\prime}(\| \theta \|) = \frac{\hat{V}^{\prime\prime}(\| \theta\|)}{\hat{V}^{\prime}(\| \theta\|)}$  can be solved as follows
$$
\begin{array}{l} \frac {\hat {V} ^ {\prime \prime} (\| \theta \|)}{\hat {V} ^ {\prime} (\| \theta \|)} = \frac {1}{\| \theta \|} \frac {1}{2 L - 1} \\ \Longrightarrow \ln \hat {V} ^ {\prime} (\| \theta \|) = \frac {1}{2 L - 1} \ln \| \theta \| \\ \Longrightarrow \hat {V} (\| \theta \|) = \frac {2 L - 1}{2 L} \| \theta \| ^ {\frac {2 L}{2 L - 1}}. \tag {52} \\ \end{array}
$$
$$
\frac {\hat {V} ^ {\prime}}{\| \theta \|} = \| \theta \| ^ {- \frac {2 (L - 1)}{2 L - 1}},
$$
Eq. (45) is automatically satisfied when  $\hat{V} (\| \theta \|)$  has the form of Eq. (51). It is now left for us to get the form of  $h$ , which can be done by noting that
$$
\begin{array}{l} \partial_ {\theta} V (\theta (0)) = 0 \\ \Rightarrow \| \theta (0) \| ^ {\frac {1}{2 L - 1}} \frac {\theta (0)}{\| \theta (0) \|} + h = 0. \\ \end{array}
$$
Thus the final form of  $V(\theta)$  is
$$
V (\theta) = \frac {1}{\Omega_ {L}} \| \theta \| ^ {\Omega_ {L}} - \theta^ {T} \frac {\theta (0)}{\| \theta (0) \| ^ {\lambda_ {L}}}. \tag {53}
$$
C.2.1 Remove the assumption of convergence to the interpolation solution
The assumption that  $X\theta(\infty) = y$  in Theorem 1 can be removed if the dimension of the span of  $X^T$  is larger than the number of samples  $n$ , i.e., when  $\dim \left( \operatorname{span}(X^T) \right) \geq n$ . This can be proved as follows.
1 Since  $\max (\dim (\operatorname {span}(X^T))) = n$  , this condition is in fact  $\dim (\operatorname {span}(X^T)) = n$
Proposition 2. For the over-parameterized regression of rank-1 linear networks Eq. (2) and the dataset  $\{(x_i, y_i)\}_{i=1}^n$  where  $x_i \in \mathbb{R}^d$  and  $n < d$ , if
$$
d i m \left(\operatorname {s p a n} \left(X ^ {T}\right)\right) \geq n,
$$
then the gradient flow solution  $\theta (\infty)$  satisfies that
$$
X \theta (\infty) = y.
$$
Proof. To prove this proposition, we study the dynamics of the loss function  $\mathcal{L} = \frac{1}{n}\sum_{i}r_{i}^{2}$ , where  $r_i = \langle \theta ,x_i\rangle -y_i$ , that is given by
$$
\begin{array}{l} \frac {d \mathcal {L}}{d t} = \frac {\partial \mathcal {L}}{\partial \theta} \frac {d \theta}{d t} \\ = \frac {2}{n} r ^ {T} X \left[ - \frac {2 \xi^ {2}}{n} \left(I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} r \right] \\ = - \frac {4 \xi^ {2}}{n ^ {2}} \left[ r X ^ {T} X r + 2 (L - 1) \frac {r ^ {T} X \theta \theta^ {T} X ^ {T} r}{\| \theta \| ^ {2}} \right] \\ = - \frac {4 \xi^ {2}}{n ^ {2}} \left[ \| X ^ {T} r \| ^ {2} + 2 (L - 1) \frac {\left(\theta^ {T} X ^ {T} r\right) ^ {2}}{\| \theta \| ^ {2}} \right] \\ \leq 0, \\ \end{array}
$$
where we have the equality in the last line when  $r = (0, \dots, 0)^T \in \mathbb{R}^d$ , i.e.,  $X\theta = y$ , or  $X^T r = 0$  which is not possible since we have assumed that  $\dim(\operatorname{span}(X^T)) \geq n$ . Therefore,  $\mathcal{L}(\theta(t))$  keeps decreasing until  $X\theta(t) = y$ , i.e., until GD finds the interpolation solution. Noting that  $\min_{\theta} \mathcal{L}(\theta) = 0$ , we complete the proof.
C.3 Proof of Theorem 2
In this section we prove Theorem 2. The techniques are similar to those in C.2, and we still need a time wrapping technique introduced in [3] to derive the form of  $V_{std}(\theta)$ , since the condition for applying the IMD approach is violated in this case. Recall that, for the standard linear networks  $f(x;W) = w_L^T W_{L - 1}\dots W_1x = \theta^T x$  where  $W_{k}\in \mathbb{R}^{d_{k}\times d_{k + 1}}$ , our purpose is to find a potential function  $V_{std}(\theta)$  such that the gradient flow solution  $\theta (\infty)$  satisfies that
$$
\theta (\infty) = \arg \min  _ {\theta} V _ {s t d} (\theta), \quad s. t. X \theta = y. \tag {54}
$$
Since we assume the balanced initialization (Definition 2), then according to [11, 6, 2], the norms of all layers grow at the same rate and are the same for any  $t > 0$ :
$$
W _ {k + 1} ^ {T} (t) W _ {k + 1} (t) = W _ {k} (t) W _ {k} ^ {T} (t).
$$
Furthermore, following the procedure of [2], we obtain that the dynamics of  $\theta$  is
$$
\frac {d \theta}{d t} = - \| \theta \| ^ {\frac {2 (L - 1)}{L}} \left[ I + (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} r. \tag {55}
$$
We now present the proof.
Proof. Eq. (55) can be written as
$$
\left\| \theta \right\| ^ {\frac {2 (1 - L)}{L}} \left[ I + (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] ^ {- 1} \frac {d \theta}{d t} = - X ^ {T} r,
$$
where, according to Lemma 2, the inverse in above equation is given by
$$
\begin{array}{l} \left[ I + (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] ^ {- 1} = I - \frac {1}{1 + \operatorname {t r} \left(\frac {(L - 1) \theta \theta^ {T}}{\| \theta \| ^ {2}}\right)} (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \\ = I - \frac {L - 1}{L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}. \tag {56} \\ \end{array}
$$
$$
\left\| \theta \right\| ^ {\frac {2 (1 - L)}{L}} \left(I - \frac {L - 1}{L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \frac {d \theta}{d t} = - X ^ {T} r. \tag {57}
$$
In this case, we still assume that  $V_{std}(\theta)$  has the following form
$$
V _ {s t d} (\theta) = \hat {V} _ {s t d} \left(\| \theta \|\right) + \beta^ {T} \theta \tag {58}
$$
for some constant vector  $\beta \in \mathbb{R}^d$ . Then following a similar procedure as in C.2, we have that
$$
\partial_ {\theta} V _ {s t d} (\theta) = \hat {V} _ {s t d} ^ {\prime} \frac {\theta}{\| \theta \|} + h ^ {T} \tag {59}
$$
$$
\partial_ {\theta} ^ {2} V _ {s t d} (\theta) = \frac {\hat {V} _ {s t d} ^ {\prime}}{\| \theta \|} \left[ I - \left(1 - \| \theta \| \frac {\hat {V} _ {s t d} ^ {\prime \prime}}{\hat {V} _ {s t d} ^ {\prime}}\right) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right]. \tag {60}
$$
In the IMD approach,  $V_{std}(\theta)$  should satisfy that
$$
\frac {d}{d t} \left(\partial_ {\theta} V _ {s t d} (\theta)\right) = - \frac {2}{n} X ^ {T} r, \tag {61}
$$
which requires that  $\partial_{\theta}^{2}V_{std}(\theta) = \| \theta \|_{L}^{\frac{2(L - 1)}{L}}\left(I - \frac{L - 1}{L}\frac{\theta\theta^{T}}{\|\theta\|^{2}}\right)\frac{d\theta}{dt}$ . But this is not possible. Therefore, we multiply a time re-scale factor  $g(\theta)$ , as long as  $g(\theta)$  is positive, to both sides of Eq. (57) and only require that  $V_{std}(\theta)$  satisfies the above relation under the new time scale  $\tau : \mathbb{R} \to \mathbb{R}$  such that  $\tau' = g(\theta)$ :
$$
g (\theta) \| \theta \| ^ {\frac {2 (1 - L)}{L}} \left(I - \frac {L - 1}{L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \frac {d \theta}{d t} = - g (\theta) X ^ {T} r. \tag {62}
$$
Then the limit point at  $t = \infty$  in Eq. (61) is also visited at the point  $\tau = \int_0^\infty g(\theta(s)) ds$  in Eq. (62). We now solve the explicit form of  $V_{std}(\theta)$  that satisfies Eq. (62). By comparing Eq. (60) and the right hand side of Eq. (57), we obtain that the following relation should be satisfied:
$$
\frac {\hat {V} _ {s t d} ^ {\prime}}{\| \theta \|} \left[ I - \left(1 - \| \theta \| \frac {\hat {V} _ {s t d} ^ {\prime \prime}}{\hat {V} _ {s t d} ^ {\prime}}\right) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] = g (\theta) \| \theta \| ^ {\frac {2 (1 - L)}{L}} \left(I - \frac {L - 1}{L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right). \tag {63}
$$
This implies that we need:
- the terms in the bracket on both sides should match:
$$
\begin{array}{l} 1 - \| \theta \| \frac {\hat {V} _ {s t d} ^ {\prime \prime}}{\hat {V} _ {s t d} ^ {\prime}} = \frac {L - 1}{L} \Rightarrow \frac {1}{x L} = \frac {\hat {V} _ {s t d} ^ {\prime \prime}}{\hat {V} _ {s t d} ^ {\prime}} \\ \Rightarrow \ln \hat {V} _ {s t d} ^ {\prime} = \frac {1}{L} \ln \| \theta \| + C \Rightarrow \hat {V} _ {s t d} = \frac {C ^ {\prime} L}{L + 1} \| \theta \| ^ {\frac {1}{L} + 1} \tag {64} \\ \end{array}
$$
for some constant  $C$  and  $C'$ , where we can simply choose  $C' = 1$ ;
- the terms outside the bracket on both sides should also match:
$$
\frac {\hat {V} _ {s t d} ^ {\prime}}{\| \theta \|} = g (\theta) \| \theta \| ^ {\frac {2 (1 - L)}{L}} \Rightarrow g (\theta) = \hat {V} _ {s t d} ^ {\prime} \| \theta \| ^ {\frac {2 (L - 1)}{L} - 1}. \tag {65}
$$
To obtain the form of the constant vector  $\beta$ , we note that  $\partial_{\theta}V(\theta (0)) = 0$ , which immediately gives us
$$
\left\| \theta (0) \right\| ^ {\frac {1}{L}} \frac {\theta (0)}{\left\| \theta (0) \right\|} + \beta = 0 \Rightarrow \beta = - \theta (0) \| \theta (0) \| ^ {\frac {1}{L} - 1}. \tag {66}
$$
Combining all these terms, we have the final form of  $V_{std}(\theta)$ :
$$
V _ {s t d} (\theta) = \frac {L}{L + 1} \| \theta \| ^ {\frac {1}{L} + 1} - \theta (0) ^ {T} \theta \| \theta (0) \| ^ {\frac {1}{L} - 1}. \tag {67}
$$
As in C.2, denoting  $\lambda_{i} = \int_{0}^{\infty}r_{i}(s)ds$  and noting that  $\partial_{\theta}V_{std}(\theta (0)) = 0$  give us
$$
\partial_ {\theta} V _ {s t d} = \sum_ {i = 1} ^ {n} x _ {i} \lambda_ {i},
$$
which is exactly the KKT stationary condition of the optimization problem (54).
Convergence to the interpolation solution. Similar to C.2.1, we can also show the convergence to the interpolation solution when  $\dim \bigl (\mathrm{span}(X^T)\bigr)\geq n$  by deriving the dynamics of  $\mathcal{L}$ :
$$
\begin{array}{l} \frac {d \mathcal {L}}{d t} = \frac {\partial \mathcal {L}}{\partial \theta} \frac {d \theta}{d t} \\ = - \frac {2 \| \theta \| ^ {\frac {2 (1 - L)}{L}}}{n} r ^ {T} X \left[ I + (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} r \\ = - \frac {2 \| \theta \| ^ {\frac {2 (1 - L)}{L}}}{n} \left[ \| r ^ {T} X \| ^ {2} + (L - 1) \left(\theta^ {T} X ^ {T} r\right) ^ {2} \right]. \\ \end{array}
$$
Since we assume that  $\dim \bigl (\mathrm{span}(X^T)\bigr)\geq n$ ,  $d\mathcal{L} / dt < 0$  until  $X\theta (t) = y$ , i.e.,  $\mathcal{L}$  keeps decreasing until GD finds the interpolation solution.
C.4 Proof of Proposition 1
In this section, we prove Proposition 1 by analyzing the gradient of model parameters. It is helpful to recall that for a matrix  $A$ , we use  $A_{ij}$  to denote its  $i$ -th row  $j$ -th column element. For weight matrices, e.g.,  $W_{k}$ , we use  $W_{k;ij}$  to denote its  $i$ -th row  $j$ -th column element.
Proof. For a standard linear network that has the initialization of Proposition 1 and  $L$  is an odd number, i.e., for an integer  $p$
note that the  $L - 1$ -th layer satisfies that  $W_{L - 1;ij} = W_{L - 1;ij}\delta_{jc_{L - 1}}$  where  $\delta_{jl} = 1$  if  $j = l$  otherwise  $\delta_{jl} = 0$ , we can write the networks at  $t = 0$  as
$$
f (x; W) = \sum_ {i} \sum_ {j} w _ {L; i} W _ {L - 1; i j} \delta_ {j c _ {L - 1}} \left(W _ {L - 2} \dots W _ {1} x\right) _ {j}.
$$
Then the gradient w.r.t  $W_{L - 1}$  at  $t = 0$  is
$$
\begin{array}{l} \left(\nabla_ {W _ {L - 1}} \mathcal {L} (\theta)\right) _ {i j} = \frac {2}{n} \sum_ {\mu = 1} ^ {n} r _ {\mu} w _ {L; i} \left(W _ {L - 2} \dots W _ {1} x\right) _ {j} \delta_ {j c _ {L - 1}} (68) \\ \Longrightarrow \left(\nabla_ {W _ {L - 1}} \mathcal {L} (\theta)\right) _ {i j} = 0 \text {i f} j \neq c _ {L - 1}. (69) \\ \end{array}
$$
This means that the parameter  $W_{L-1;ij}$  will be updated only when  $j = c_{L-1}$ . As a result, the initialization shape for  $W_{L-1}$  will be maintained, i.e., only the non-zero column of  $W_{L-1}$  at  $t = 0$  will be updated and all other elements of  $W_{L-1}$  will be zero for any  $t > 0$  since the corresponding gradients vanish. Similarly, for  $W_{L-1}$ , we can write the network at  $t = 0$  as:
$$
f (x; W) = \sum_ {j} \sum_ {l} (w _ {L} ^ {T} W _ {L - 1}) _ {j} \delta_ {j c _ {L - 1}} W _ {L - 2; j l} (W _ {L - 3} \dots W _ {1} x) _ {l},
$$
$$
\left(\nabla_ {W _ {L - 2}} \mathcal {L} (\theta)\right) _ {i j} = 0 \text {i f} i \neq c _ {L - 1}
$$
and the initialization shape for  $W_{L - 2}$  will also be maintained. Following a similar procedure, we conclude that all the initialization shapes will be maintained for any  $t > 0$ .
We now consider the diagonal initialization for weight matrices, namely that
$$
W _ {k; i j} = 0 \text {i f} i \neq j, \quad \forall k \in \{1, \dots , L - 1 \}.
$$
For any  $k \in \{1, \dots, L - 1\}$ , the gradient w.r.t  $W_{k}$  at  $t = 0$  is
$$
\left(\nabla_ {W _ {k}} \mathcal {L} (\theta)\right) _ {i j} = \frac {2}{n} \sum_ {\mu = 1} ^ {n} r _ {\mu} \left(w _ {L} ^ {T} \dots W _ {k + 1}\right) _ {i} \left(W _ {k - 1} \dots W _ {1} x\right) _ {j},
$$
where, clearly, there is not any constraint on which elements of  $(w_{L}^{T}\dots W_{k + 1})^{T}\in \mathbb{R}^{d}$  and  $(W_{k - 1}\dots W_1x)\in \mathbb{R}^d$  are zeros if we do not require that many diagonal elements of the initialization
matrices are zeros, which clearly violates the diagonal initialization requirements. Thus we can not conclude that
$$
(\nabla_ {W _ {k}} \mathcal {L} (\theta)) _ {i j} = 0 \text {i f} i \neq j,
$$
i.e., the off-diagonal elements of  $W_{k}$  will also be updated to be non-zeros. Thus the diagonal initialization of weight matrices can not be maintained. The conclusions for other layers can be easily derived by following similar arguments.
In this section, we present the technical details for Section 4. In particular, we introduce our modelling details of the SGD dynamics in D.1, derive the SDE of the model parameter  $\theta$  in D.2, and, finally, prove Theorem 3 in D.3.
D.1 SDE modelling details
Similar to the case of GD, to derive the implicit bias of SGD for rank-1 linear networks, we need to first derive the dynamics of the overall model parameter  $\theta$ . For this purpose, the structure of the noise of SGD is crucial. For convenience, we first discuss the basic SGD where only one data is randomly sampled at each step, while the generalization to batch-size SGD is straightforward.
Structure of the noise. We present the details for  $u_{k}$ , and the case for  $v_{k + 1}$  is similar. Recall that the empirical loss is  $\mathcal{L} = \sum_{i}\ell_{i} / \mathrm{n},\eta$  is the learning rate and that the network can be written as  $f(x;u,v) = \rho_k\rho_{-k}v_{k + 1}^T u_kv_1^Tx$ , we start with the SGD update equation for  $u_{k}$  where we let  $j_{t}$  denote the index of the sampled data at the  $t$ -th step:
$$
\begin{array}{l} u _ {k} (t + 1) = u _ {k} (t) - \eta \frac {\partial \ell_ {j _ {t}}}{\partial u _ {k}} = u _ {k} (t) - \eta \frac {\partial \mathcal {L}}{\partial u _ {k}} + \eta \left(\frac {\partial \mathcal {L}}{\partial u _ {k}} - \frac {\partial \ell_ {j _ {t}}}{\partial u _ {k}}\right) \\ = u _ {k} (t) - \frac {2 \eta}{n} v _ {1} ^ {T} X ^ {T} r \rho_ {k} \rho_ {- k} v _ {k + 1} \\ + 2 \eta \left(\frac {1}{n} \sum_ {i} r _ {i} \rho_ {k} \rho_ {- k} v _ {1} ^ {T} x _ {i} v _ {k + 1} - r _ {j _ {t}} \rho_ {k} \rho_ {- k} v _ {1} ^ {T} x _ {j _ {t}} v _ {k + 1}\right) \\ = u _ {k} (t) - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \rho_ {k} \rho_ {- k} v _ {k + 1} + 2 \eta \rho_ {k} \rho_ {- k} v _ {k + 1} v _ {1} ^ {T} X ^ {T} Z _ {j _ {t}} \tag {70} \\ \end{array}
$$
where we let  $\vec{e_{j_t}}$  be the basis vector in  $\mathbb{R}^n$  such that the  $j_{t}$ -th element is 1 while all other elements are 0 and
$$
Z _ {j _ {t}} = \mathrm {E} _ {j _ {t}} \left[ r _ {j _ {t}} \vec {e _ {j _ {t}}} \right] - r _ {j _ {t}} \vec {e _ {j _ {t}}}, \quad \operatorname {c o v} \left[ Z _ {j _ {t}} \right] \sim \frac {\mathcal {L}}{n} \mathbf {I} _ {n}. \tag {71}
$$
As a result of this, the noise of SGD for  $u_{k}$  is now
$$
\Sigma \left(u _ {k} (t)\right) = \frac {4 \eta^ {2} \mathcal {L}}{n} \left(\rho_ {k} \rho_ {- k}\right) ^ {2} v _ {1} ^ {T} X ^ {T} X v _ {1} v _ {k + 1} v _ {k + 1} ^ {T}. \tag {72}
$$
Continuous Modelling of SGD. The continuous modelling techniques for SGD have been widely applied in recent works [1, 9, 23, 22] to study the dynamics of SGD. In our setting, the continuous counterpart of SGD is established as follows. First, the discrete SGD updating equations Eq. (8) can be equivalently written as
$$
u _ {k} (t + 1) = u _ {k} (t) - \eta \nabla_ {u _ {k}} \mathcal {L} (\theta) + \eta \left[ \nabla_ {u _ {k}} \mathcal {L} (\theta) - \nabla_ {u _ {k}} \ell_ {i (t)} (\theta) \right],
$$
$$
v _ {k + 1} (t + 1) = v _ {k + 1} (t) - \eta \nabla_ {v _ {k + 1}} \mathcal {L} (\theta) + \eta [ \nabla_ {v _ {k + 1}} \mathcal {L} (\theta) - \nabla_ {v _ {k + 1}} \ell_ {i (t)} (\theta) ]
$$
for  $k \in \{1, \dots, L - 1\}$  where both  $\nabla_{u_k}\mathcal{L}(\theta) - \nabla_{u_k}\ell_{i(t)}(\theta)$  and  $\nabla_{v_{k + 1}}\mathcal{L}(\theta) - \nabla_{v_{k + 1}}\ell_{i(t)}(\theta)$  are zero-mean noises with covariance matrices in  $\mathbb{R}^{d_k\times d_k}$
$$
\begin{array}{l} \Sigma (u _ {k}) = \frac {4 \mathcal {L} (\rho_ {k} \rho_ {- k}) ^ {2}}{n} v _ {1} ^ {T} X ^ {T} X v _ {1} v _ {k + 1} v _ {k + 1} ^ {T}, \\ \Sigma (v _ {k + 1}) = \frac {4 \mathcal {L} (\rho_ {k} \rho_ {- k}) ^ {2}}{n} v _ {1} ^ {T} X ^ {T} X v _ {1} u _ {k} u _ {k} ^ {T}, \\ \end{array}
$$
with  $\rho_{k}$  and  $\rho_{-k}$  defined in Eq. (3). Second, as we identify the noise covariance, letting  $\eta \rightarrow 0^2$ , then we obtain the continuous counterpart of the discrete SGD that is a set of stochastic differential equations (SDE):
$$
d u _ {k} = - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \rho_ {k} \rho_ {- k} v _ {k + 1} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} (\rho_ {k} \rho_ {- k}) v _ {k + 1} v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t} \tag {73}
$$
$$
d v _ {k + 1} = - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \rho_ {k} \rho_ {- k} u _ {k} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} (\rho_ {k} \rho_ {- k}) u _ {k} v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t} \tag {74}
$$
where we let  $r = (f(x_1;u,v) - y_1,\dots ,f(x_n;u,v) - y_n)^T\in \mathbb{R}^n$  be the residuals and  $\mathcal{W}_t$  is a standard Brownian motion in  $\mathbb{R}^n$ . Similarly, recalling the definition of  $\xi = w_L^T W_{L - 1}\dots W_2u_1$ , the SDE of  $v_{1}$  is
$$
d v _ {1} = - \frac {2}{n} \xi X ^ {T} r d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \xi X ^ {T} d \mathcal {W} _ {t}. \tag {75}
$$
Generalization to batch-SGD. When  $b$  (a positive constant) data points  $\mathcal{B}_t$  are sampled in each iteration of SGD, i.e., batch-SGD, we can change the SGD update equation as follows (taking  $u_k$  as an example)
$$
u _ {k} (t + 1) = u _ {k} (t) - \frac {\eta}{b} \sum_ {j _ {t} \in \mathcal {B} _ {t}} \frac {\partial \ell_ {j _ {t}}}{\partial u _ {k}}.
$$
This only changes the noise  $\Sigma (u_k(t))$  to a batch version
$$
\Sigma_ {b} (u _ {k} (t)) = \frac {1}{b} \Sigma (u _ {k} (t)),
$$
which only affects the noise part of the SDE of  $u_{k}$  and leads it to become a batch version SDE:
$$
d u _ {k} = - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \rho_ {k} \rho_ {- k} v _ {k + 1} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n b}} (\rho_ {k} \rho_ {- k}) v _ {k + 1} v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t}.
$$
This is equivalent to re-scale the learning rate  $\eta$  to
$$
\eta_ {b} = \frac {\eta}{b},
$$
and leaving other parts unchanged. Thus the generalization to batch-SGD is straightforward—simply replacing all  $\eta$  with  $\eta_{b}$ .
D.2 The SDE of  $\theta$
In this section, we carefully derive the continuous dynamics of SGD for the parameterization of our rank-1 linear networks. We first discuss the balanced initialization condition.
Balanced initialization. Similar to the case for GD (Definition 2), we also assume the balanced initialization Eq. (26) across layers. Although the dynamics of SGD is different with that of GD, it still applies the gradient to update the parameters at every step that will maintain the balanced property thus the dynamics of SGD will also maintain the balanced property, i.e.,
$$
\frac {\left\langle v _ {k + 1} (t) , u _ {k} (t) \right\rangle^ {2}}{\| v _ {k + 1} (t) \| ^ {2} \| u _ {k} (t) \| ^ {2}} = 1
$$
$$
\left\| v _ {k + 1} (t) \right\| = \left\| u _ {k} (t) \right\| = \left\| v _ {1} (t) \right\|
$$
for  $t > 0$  and  $k \in \{1, \dots, L - 1\}$ , during the training of rank-1 linear networks.
With the equations for  $u_{k}$  and  $v_{k + 1}$ , we now derive the SDE of  $\theta$  summarized in the following lemma.
2More details of this modelling technique can be found in [14].
Lemma 3. For an  $L$ -layer rank-1 linear network Eq. (2), if we assume balanced initialization, then the stochastic gradient flow of  $\theta$  is
$$
\begin{array}{l} d \theta = - \frac {2 \xi^ {2}}{n} H (\theta) X ^ {T} r d t + 2 \xi^ {2} \sqrt {\frac {\eta \mathcal {L}}{n}} H (\theta) X ^ {T} d \mathcal {W} _ {t} \\ + \frac {8 \eta \mathcal {L} (L - 1)}{n \| \theta \| ^ {\frac {2}{2 L - 1}}} \left[ I + \frac {2 L - 3}{2} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} X \theta d t \\ \end{array}
$$
where  $H(\theta) = I + 2(L - 1)\theta \theta^T /\| \theta \| ^2$
Proof. According to the Ito's Lemma, we have
$$
d \theta = d \left(\xi v _ {1}\right) = d \xi v _ {1} + \xi d v _ {1} + d \xi d v _ {1}. \tag {76}
$$
Thus to obtain the SDE of  $\theta$ , we need to analyze every term of the above equation. We first give  $d\xi$ .
The form of  $d\xi$ . Let  $\omega_{k} = v_{k + 1}^{T}u_{k}$  and  $\psi_{k} = \|v_{k+1}\|^{2} + \|u_{k}\|^{2}$ , we first characterize the SDE of  $\omega_{k}$ . According to the Ito's calculus, we obtain that
$$
d \omega_ {k} = d \left(v _ {k + 1} ^ {T} u _ {k}\right) = \underbrace {u _ {k} ^ {T} d v _ {k + 1} + v _ {k + 1} ^ {T} d u _ {k}} _ {\text {串}} + \underbrace {d u _ {k} ^ {T} d v _ {k + 1}} _ {\diamond}, \tag {77}
$$
where, by applying Eq. (73) and Eq. (74) and noting that  $(d\mathcal{W}_t)^2 = dt$