SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \frac {\partial}{\partial t} \mathbb {E} _ {p _ {t}} [ \Phi (\boldsymbol {W}) ] = \lim  _ {\eta \rightarrow 0} \frac {\mathbb {E} _ {p _ {t}} [ \langle \nabla \Phi (\boldsymbol {W}) , - \eta \nabla \mathcal {L} (\boldsymbol {W} ; \mathcal {D}) + \sqrt {2 \sigma^ {2} \eta} Z \rangle + \sigma^ {2} \eta Z ^ {\top} \nabla^ {2} \Phi (\boldsymbol {W}) Z + o (\eta) ]}{\eta} \\ = - \mathbb {E} _ {p _ {t}} [ \langle \nabla \Phi (\boldsymbol {W}), \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D}) \rangle ] + \sigma^ {2} \mathbb {E} _ {p _ {t}} [ \operatorname {T r} (\nabla^ {2} \Phi (\boldsymbol {W})) ]. \tag {24} \\ \end{array}
$$
By plugging  $\Phi (\pmb {W}) = \| \pmb {W} - \pmb{W}_0\|_{M_T}^2$  into Eq. (24), we have that
$$
\begin{array}{l} \frac {\partial}{\partial t} \mathbb {E} _ {p _ {t}} [ \| \boldsymbol {W} - \boldsymbol {W} _ {0} \| _ {M _ {T}} ^ {2} ] \leq - 2 \mathbb {E} _ {p _ {t}} [ \langle \boldsymbol {W} - \boldsymbol {W} _ {0}, \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D}) \rangle ] + 2 \sigma^ {2} \operatorname {r a n k} (M _ {T}) (25) \\ = - 2 \mathbb {E} _ {p _ {t}} [ \langle \boldsymbol {W} - \boldsymbol {W} _ {0}, \nabla \mathcal {L} (\boldsymbol {W} _ {0}; \mathcal {D}) \rangle ] + 2 \sigma^ {2} \operatorname {r a n k} (M _ {T}) \\ - 2 \mathbb {E} _ {p _ {t}} [ \langle \boldsymbol {W} - \boldsymbol {W} _ {0}, \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D}) - \nabla \mathcal {L} (\boldsymbol {W} _ {0}; \mathcal {D}) \rangle ] \\ \leq \mathbb {E} [ \| \nabla \mathcal {L} (\boldsymbol {W} _ {0}; \mathcal {D}) \| _ {2} ^ {2} ] + \mathbb {E} _ {p _ {t}} [ \| \boldsymbol {W} - \boldsymbol {W} _ {0} \| _ {M _ {t}} ^ {2} ] + 2 \sigma^ {2} \operatorname {r a n k} (M _ {T}) (26) \\ + \mathbb {E} _ {p _ {t}} [ \| \boldsymbol {W} - \boldsymbol {W} _ {0} \| _ {M _ {T}} ^ {2} ] + \mathbb {E} _ {p _ {t}} [ \| \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D}) - \nabla \mathcal {L} (\boldsymbol {W} _ {0}; \mathcal {D}) \| _ {2} ^ {2} ], (27) \\ \end{array}
$$
where the last inequality is by using the Cauchy-schwartz inequality. By plugging the assumption that  $\| \nabla \ell (f_{\pmb{W}}(\pmb {x});\pmb {y})) - \nabla \ell (f_{\pmb{W}^{\prime}}(\pmb {x});\pmb {y})\|_{2} <   \max \{c,\beta \| \pmb {W} - \pmb{W}^{\prime}\|_{2}\}$  into the above inequality, we have that
$$
\frac {\partial}{\partial t} \mathbb {E} _ {p _ {t}} [ \| \boldsymbol {W} - \boldsymbol {W} _ {0} \| _ {M _ {T}} ^ {2} ] \leq (2 + \beta^ {2}) \mathbb {E} _ {p _ {t}} [ \| \boldsymbol {W} - \boldsymbol {W} _ {0} \| _ {M _ {T}} ^ {2} ] + \mathbb {E} [ \| \nabla \mathcal {L} (\boldsymbol {W} _ {0}; \mathcal {D}) \| _ {2} ^ {2} ] + 2 \sigma^ {2} \operatorname {r a n k} (M _ {T}) + c ^ {2}. \tag {28}
$$
By solving the above ordinary differential inequality on  $t \in [0, T]$ , we have that
$$
\mathbb {E} _ {p _ {t}} \left[ \| \boldsymbol {W} - \boldsymbol {W} _ {0} \| _ {M _ {T}} ^ {2} \right] \leq \frac {e ^ {(2 + \beta^ {2}) t} - 1}{2 + \beta^ {2}} \left(\mathbb {E} [ \| \nabla \mathcal {L} (\boldsymbol {W} _ {0}; \mathcal {D}) \| _ {2} ^ {2} ] + 2 \sigma^ {2} \operatorname {r a n k} (M _ {T}) + c ^ {2}\right). \tag {29}
$$
By plugging Eq. (29) into Eq. (21) and Eq. (20), followed by integration over time  $t \in [0, T]$ , we have that
$$
\begin{array}{l} \int_ {0} ^ {T} \mathbb {E} _ {p _ {t}} \left[ \| \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D}) - \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D} ^ {\prime}) \| _ {2} ^ {2} \right] \mathrm {d} t \leq 2 T \cdot \mathbb {E} _ {p _ {0}} \left[ \| \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D}) - \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D} ^ {\prime}) \| _ {2} ^ {2} \right] (30) \\ + \frac {2 \beta^ {2}}{n ^ {2} (2 + \beta^ {2})} \left(\frac {e ^ {(2 + \beta^ {2}) T} - 1}{2 + \beta^ {2}} - T\right) \cdot \left(\mathbb {E} _ {p _ {0}} \left[ \| \nabla \mathcal {L} (\boldsymbol {W}; \mathcal {D}) \| _ {2} ^ {2} \right] + 2 \sigma^ {2} \operatorname {r a n k} \left(M _ {T}\right) + c ^ {2}\right) + \frac {2 c ^ {2} T}{n ^ {2}}, (31) \\ \end{array}
$$
C Deferred proofs for Section 4
C.1 Bounding the gradient norm at initialization
To bound the moment of  $\ell_2$  norm of the gradient  $\frac{\partial f(x)}{\partial W_l}$  of network output function, we need the following (extended) lemmas from Zhu et al. [54].
Lemma C.1 ([54, Lemma 1]). Let  $\pmb{w} \sim \mathcal{N}(0, \sigma^2\mathbb{I}_n)$ , then for two fixed non-zero vectors  $\pmb{h}_1, \pmb{h}_2 \in \mathbb{R}^n$  whose correlation is unknown, define two random variables  $X = (\pmb{w}^\top \pmb{h}_11_{\{\pmb{w}^\top \pmb{h}_2 \geq 0\}})^2$  and  $Y = s(\pmb{w}^\top \pmb{h}_1)^2$ , where  $s \sim \text{Ber}(1, 1/2)$  follows a Bernoulli distribution with 1 trial and  $\frac{1}{2}$  success rate, and  $s$  and  $\pmb{w}$  are independent random variables. Then  $X$  and  $Y$  have the same distribution.
Lemma C.2 (Extension of [54, Lemma 2]). Given a fixed non-zero matrix  $\pmb{H}_1 \in \mathbb{R}^{p \times r}$  and a fixed non-zero vector  $\pmb{h}_2 \in \mathbb{R}^p$  and, let  $\pmb{W} \in \mathbb{R}^{q \times p}$  be a random matrix with i.i.d. entries  $W_{ij} \sim \mathcal{N}(0, \beta)$  and a matrix (or vector)  $V = \phi'(W\pmb{h}_2)W\pmb{H}_1 \in \mathbb{R}^{q \times r}$ , then, we have  $\mathbb{E}\left[\frac{\|\pmb{V}\|_F^2}{\|\pmb{H}_1\|_F^2}\right] = \frac{q\beta}{2}$ .
Proof. According to the definition of  $\mathbf{V} = \phi^{\prime}(\mathbf{W}\mathbf{h}_{2})\mathbf{W}\mathbf{H}_{1}\in \mathbb{R}^{q\times r}$ , we have:
$$
\| \boldsymbol {V} \| _ {F} ^ {2} = \sum_ {i = 1} ^ {q} \sum_ {j = 1} ^ {r} \left(\boldsymbol {D} _ {i, i} \langle \boldsymbol {W} ^ {[ i; ]}, \boldsymbol {H} _ {1} ^ {[; j ]} \rangle\right) ^ {2},
$$
where  $D_{i,i} = 1_{\{\langle \pmb{W}^{[i]},\pmb{h}_2\rangle \geq 0\}}$ $\pmb{W}^{[i]}$  is the  $i$ -th row of  $\pmb{W}$ , and  $\pmb{H}_1^{[;j]}$  is the  $j$ -th column vector of  $\pmb{H}_1$ . Therefore by Lemma C.1, with i.i.d. Bernoulli random variable  $\rho_1,\dots ,\rho_q\sim \mathrm{Ber}(1,1 / 2)$ , we have
$$
\left\| \boldsymbol {V} \right\| _ {F} ^ {2} = \sum_ {i = 1} ^ {q} \sum_ {j = 1} ^ {r} \rho_ {i} \left\langle \boldsymbol {W} ^ {[ i; ]}, \boldsymbol {H} _ {1} ^ {[; j ]} \right\rangle^ {2} = \sum_ {i = 1} ^ {q} \sum_ {j = 1} ^ {r} \rho_ {i} \beta \left\| \boldsymbol {H} _ {1} ^ {[; j ]} \right\| _ {2} ^ {2} \tilde {w} _ {i j} ^ {2}, \tag {32}
$$
where  $\tilde{w}_{ij} = \langle \pmb{W}^{[i;]},\pmb{H}_1^{[j;]}\rangle /\left(\sqrt{\beta\|\pmb{H}_1^{[j;]} \|_2^2}\right)$ . By the fact that  $\pmb{W}^{[i;]}$  has i.i.d. Gaussian entries, for any fixed  $j$ , we have that  $\tilde{w}_{ij}\sim \mathcal{N}(0,1),i = 1,\dots ,q$  independently. Therefore, we have
$$
\mathbb {E} \left[ \| \boldsymbol {V} \| _ {F} ^ {2} \right] = \sum_ {i = 1} ^ {q} \sum_ {j = 1} ^ {r} \mathbb {E} [ \rho_ {i} ] \beta \| \boldsymbol {H} _ {1} ^ {[; j ]} \| _ {2} ^ {2} \mathbb {E} [ \tilde {w} _ {i j} ^ {2} ] = \frac {q \beta}{2} \mathbb {E} [ \| \boldsymbol {H} _ {1} \| _ {F} ^ {2} ].
$$
Now, we are ready to prove output gradient expectation at random initialization as follows.
Lemma C.3 (Output Gradient Expectation Bound at Random Initialization). Fix any data record  $\mathbf{x}$  then over the randomness of the initialization distributions for  $W_{1},\dots ,W_{L}$ , i.e.,  $W_{l}\sim \mathcal{N}(0,\beta_{l}\mathbb{I})$  for  $l = 1,\dots ,L - 1$ , it satisfies that
$$
\mathbb {E} _ {\boldsymbol {W}} \left[ \| \frac {\partial f (\boldsymbol {x})}{\partial \operatorname {V e c} (\boldsymbol {W})} \| _ {F} ^ {2} \right] = \| \boldsymbol {x} \| _ {2} ^ {2} o \left(\prod_ {i = 1} ^ {L - 1} \frac {\beta_ {i} m _ {i}}{2}\right) \sum_ {l = 1} ^ {L} \frac {\beta_ {L}}{\beta_ {l}}. \tag {33}
$$
Proof. We use  $\operatorname{Vec}(\mathbf{W}_l)$  to denote the concatenation of all row vector of the parameter matrix  $\mathbf{W}_l$ . By chain rule, for  $l = 1, \dots, L - 1$ , we have that
$$
\begin{array}{l} \frac {\partial \boldsymbol {f} (\boldsymbol {x})}{\partial \operatorname {V e c} \left(\boldsymbol {W} _ {l}\right)} = \frac {\partial h _ {L} (\boldsymbol {x})}{\partial h _ {L - 1} (\boldsymbol {x})} \left(\prod_ {i = 1} ^ {L - 1 - l} \frac {\partial h _ {L - i} (\boldsymbol {x})}{\partial h _ {L - 1 - i} (\boldsymbol {x})}\right) \frac {\partial h _ {l}}{\operatorname {V e c} \left(\boldsymbol {W} _ {l}\right)} (34) \\ = \boldsymbol {W} _ {L} \left(\prod_ {i = 1} ^ {L - 1 - l} \sigma_ {L - i} ^ {\prime} \boldsymbol {W} _ {L - i}\right) \sigma_ {l} ^ {\prime} \left( \begin{array}{c c c} h _ {l - 1} ^ {\top} & 0 & \dots \\ \vdots & \vdots & \vdots \\ 0 & 0 & h _ {l - 1} ^ {\top} \end{array} \right) _ {m _ {l} \times m _ {l} m _ {l - 1}}. (35) \\ \end{array}
$$
Similarly, for the  $L$ -th layer, we have that
$$
\frac {\partial \boldsymbol {f} (\boldsymbol {x})}{\partial \operatorname {V e c} \left(\boldsymbol {W} _ {L}\right)} = \left( \begin{array}{c c c} h _ {L - 1} ^ {\top} & 0 & \dots \\ \vdots & \vdots & \vdots \\ 0 & 0 & h _ {L - 1} ^ {\top} \end{array} \right) _ {o \times o m _ {L - 1}}. \tag {36}
$$
By properties of ReLU activation  $\phi$ , we have  $\phi_{L - i}^{\prime} = \mathrm{diag}[\mathrm{sgn}(W_{L - i}h_{L - 1 - i})]$ , where  $\mathrm{sgn}(x) = \left\{ \begin{array}{ll}1 & x > 0\\ 0 & x\leq 0 \end{array} \right.$  operates coordinate-wise with regard to the input matrix. Therefore, we have that for  $l = 1,\dots ,L - 1$
$$
\frac{\partial\boldsymbol{f}(\boldsymbol{x})}{\partial\operatorname{Vec}(\boldsymbol{W}_{l})} = \boldsymbol{W}_{L}\left(\prod_{i = 1}^{L - 1 - l}\operatorname {diag}[\operatorname {sgn}(W_{L - i}h_{L - 1 - i})]\boldsymbol{W}_{L - i}\right)\cdot \operatorname {diag}[\operatorname {sgn}(W_{l}h_{l - 1})]\left( \begin{array}{ccc}h_{l - 1}^{\top} & 0 & \dots \\ \vdots & \vdots & \vdots \\ 0 & 0 & h_{l - 1}^{\top} \end{array} \right)_{m_{l}\times m_{l}m_{l - 1}}.
$$
For notational simplicity, we introduce the notation of  $\pmb{t}_l^{l'}$  for  $l = 1, \dots, L - 1$  and  $l \leq l' < L$  as follows.
$$
\boldsymbol{t}_{l}^{l^{\prime}}:= \left(\prod_{i = L - l^{\prime}}^{L - 1 - l}\operatorname {diag}[\operatorname {sgn}(W_{L - i}h_{L - 1 - i})]\boldsymbol{W}_{L - i}\right)\cdot \operatorname {diag}[\operatorname {sgn}(W_{l}h_{l - 1})]\left( \begin{array}{ccc}h_{l - 1}^{\top} & 0 & \dots \\ \vdots & \vdots & \vdots \\ 0 & 0 & h_{l - 1}^{\top} \end{array} \right)_{m_{l}\times m_{l}m_{l - 1}}
$$
Then by definition, we have that
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {W}} \left[ \| \frac {\partial f (\boldsymbol {x})}{\partial \operatorname {V e c} (\boldsymbol {W} _ {l})} \| _ {F} ^ {2} \right] = \mathbb {E} _ {\boldsymbol {W}} \left[ \| \boldsymbol {W} _ {L} \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2} \right] = \mathbb {E} _ {\boldsymbol {W}} \left[ \frac {\| \boldsymbol {W} _ {L} \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2}}{\| \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2}} \cdot \frac {\| \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2}}{\| \boldsymbol {t} _ {l} ^ {L - 2} \| _ {F} ^ {2}} \dots \frac {\| \boldsymbol {t} _ {l} ^ {l + 1} \| _ {F} ^ {2}}{\| \boldsymbol {t} _ {l} ^ {l} \| _ {F} ^ {2}} \cdot \| \boldsymbol {t} _ {l} ^ {l} \| _ {F} ^ {2} \right] \\ = \mathbb {E} _ {\boldsymbol {W} _ {1}, \dots , \boldsymbol {W} _ {l}} \left[ \| \boldsymbol {t} _ {l} ^ {l} \| _ {F} ^ {2} \cdot \mathbb {E} _ {\boldsymbol {W} _ {l + 1}} \left[ \frac {\| \boldsymbol {t} _ {l} ^ {l + 1} \| _ {F} ^ {2}}{\| \boldsymbol {t} _ {l} ^ {l} \| _ {F} ^ {2}} \dots \mathbb {E} _ {\boldsymbol {W} _ {L}} \left[ \frac {\| \boldsymbol {W} _ {L} \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2}}{\| \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2}} \right] \right] \right]. \\ \end{array}
$$
By rotational invariance of Gaussian column vectors, we prove that for any possible value of  $t_l^{L-1}$  (which is completely determined by  $W_1, \dots, W_{L-1}$  and  $x$ ), for any  $l = 1, \dots, L-1$ , we have that
$$
\mathbb {E} _ {\boldsymbol {W} _ {L}} \left[ \frac {\| \boldsymbol {W} _ {L} \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2}}{\| \boldsymbol {t} _ {l} ^ {L - 1} \| _ {F} ^ {2}} \right] = \mathbb {E} _ {\boldsymbol {W} _ {L}} \left[ \frac {\| \boldsymbol {W} _ {L} \boldsymbol {e} _ {1} \| _ {2} ^ {2}}{\| \boldsymbol {e} _ {1} \| _ {2} ^ {2}} \right] = \beta_ {L o}. \tag {37}
$$
By Lemma C.2, for any  $l = 1,\dots ,L - 2$  and  $l\leq l^{\prime}\leq L - 2$ , we have that
$$
\mathbb {E} _ {\boldsymbol {W} _ {l + 1}} \left[ \frac {\| \boldsymbol {t} _ {l} ^ {l ^ {\prime} + 1} \| _ {F} ^ {2}}{\| \boldsymbol {t} _ {l} ^ {l ^ {\prime}} \| _ {F} ^ {2}} \right] = \frac {\beta_ {l ^ {\prime} + 1}}{2} m _ {l ^ {\prime} + 1}. \tag {38}
$$
We now bound the last term  $\mathbb{E}_{\pmb{W}_1,\dots ,\pmb{W}_l}\left[\| \pmb {t}_l^l\| _F^2\right]$  for  $l = 1,\dots ,L - 1$  . By definition, we have that
$$
\mathbb {E} _ {\boldsymbol {W} _ {1}, \dots , \boldsymbol {W} _ {l}} \left[ \| \boldsymbol {t} _ {l} ^ {l} \| _ {F} ^ {2} \right] = \mathbb {E} _ {\boldsymbol {W} _ {1}, \dots , \boldsymbol {W} _ {l}} \left[ \sum_ {i = 1} ^ {m _ {l}} 1 _ {\{\boldsymbol {W} _ {l} ^ {[ i: ]} h _ {l - 1} \leq 0 \}} \cdot \| h _ {l - 1} \| _ {2} ^ {2} \right] = \frac {m _ {l}}{2} \mathbb {E} _ {\boldsymbol {W} _ {1}, \dots , \boldsymbol {W} _ {l - 1}} \mathbb {E} [ \| h _ {l - 1} \| _ {2} ^ {2} ]. \tag {39}
$$
To bound the  $\mathbb{E}_{\pmb{W}_1,\dots ,\pmb{W}_l}[\| h_{l - 1}\| _2^2 ]$  note that by Lemma C.2, we prove that for any  $l = 1,\dots ,L - 1$
$$
\mathbb {E} _ {\boldsymbol {W} _ {l - 1}} \left[ \frac {\| h _ {l - 1} (\boldsymbol {x}) \| _ {2} ^ {2}}{\| h _ {l - 2} (\boldsymbol {x}) \| _ {2} ^ {2}} \right] = \frac {\beta_ {l - 1}}{2} m _ {l - 1}. \tag {40}
$$
Therefore, for any  $l = 1,\dots ,L$  , we have that
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {W} _ {1}, \dots , \boldsymbol {W} _ {l - 1}} \left[ \| h _ {l - 1} (\boldsymbol {x}) \| _ {2} ^ {2} \right] = \mathbb {E} _ {\boldsymbol {W} _ {1}, \dots , \boldsymbol {W} _ {l - 1}} \left[ \frac {\| h _ {l - 1} (\boldsymbol {x}) \| _ {2} ^ {2}}{\| h _ {l - 2} (\boldsymbol {x}) \| _ {2} ^ {2}} \dots \frac {\| h _ {1} (\boldsymbol {x}) \| _ {2} ^ {2}}{\| \boldsymbol {x} \| _ {2} ^ {2}} \right] \cdot \| \boldsymbol {x} \| _ {2} ^ {2} (41) \\ = \left(\prod_ {i = 1} ^ {l - 1} \frac {\beta_ {i}}{2} m _ {i}\right) \| \boldsymbol {x} \| _ {2} ^ {2}. (42) \\ \end{array}
$$
By plugging (42) into (39), we have that
$$
\mathbb {E} _ {\boldsymbol {W} _ {1}, \dots , \boldsymbol {W} _ {l}} \left[ \| \boldsymbol {t} _ {l} ^ {l} \| _ {2} ^ {2} \right] = \frac {m _ {l}}{2} \left(\prod_ {i = 1} ^ {l - 1} \frac {\beta_ {i}}{2} m _ {i}\right) \| \boldsymbol {x} \| _ {2} ^ {2}. \tag {43}
$$
By combining (37), (38) and (43), we have for any  $l = 1, \dots, L - 1$
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {W}} \left[ \| \frac {\partial f (\boldsymbol {x})}{\partial \operatorname {V e c} (\boldsymbol {W} _ {l})} \| _ {F} ^ {2} \right] \\ = \frac {m _ {l}}{2} \left(\prod_ {i = 1} ^ {l - 1} \frac {\beta_ {i}}{2} m _ {i}\right) \cdot \left(\prod_ {i = l + 1} ^ {L - 1} \frac {\beta_ {i}}{2} m _ {i}\right) \cdot \beta_ {L} o \cdot \| \boldsymbol {x} \| _ {2} ^ {2} = \frac {\beta_ {L}}{\beta_ {l}} \| \boldsymbol {x} \| _ {2} ^ {2} o \left(\prod_ {i = 1} ^ {L - 1} \frac {\beta_ {i} m _ {i}}{2}\right). \\ \end{array}
$$
On the other hand, by plugging Eq. (42) (under  $\ell = L$ ) into Eq. (36), we have that
$$
\mathbb {E} _ {\boldsymbol {W}} \left[ \| \frac {\partial f (\boldsymbol {x})}{\partial \operatorname {V e c} \left(\boldsymbol {W} _ {L}\right)} \| _ {2} ^ {2} \right] = o \left(\prod_ {i = 1} ^ {L - 1} \frac {\beta_ {i}}{2} m _ {i}\right) \| \boldsymbol {x} \| _ {2} ^ {2}
$$
$$
\mathbb {E} _ {\boldsymbol {W}} \left[ \| \frac {\partial f (\boldsymbol {x})}{\partial \operatorname {V e c} (\boldsymbol {W})} \| _ {F} ^ {2} \right] = \sum_ {l = 1} ^ {L} \| \frac {\partial f (\boldsymbol {x})}{\partial \boldsymbol {W} _ {l}} \| _ {F} ^ {2} = \| \boldsymbol {x} \| _ {2} ^ {2} o \left(\prod_ {i = 1} ^ {L - 1} \frac {\beta_ {i} m _ {i}}{2}\right) \sum_ {l = 1} ^ {L} \frac {\beta_ {L}}{\beta_ {l}},
$$
which suffices to prove Eq. (33).
C.2 Deferred proof for Theorem 4.1
Finally, we prove that the gradient difference between two training datasets under linearized network is bounded by a constant throughout training (which only depends on the network width, depth and initialization distribution).
Theorem 4.1. Under Assumption 2.2, taking over the randomness of the random initialization and the Brownian motion, for any  $t \in [0,T]$ , running Langevin diffusion on linearized network Eq. (3) satisfies that
$$
\mathbb {E} \left[ \| \nabla \mathcal {L} \left(\boldsymbol {W} _ {t} ^ {\text {l i n}}; \mathcal {D}\right) - \mathcal {L} \left(\boldsymbol {W} _ {t} ^ {\text {l i n}}; \mathcal {D} ^ {\prime}\right) \| _ {2} ^ {2} \right] \leq \frac {4 B}{n ^ {2}}, \tag {44}
$$
where  $n$  is the training dataset size, and  $B$  is a constant that only depends on the data dimension  $d$ , the number of classes  $o$ , the network depth  $L$ , the per-layer network width  $\{m_i\}_{i=1}^L$ , and the per-layer variances  $\{\beta_i\}_{i=1}^L$  of the Gaussian initialization distribution as follows.
$$
B := d \cdot o \cdot \left(\prod_ {i = 1} ^ {L - 1} \frac {\beta_ {i} m _ {i}}{2}\right) \sum_ {l = 1} ^ {L} \frac {\beta_ {L}}{\beta_ {l}}, \tag {45}
$$
Proof. Denote  $\mathbf{W}$  as the initialization parameters and denote  $\mathbf{W}_t^{lin}$  as the parameters for linearized network after training time  $t$ . Then the gradient difference under linearized network and cross-entropy loss function is as follows.
$$
\begin{array}{l} \| \nabla \mathcal {L} \left(W _ {t}; \mathcal {D}\right) - \nabla \mathcal {L} \left(W _ {t}; \mathcal {D} ^ {\prime}\right) \| _ {F} ^ {2} \\ = \left\| \frac {\nabla f _ {W} (\boldsymbol {x}) ^ {\top} \left(\operatorname {s o f t m a x} \left(\boldsymbol {f} _ {W _ {t}} (\boldsymbol {x})\right) - \boldsymbol {y}\right)}{n} - \frac {\nabla f _ {W} \left(\boldsymbol {x} ^ {\prime}\right) ^ {\top} \left(\operatorname {s o f t m a x} \left(\boldsymbol {f} _ {W _ {t}} \left(\boldsymbol {x} ^ {\prime}\right)\right) - \boldsymbol {y} ^ {\prime}\right)}{n} \right\| _ {F} ^ {2} \\ \leq \frac {2}{n ^ {2}} \left(\| \nabla \boldsymbol {f} _ {\boldsymbol {W}} (\boldsymbol {x}) \| _ {F} ^ {2} + \| \nabla \boldsymbol {f} _ {\boldsymbol {W}} (\boldsymbol {x} ^ {\prime}) \| _ {F} ^ {2}\right). \\ \end{array}
$$
Plugging Lemma C.3 into the above equation with data Assumption 2.2 suffice to prove the result.
D Deferred proofs for Section 6
D.1 Deferred proof for Lemma 6.1 on convergence of training linearized network
In this section, we prove empirical risk bound for the average of all iterates of Langevin diffusion, building on standard results [42, Theorem 2] and [45, Theorem 3.1] for the (discrete time) stochastic gradient descent algorithm.
Lemma 6.1. (Extension of [42, Theorem 2] and [45, Theorem 3.1]) Let  $\mathcal{L}_0^{lin}(\pmb{W};\mathcal{D})$  be the empirical risk function of linearized newtork Eq. (3) expanded at initialization vector  $\pmb{W}_0^{lin}$ . Let  $\pmb{W}_0^*$  be an  $\alpha$ -near-optimal solution for the ERM problem such that  $\mathcal{L}_0^{lin}(\pmb{W}_0^*;\mathcal{D}) - \min_{\pmb{W}}\mathcal{L}_0^{lin}(\pmb{W};\mathcal{D}) \leq \alpha$ . Let  $\mathcal{D} = \pmb{x}_1,\dots,\pmb{x}_n$  be an arbitrary training dataset of size  $n$ , and denote  $M_0 = (\nabla f_{\pmb{W}_0^{lin}}(\pmb{x}_1),\dots,\nabla f_{\pmb{W}_0^{lin}}(\pmb{x}_n))^{\top}$  as the NTK feature matrix at initialization. Then running Langevin diffusion (4) on  $\mathcal{L}_0^{lin}(\pmb{W})$  with time  $T$  and initialization vector  $\pmb{W}_0^{lin}$  satisfies
where the expectation is over Brownian motion  $B_{T}$  in Langevin diffusion Eq. (4),  $\bar{\mathbf{W}}_T^{lin} = \frac{1}{T}\int \bar{\mathbf{W}}_t^{lin}\mathrm{d}t$  is the average of all iterates, and  $R = \| \pmb{W}_0^{lin} - \pmb{W}_0^*\|_{M_0}^2$  is the gap between initialization parameters  $\pmb{W}_0^{lin}$  and solution  $\pmb{W}_0^*$ .
Proof. Our proofs are heavily based on the idea in [45, Theorem 3.1] to work only in the parameter space spanned by the input feature vectors. And our proof serves as an extension of their bound to the continuous-time Langevin diffusion algorithm. We begin by using convexity of the empirical loss function  $\mathcal{L}_0^{lin}(\boldsymbol{W};\mathcal{D})$  for linearized network to prove the following standard results
$$
\mathcal {L} _ {0} ^ {l i n} \left(\bar {\boldsymbol {W}} _ {T} ^ {l i n}; \mathcal {D}\right) - \mathcal {L} _ {0} ^ {l i n} \left(\boldsymbol {W} _ {0} ^ {*}; \mathcal {D}\right) \leq \left\langle \bar {\boldsymbol {W}} _ {T} ^ {l i n} - \boldsymbol {W} _ {0} ^ {*}, \nabla \mathcal {L} _ {0} ^ {l i n} \left(\bar {\boldsymbol {W}} _ {T} ^ {l i n}; \mathcal {D}\right) \right\rangle . \tag {46}
$$
Denote  $M_0 = \left( \nabla f_{\mathbf{W}_0^{lin}}(\mathbf{x}_1) \dots \nabla f_{\mathbf{W}_0^{lin}}(\mathbf{x}_n) \right)$  and compute the gradient under cross entropy loss and linearized network, we have  $\nabla \mathcal{L}_0^{lin}(\mathbf{W}_T^{lin}; \mathcal{D})$  lies in the column space of  $M_0$ . Denote  $\Pi_{M_0}$  as the projection operator to the column space of  $M_0$ , then (46) can be rewritten as
$$
\mathcal {L} _ {0} ^ {l i n} \left(\bar {\boldsymbol {W}} _ {T} ^ {l i n}; \mathcal {D}\right) - \mathcal {L} _ {0} ^ {l i n} \left(\boldsymbol {W} _ {0} ^ {*}; \mathcal {D}\right) \leq \left\langle \Pi_ {M _ {0}} \left(\bar {\boldsymbol {W}} _ {T} ^ {l i n} - \boldsymbol {W} _ {0} ^ {*}\right), \nabla \mathcal {L} _ {0} ^ {l i n} \left(\bar {\boldsymbol {W}} _ {T} ^ {l i n}; \mathcal {D}\right) \right\rangle . \tag {47}
$$
By taking expectation over the randomness of Brownian motion in Langevin diffusion, we have
$$
\mathbb {E} \left[ \mathcal {L} _ {0} ^ {\text {l i n}} \left(\bar {\boldsymbol {W}} _ {T} ^ {\text {l i n}}\right) \right] - \mathcal {L} _ {0} ^ {\text {l i n}} \left(\boldsymbol {W} _ {0} ^ {*}; \mathcal {D}\right) \leq \frac {1}{T} \int_ {0} ^ {T} \mathbb {E} \left[ \left\langle \Pi_ {M _ {0}} \left(\boldsymbol {W} _ {t} ^ {\text {l i n}} - \boldsymbol {W} _ {0} ^ {*}\right), \nabla \mathcal {L} _ {0} ^ {\text {l i n}} \left(\boldsymbol {W} _ {t} ^ {\text {l i n}}; \mathcal {D}\right) \right\rangle \right] d t. \tag {48}
$$
We now rewrite  $\mathbb{E}\left[\langle \Pi_{M_0}(\bar{\boldsymbol{W}}_t^{lin} - \boldsymbol{W}_0^*),\nabla \mathcal{L}_0^{lin}(\boldsymbol{W}_t^{lin};\mathcal{D})\rangle \right]$  by computing  $\frac{\partial}{\partial t}\mathbb{E}[\| \boldsymbol{W}_t^{lin} - \boldsymbol{W}_0^*\|_{M_0}^2 ]$  where  $\| \boldsymbol {W}_t^{lin} - \boldsymbol {W}_0^*\|_{M_0}^2 = \Pi_{M_0}\left(\boldsymbol {W}_t^{lin} - \boldsymbol {W}_0^*\right)^\top \Pi_{M_0}\left(\boldsymbol {W}_t^{lin} - \boldsymbol {W}_0^*\right)$ . By applying (24) with  $p_t$  being the distribution for  $\boldsymbol{W}_{t}^{lin}$  in Langevin diffusion for linearized network starting from point initialization  $\boldsymbol{W}_0^{lin}$ , and with function  $d(\boldsymbol {W}) = \| \boldsymbol {W} - \boldsymbol{W}_0^*\|_{M_0}^2$ , we have that
$$
\frac {\partial}{\partial t} \mathbb {E} [ \| \boldsymbol {W} _ {t} ^ {l i n} - \boldsymbol {W} _ {0} ^ {*} \| _ {M _ {0}} ^ {2} ] \leq - 2 \mathbb {E} [ \left\langle \Pi_ {M _ {0}} \left(\boldsymbol {W} _ {t} ^ {l i n} - \boldsymbol {W} _ {0} ^ {*}\right), \nabla \mathcal {L} _ {0} ^ {l i n} \left(\boldsymbol {W} _ {t} ^ {l i n}; \mathcal {D}\right) \right\rangle ] + \sigma^ {2} r a n k (M _ {0}). \tag {49}
$$
Therefore by plugging (49) into (48), we have that
$$
\begin{array}{l} \mathbb {E} \left[ \mathcal {L} _ {0} ^ {\text {l i n}} \left(\bar {\boldsymbol {W}} _ {T} ^ {\text {l i n}}\right) \right] - \mathcal {L} _ {0} ^ {\text {l i n}} \left(\boldsymbol {W} _ {0} ^ {*}; \mathcal {D}\right) \leq - \frac {1}{2 T} \int_ {0} ^ {T} \frac {\partial}{\partial t} \mathbb {E} \left[ \| \boldsymbol {W} _ {t} - \boldsymbol {W} _ {0} ^ {*} \| _ {M _ {0}} ^ {2} \right] d t + \frac {1}{2} \sigma^ {2} \operatorname {r a n k} (M _ {0}) (50) \\ \leq \frac {1}{2 T} \| \boldsymbol {W} _ {0} ^ {\text {l i n}} - \boldsymbol {W} _ {0} ^ {*} \| _ {M _ {0}} ^ {2} + \frac {1}{2} \sigma^ {2} \operatorname {r a n k} (M _ {0}) (51) \\ \end{array}
$$
D.2 Deferred proof for Lemma 6.3