SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} \geq - \zeta_ {\phi , g} (\left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle) \| \overline {{\boldsymbol {u}}} _ {\perp} ^ {t} \| ^ {2} - \tilde {\mathcal {O}} (\sqrt {d / n}) \\ \geq c \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {s - 1} \left(1 - \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {2}\right) - \tilde {\mathcal {O}} (\sqrt {d / n}) \quad \text {(B y A s s u m p t i o n 2)}. \\ \end{array}
$$
We need to ensure the noise term, i.e. the second term on the RHS remains smaller than the signal, i.e. the first term. The signal term attains its minimum at either initialization  $t = 0$  or at the end of the trajectory  $t = T$  where  $\langle \overline{\boldsymbol{w}}^T, \overline{\boldsymbol{u}} \rangle = 1 - \varepsilon$ , which imposes the following sufficient conditions on  $n$ . Namely, at initialization we require
$$
n \gtrsim C d \langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \rangle^ {2 (1 - s)},
$$
while at  $t = T$  we require
$$
n \gtrsim C d / \varepsilon^ {2},
$$
where  $C$  hides constant that only depend on  $s$  and polylogarithmic factors of  $d$ . Hence, we obtain
$$
\frac {\mathrm {d} \left\langle \overline {{\boldsymbol {w}}} ^ {t} , \overline {{\boldsymbol {u}}} \right\rangle}{\mathrm {d} t} \geq c ^ {\prime} \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {s - 1} (1 - \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle^ {2}).
$$
for some universal constant  $c' > 0$ . Via integration (similar to the proof of Proposition 4 in Appendix D.1), for
$$
T _ {1} := \sup  \left\{t > 0: \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle <   1 / 2 \right\}
$$
$$
T _ {1} \lesssim \tau_ {s} (\langle \overline {{\boldsymbol {w}}} ^ {0}, \overline {{\boldsymbol {u}}} \rangle),
$$
and for
$$
T _ {2} := \sup  \left\{t > 0: \left\langle \overline {{\boldsymbol {w}}} ^ {t}, \overline {{\boldsymbol {u}}} \right\rangle <   1 - \varepsilon \right\}
$$
we obtain  $T_{2} - T_{1} \lesssim \ln(1 / \varepsilon)$ , which completes the proof. We conclude by remarking that the proof of Corollary 8 is immediate given Theorem 7 and Lemma 19.
E.2 Preliminary Lemmas for Proving Theorem 9
We will adapt the following lemma from  $\mathrm{[MHPG^{+}23]}$ , which provides a non-parametric approximation of  $g$  via random biases.
Lemma 20.  $\left[\mathrm{MHPG}^{+}23\right.$  , Lemma 22] For any smooth  $g:\mathbb{R}\to \mathbb{R}$  and  $\Delta >0$  , let  $\tilde{g}:\mathbb{R}\rightarrow \mathbb{R}$  be a smooth function such that  $\tilde{g} (z) = g(z)$  for  $|z|\leq \Delta$  and  $\tilde{g} (-2\Delta) = \tilde{g} '(-2\Delta) = 0$  . Suppose  $\{b_j\}_{j = 1}^m\stackrel {i.i.d.}{\sim}\mathrm{Unif}(-2\Delta ,2\Delta)$  , and let  $\Delta_{*}\coloneqq \Delta \sup_{|z|\leq 2\Delta}|\tilde{g}''(z)|$  . Then, there exist second layer weights  $\{a_j(b_j)\}_{j = 1}^m$  with  $\| \pmb {a}\| \lesssim \Delta_{*} / \sqrt{m}$  , such that for any fixed  $z\in [-\Delta ,\Delta ]$  and any  $\delta >0$  , with probability at least  $1 - \delta$  over the random biases,
$$
\left| \sum_ {j = 1} ^ {m} a (b _ {j}) \phi (z + b _ {j}) - g (b _ {j}) \right| \lesssim \Delta \Delta_ {*} \sqrt {\frac {\ln (1 / \delta)}{m}},
$$
where  $\phi$  is the ReLU activation.
We use the above lemma to show the existence of a second layer with  $\tilde{\mathcal{O}}(1/\sqrt{m})$  norm with training error of order  $\tilde{\mathcal{O}}(1/m)$ .
Lemma 21. For any  $\varepsilon < 1$ , suppose  $\langle \overline{\mathbf{w}}, \overline{\mathbf{u}} \rangle \geq 1 - \varepsilon$ . Then for any  $q > 0$ , sufficiently large  $d$ ,  $n \gtrsim d / \varepsilon^2$ , with probability at least  $1 - \mathcal{O}(d^{-q})$  over the random biases and the dataset, there exists a second layer  $\mathbf{a}$  with  $\| \mathbf{a} \| \leq \tilde{\mathcal{O}}(1 / \sqrt{m})$  described by Lemma 20 such that
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \left(\sum_ {j = 1} ^ {m} a _ {j} \phi (\left\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle + b _ {j}) - y ^ {(i)}\right) ^ {2} \lesssim \mathbb {E} [ \epsilon^ {2} ] + \tilde {\mathcal {O}} (1 / m + \varepsilon),
$$
Proof. We begin by replacing  $\overline{w}$  and  $\tilde{z}^{(i)}$  with  $\overline{u}$  and  $z^{(i)}$ . Specifically, via Jensen's inequality,
$$
\begin{array}{l} \frac {1}{n} \sum_ {i = 1} ^ {n} \left\{\sum_ {j = 1} ^ {m} a _ {j} \phi \left(\left\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle + b _ {j}\right) - y ^ {(i)} \right\} ^ {2} \leq \underbrace {\frac {4}{n} \sum_ {i = 1} ^ {n} \left\{\sum_ {j = 1} ^ {m} a _ {j} \phi \left(\left\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} ^ {(i)} \right\rangle + b _ {j}\right) - g \left(\left\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} ^ {(i)} \right\rangle\right) \right\}} _ {=: \mathcal {E} _ {1}} ^ {2} \\ \leq \underbrace {\frac {4}{n} \sum_ {i = 1} ^ {n} \left\{y ^ {(i)} - g \left(\left\langle \overline {{u}} , z ^ {(i)} \right\rangle\right) \right\} ^ {2}} _ {=: \mathcal {E} _ {2}} \\ \leq \underbrace {\frac {4}{n} \sum_ {i = 1} ^ {n} \left(\sum_ {j = 1} ^ {m} a _ {j} \left\{\phi \left(\left\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle + b _ {j}\right) - \phi \left(\left\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} ^ {(i)} \right\rangle + b _ {j}\right) \right\}\right) ^ {2}} _ {=: \mathcal {E} _ {3}} \\ \leq \underbrace {\frac {4}{n} \sum_ {i = 1} ^ {n} \left(\sum_ {j = 1} ^ {m} a _ {j} \left\{\phi \left(\left\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} ^ {(i)} \right\rangle + b _ {j}\right) - \phi \left(\left\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} ^ {(i)} \right\rangle + b _ {j}\right) \right\}\right) ^ {2}} _ {=: \mathcal {E} _ {4}}. \\ \end{array}
$$
We bound each term separately. For  $\mathcal{E}_1$ , we can invoke Lemma 20 which implies that each term in the sum can be bounded by  $\tilde{\mathcal{O}}(1/m)$  with probability at least  $1 - 1/(nd^q)$ , thus by a union bound, with probability at least  $1 - d^{-q}$  over the random biases,
$$
\mathcal {E} _ {1} \leq \tilde {\mathcal {O}} (1 / m).
$$
By sub-Gaussianity of  $\epsilon^{(i)}$  (hence sub-exponentiality of  $\epsilon^{(i)^2}$ ), for  $n \gtrsim d$  (with a sufficiently large constant) we have
$$
\mathcal {E} _ {2} \lesssim \mathbb {E} \left[ \epsilon^ {2} \right] + \sqrt {d / n},
$$
with probability at least  $1 - e^{-d}$ .
For  $\mathcal{E}_3$ , via the Lipschitzness of ReLU and the Cauchy-Schwartz inequality we can write
$$
\mathcal {E} _ {3} \leq \frac {\tilde {\mathcal {O}} (1)}{n} \sum_ {i = 1} ^ {n} \left\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} ^ {(i)} - \boldsymbol {z} ^ {(i)} \right\rangle^ {2} \leq \tilde {\mathcal {O}} (1) \| \mathbf {I} - \hat {\boldsymbol {\Sigma}} ^ {- 1 / 2} \boldsymbol {\Sigma} ^ {1 / 2} \| ^ {2} \leq \tilde {\mathcal {O}} (d / n ^ {\prime}),
$$
where we used the event of Lemma 24 which happens with probability at least  $1 - 2e^{-d}$ , and  $\tilde{\mathcal{O}}(1)$  represents a constant that depends at most polylogarithmically on  $d$ .
Finally, we bound the last term. Once again via the Lipschitzness of the ReLU activation and the Cauchy-Schwartz inequality
$$
\mathcal {E} _ {4} \leq \frac {\tilde {\mathcal {O}} (1)}{n} \sum_ {i = 1} ^ {n} \left\langle \overline {{\boldsymbol {w}}} - \overline {{\boldsymbol {u}}}, \boldsymbol {z} ^ {(i)} \right\rangle^ {2} \leq \tilde {\mathcal {O}} (\| \overline {{\boldsymbol {w}}} - \overline {{\boldsymbol {u}}} \| ^ {2}) \leq \tilde {\mathcal {O}} (\varepsilon),
$$
where once again we used the event of Lemma 24. On the intersection of all desired events, we have
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \left(\sum_ {j = 1} ^ {m} a _ {j} \phi (\left\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle + b _ {j}) - y ^ {(i)}\right) \lesssim \mathbb {E} \left[ \epsilon^ {2} \right] + \tilde {\mathcal {O}} (1 / m + \sqrt {d / n} + d / n ^ {\prime} + \varepsilon).
$$
We conclude the proof by noticing that  $n' \gtrsim n^2$  and  $n \gtrsim d\varepsilon^{-2}$ .
Additionally, we will use the following standard Lemma on the Rademacher complexity of two-layer neural networks, which in particular is a restatement of  $\left[\mathrm{MHPG}^{+}23\right.$ , Lemma 18] in a way suitable for our analysis.
Lemma 22. Let  $\mathcal{F}$  be a class of real-valued functions on  $(z,y)$ . Given  $n$  samples  $\{z^{(i)},y\}_{i = 1}^{n}$ , define the empirical Rademacher complexity of  $\mathcal{F}$  as
$$
\hat {\mathfrak {N}} _ {n} (\mathcal {F}) := \mathbb {E} _ {(\varsigma_ {i}) _ {i = 1} ^ {n}} \left[ \sup _ {f \in \mathcal {F}} \frac {1}{n} \sum_ {i = 1} ^ {n} \varsigma_ {i} f (z ^ {(i)}, y ^ {(i)}) \right],
$$
where  $(\varsigma_{i})$  are i.i.d. Rademacher random variables (i.e.  $\pm 1$  with equal probability). Suppose  $\mathcal{F}$  is given by
$$
\mathcal {F} := \left\{\left(\boldsymbol {z}, y\right) \mapsto \left(\sum_ {j = 1} ^ {m} a _ {j} \phi \left(\left\langle \overline {{\boldsymbol {u}}}, \boldsymbol {z} \right\rangle + b _ {j}\right) - y\right) ^ {2} \wedge C: \| \boldsymbol {a} \| \leq r _ {a} / \sqrt {m}, | b _ {j} | \leq r _ {b}, \forall 1 \leq j \leq m \right\},
$$
for some fixed  $\overline{\mathbf{u}}\in \mathbb{S}^{d - 1}$ . Suppose  $\{\pmb {z}^{(i)}\}_{i = 1}^{n}\stackrel {i.i.d.}{\sim}\mathcal{N}(0,\mathbf{I}_d)$ , and suppose  $|\phi^{\prime}|\leq 1$ . Then,
$$
\mathbb {E} _ {(\boldsymbol {z} ^ {(i)}, \boldsymbol {y} ^ {(i)}) _ {i = 1} ^ {n}} \left[ \hat {\Re} _ {n} (\mathcal {F}) \right] \leq \frac {2 \sqrt {2 C} (1 + r _ {b}) r _ {a}}{\sqrt {n}}.
$$
Proof. See the proof of  $\left[\mathrm{MHPG}^{+}23\right.$  , Lemma 18].
E.3 Proof of Theorem 9
Throughout the proof, we will assume  $\langle \overline{\boldsymbol{w}},\overline{\boldsymbol{u}}\rangle \geq 1 - \varepsilon$  where we recall
$$
\overline {{\boldsymbol {w}}} := \frac {\hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w}}{\| \hat {\boldsymbol {\Sigma}} ^ {1 / 2} \boldsymbol {w} \|} \quad \text {a n d} \quad \overline {{\boldsymbol {u}}} := \frac {\boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u}}{\| \boldsymbol {\Sigma} ^ {1 / 2} \boldsymbol {u} \|}.
$$
From either Theorem 5 or Theorem 7, we can assume  $\langle \overline{w},\overline{u}\rangle \geq 1 - \varepsilon$  with probability at least  $1 - \mathcal{O}(d^{-q})$  for any fixed  $q > 0$ . For simplicity, let
$$
\hat {y} (\tilde {\boldsymbol {z}}; \overline {{\boldsymbol {w}}}) = \sum_ {j = 1} ^ {m} a _ {j} \phi (\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} \rangle + b _ {j}),
$$
and similarly define  $\hat{y} (z;\overline{u})$ . We define the following quantities,
$$
\mathcal {R} (\overline {{\boldsymbol {w}}}) := \mathbb {E} _ {\boldsymbol {z}, y} \left[ \left(\hat {y} (\tilde {\boldsymbol {z}}; \overline {{\boldsymbol {w}}}) - y\right) ^ {2} \right] \quad \text {a n d} \quad R (\overline {{\boldsymbol {u}}}) := \mathbb {E} _ {\boldsymbol {z}, y} \left[ \left(\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y\right) ^ {2} \right], \tag {E.1}
$$
and similarly define their empirical counterparts,
$$
\hat {\mathcal {R}} (\overline {{\boldsymbol {w}}}) := \frac {1}{n} \sum_ {i = 1} ^ {n} \left(\hat {y} \left(\tilde {\boldsymbol {z}} ^ {(i)}; \overline {{\boldsymbol {w}}}\right) - y ^ {(i)}\right) ^ {2} \quad \text {a n d} \quad \hat {R} (\overline {{\boldsymbol {u}}}) := \frac {1}{n} \sum_ {i = 1} ^ {n} \left(\hat {y} \left(\boldsymbol {z} ^ {(i)}; \overline {{\boldsymbol {u}}}\right) - y ^ {(i)}\right) ^ {2}. \tag {E.2}
$$
Notice that ultimately, we are interested in bounding  $\mathcal{R}(\overline{\boldsymbol{w}})$ . We break down the proof into three steps. In the first step, we show that  $\mathcal{R}(\overline{\boldsymbol{w}})$  can be upper bounded by  $R(\overline{\boldsymbol{u}})$ . Then, via a generalization bound, we show that the  $R(\overline{\boldsymbol{u}})$  can be upper bounded by  $\hat{R}(\overline{\boldsymbol{u}})$ . Finally, we show that  $\hat{R}(\overline{\boldsymbol{u}})$  can be upper bounded by the training error, i.e.  $\hat{\mathcal{R}}(\overline{\boldsymbol{w}})$ , and convex optimization of the last layer can attain the near-optimal value of this training error which is bounded by Lemma 21.
Step 1. Bounding  $\mathcal{R}(\overline{w})$  via  $R(\overline{u})$ . By Jensen's inequality,
$$
\mathbb {E} _ {\boldsymbol {z}, y} \Big [ (\hat {y} (\tilde {\boldsymbol {z}}; \overline {{\boldsymbol {w}}}) - y) ^ {2} \Big ] \leq 3 \mathbb {E} _ {\boldsymbol {z}} \Big [ (\hat {y} (\tilde {\boldsymbol {z}}; \overline {{\boldsymbol {w}}}) - \hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {w}}})) ^ {2} \Big ] + 3 \mathbb {E} _ {\boldsymbol {z}} \Big [ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {w}}}) - \hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}})) ^ {2} \Big ] + 3 \mathbb {E} _ {\boldsymbol {z}, y} \Big [ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \Big ].
$$
Suppose  $\| \mathbf{a} \| \leq r_a / \sqrt{m}$ . For the first term, by Lipschitzness of  $\phi$  and the Cauchy-Schwartz inequality
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {z}} \left[ (\hat {y} (\tilde {\boldsymbol {z}}; \overline {{\boldsymbol {w}}}) - \hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {w}}})) ^ {2} \right] = \mathbb {E} _ {\boldsymbol {z}} \left[ \left(\sum_ {j = 1} ^ {m} a _ {j} \left\{\phi \left(\left\langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} ^ {(i)} \right\rangle + b _ {j}\right) - \phi \left(\left\langle \overline {{\boldsymbol {w}}}, \boldsymbol {z} ^ {(i)} \right\rangle + b _ {j}\right) \right\}\right) ^ {2} \right] \\ \leq r _ {a} ^ {2} \mathbb {E} _ {\boldsymbol {z}} \left[ \langle \overline {{\boldsymbol {w}}}, \tilde {\boldsymbol {z}} - \boldsymbol {z} \rangle^ {2} \right] \\ \leq r _ {a} ^ {2} \left\| \mathbf {I} _ {d} - \hat {\boldsymbol {\Sigma}} ^ {- 1 / 2} \boldsymbol {\Sigma} ^ {1 / 2} \right\| ^ {2} \lesssim r _ {a} ^ {2} d / n ^ {\prime}, \\ \end{array}
$$
where the last inequality holds with probability at least  $1 - 2e^{-d}$  on the event of Lemma 24.
For the middle term, via a similar argument,
$$
\mathbb {E} _ {\boldsymbol {z}} \left[ \left(\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {w}}}) - \hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}})\right) ^ {2} \right] \leq r _ {a} ^ {2} \mathbb {E} _ {\boldsymbol {z}} \left[ \left\langle \overline {{\boldsymbol {w}}} - \overline {{\boldsymbol {u}}}, \boldsymbol {z} \right\rangle^ {2} \right] \leq 2 r _ {a} \varepsilon .
$$
In what follows, we will restrict the analysis to the case where  $r_a = \tilde{\mathcal{O}}(1)$ . Therefore, we have
$$
\mathbb {E} _ {\boldsymbol {z}, \boldsymbol {y}} \left[ (\hat {y} (\tilde {\boldsymbol {z}}; \overline {{\boldsymbol {w}}}) - \boldsymbol {y}) ^ {2} \right] \leq 3 \mathbb {E} _ {\boldsymbol {z}, \boldsymbol {y}} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - \boldsymbol {y}) ^ {2} \right] + \tilde {\mathcal {O}} (d / n ^ {\prime} + \varepsilon).
$$
Step 2. Generalization: Bounding  $R(\overline{u})$  via  $\hat{R}(\overline{u})$ .
$$
E := \left\{\left| \langle \overline {{u}}, z \rangle \right| \vee | \epsilon | \leq \sqrt {2 \ln (n d ^ {q})} \right\}.
$$
and similarly define  $E^{(i)}$  by replacing  $z$  and  $\epsilon$  with  $z^{(i)}$  and  $\epsilon^{(i)}$  respectively. Via the Cauchy-Schwartz and Jensen inequalities
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {z}, y} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \right] = \mathbb {E} _ {\boldsymbol {z}, y} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \mathbf {1} (E) \right] + \mathbb {E} _ {\boldsymbol {z}, y} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \mathbf {1} (E ^ {C}) \right] \\ \leq \mathbb {E} _ {\boldsymbol {z}, y} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \mathbf {1} (E) \right] + \sqrt {8} \left(\mathbb {E} \left[ \hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) ^ {4} \right] + \mathbb {E} \left[ y ^ {4} \right]\right) ^ {1 / 2} \mathbb {P} \left(E ^ {C}\right) ^ {1 / 2} \\ \end{array}
$$
Moreover,  $\mathbb{E}\big[y^4\big] \lesssim 1$ ,  $\mathbb{E}_{z,y}\big[\hat{y}(z; \overline{\boldsymbol{u}})^4\big] \leq \tilde{\mathcal{O}}(1)$ , and  $\mathbb{P}\big(E^C\big) \leq 4/(nd^q)$  (via a standard sub-Gaussian tail bound). Consequently,
$$
\mathbb {E} _ {\boldsymbol {z}, y} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \right] \leq \mathbb {E} _ {\boldsymbol {z}, y} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \mathbf {1} (E) \right] + \tilde {\mathcal {O}} (n ^ {- 1} d ^ {- q}),
$$
$$
\ell \left(\boldsymbol {z} ^ {(i)}, y ^ {(i)}; \boldsymbol {a}, \boldsymbol {b}\right) := \left(\sum_ {j = 1} ^ {m} a _ {j} \phi \left(\left\langle \bar {\boldsymbol {u}}, \boldsymbol {z} ^ {(i)} \right\rangle + b _ {j}\right) - y ^ {(i)}\right) ^ {2} \mathbf {1} (E).
$$
Notice that  $\overline{\pmb{u}}$  is fixed. Then, by a standard symmetrization argument (see e.g. [VH16, Lemma 7.4]) and Lemma 22
$$
\begin{array}{l} \mathbb {E} \left[ \sup  _ {\| \boldsymbol {a} \| \leq r _ {a} / \sqrt {m}, | b _ {j} | \leq r _ {b}} \mathbb {E} _ {\boldsymbol {z}, y} [ \ell (\boldsymbol {z}, y; \boldsymbol {a}, \boldsymbol {b}) ] - \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (\boldsymbol {z} ^ {(i)}, y ^ {(i)}; \boldsymbol {a}, \boldsymbol {b}) \right] \leq 2 \mathbb {E} \left[ \hat {\Re} _ {n} (\mathcal {F}) \right] \\ \leq \tilde {\mathcal {O}} (\sqrt {1 / n}). \\ \end{array}
$$
where  $\hat{\Re}_n(\mathcal{F})$ . As the loss is bounded, we can apply McDiarmid's inequality to turn the above bound in expectation into a bound in probability, in particular
$$
\sup  _ {\| \boldsymbol {a} \| \leq r _ {a} / \sqrt {m}, | b _ {j} | \leq r _ {b}} \mathbb {E} _ {\boldsymbol {z}, y} [ \ell (\boldsymbol {z}, y; \boldsymbol {a}, \boldsymbol {b}) ] - \frac {1}{n} \sum_ {i = 1} ^ {n} \ell (\boldsymbol {z} ^ {(i)}, y ^ {(i)}; \boldsymbol {a}, \boldsymbol {b}) \leq \tilde {\mathcal {O}} (\sqrt {d / n})
$$
with probability at least  $1 - 2e^{-d}$ . Therefore, we conclude this step by noticing that
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {z}, y} \left[ (\hat {y} (\boldsymbol {z}; \overline {{\boldsymbol {u}}}) - y) ^ {2} \right] \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \left(\hat {y} (\boldsymbol {z} ^ {(i)}; \overline {{\boldsymbol {u}}}) - y ^ {(i)}\right) ^ {2} \mathbf {1} (E ^ {(i)}) + \tilde {\mathcal {O}} (\sqrt {d / n} + n ^ {- 1} d ^ {- q}) \\ \leq \frac {1}{n} \sum_ {i = 1} ^ {n} \Big (\hat {y} (\boldsymbol {z} ^ {(i)}; \overline {{\boldsymbol {u}}}) - y ^ {(i)} \Big) ^ {2} + \tilde {\mathcal {O}} (\sqrt {d / n}), \\ \end{array}
$$
with probability at least  $1 - 2e^{-d}$ .
Step 3. Bounding the training error and finishing the proof. This step is similar to the proof of  $\mathrm{[MHPG^{+}23, Theorem 4]}$ . For conciseness, define
$$
\hat {\mathcal {R}} (\boldsymbol {a}) := \frac {1}{n} \sum_ {i = 1} ^ {n} \left(\hat {y} (\boldsymbol {x} ^ {(i)}; \boldsymbol {W}, \boldsymbol {a}, \boldsymbol {b}) - y ^ {(i)}\right) ^ {2}.
$$
and  $\hat{\mathcal{R}}_{\lambda}(\pmb {a})\coloneqq \hat{\mathcal{R}} (\pmb {a}) + \lambda \| \pmb {a}\| ^2 /2$  . Our goal is to choose suitable  $\lambda$  such that the minimizer
$$
\boldsymbol{a}^{*}:= \operatorname *{arg  min}_{\boldsymbol {a}\in \mathbb{R}^{m}}\hat{\mathcal{R}} (\boldsymbol {a}) + \lambda \| \boldsymbol {a}\|^{2} / 2,
$$
satisfies  $\| \pmb{a}^{*}\| \leq r_{a} / \sqrt{m}$  while the value of the above minimization problem which we denote with  $\mathcal{R}_{\lambda}^{*}$  does not significantly exceed
$$
\min_{\| \boldsymbol {a}\| \leq r_{a} / \sqrt{m}}\hat{\mathcal{R}} (\boldsymbol {a}).
$$
We argue that the suitable choice for  $\lambda$  is
$$
\lambda \asymp \frac {m \mathbb {E} \left[ \epsilon^ {2} \right] + m \varepsilon + 1}{r _ {a} ^ {2}} = \tilde {\Theta} (m \mathbb {E} [ \epsilon^ {2} ] + m \varepsilon + 1). \tag {E.3}
$$
Let  $\hat{\mathcal{R}}^*$  denote the minimizer of the regularized problem and  $\tilde{\pmb{a}} := \arg \min_{\| \pmb{a} \| \leq r_a / \sqrt{m}} \hat{\mathcal{R}}(\pmb{a})$ . From Lemma 21, with a proper choice of  $r_a = \tilde{\Theta}(1)$ , we have
$$
\hat {\mathcal {R}} (\tilde {\boldsymbol {a}}) \lesssim \mathbb {E} \left[ \epsilon^ {2} \right] + \tilde {\mathcal {O}} (1 / m + \varepsilon).
$$
with probability at least  $1 - \mathcal{O}(d^{-q})$  over the biases and the dataset. Note that as  $\pmb{a}^{*}$  is the minimizer of  $\hat{\mathcal{R}}_{\lambda}$ , we have
$$
\hat {\mathcal {R}} (\boldsymbol {a} ^ {*}) + \frac {\lambda \| \boldsymbol {a} ^ {*} \| ^ {2}}{2} \leq \hat {\mathcal {R}} (\tilde {\boldsymbol {a}}) + \frac {\lambda \| \tilde {\boldsymbol {a}} \| ^ {2}}{2},
$$
and in particular
$$
\frac {\lambda \| \boldsymbol {a} ^ {*} \| ^ {2}}{2} \leq \hat {\mathcal {R}} (\tilde {\boldsymbol {a}}) + \frac {\lambda \| \tilde {\boldsymbol {a}} \| ^ {2}}{2} \Rightarrow \| \boldsymbol {a} ^ {*} \| \leq \tilde {\mathcal {O}} (1 / \sqrt {m}).
$$
$$
\hat {\mathcal {R}} (\boldsymbol {a} ^ {*}) \leq \hat {\mathcal {R}} (\tilde {\boldsymbol {a}}) + \frac {\lambda \| \tilde {\boldsymbol {a}} \| ^ {2}}{2} \lesssim \mathbb {E} [ \epsilon ] ^ {2} + \tilde {\mathcal {O}} (1 / m + \varepsilon).
$$
Let  $\{\pmb{a}^t\}_{t\geq 0}$  be the solution to the gradient flow of  $\pmb{a}$ . Then,
$$
\frac {\mathrm {d} \left\| \boldsymbol {a} ^ {t} - \boldsymbol {a} ^ {*} \right\| ^ {2}}{\mathrm {d} t} = - 2 \left\langle \boldsymbol {a} ^ {t} - \boldsymbol {a} ^ {*}, \nabla \hat {\mathcal {R}} _ {\lambda} (\boldsymbol {a} ^ {t}) \right\rangle ,
$$
and by the first-order condition of strong convexity
$$
\left\langle \boldsymbol {a} ^ {t} - \boldsymbol {a} ^ {*}, \nabla \hat {\mathcal {R}} _ {\lambda} (\boldsymbol {a} ^ {t}) \right\rangle \geq \lambda \| \boldsymbol {a} ^ {t} - \boldsymbol {a} ^ {*} \| ^ {2},
$$
$$
\left\| \boldsymbol {a} ^ {T ^ {\prime}} - \boldsymbol {a} ^ {*} \right\| ^ {2} \leq e ^ {- 2 \lambda T ^ {\prime}} \left\| \boldsymbol {a} ^ {0} - \boldsymbol {a} ^ {*} \right\| ^ {2}.
$$
As the training error (of the regularized problem) is  $\lambda$ -strongly convex in  $a$ , by applying the standard Polyak-Lojasiewicz condition, gradient flow for training  $a$  obtains
$$
\hat {\mathcal {R}} _ {\lambda} \left(\boldsymbol {a} ^ {T ^ {\prime}}\right) - \hat {\mathcal {R}} _ {\lambda} ^ {*} \leq \left(\hat {\mathcal {R}} _ {\lambda} \left(\boldsymbol {a} ^ {0}\right) - \hat {\mathcal {R}} _ {\lambda} ^ {*}\right) e ^ {- 2 \lambda T ^ {\prime}}.
$$
Furthermore, since
$$
\left\| \boldsymbol {a} ^ {*} \right\| ^ {2} - \left\| \boldsymbol {a} ^ {T ^ {\prime}} \right\| ^ {2} \leq 2 \left\| \boldsymbol {a} ^ {*} \right\| \left\| \boldsymbol {a} ^ {T ^ {\prime}} - \boldsymbol {a} ^ {*} \right\| - \left\| \boldsymbol {a} ^ {T ^ {\prime}} - \boldsymbol {a} ^ {*} \right\| ^ {2} \leq 2 \left\| \boldsymbol {a} ^ {T ^ {\prime}} - \boldsymbol {a} ^ {*} \right\| \left\| \boldsymbol {a} ^ {*} \right\|,
$$
$$
\hat {\mathcal {R}} (\boldsymbol {a} ^ {T ^ {\prime}}) - \hat {\mathcal {R}} ^ {*} \leq 2 \| \boldsymbol {a} ^ {0} - \boldsymbol {a} ^ {*} \| \| \boldsymbol {a} ^ {*} \| e ^ {- \lambda T ^ {\prime}} + \left(\hat {\mathcal {R}} _ {\lambda} (\boldsymbol {a} ^ {0}) - \hat {\mathcal {R}} _ {\lambda} ^ {*}\right) e ^ {- 2 \lambda T ^ {\prime}}.
$$
Consequently, choosing
$$
T ^ {\prime} \geq \frac {\ln \left(\frac {\| \boldsymbol {a} ^ {0} - \boldsymbol {a} ^ {*} \|}{\| \boldsymbol {a} ^ {*} \|}\right)}{\lambda} \vee \frac {\ln \left(\frac {4 \| \boldsymbol {a} ^ {0} - \boldsymbol {a} ^ {*} \| \| \boldsymbol {a} ^ {*} \|}{\varepsilon}\right)}{\lambda} \vee \frac {\ln \left(\frac {2 \left(\hat {\mathcal {R}} _ {\lambda} (\boldsymbol {a} ^ {0}) - \hat {\mathcal {R}} _ {\lambda} ^ {*}\right)}{\varepsilon}\right)}{2 \lambda}, \tag {E.4}
$$