SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\mathbf {a} ^ {(t)} = \mathbf {a} ^ {(t - 1)} - \eta^ {(t)} \nabla_ {\mathbf {a}} \mathcal {L} _ {\mathcal {D}} ^ {\lambda (t)} \left(f _ {\Xi^ {(t - 1)}}\right)
$$
$$
\mathbf {W} ^ {(t)} = \mathbf {W} ^ {(t - 1)} - \eta^ {(t)} \nabla_ {\mathbf {W}} \mathcal {L} _ {\mathcal {D}} ^ {\lambda^ {(t)}} \left(f _ {\Xi^ {(t - 1)}}\right)
$$
Given an input distribution, we can get a Gradient Feature set  $S_{p,\gamma,B_G}$  and  $f^*(\mathbf{x}) = \sum_{j=1}^{r} \mathbf{a}_j^* \sigma(\langle \mathbf{w}_j^*, \mathbf{x} \rangle - \mathbf{b}_j^*)$ , where  $f^* \in \mathcal{F}_{d,r,B_F,S_{p,\gamma,B_G}}$  is a Gradient Feature Induced networks defined in Definition 3.11. Considering training by Algorithm 3, we have the following results.
Theorem D.1 (Gradient Feature Learning Framework under Expected Risk). Assume Assumption 3.1. For any  $\epsilon, \delta \in (0,1)$ , if  $m \leq e^d$  and
$$
m = \Omega \left(\frac {1}{p} \left(\frac {r B _ {a 1} B _ {x 1}}{\epsilon} \sqrt {\frac {B _ {b}}{B _ {G}}}\right) ^ {4} + \frac {1}{\sqrt {\delta}} + \frac {1}{p} \left(\log \left(\frac {r}{\delta}\right)\right) ^ {2}\right), \tag {25}
$$
$$
T = \Omega \left(\frac {1}{\epsilon} \left(\frac {\sqrt {r} B _ {a 2} B _ {b} B _ {x 1}}{(m p) ^ {\frac {1}{4}}} + m \tilde {b}\right) \left(\frac {\sqrt {\log m}}{\sqrt {B _ {b} B _ {G}}} + \frac {1}{B _ {x 1} (m p) ^ {\frac {1}{4}}}\right)\right), \tag {26}
$$
then with proper hyper-parameter values, we have with probability  $\geq 1 - \delta$ , there exists  $t \in [T]$  in Algorithm 3 with
$$
\Pr \left[ \operatorname {s i g n} \left(f _ {\Xi^ {(t)}} (\mathbf {x})\right) \neq y \right] \leq \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(t)}}\right) \leq \mathrm {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}} + r B _ {a 1} B _ {x 1} \sqrt {2 \gamma} + \epsilon . \tag {27}
$$
See the full statement and proof in Theorem D.9. Below, we show some lemmas used in the analysis of population loss.
D.2.1 Feature Learning
We first show that a large subset of neurons has gradients at the first step as good features.
Definition D.2 (Nice Gradients Set. Equivalent to Equation (9)). We define
$$
G _ {(D, + 1), N i c e} := \left\{i \in [ m ]: \left\langle \mathbf {w} _ {i} ^ {(1)}, D \right\rangle > (1 - \gamma) \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}, \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} \geq \left| \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i} ^ {(0)} \right| B _ {G} \right\}
$$
$$
G _ {(D, - 1), N i c e} := \left\{i \in [ 2 m ] \backslash [ m ]: \left\langle \mathbf {w} _ {i} ^ {(1)}, D \right\rangle > (1 - \gamma) \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}, \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} \geq \left| \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i} ^ {(0)} \right| B _ {G} \right\}
$$
where  $\gamma, B_G$  is the same in the Definition 3.7.
Lemma D.3 (Feature Emergence. Full Statement of Lemma 3.13). Let  $\lambda^{(1)} = \frac{1}{\eta^{(1)}}$ . For any  $r$  size subset  $\{(D_1, s_1), \ldots, (D_r, s_r)\} \subseteq S_{p,\gamma, B_G}$ , with probability at least  $1 - 2re^{-cmp}$  where  $c > 0$  is a universal constant, we have that for all  $j \in [r]$ ,  $|G_{(D_j, s_j), Nice}| \geq \frac{mp}{4}$ .
Proof of Lemma D.3. By symmetric initialization and Lemma F.1, we have for all  $i \in [2m]$
$$
\mathbf {w} _ {i} ^ {(1)} = - \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i} ^ {(0)} \mathbb {E} _ {(\mathbf {x}, y)} \left[ y \sigma^ {\prime} \left[ \left\langle \mathbf {w} _ {i} ^ {(0)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i} \right] \mathbf {x} \right] \tag {28}
$$
$$
= - \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i} ^ {(0)} G \left(\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}\right). \tag {29}
$$
For all  $j\in [r]$  , as  $(D_j,s_j)\in S_{p,\gamma ,B_G}$  , by Lemma F.3,
(1) if  $s_j = +1$ , for all  $i \in [m]$ , we have
$$
\Pr \left[ i \in G _ {\left(D _ {j}, s _ {j}\right), N i c e} \right] \tag {30}
$$
$$
\begin{array}{l} = \Pr \left[ \frac {\left\langle \mathbf {w} _ {i} ^ {(1)} , D _ {j} \right\rangle}{\left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}} > (1 - \gamma), \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} \geq \left| \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i} ^ {(0)} \right| B _ {G} \right] (31) \\ = \Pr \left[ \frac {\left\langle \mathbf {w} _ {i} ^ {(1)} , D _ {j} \right\rangle}{\left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}} > (1 - \gamma), \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} \geq \left| \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i} ^ {(0)} \right| B _ {G}, \frac {\mathbf {b} _ {i}}{| \mathbf {b} _ {i} |} = s _ {j} \right] (32) \\ \geq \Pr \left[ G (\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}) \in \mathcal {C} _ {D _ {j}, \gamma}, \| G (\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}) \| _ {2} \geq B _ {G}, \frac {\mathbf {b} _ {i}}{| \mathbf {b} _ {i} |} = s _ {j}, \mathbf {a} _ {i} ^ {(0)} \left\langle G (\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}), D _ {j} \right\rangle > 0 \right] \\ \geq \frac {p}{2}, (33) \\ \end{array}
$$
(2) if  $s_j = -1$ , for all  $i \in [2m] \setminus [m]$ , similarly we have
$$
\Pr \left[ i \in G _ {\left(D _ {j}, s _ {j}\right), N i c e} \right] \geq \frac {p}{2}. \tag {34}
$$
By concentration inequality, (Chernoff's inequality under small deviations), we have
$$
\Pr \left[ | G _ {(D _ {j}, s _ {j}), N i c e} | <   \frac {m p}{4} \right] \leq 2 e ^ {- c m p}. \tag {35}
$$
We complete the proof by union bound.
D.2.2 Good Network Finds
Then, the gradients allow for obtaining a set of neurons approximating the "ground-truth" network with comparable loss.
Lemma D.4 (Existence of Good Networks. Full Statement of Lemma 3.14). Let  $\lambda^{(1)} = \frac{1}{\eta^{(1)}}$ . For any  $B_{\epsilon} \in (0, B_{b})$ , let  $\sigma_{a} = \Theta \left(\frac{\tilde{b}}{-\ell'(0)\eta^{(1)}B_{G}B_{\epsilon}}\right)$  and  $\delta = 2re^{-\sqrt{mp}}$ . Then, with probability at least  $1 - \delta$  over the initialization, there exists  $\tilde{\mathbf{a}}_i$ 's such that  $f_{(\tilde{\mathbf{a}},\mathbf{W}^{(1)},\mathbf{b})}(\mathbf{x}) = \sum_{i=1}^{4m}\tilde{\mathbf{a}}_i\sigma\left(\left\langle\mathbf{w}_i^{(1)},\mathbf{x}\right\rangle - \mathbf{b}_i\right)$  satisfies
$$
\mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) \leq r B _ {a 1} \left(\frac {B _ {x 1} ^ {2} B _ {b}}{\sqrt {m p} B _ {G} B _ {\epsilon}} + B _ {x 1} \sqrt {2 \gamma} + B _ {\epsilon}\right) + \mathrm {O P T} _ {d, r, B _ {F}, S _ {p}, \gamma , B _ {G}}, \tag {36}
$$
and  $\| \tilde{\mathbf{a}} \|_0 = O\left(r(mp)^{\frac{1}{2}}\right),\| \tilde{\mathbf{a}} \|_2 = O\left(\frac{B_{a2}B_b}{\tilde{b}(mp)^{\frac{1}{4}}}\right),\| \tilde{\mathbf{a}} \|_\infty = O\left(\frac{B_{a1}B_b}{\tilde{b}(mp)^{\frac{1}{2}}}\right).$
Proof of Lemma D.4. Recall  $f^{*}(\mathbf{x}) = \sum_{j=1}^{r} \mathbf{a}_{j}^{*} \sigma(\langle \mathbf{w}_{j}^{*}, \mathbf{x} \rangle - \mathbf{b}_{j}^{*})$ , where  $f^{*} \in \mathcal{F}_{d,r,B_{F},S_{p,\gamma},B_{G}}$  is defined in Definition 3.11 and let  $s_{j}^{*} = \frac{\mathbf{b}_{j}^{*}}{|\mathbf{b}_{j}^{*}|}$ . By Lemma D.3, with probability at least  $1 - \delta_{1}$ ,  $\delta_{1} = 2re^{-cmp}$ , for all  $j \in [r]$ , we have  $|G_{(\mathbf{w}_{j}^{*}, s_{j}^{*}), Nice}| \geq \frac{mp}{4}$ . Then for all  $i \in G_{(\mathbf{w}_{j}^{*}, s_{j}^{*}), Nice} \subseteq [2m]$ , we have  $-\ell'(0)\eta^{(1)} G(\mathbf{w}_{i}^{(0)}, \mathbf{b}_{i})\frac{\mathbf{b}_{j}^{*}}{b}$  only depend on  $\mathbf{w}_{i}^{(0)}$  and  $\mathbf{b}_{i}$ , which is independent of  $\mathbf{a}_{i}^{(0)}$ . Given Definition 3.7, we have
$$
- \ell^ {\prime} (0) \eta^ {(1)} \| G \left(\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}\right) \| _ {2} \frac {\mathbf {b} _ {j} ^ {*}}{\tilde {b}} \in \left[ \ell^ {\prime} (0) \eta^ {(1)} B _ {x 1} \frac {B _ {b}}{\tilde {b}}, - \ell^ {\prime} (0) \eta^ {(1)} B _ {x 1} \frac {B _ {b}}{\tilde {b}} \right]. \tag {37}
$$
We split  $[r]$  into  $\Gamma = \{j\in [r]:|\mathbf{b}_j^* | < B_\epsilon \}$ ,  $\Gamma_{-} = \{j\in [r]:\mathbf{b}_{j}^{*}\leq -B_{\epsilon}\}$  and  $\Gamma_{+} = \{j\in [r]:$ $\mathbf{b}_j^*\geq B_\epsilon \}$ . Let  $\epsilon_{a} = \frac{B_{x1}B_{b}}{\sqrt{mp}B_{G}B_{\epsilon}}$ . Then we know that for all  $j\in \Gamma_{+}\cup \Gamma_{-}$ , for all  $i\in G_{(\mathbf{w}_j^*,s_j^*)},Nice,$  we have
$$
\begin{array}{l} \Pr_ {\mathbf {a} _ {i} ^ {(0)} \sim \mathcal {N} \left(0, \sigma_ {a} ^ {2}\right)} \left[ \left| - \mathbf {a} _ {i} ^ {(0)} \ell^ {\prime} (0) \eta^ {(1)} \| G \left(\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}\right) \| _ {2} \frac {\left| \mathbf {b} _ {j} ^ {*} \right|}{\tilde {b}} - 1 \right| \leq \epsilon_ {a} \right] (38) \\ = \Pr_ {\mathbf {a} _ {i} ^ {(0)} \sim \mathcal {N} \left(0, \sigma_ {a} ^ {2}\right)} \left[ 1 - \epsilon_ {a} \leq - \mathbf {a} _ {i} ^ {(0)} \ell^ {\prime} (0) \eta^ {(1)} \| G \left(\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}\right) \| _ {2} \frac {\left| \mathbf {b} _ {j} ^ {*} \right|}{\tilde {b}} \leq 1 + \epsilon_ {a} \right] (39) \\ = \Pr_ {g \sim \mathcal {N} (0, 1)} \left[ 1 - \epsilon_ {a} \leq g \Theta \left(\frac {\| G \left(\mathbf {w} _ {i} ^ {(0)} , \mathbf {b} _ {i}\right) \| _ {2} \left| \mathbf {b} _ {j} ^ {*} \right|}{B _ {G} B _ {\epsilon}}\right) \leq 1 + \epsilon_ {a} \right] (40) \\ = \operatorname * {P r} _ {g \sim \mathcal {N} (0, 1)} \left[ (1 - \epsilon_ {a}) \Theta \left(\frac {B _ {G} B _ {\epsilon}}{\| G (\mathbf {w} _ {i} ^ {(0)} , \mathbf {b} _ {i}) \| _ {2} | \mathbf {b} _ {j} ^ {*} |}\right) \leq g \leq (1 + \epsilon_ {a}) \Theta \left(\frac {B _ {G} B _ {\epsilon}}{\| G (\mathbf {w} _ {i} ^ {(0)} , \mathbf {b} _ {i}) \| _ {2} | \mathbf {b} _ {j} ^ {*} |}\right) \right] \\ = \Theta \left(\frac {\epsilon_ {a} B _ {G} B _ {\epsilon}}{\| G \left(\mathbf {w} _ {i} ^ {(0)} , \mathbf {b} _ {i}\right) \| _ {2} \left| \mathbf {b} _ {j} ^ {*} \right|}\right) (41) \\ \geq \Omega \left(\frac {\epsilon_ {a} B _ {G} B _ {\epsilon}}{B _ {x 1} B _ {b}}\right) (42) \\ = \Omega \left(\frac {1}{\sqrt {m p}}\right). (43) \\ \end{array}
$$
Thus, with probability  $\Omega\left(\frac{1}{\sqrt{mp}}\right)$  over  $\mathbf{a}_i^{(0)}$ , we have
$$
\left| - \mathbf {a} _ {i} ^ {(0)} \ell^ {\prime} (0) \eta^ {(1)} \| G \left(\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}\right) \| _ {2} \frac {\left| \mathbf {b} _ {j} ^ {*} \right|}{\tilde {b}} - 1 \right| \leq \epsilon_ {a}, \quad \left| \mathbf {a} _ {i} ^ {(0)} \right| = O \left(\frac {\tilde {b}}{- \ell^ {\prime} (0) \eta^ {(1)} B _ {G} B _ {\epsilon}}\right). \tag {44}
$$
Similarly, for  $j \in \Gamma$ , for all  $i \in G_{(\mathbf{w}_j^*, s_j^*), Nice}$ , with probability  $\Omega\left(\frac{1}{\sqrt{mp}}\right)$  over  $\mathbf{a}_i^{(0)}$ , we have
$$
\left| - \mathbf {a} _ {i} ^ {(0)} \ell^ {\prime} (0) \eta^ {(1)} \| G \left(\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}\right) \| _ {2} \frac {B _ {\epsilon}}{\tilde {b}} - 1 \right| \leq \epsilon_ {a}, \quad \left| \mathbf {a} _ {i} ^ {(0)} \right| = O \left(\frac {\tilde {b}}{- \ell^ {\prime} (0) \eta^ {(1)} B _ {G} B _ {\epsilon}}\right). \tag {45}
$$
For all  $j \in [r]$ , let  $\Lambda_j \subseteq G_{(\mathbf{w}_j^*, s_j^*)}$ . Nice be the set of  $i$ 's such that condition Equation (44) or Equation (45) are satisfied. By Chernoff bound and union bound, with probability at least  $1 - \delta_2$ ,  $\delta_2 = re^{-\sqrt{mp}}$ , for all  $j \in [r]$  we have  $|\Lambda_j| \geq \Omega(\sqrt{mp})$ .
We have for  $\forall j\in \Gamma_{+}\cup \Gamma_{-},\forall i\in \Lambda_{j}$
$$
\begin{array}{l} \left| \frac {\left| \mathbf {b} _ {j} ^ {*} \right|}{\tilde {b}} \left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle \right| (46) \\ \leq \left\| - \mathbf {a} _ {i} ^ {(0)} \ell^ {\prime} (0) \eta^ {(1)} \| G \left(\mathbf {w} _ {i} ^ {(0)}, \mathbf {b} _ {i}\right) \| _ {2} \frac {\left| \mathbf {b} _ {j} ^ {*} \right|}{\tilde {b}} \frac {\mathbf {w} _ {i} ^ {(1)}}{\left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}} - \frac {\mathbf {w} _ {i} ^ {(1)}}{\left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}} + \frac {\mathbf {w} _ {i} ^ {(1)}}{\left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}} - \mathbf {w} _ {j} ^ {*} \right\| \| \mathbf {x} \| _ {2} (47) \\ \leq \left(\epsilon_ {a} + \sqrt {2 \gamma}\right) \| \mathbf {x} \| _ {2}. (48) \\ \end{array}
$$
Similarly, for  $\forall j\in \Gamma ,\forall i\in \Lambda_{j}$
$$
\left| \frac {B _ {\epsilon}}{\tilde {b}} \left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle \right| \leq \left(\epsilon_ {a} + \sqrt {2 \gamma}\right) \| \mathbf {x} \| _ {2}. \tag {49}
$$
If  $i\in \Lambda_j,j\in \Gamma_+\cup \Gamma_-$  , set  $\tilde{\mathbf{a}}_i = \mathbf{a}_j^*\frac{|\mathbf{b}_j^*|}{|\Lambda_j|\tilde{b}}$  , if  $i\in \Lambda_j,j\in \Gamma$  , set  $\tilde{\mathbf{a}}_i = \mathbf{a}_j^*\frac{B_\epsilon}{|\Lambda_j|\tilde{b}}$  otherwise set  $\tilde{\mathbf{a}}_i = 0$  we have  $\| \tilde{\mathbf{a}}\| _0 = O\left(r(mp)^{\frac{1}{2}}\right),\| \tilde{\mathbf{a}}\| _2 = O\left(\frac{B_{a2}B_b}{\tilde{b}(mp)^{\frac{1}{4}}}\right),\| \tilde{\mathbf{a}}\|_{\infty} = O\left(\frac{B_{a1}B_b}{\tilde{b}(mp)^{\frac{1}{2}}}\right).$
$$
\begin{array}{l} \mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) (50) \\ = \mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) - \mathcal {L} _ {\mathcal {D}} \left(f ^ {*}\right) + \mathcal {L} _ {\mathcal {D}} \left(f ^ {*}\right) (51) \\ \leq \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left| f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})} (\mathbf {x}) - f ^ {*} (\mathbf {x}) \right| \right] + \mathcal {L} _ {\mathcal {D}} \left(f ^ {*}\right) (52) \\ \end{array}
$$
$$
\begin{array}{l} \leq \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left| \sum_ {i = 1} ^ {m} \tilde {\mathbf {a}} _ {i} \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \tilde {b}\right) + \sum_ {i = m + 1} ^ {2 m} \tilde {\mathbf {a}} _ {i} \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle + \tilde {b}\right) - \sum_ {j = 1} ^ {r} \mathbf {a} _ {j} ^ {*} \sigma \left(\left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle - \mathbf {b} _ {j} ^ {*}\right) \right| \right] \\ + \mathcal {L} _ {\mathcal {D}} \left(f ^ {*}\right) \tag {53} \\ \end{array}
$$
$$
\begin{array}{l} \leq \mathbb {E} _ {(\mathbf {x}, y)} \left[ \right.\left| \sum_ {j \in \Gamma_ {+}} \sum_ {i \in \Lambda_ {j}} \mathbf {a} _ {j} ^ {*} \frac {1}{| \Lambda_ {j} |} \right| \frac {| \mathbf {b} _ {j} ^ {*} |}{\tilde {b}} \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \tilde {b}\right) - \sigma \left(\left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle - \mathbf {b} _ {j} ^ {*}\right)\left. \right|\left. \right] (54) \\ + \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left| \sum_ {j \in \Gamma_ {-}} \sum_ {i \in \Lambda_ {j}} \mathbf {a} _ {j} ^ {*} \frac {1}{| \Lambda_ {j} |} \left| \frac {| \mathbf {b} _ {j} ^ {*} |}{\tilde {b}} \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle + \tilde {b}\right) - \sigma \left(\left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle - \mathbf {b} _ {j} ^ {*}\right) \right| \right| \right] (55) \\ + \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left| \sum_ {j \in \Gamma} \sum_ {i \in \Lambda_ {j}} \mathbf {a} _ {j} ^ {*} \frac {1}{| \Lambda_ {j} |} \left| \frac {B _ {\epsilon}}{\tilde {b}} \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \tilde {b}\right) - \sigma \left(\left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle - \mathbf {b} _ {j} ^ {*}\right) \right| \right] + \mathcal {L} _ {\mathcal {D}} (f ^ {*}) \right. (56) \\ \end{array}
$$
$$
\begin{array}{l} \leq \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left| \sum_ {j \in \Gamma_ {+}} \sum_ {i \in \Lambda_ {j}} \mathbf {a} _ {j} ^ {*} \frac {1}{| \Lambda_ {j} |} \left| \frac {| \mathbf {b} _ {j} ^ {*} |}{\tilde {b}} \left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle \right| \right| \right] (57) \\ + \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left| \sum_ {j \in \Gamma_ {-}} \sum_ {i \in \Lambda_ {j}} \mathbf {a} _ {j} ^ {*} \frac {1}{| \Lambda_ {j} |} \left| \frac {| \mathbf {b} _ {j} ^ {*} |}{\tilde {b}} \left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle \right| \right| \right] (58) \\ + \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left| \sum_ {j \in \Gamma} \sum_ {i \in \Lambda_ {j}} \mathbf {a} _ {j} ^ {*} \frac {1}{| \Lambda_ {j} |} \left| \frac {B _ {\epsilon}}{\tilde {b}} \left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle + B _ {\epsilon} - \left\langle \mathbf {w} _ {j} ^ {*}, \mathbf {x} \right\rangle \right| \right| \right] + \mathcal {L} _ {\mathcal {D}} (f ^ {*}) (59) \\ \end{array}
$$
$$
\begin{array}{l} \leq r \| \mathbf {a} ^ {*} \| _ {\infty} \left(\epsilon_ {a} + \sqrt {2 \gamma}\right) \mathbb {E} _ {(\mathbf {x}, y)} \| \mathbf {x} \| _ {2} + | \Gamma | \| \mathbf {a} ^ {*} \| _ {\infty} B _ {\epsilon} + \mathcal {L} _ {\mathcal {D}} (f ^ {*}) (60) \\ \leq r B _ {x 1} B _ {a 1} \left(\epsilon_ {a} + \sqrt {2 \gamma}\right) + \left| \Gamma \right| B _ {a 1} B _ {\epsilon} + \operatorname {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}}. (61) \\ \end{array}
$$
We finish the proof by union bound and  $\delta \geq \delta_1 + \delta_2$ .
D.2.3 Learning an Accurate Classifier
We will use the following theorem from existing work to prove that gradient descent learns a good classifier (Theorem D.9). Theorem D.1 is simply a direct corollary of Theorem D.9.
Theorem D.5 (Theorem 13 in [33]). Fix some  $\eta$ , and let  $f_{1},\ldots ,f_{T}$  be some sequence of convex functions. Fix some  $\theta_{1}$ , and assume we update  $\theta_{t + 1} = \theta_t - \eta \nabla f_t(\theta_t)$ . Then for every  $\theta^{*}$  the following holds:
$$
\frac {1}{T} \sum_ {t = 1} ^ {T} f _ {t} (\theta_ {t}) \leq \frac {1}{T} \sum_ {t = 1} ^ {T} f _ {t} (\theta^ {*}) + \frac {1}{2 \eta T} \| \theta^ {*} \| _ {2} ^ {2} + \| \theta_ {1} \| _ {2} \frac {1}{T} \sum_ {t = 1} ^ {T} \| \nabla f _ {t} (\theta_ {t}) \| _ {2} + \eta \frac {1}{T} \sum_ {t = 1} ^ {T} \| \nabla f _ {t} (\theta_ {t}) \| _ {2} ^ {2}.
$$
To apply the theorem we first present a few lemmas bounding the change in the network during steps.
Lemma D.6 (Bound of  $\Xi^{(0)},\Xi^{(1)})$  . Assume the same conditions as in Lemma D.4, and  $d\geq \log m$  with probability at least  $1 - \delta -\frac{1}{m^2}$  over the initialization,  $\| \mathbf{a}^{(0)}\|_{\infty} = O\left(\frac{\tilde{b}\sqrt{\log m}}{-\ell'(0)\eta^{(1)}B_{G}B_{\epsilon}}\right)$  , and for all  $i\in [4m]$  , we have  $\| \mathbf{w}_i^{(0)}\| _2 = O\left(\sigma_w\sqrt{d}\right)$  Finally,  $\| \mathbf{a}^{(1)}\|_{\infty} = O\left(-\eta^{(1)}\ell '(0)(B_{x1}\sigma_w\sqrt{d} +\tilde{b})\right)$  and for all  $i\in [4m]$ $\| \mathbf{w}_i^{(1)}\| _2 = O\left(\frac{\tilde{b}\sqrt{\log m}B_{x1}}{B_GB_\epsilon}\right)$
Proof of Lemma D.6. By Lemma F.4, we have  $\| \mathbf{a}^{(0)}\|_{\infty} = O\left(\frac{\tilde{b}\sqrt{\log m}}{-\ell'(0)\eta^{(1)}B_{G}B_{\epsilon}}\right)$  with probability at least  $1 - \frac{1}{2m^2}$  by property of maximum i.i.d Gaussians. For any  $i\in [4m]$ , by Lemma F.5 and  $d\geq \log m$ , we have
$$
\Pr \left(\frac {1}{\sigma_ {w} ^ {2}} \left\| \mathbf {w} _ {i} ^ {(0)} \right\| _ {2} ^ {2} \geq d + 2 \sqrt {4 d \log (m)} + 8 \log (m)\right) \leq O \left(\frac {1}{m ^ {4}}\right). \tag {62}
$$
Thus, by union bound, with probability at least  $1 - \frac{1}{2m^2}$ , for all  $i \in [4m]$ , we have  $\| \mathbf{w}_i^{(0)}\|_2 = O\left(\sigma_w\sqrt{d}\right)$ .
For all  $i\in [4m]$  , we have
$$
\begin{array}{l} \left. \left| \mathbf {a} _ {i} ^ {(1)} \right| = - \eta^ {(1)} \ell^ {\prime} (0) \left| \mathbb {E} _ {(\mathbf {x}, y)} \left[ y \left[ \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(0)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i}\right) \right] \right] \right| \right. (63) \\ \leq - \eta^ {(1)} \ell^ {\prime} (0) \left(\| \mathbf {w} _ {i} ^ {(0)} \| _ {2} \mathbb {E} _ {(\mathbf {x}, y)} [ \| \mathbf {x} \| _ {2} ] + \tilde {b}\right) (64) \\ \leq O \left(- \eta^ {(1)} \ell^ {\prime} (0) \left(B _ {x 1} \sigma_ {w} \sqrt {d} + \tilde {b}\right)\right). (65) \\ \end{array}
$$
$$
\begin{array}{l} \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} = - \eta^ {(1)} \ell^ {\prime} (0) \left\| \mathbf {a} _ {i} ^ {(0)} \mathbb {E} _ {(\mathbf {x}, y)} \left[ y \sigma^ {\prime} \left[ \left\langle \mathbf {w} _ {i} ^ {(0)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i} \right] \mathbf {x} \right] \right\| _ {2} (66) \\ \leq O \left(\frac {\tilde {b} \sqrt {\log m} B _ {x 1}}{B _ {G} B _ {\epsilon}}\right). (67) \\ \end{array}
$$
Lemma D.7 (Bound of  $\Xi^{(t)}$ ). Assume the same conditions as in Lemma D.6, and let  $\eta = \eta^{(t)}$  for all  $t \in \{2, 3, \ldots, T\}$ ,  $0 < T\eta B_{x1} \leq o(1)$ , and  $0 = \lambda = \lambda^{(t)}$  for all  $t \in \{2, 3, \ldots, T\}$ , for all  $i \in [4m]$ , we have
$$
\left| \mathbf {a} _ {i} ^ {(t)} \right| \leq O \left(\left| \mathbf {a} _ {i} ^ {(1)} \right| + \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} + \frac {\tilde {b}}{B _ {x 1}} + \eta \tilde {b}\right) \tag {68}
$$
$$
\left\| \mathbf {w} _ {i} ^ {(t)} - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} \leq O \left(t \eta B _ {x 1} \left| \mathbf {a} _ {i} ^ {(1)} \right| + t \eta^ {2} B _ {x 1} ^ {2} \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} + t \eta^ {2} B _ {x 1} \tilde {b}\right). \tag {69}
$$
Proof of Lemma D.7. For all  $i\in [4m]$ , by Lemma D.6,
$$
\begin{array}{l} \left. \left| \mathbf {a} _ {i} ^ {(t)} \right| = \left| (1 - \eta \lambda) \mathbf {a} _ {i} ^ {(t - 1)} - \eta \mathbb {E} _ {(\mathbf {x}, y)} \left[ \ell^ {\prime} \left(y f _ {\Xi^ {(t - 1)} (\mathbf {x})}\right) y \left[ \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(t - 1)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i}\right) \right] \right] \right| \right| (70) \\ \leq \left| (1 - \eta \lambda) \mathbf {a} _ {i} ^ {(t - 1)} \right| + \eta \left| \mathbb {E} _ {(\mathbf {x}, y)} \left[ \left[ \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(t - 1)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i}\right) \right] \right] \right| (71) \\ \leq \left| \mathbf {a} _ {i} ^ {(t - 1)} \right| + \eta \left(B _ {x 1} \| \mathbf {w} _ {i} ^ {(t - 1)} \| _ {2} + \tilde {b}\right) (72) \\ \leq \left| \mathbf {a} _ {i} ^ {(t - 1)} \right| + \eta B _ {x 1} \| \mathbf {w} _ {i} ^ {(t - 1)} - \mathbf {w} _ {i} ^ {(1)} \| _ {2} + \eta B _ {x 1} \| \mathbf {w} _ {i} ^ {(1)} \| _ {2} + \eta \tilde {b} (73) \\ = \left| \mathbf {a} _ {i} ^ {(t - 1)} \right| + \eta B _ {x 1} \| \mathbf {w} _ {i} ^ {(t - 1)} - \mathbf {w} _ {i} ^ {(1)} \| _ {2} + \eta Z _ {i}, (74) \\ \end{array}
$$
where we denote  $Z_{i} = B_{x1}\| \mathbf{w}_{i}^{(1)}\|_{2} + \tilde{b}$ . Then we give a bound of the first layer's weights change,
$$
\begin{array}{l} \left\| \mathbf {w} _ {i} ^ {(t)} - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} (75) \\ = \left\| (1 - \eta \lambda) \mathbf {w} _ {i} ^ {(t - 1)} - \eta \mathbf {a} _ {i} ^ {(t - 1)} \mathbb {E} _ {(\mathbf {x}, y)} \left[ \ell^ {\prime} \left(y f _ {\Xi^ {(t - 1)} (\mathbf {x})}\right) y \sigma^ {\prime} \left[ \left\langle \mathbf {w} _ {i} ^ {(t - 1)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i} \right] \mathbf {x} \right] - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} (76) \\ \leq \left\| \mathbf {w} _ {i} ^ {(t - 1)} - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} + \eta B _ {x 1} \left| \mathbf {a} _ {i} ^ {(t - 1)} \right|. (77) \\ \end{array}
$$
Combine two bounds, we can get