SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\clubsuit = - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \rho_ {k} \rho_ {- k} \psi_ {k} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \rho_ {k} \rho_ {- k} \psi_ {k} v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t}
$$
$$
\diamondsuit = \frac {4 \eta \mathcal {L}}{n} \left(\rho_ {k} \rho_ {- k}\right) ^ {2} \omega_ {k} v _ {1} ^ {T} X ^ {T} X v _ {1} d t.
$$
Combining the above two terms gives us the SDE of  $\omega_{k}$
$$
d \omega_ {k} = \left[ - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \rho_ {k} \rho_ {- k} \psi_ {k} + \frac {4 \eta (\rho_ {k} \rho_ {- k}) ^ {2} \mathcal {L}}{n} \omega_ {k} v _ {1} ^ {T} X ^ {T} X v _ {1} \right] d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \rho_ {k} \rho_ {- k} \psi_ {k} v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t}.
$$
Since  $\xi = d(\prod_{k=1}^{L-1} \omega_k)$ , its SDE can be done by repeatedly applying the Ito's Lemma and the SDE of  $\omega_k$ :
$$
\begin{array}{l} d \xi = d (\prod_ {k = 1} ^ {L - 1} \omega_ {k}) \\ = \underbrace {\sum_ {k = 1} ^ {L - 1} \frac {\xi}{\omega_ {k}} d \omega_ {k}} _ {\text {♣}} + \frac {1}{2} \sum_ {k ^ {\prime}, k = 1, k \neq k ^ {\prime}} ^ {L - 1} \underbrace {\frac {\xi}{\omega_ {k} \omega_ {k ^ {\prime}}} d \omega_ {k} d \omega_ {k ^ {\prime}}} _ {\text {♥}}. \tag {78} \\ \end{array}
$$
For convenience, we first define several helper notations:
$$
\Phi_ {1} = \sum_ {k = 1} ^ {L - 1} \phi_ {k} = \sum_ {k = 1} ^ {L - 1} \frac {\psi_ {k}}{\omega_ {k} ^ {2}},
$$
$$
\Phi_ {2} = \sum_ {k = 1} ^ {L - 1} \frac {1}{\omega_ {k} ^ {2}},
$$
$$
\Phi_ {3} = \frac {1}{2} \sum_ {k, k ^ {\prime} = 1, k \neq k ^ {\prime}} ^ {L - 1} \frac {\psi_ {k} \psi_ {k ^ {\prime}}}{\omega_ {k} ^ {2} \omega_ {k ^ {\prime}} ^ {2}}.
$$
Now plugging the form of  $d\omega_{k}$  into  $\spadesuit$  gives us the first term of  $d\xi$ :
$$
\begin{array}{l} \spadesuit = - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \xi^ {2} \left(\sum_ {k = 1} ^ {L - 1} \frac {\psi_ {k}}{\omega_ {k} ^ {2}}\right) d t + \frac {4 \eta \mathcal {L} \xi^ {3} v _ {1} ^ {T} X ^ {T} X v _ {1}}{n} \left(\sum_ {k = 1} ^ {L - 1} \frac {1}{\omega_ {k} ^ {2}}\right) d t \\ + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \xi^ {2} \left(\sum_ {k = 1} ^ {L - 1} \frac {\psi_ {k}}{\omega_ {k} ^ {2}}\right) \left(v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t}\right) \\ = - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \xi^ {2} \Phi_ {1} d t + \frac {4 \eta \mathcal {L} \xi^ {3} v _ {1} ^ {T} X ^ {T} X v _ {1}}{n} \Phi_ {2} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \xi^ {2} \Phi_ {1} \left(v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t}\right), \tag {79} \\ \end{array}
$$
and applying again  $(d\mathcal{W}_t)^2 = dt$  and the form of  $d\omega_{k}$  gives us each term of the second sum of  $d\xi$
$$
\heartsuit = \frac {\xi}{\omega_ {k} \omega_ {k ^ {\prime}}} \frac {4 \eta \xi^ {2} \mathcal {L}}{n \omega_ {k} \omega_ {k ^ {\prime}}} \psi_ {k} \psi_ {k ^ {\prime}} v _ {1} ^ {T} X ^ {T} X v _ {1} d t.
$$
Summing all  $\hat{\mathbf{\Omega}}$  and  $\spadesuit$ , we obtain the SDE of  $\xi$ :
$$
\begin{array}{l} d \xi = \left[ - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \xi^ {2} \Phi_ {1} + \frac {4 \eta \mathcal {L} \xi^ {3} v _ {1} ^ {T} X ^ {T} X v _ {1}}{n} \Phi_ {2} + \frac {4 \eta \xi^ {3} \mathcal {L} v _ {1} ^ {T} X ^ {T} X v _ {1}}{n} \Phi_ {3} \right] d t \\ + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \xi^ {2} \Phi_ {1} \left(v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t}\right) \\ = \left[ - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \xi^ {2} \Phi_ {1} + \frac {4 \eta \mathcal {L} \xi^ {3} v _ {1} ^ {T} X ^ {T} X v _ {1}}{n} \left(\Phi_ {2} + \Phi_ {3}\right) \right] d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \xi^ {2} \Phi_ {1} \left(v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t}\right). \tag {80} \\ \end{array}
$$
The SDE of  $v_{1}$  is much simpler. To get this, we start with the SGD update equation for  $v_{1}$ :
$$
\begin{array}{l} v _ {1} (t + 1) = v _ {1} (t) - \eta \frac {\partial \mathcal {L}}{\partial v _ {1}} + \eta \left(\frac {\partial \mathcal {L}}{\partial v _ {1}} - \frac {\partial \ell_ {j _ {t}}}{\partial v _ {1}}\right) \\ = v _ {1} (t) - \frac {2 \eta}{n} \xi X ^ {T} r + 2 \eta \left(\frac {1}{n} \sum_ {i} \xi x _ {i} r _ {i} - r _ {j _ {t}} \xi x _ {j _ {t}}\right) \\ = v _ {1} (t) - \frac {2 \eta}{n} \xi X ^ {T} r + 2 \eta \xi X ^ {T} Z _ {j _ {t}}, \tag {81} \\ \end{array}
$$
which implies that the noise covariance in this case is
$$
\Sigma (v _ {1} (t)) = \frac {4 \eta^ {2} \xi^ {2} \mathcal {L}}{n} X ^ {T} X.
$$
Then using a similar approach as that of  $u_{k}$ , we get the SDE of  $v_{1}$
$$
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 {82}
$$
Now it is sufficient for us to derive the form of  $d\theta$
The form of  $d\theta$ . Combined with the SDE of  $\xi$ , we now have
$$
d \theta = d \left(\xi v _ {1}\right) = \underbrace {\xi d v _ {1} + v _ {1} d \xi} _ {\clubsuit} + \underbrace {d \xi d v _ {1}} _ {\clubsuit}. \tag {83}
$$
For the  $\clubsuit$  term, as we already have the form of  $dv_{1}$  in Eq. (82) and  $d\xi$  in Eq. (80), we simply plug them into  $\clubsuit$  and obtain that:
$$
\begin{array}{l} \clubsuit = - \frac {2}{n} \xi^ {2} X ^ {T} r d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \xi^ {2} (X ^ {T} d \mathcal {W} _ {t}) + v _ {1} \left[ - \frac {2}{n} v _ {1} ^ {T} X ^ {T} r \xi^ {2} \Phi_ {1} + \frac {4 \eta L \xi^ {3} v _ {1} ^ {T} X ^ {T} X v _ {1}}{n} (\Phi_ {2} + \Phi_ {3}) \right] d t \\ + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \Phi_ {1} \xi^ {2} v _ {1} v _ {1} ^ {T} X ^ {T} d \mathcal {W} _ {t} \\ = - \frac {2}{n} \left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) X ^ {T} r d t + \frac {4 \eta \mathcal {L} \xi^ {3} v _ {1} v _ {1} ^ {T} X ^ {T} X v _ {1}}{n} \left(\Phi_ {2} + \Phi_ {3}\right) d t \\ + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) X ^ {T} d \mathcal {W} _ {t}. \tag {84} \\ \end{array}
$$
For the  $\spadesuit$  term, we only need to consider the  $d\mathcal{W}_t$  terms of Eq. (82) and (80):
$$
\clubsuit = 4 \frac {\eta \mathcal {L} \xi^ {3}}{n} \Phi_ {1} X ^ {T} X v _ {1} d t.
$$
Combining the above two equations gives us the final SDE of  $d\theta$ :
$$
\begin{array}{l} d \theta = - \frac {2}{n} \left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) X ^ {T} r d t + \frac {4 \eta \mathcal {L}}{n} \left[ \xi^ {3} \Phi_ {1} I + \xi^ {3} (\Phi_ {2} + \Phi_ {3}) v _ {1} v _ {1} ^ {T} \right] X ^ {T} X v _ {1} d t \\ + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) X ^ {T} d \mathcal {W} _ {t}. \tag {85} \\ \end{array}
$$
On the other hand, recall that our assumptions regarding the initial conditions of  $f(x; u, v)$  in the Lemma (Definition 1) and following similar techniques as in the case for GD, we have that for any  $t > 0$ :
1.  $\forall k: \| u_{k}(t)\|^{2} = \| v_{k}(t)\|^{2} = \| v_{1}(t)\|^{2}$
2.  $\forall k:\omega_k^2 = \langle u_k(t),v_{k + 1}(t)\rangle^2 = \| u_k(t)\| ^4 = \| v_1(t)\| ^4.$
Plugging these terms back to the definitions of  $\Phi_1$ ,  $\Phi_2$  and  $\Phi_3$ , we obtain that
$$
\psi_ {k} = 2 \| v _ {1} \| ^ {2} \text {a n d} \xi^ {2} = \| \theta \| ^ {\frac {4 (L - 1)}{2 L - 1}} \tag {86}
$$
as in the case for GD and
$$
\Phi_ {1} = \frac {2 (L - 1)}{\| v _ {1} \| ^ {2}} = \frac {2 (L - 1) \xi^ {2}}{\| \theta \| ^ {2}}, \tag {87}
$$
$$
\Phi_ {2} = \frac {L - 1}{\| v _ {1} \| ^ {4}} = \frac {\xi^ {4} (L - 1)}{\| \theta \| ^ {4}}, \tag {88}
$$
$$
\Phi_ {3} = \frac {2 (L - 1) (L - 2)}{\| v _ {1} \| ^ {4}} = \frac {2 \xi^ {4} (L - 1) (L - 2)}{\| \theta \| ^ {4}}, \tag {89}
$$
where we use that  $\| \theta \| = \| \xi v_{1}\| = |\xi ||v_{1}|$ . Therefore, the final form of  $d\theta$  is now
$$
H (\theta) = I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}.
$$
D.3 Proof for Theorem 3
In this section, we determine the form of  $V^{\mathcal{S}}(\theta)$  such that  $\theta$  follows a stochastic mirror flow
$$
d \partial_ {\theta} V ^ {\mathcal {S}} (\theta , t) = - \frac {\partial \mathcal {L}}{\partial \theta} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \left(X ^ {T} d \mathcal {W} _ {t}\right), \tag {90}
$$
which then proves the claims of Theorem 3. For this purpose, we start with manipulating the SDE of  $\theta$  derived in Lemma 3. Note that Eq. (85) can be written as
$$
\left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) ^ {- 1} d \theta - \frac {4 \eta \mathcal {L}}{n} \mathcal {P} d t = - \frac {2}{n} X ^ {T} r d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} X ^ {T} d \mathcal {W} _ {t}, \tag {91}
$$
$$
\mathcal {P} = \left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) ^ {- 1} \left[ \xi^ {3} \Phi_ {1} I + \xi^ {3} (\Phi_ {2} + \Phi_ {3}) v _ {1} v _ {1} ^ {T} \right] X ^ {T} X v _ {1}.
$$
To solve the inverse appeared in the above equation, we apply Lemma 2 and noting that
$$
\mathfrak {t r} \left(\xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T} \xi^ {- 2} I\right) = \Phi_ {1} \| v _ {1} \| ^ {2},
$$
$$
\begin{array}{l} \left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) ^ {- 1} = \frac {1}{\xi^ {2}} I - \frac {1}{1 + \Phi_ {1} \| v _ {1} \| ^ {2}} \frac {\xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}}{\xi^ {4}} \\ = \frac {1}{\xi^ {2}} \left(I - \frac {\Phi_ {1} v _ {1} v _ {1} ^ {T}}{1 + \Phi_ {1} \| v _ {1} \| ^ {2}}\right) \\ = \frac {1}{\xi^ {2}} \left(I - \frac {1}{1 + \frac {\xi^ {2}}{\Phi_ {1} \| \theta \| ^ {2}}} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right), \\ \end{array}
$$
where we use that  $\theta = \xi v_{1}$  in the last equality, which enables us to simplify  $\mathcal{P}$ :
$$
\begin{array}{l} \mathcal {P} = \left(\xi^ {2} I + \xi^ {2} \Phi_ {1} v _ {1} v _ {1} ^ {T}\right) ^ {- 1} \left[ \xi^ {3} \Phi_ {1} I + \xi^ {3} \left(\Phi_ {2} + \Phi_ {3}\right) v _ {1} v _ {1} ^ {T} \right] X ^ {T} X v _ {1} \\ = \Phi_ {1} \left(I - \frac {1}{1 + \frac {\xi^ {2}}{\Phi_ {1} \| \theta \| ^ {2}}} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \left[ I + \frac {\Phi_ {2} + \Phi_ {3}}{\xi^ {2} \Phi_ {1}} \theta \theta^ {T} \right] X ^ {T} X \theta \\ = \Phi_ {1} \left[ I - \left(\frac {1}{\| \theta \| ^ {2} + \frac {\xi^ {2}}{\Phi_ {1}}} - \frac {\Phi_ {2} + \Phi_ {3}}{\xi^ {2} \Phi_ {1}} + \frac {\| \theta \| ^ {2}}{\| \theta \| ^ {2} + \frac {\xi^ {2}}{\Phi_ {1}}} \frac {\Phi_ {2} + \Phi_ {3}}{\xi^ {2} \Phi_ {1}}\right) \theta \theta^ {T} \right] X ^ {T} X \theta \\ = \Phi_ {1} \left[ I - \frac {\Phi_ {1} ^ {2} - \Phi_ {2} - \Phi_ {3}}{\Phi_ {1} ^ {2} \| \theta \| ^ {2} + \xi^ {2} \Phi_ {1}} \theta \theta^ {T} \right] X ^ {T} X \theta . \\ \end{array}
$$
Thus Eq. (91) for the overall SDE of  $\theta$  now becomes
$$
\frac {1}{\xi^ {2}} \left(I - \frac {1}{1 + \frac {\xi^ {2}}{\Phi_ {1} \| \theta \| ^ {2}}} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) d \theta - \frac {4 \eta \mathcal {L}}{n} \mathcal {P} d t = - \frac {\partial \mathcal {L}}{\partial \theta} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} (X ^ {T} d \mathcal {W} _ {t}). \tag {92}
$$
The balanced initialization gives us Eq. (87), (88), and (89), and recall that
$$
\lambda_ {L} = \frac {2 (L - 1)}{2 L - 1},
$$
thus we can further rewrite
$$
\begin{array}{l} \frac {1}{\xi^ {2}} \left(I - \frac {1}{1 + \frac {\xi^ {2}}{\Phi_ {1} \| \theta \| ^ {2}}} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) = \frac {1}{\xi^ {2}} \left(I - \frac {1}{1 + \frac {1}{2 (L - 1)}} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \\ = \frac {1}{\xi^ {2}} \left(I - \lambda_ {L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \tag {93} \\ \end{array}
$$
$$
\begin{array}{l} \mathcal {P} = \frac {1}{\xi^ {2}} \left(I - \lambda_ {L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \left[ \xi^ {2} \Phi_ {1} I + (\Phi_ {2} + \Phi_ {3}) \theta \theta^ {T} \right] X ^ {T} X \theta \\ = \frac {1}{\xi^ {2}} \left[ \xi^ {2} \Phi_ {1} I + \left((\Phi_ {2} + \Phi_ {3}) (1 - \lambda_ {L}) - \frac {\lambda_ {L} \xi^ {2} \Phi_ {1}}{\| \theta \| ^ {2}}\right) \theta \theta^ {T} \right] X ^ {T} X \theta \\ = \frac {1}{\xi^ {2}} \left[ \frac {2 (L - 1) \xi^ {4}}{\| \theta \| ^ {2}} I + \left((L - 1 + 4 (L - 1) (L - 2)) (1 - \lambda_ {L}) - 2 (L - 1) \lambda_ {L}\right) \frac {\xi^ {4}}{\| \theta \| ^ {4}} \theta \theta^ {T} \right] X ^ {T} X \theta \\ = \frac {\xi^ {2}}{\| \theta \| ^ {2}} \left[ 2 (L - 1) I - \frac {3 (L - 1)}{2 L - 1} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} X \theta \\ = \frac {2 (L - 1) \xi^ {2}}{\| \theta \| ^ {2}} \left(I - \frac {\theta \theta^ {T}}{2 \| \theta \| ^ {2}}\right) X ^ {T} X \theta . \tag {94} \\ \end{array}
$$
Thus, Eq. (91) now<sup>3</sup> can be rewritten as
$$
\begin{array}{l} \frac {1}{\xi^ {2}} \left(I - \lambda_ {L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) d \theta - \frac {8 \eta \mathcal {L}}{n} \frac {(L - 1) \xi^ {2}}{\| \theta \| ^ {2}} \left(I - \frac {1}{2} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} X \theta d t \\ = - \frac {\partial \mathcal {L}}{\partial \theta} d t + 2 \sqrt {\frac {\eta \mathcal {L}}{n}} \left(X ^ {T} d \mathcal {W} _ {t}\right). \tag {95} \\ \end{array}
$$
Finding the form of  $V^S(\theta)$ . We now proceed to find the form of  $V^S(\theta)$ . For convenience, we apply a time re-scaling technique such that  $d\tau = |\xi|dt$  ( $d\mathcal{W}_{\tau} = \sqrt{|\xi|}d\mathcal{W}_{t}$  according to [12]) to the above equation. For convenience, we still use  $t$  to represent the time after re-scaling. Then the above equation becomes
$$
\begin{array}{l} \frac {1}{\| \theta \| ^ {\frac {2 (L - 1)}{2 L - 1}}} \left(I - \lambda_ {L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) d \theta - \frac {8 \eta \mathcal {L}}{n} \frac {(L - 1) \xi^ {2}}{\| \theta \| ^ {2}} \left(I - \frac {1}{2} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} X \theta d t \\ = - \frac {\partial \mathcal {L}}{\partial \theta} d t + 2 \sqrt {\frac {\eta | \xi | \mathcal {L}}{n}} \left(X ^ {T} d \mathcal {W} _ {t}\right). \tag {96} \\ \end{array}
$$
Recall that  $\theta(0) \in \mathbb{R}^d$  is the initialization of  $\theta$ , since  $V^S(\theta)$  for SGD should have similar form as that for GD, we borrow from the GD results and first define a constant vector  $\gamma \in \mathbb{R}^d$
$$
\gamma = - \| \theta (0) \| ^ {- \frac {2 L - 2}{2 L - 1}} \theta (0).
$$
Suppose now that  $V^{\mathcal{S}}$  has the following form:
$$
V ^ {\mathcal {S}} (\theta , t) = \frac {2 L - 1}{2 L} \| \theta \| ^ {\frac {2 L}{2 L - 1}} + \gamma^ {T} \theta + g (t) ^ {T} \theta \text {f o r} g \in \mathbb {R} ^ {d}, \tag {97}
$$
by similar techniques as in the case of GD, we obtain the first and second derivatives of  $V^{\mathcal{S}}(\theta)$  w.r.t  $\theta$ :
$$
\partial_ {\theta} V ^ {\mathcal {S}} (\theta , t) = \| \theta \| ^ {- \frac {2 (L - 1)}{2 L - 1}} \theta + \gamma + g (t),
$$
$$
\partial_ {\theta} ^ {2} V ^ {\mathcal {S}} (\theta , t) = \frac {1}{\| \theta \| ^ {\frac {2 (L - 1)}{2 L - 1}}} \left[ I - \lambda_ {L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right].
$$
There exists a function of  $\theta$ ,  $G(\theta)$  whose exact form can be derived but not necessary for us, corresponding to our  $V^{\mathcal{S}}(\theta, t)$  defined above, such that  $d\partial_{\theta}V^{\mathcal{S}}(\theta, t)$  can be written as
$$
\begin{array}{l} d \partial_ {\theta} V ^ {\mathcal {S}} (\theta) = \partial_ {\theta} \left(\partial_ {\theta} V ^ {\mathcal {S}} (\theta , t)\right) d \theta + \partial_ {t} \left(\partial_ {\theta} V ^ {\mathcal {S}} (\theta , t)\right) d t + \frac {\partial^ {2}}{\partial \theta \partial \theta} \partial_ {\theta} V ^ {\mathcal {S}} (\theta , t) (d \theta) ^ {2} \\ = \partial_ {\theta} \left(\partial_ {\theta} V ^ {\mathcal {S}} (\theta , t)\right) d \theta + \left[ \partial_ {t} \left(\partial_ {\theta} V ^ {\mathcal {S}} (\theta , t)\right) + G (\theta) \frac {\partial^ {2}}{\partial \theta \partial \theta} \partial_ {\theta} V ^ {\mathcal {S}} (\theta , t) \right] d t \\ \end{array}
$$
Now suppose that we can choose a particular  $g(t)$ , thus a particular  $G(\theta)$ , such that the following relation is satisfied:
$$
\begin{array}{l} \frac {\partial}{\partial t} \partial_ {\theta} V ^ {S} (\theta , t) + G (\theta) \frac {\partial^ {2}}{\partial \theta \partial \theta} \partial_ {\theta} V ^ {S} (\theta , t) \\ = - \frac {8 \eta \mathcal {L}}{n} \frac {(L - 1) \xi^ {2}}{\| \theta \| ^ {2}} \left(I - \frac {1}{2} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} X \theta , \tag {98} \\ \end{array}
$$
then we can rewrite the SDE of  $\partial_{\theta}V^{\mathcal{S}}(\theta ,t)$  as follows:
$$
\begin{array}{l} d \partial_ {\theta} V ^ {S} (\theta) = \partial_ {\theta} \left(\partial_ {\theta} V ^ {S} (\theta , t)\right) d \theta + \partial_ {t} \left(\partial_ {\theta} V ^ {S} (\theta , t)\right) d t + G (\theta) \frac {\partial^ {2}}{\partial \theta \partial \theta} \partial_ {\theta} V ^ {S} (\theta , t) d t \\ = \frac {1}{\left\| \theta \right\| ^ {\frac {2 (L - 1)}{2 L - 1}}} \left(I - \lambda_ {L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) d \theta - \frac {8 \eta \mathcal {L}}{n} \frac {(L - 1) \xi^ {2}}{\| \theta \| ^ {2}} \left(I - \frac {1}{2} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} X \theta d t \tag {99} \\ \end{array}
$$
Note that when  $L = 1$  we do not have the second  $dt$  term in the LHS of the equation, this term is brought by adding layers to the model.
which will give us the desired "mirror flow" equation:
$$
d \partial_ {\theta} V ^ {\mathcal {S}} (\theta) = - \nabla_ {\theta} \mathcal {L} d t + 2 \sqrt {\frac {\eta | \xi | \mathcal {L}}{n}} X ^ {T} d \mathcal {W} _ {t}.
$$
It is now suffice for us to find the particular  $g(t)$  that makes Eq. (98) satisfied. For convenience, recall that for a vector  $a$  we use  $a_{\mu}$  to denote its  $\mu$ -th component, we define the following helper notations:
$$
k = 2 \xi^ {2} \sqrt {\frac {\eta \mathcal {L}}{n}} \in \mathbb {R}, \tag {100}
$$
$$
B = \left[ I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} \in \mathbb {R} ^ {d \times n}, \tag {101}
$$
$$
d \theta_ {\mu} = O _ {\mu} d t + \sum_ {i} B _ {\mu i} d \mathcal {W} _ {t, i}, \tag {102}
$$
$$
D _ {\mu} (\theta , t) = \left(\frac {\partial V ^ {S} (\theta , t)}{\partial \theta}\right) _ {\mu} = \| \theta \| ^ {- \lambda_ {L}} \theta_ {\mu} + \gamma_ {\mu} + g _ {\mu} (t), \quad D (\theta , t) \in \mathbb {R} ^ {d}, \tag {103}
$$
$$
\delta_ {\rho \sigma} = 1 \text {i f} \rho = \sigma \text {o t h e r w i s e} \delta_ {\rho \sigma} = 0, \tag {104}
$$
where the exact form of  $O$  can be obtained from Lemma 3 but is not necessary for us. In the following, we use  $D$  to represent  $D(\theta, t)$  for convenience. For our purpose of choosing the particular  $g(t)$ ,  $dD(\theta, t) = d\partial_{\theta}V^{S}(\theta, t)$  should match the R.H.S of Eq. (99). To make this relation clear, according to the Ito's Lemma, we obtain that
$$
d D = \partial_ {t} D d t + \partial_ {\theta} D d \theta + \underbrace {\frac {1}{2} \sum_ {\rho , \sigma} \frac {\partial^ {2} D}{\partial \theta_ {\rho} \partial \theta_ {\sigma}} d \theta_ {\rho} d \theta_ {\sigma}} _ {\spadesuit},
$$
where the last term is crucial. From the SDE of  $\theta$  (Lemma 3), we have that the  $d\theta_{\rho}d\theta_{\sigma}$  appeared in  $\spadesuit$  can be written as
$$
d \theta_ {\rho} d \theta_ {\sigma} = k ^ {2} \sum_ {j, i} B _ {j \sigma} B _ {\rho i} d \mathcal {W} _ {t, i} d \mathcal {W} _ {t, j} = k ^ {2} \sum_ {i} B _ {i \sigma} B _ {\rho i} d t.
$$
On the other hand, according to
$$
\partial_ {\theta} D = \partial_ {\theta} ^ {2} V ^ {\mathcal {S}} (\theta , t) = \frac {1}{\| \theta \| ^ {\lambda_ {L}}} \left(I - \lambda_ {L} \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right),
$$
we have that, with explicit subscripts of vectors,
$$
\begin{array}{l} \frac {\partial^ {2} D _ {\mu}}{\partial \theta_ {\rho} \partial \theta_ {\sigma}} = \frac {\partial}{\partial \theta_ {\sigma}} \left[ \| \theta \| ^ {- \lambda_ {L}} \left(\delta_ {\mu \rho} - \lambda_ {L} \frac {\theta_ {\rho} \theta_ {\mu}}{\| \theta \| ^ {2}}\right) \right] \\ = - \lambda_ {L} \| \theta \| ^ {- \lambda_ {L} - 2} \delta_ {\mu \rho} \theta_ {\sigma} - \lambda_ {L} \frac {\partial}{\partial \theta_ {\sigma}} \left(\| \theta \| ^ {- \lambda_ {L} - 2} \theta_ {\rho} \theta_ {\mu}\right) \\ = - \lambda_ {L} \| \theta \| ^ {- \lambda_ {L} - 2} \delta_ {\mu \rho} \theta_ {\sigma} \\ - \lambda_ {L} \left[ - (\lambda_ {L} + 2) \| \theta \| ^ {- \lambda_ {L} - 4} \theta_ {\sigma} \theta_ {\mu} \theta_ {\rho} + \| \theta \| ^ {- \lambda_ {L} - 2} \left(\delta_ {\rho \sigma} \theta_ {\mu} + \delta_ {\mu \sigma} \theta_ {\rho}\right) \right] \\ = - \frac {\lambda_ {L}}{\| \theta \| ^ {\lambda_ {L} + 2}} \left(\delta_ {\mu \rho} \theta_ {\sigma} + \delta_ {\mu \sigma} \theta_ {\rho}\right) - \frac {\lambda_ {L}}{\| \theta \| ^ {\lambda_ {L} + 2}} \theta_ {\mu} \left(\delta_ {\rho \sigma} - \left(\lambda_ {L} + 2\right) \frac {\theta_ {\rho} \theta_ {\sigma}}{\| \theta \| ^ {2}}\right). \tag {105} \\ \end{array}
$$
Note that the above expression implies that  $\frac{\partial^2D}{\partial\theta\partial\theta}$  is in fact a rank-3 tensor. Combined with the expression of  $d\theta_{\rho}d\theta_{\sigma}$  derived above, we have that
$$
\spadesuit = - \frac {k ^ {2} \lambda_ {L}}{2 \| \theta \| ^ {\lambda_ {L} + 2}} \sum_ {\rho , \sigma} \sum_ {i} \left[ \underbrace {\delta_ {\mu \rho} \theta_ {\sigma} + \delta_ {\mu \sigma} \theta_ {\rho}} _ {H _ {\mu \sigma} ^ {\rho}} + \theta^ {\mu} \underbrace {\left(\delta_ {\rho \sigma} - (\lambda_ {L} + 2) \frac {\theta_ {\rho} \theta_ {\sigma}}{\| \theta \| ^ {2}}\right)} _ {P _ {\rho \sigma}} \right] B _ {\sigma i} B _ {i \rho} d t \tag {106}
$$
where  $H \in \mathbb{R}^{d \times d \times d}$  is a rank-3 tensor. To get the exact form of  $\spadesuit$ , we need to have the exact forms of
1.  $\sum_{\rho, \sigma, i} H_{\mu \sigma}^{\rho} B_{\sigma i} B_{i\rho}$ . Note that  $HBB \in \mathbb{R}^{d}$  and there are two different terms, both of which will induce a same vector. In particular,
$$
\sum_ {\rho , \sigma , i} \delta_ {\rho \mu} \theta_ {\sigma} B _ {\sigma i} B _ {i \rho} = \sum_ {\rho \sigma , i} B _ {\mu i} B _ {i \sigma} \theta_ {\sigma} = (B B ^ {T} \theta) _ {\mu}
$$
$$
\sum_ {\rho , \sigma , i} \delta_ {\mu \sigma} \theta_ {\rho} B _ {\sigma i} B _ {i \rho} = ((\boldsymbol {\theta} ^ {T} B B ^ {T}) ^ {T}) _ {\mu} = (B B ^ {T} \boldsymbol {\theta}) _ {\mu}.
$$
$$
\sum_ {\rho , \sigma , i} H _ {\mu \sigma} ^ {\rho} B _ {\sigma i} B _ {i \rho} = 2 (B ^ {T} B \theta) _ {\mu},
$$
$$
\begin{array}{l} B B ^ {T} \theta = \left(I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} X \left(I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \theta \\ = \left(I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} X (\theta + 2 (L - 1) \theta) \\ = (2 L - 1) \left[ I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} X \theta . \tag {107} \\ \end{array}
$$
2.  $\sum_{\rho, \sigma, i} P_{\rho}^{\sigma} B_{\sigma}^{i} B_{i}^{\rho}$ . Using the matrix notation, we can easily find that
$$
\sum_ {\rho , \sigma , i} P _ {\rho} ^ {\sigma} B _ {\sigma} ^ {i} B _ {i} ^ {\rho} = \operatorname {t r} \left(B ^ {T} P B\right),
$$
where, recall the definition of  $B$  in Eq. (101),
$$
\begin{array}{l} B ^ {T} P B = X \left(I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \left(I - (\lambda_ {L} + 2) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) \left(I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}}\right) X ^ {T} \\ = X \left[ I + (2 (L - 1) - (\lambda_ {L} + 2) - 2 (L - 1) (\lambda_ {L} + 2)) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] \\ \times \left[ I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} \\ = X \left[ I - 2 (2 L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] \left[ I + 2 (L - 1) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T} \\ = X \left[ I - (8 L ^ {2} - 1 0 L + 4) \frac {\theta \theta^ {T}}{\| \theta \| ^ {2}} \right] X ^ {T}. \tag {108} \\ \end{array}
$$
Taking the trace of the above equation gives us the second term of  $\spadesuit$ :