SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} f ^ {*} (\mathbf {x}) = \sum_ {i = 0} ^ {k} (- 1) ^ {i} \sqrt {k} \tag {503} \\ \cdot \left[ \sigma \left(\langle D, \mathbf {x} \rangle - \frac {2 i - k - 1}{\sqrt {k}}\right) - 2 \sigma \left(\langle D, \mathbf {x} \rangle - \frac {2 i - k}{\sqrt {k}}\right) + \sigma \left(\langle D, \mathbf {x} \rangle - \frac {2 i - k + 1}{\sqrt {k}}\right) \right]. \\ \end{array}
$$
For  $\mathcal{D}_{\text{parity-uniform setting}}$ , we have  $f^{*} \in \mathcal{F}_{d,3(k + 1),B_{F},S_{p,\gamma_{\infty},B_{G},B_{G1}}^{\infty}}$  where  $B_{F} = (B_{a1},B_{a2},B_{b}) = \left(2\sqrt{k},2\sqrt{(k(k + 1))},\frac{k + 1}{\sqrt{k}}\right)$ ,  $p = \Theta \left(\frac{1}{2^{k - 1}}\right)$ ,  $\gamma_{\infty} = O\left(\frac{\sqrt{k}}{d - k}\right)$ ,  $B_{G} = \Theta(B_{G1}) = \Theta(d^{-k})$  and  $B_{x1} = \sqrt{d}$ ,  $B_{x2} = d$ . We also have  $\mathrm{OPT}_{d,3(k + 1),B_{F},S_{p,\gamma_{\infty},B_{G},B_{G1}}^{\infty}} = 0$ .
Proof of Lemma E.39. Fix index  $i$ , with probability  $p_1 = \Theta(2^{-k})$ , we will have  $\mathbf{w}_{i,j}^{(0)} = \mathrm{sign}(\mathbf{a}_i^{(0)}) \cdot \mathrm{sign}(\xi_{k-1})$ , for  $\forall j$ . For  $\mathbf{w}_i^{(0)}$  that satisfy these conditions, we will have:
$$
\operatorname {s i g n} \left(\mathbf {a} _ {i} ^ {(0)}\right) g _ {i, j} = | \xi_ {k - 1} |, \forall j \in A \tag {504}
$$
$$
\operatorname {s i g n} \left(\mathbf {a} _ {i} ^ {(0)}\right) g _ {i, j} = \left| \xi_ {k + 1} \right|, \forall j \notin A. \tag {505}
$$
Then by Lemma 4 in [15], we have
$$
\begin{array}{l} \left\| \frac {\operatorname {s i g n} \left(\mathbf {a} _ {i} ^ {(0)}\right) G \left(\mathbf {w} _ {i} ^ {(0)} , \tilde {b}\right)}{\| G \left(\mathbf {w} _ {i} ^ {(0)} , \tilde {b}\right) \|} - D \right\| _ {\infty} \leq \max  \left\{\left| \frac {1}{k \sqrt {\frac {1}{k} + \frac {1}{d - k}}} - \frac {1}{\sqrt {k}} \right|, \left| \frac {1}{(d - k) \sqrt {\frac {1}{k} + \frac {1}{d - k}}} \right| \right\} (506) \\ \leq \frac {\sqrt {k}}{d - k} (507) \\ \end{array}
$$
$$
\left\| \operatorname {s i g n} \left(\mathbf {a} _ {i} ^ {(0)}\right) G \left(\mathbf {w} _ {i} ^ {(0)}, \tilde {b}\right) \right\| _ {2} = \sqrt {k \left| \xi_ {k - 1} \right| ^ {2} + (d - k) \left| \xi_ {k + 1} \right| ^ {2}} = \Theta \left(d ^ {\Theta (k)}\right). \tag {508}
$$
From here, we can see that if we set  $\gamma_{\infty} = \frac{\sqrt{k}}{d - k}$ ,  $B_{G} = B_{G1} = \sqrt{k|\xi_{k - 1}|^{2} + (d - k)|\xi_{k + 1}|^{2}}$ ,  $p = p_1$ , we will have  $(D, + 1)$ ,  $(D, - 1) \in S_{p,\gamma_{\infty},B_{G},B_{G1}}^{\infty}$  by our symmetric initialization. As a result, we have  $f^{*} \in \mathcal{F}_{d,3(k + 1),B_{F},S_{p,\gamma_{\infty},B_{G},B_{G1}}^{\infty}}$ . Finally, it is easy to verify that  $f^{*}(\mathbf{x}) = \mathrm{XOR}(\mathbf{x}_{A})$ , thus  $\mathrm{OPT}_{d,3(k + 1),B_F,S_{p,\gamma_\infty,B_G,B_{G1}}^\infty} = 0$ .
Theorem E.40 (Uniform Parity Functions: Main Result (Alternative)). For  $\mathcal{D}_{\text{parity-uniform setting}}$ , for any  $\delta \in (0,1)$  satisfying  $\delta \leq O\left(\frac{1}{d^2}\right)$  and for any  $\epsilon \in (0,1)$  when
$$
m = \operatorname {p o l y} \left(\log \left(\frac {1}{\delta}\right), \frac {1}{\epsilon}, 2 ^ {\Theta (k)}, d\right), T = \Theta \left(d ^ {\Theta (k)}\right), n = \Theta \left(d ^ {\Theta (k)}\right) \tag {509}
$$
$$
\Pr \left[ \operatorname {s i g n} \left(f _ {\Xi^ {(t)}} (\mathbf {x})\right) \neq y \right] \leq \frac {k ^ {2} \sqrt {d \log (d)}}{d - k} + \epsilon . \tag {510}
$$
Proof of Theorem E.40. Plug the values of parameters into Theorem E.37 and directly get the result.
E.7 Multiple Index Model with Low Degree Polynomial
E.7.1 Problem Setup
The multiple-index data problem has been used for studying network learning [18, 32]. We consider proving guarantees for the setting in [32], following our framework. We use the properties of the problem to prove the key lemma (i.e., the existence of good networks) in our framework and then derive the final guarantee from our theorem of the simple setting (Theorem 3.4).
Data Distributions. We draw input from the distribution  $\mathcal{D}_{\mathcal{X}} = \mathcal{N}(0, I_{d \times d})$ , and we assume the target function is  $g^{*}(\mathbf{x}): \mathbb{R}^{d} \to \mathbb{R}$ , where  $g^{*}$  is a degree  $\tau$  polynomial normalized so that  $\mathbb{E}_{\mathbf{x} \sim \mathcal{D}_{\mathcal{X}}} [g^{*}(\mathbf{x})^{2}] = 1$ .
Assumption E.41. There exists linearly independent vectors  $u_{1}, \ldots, u_{r}$  such that  $g^{*}(\mathbf{x}) = g(\langle \mathbf{x}, u_{1} \rangle, \ldots, \langle \mathbf{x}, u_{r} \rangle)$ .  $H := \mathbb{E}_{\mathbf{x} \sim \mathcal{D}_{\mathcal{X}}}[\nabla^{2} g^{*}(\mathbf{x})]$  has rank  $r$ , where  $H$  is a Hessian matrix.
Definition E.42. Denote the normalized condition number of  $H$  by
$$
\kappa := \frac {\| H ^ {\dagger} \|}{\sqrt {r}}. \tag {511}
$$
Initialization and Loss. For  $\forall i\in [m]$ , we use the following initialization:
$$
\mathbf {a} _ {i} ^ {(0)} \sim \{- 1, 1 \}, \quad \mathbf {w} _ {i} ^ {(0)} \sim \mathcal {N} \left(0, \frac {1}{d} I _ {d \times d}\right) \text {a n d} \mathbf {b} _ {i} = 0. \tag {512}
$$
For this regression problem, we use mean square loss:
$$
\mathcal {L} _ {\mathcal {D} _ {\chi}} \left(f _ {\Xi}\right) = \mathbb {E} _ {\mathbf {x} \sim \mathcal {D} _ {\chi}} \left[ \left(f _ {\Xi} (\mathbf {x}) - g ^ {*} (\mathbf {x})\right) ^ {2} \right]. \tag {513}
$$
Algorithm 5 Network Training via Gradient Descent [32]. Special case of Algorithm 2
Initialize  $(\mathbf{a}^{(0)},\mathbf{W}^{(0)},\mathbf{b})$  as in Equation (8) and Equation (512); Sample  $\mathcal{Z}\sim \mathcal{D}_{\mathcal{X}}^n$ $\rho = \frac{1}{n}\sum_{\mathbf{x}\in \mathcal{Z}}g^{*}(\mathbf{x}),\beta = \frac{1}{n}\sum_{\mathbf{x}\in \mathcal{Z}}g^{*}(\mathbf{x})\mathbf{x}$ $y = g^{*}(\mathbf{x}) - \rho -\beta \cdot \mathbf{x}$ $\mathbf{W}^{(1)} = \mathbf{W}^{(0)} - \eta^{(1)}(\nabla_{\mathbf{W}}\widetilde{\mathcal{L}}_{\mathcal{Z}}(f_{\Xi^{(0)}}) + \lambda^{(1)}\mathbf{W}^{(0)})$    
Re-initialize  $\mathbf{b}_i\sim \mathcal{N}(0,1)$    
for  $t = 2$  to  $T$  do  $\mathbf{a}^{(t)} = \mathbf{a}^{(t - 1)} - \eta^{(t)}\nabla_{\mathbf{a}}\widetilde{\mathcal{L}}_{\mathcal{Z}}(f_{\Xi^{(t - 1)}})$
Lemma E.43 (Multiple Index Model with Low Degree Polynomial: Existence of Good Networks. Rephrase of Lemma 25 in [32]). Assume  $n \geq d^2 r\kappa^2 (C_l\log (nmd))^{\tau +1}$ ,  $d \geq C_d\kappa r^{3 / 2}$ , and  $m \geq r^{\tau}\kappa^{2\tau}(C_l\log (nmd))^{6\tau +1}$  for sufficiently large constants  $C_d, C_l$ , and let  $\eta^{(1)} = \sqrt{\frac{d}{(C_l\log(nmd))^3}}$  and  $\lambda^{(1)} = \frac{1}{\eta^{(1)}}$ . Then with probability  $1 - \frac{1}{\mathrm{poly}(m,d)}$ , there exists  $\tilde{\mathbf{a}} \in \mathbb{R}^m$  such that  $f_{(\tilde{\mathbf{a}},\mathbf{W}^{(1)},\mathbf{b})}$  satisfies
$$
\mathcal {L} _ {\mathcal {D} _ {\mathcal {X}}} \left(f _ {(\hat {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) \leq O \left(\frac {1}{n} + \frac {r ^ {\tau} \kappa^ {2 \tau} \left(C _ {l} \log (n m d)\right) ^ {6 \tau + 1})}{m}\right) \tag {514}
$$
$$
\left\| \tilde {\mathbf {a}} \right\| _ {2} ^ {2} \leq O \left(\frac {r ^ {\tau} \kappa^ {2 \tau} \left(C _ {l} \log (n m d)\right) ^ {6 \tau}}{m}\right). \tag {515}
$$
E.7.2 Multiple Index Model: Final Guarantee
Considering training by Algorithm 5, we have the following results.
Theorem E.44 (Multiple Index Model with Low Degree Polynomial: Main Result). Assume  $n \geq \Omega \left(d^{2}r\kappa^{2}(C_{l}\log(nmd))^{\tau+1} + m\right)$ ,  $d \geq C_{d}\kappa r^{3/2}$ , and  $m \geq \Omega \left(\frac{1}{\epsilon}r^{\tau}\kappa^{2\tau}(C_{l}\log(nmd))^{6\tau+1}\right)$  for sufficiently large constants  $C_{d}, C_{l}$ . Let  $\eta^{(1)} = \sqrt{\frac{d}{(C_{l}\log(nmd))^{3}}}$  and  $\lambda^{(1)} = \frac{1}{\eta^{(1)}}$ , and  $\eta = \eta^{(t)} = \Theta(m^{-1})$ , for all  $t \in \{2,3,\ldots,T\}$ . For any  $\epsilon \in (0,1)$ , if  $T \geq \Omega\left(\frac{m^{2}}{\epsilon}\right)$ , then with properly set parameters and Algorithm 5, with high probability that there exists  $t \in [T]$  such that
$$
\mathcal {L} _ {\mathcal {D} _ {\mathcal {X}}} f _ {(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b})} \leq \epsilon . \tag {516}
$$
Proof of Theorem E.44. By Lemma E.43, we have for properly chosen hyper-parameters,
$$
\begin{array}{l} \mathrm {O P T} _ {\mathbf {W} ^ {(1)}, \mathbf {b}, B _ {a 2}} \leq \mathcal {L} _ {\mathcal {D} _ {\chi}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) \leq O \left(\frac {1}{n} + \frac {r ^ {\tau} \kappa^ {2 \tau} \left(C _ {l} \log (n m d)\right) ^ {6 \tau + 1})}{m}\right) (517) \\ \leq \frac {\epsilon}{3}. (518) \\ \end{array}
$$
We compute the  $L$ -smooth constant of  $\widetilde{\mathcal{L}}_{\mathcal{Z}}\left(f_{(\mathbf{a},\mathbf{W}^{(1)},\mathbf{b})}\right)$  to a.
$$
\begin{array}{l} \left\| \nabla_ {\mathbf {a}} \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} _ {1}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \nabla_ {\mathbf {a}} \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} _ {2}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) \right\| _ {2} (519) \\ = \left\| \frac {1}{n} \sum_ {\mathbf {x} \in \mathcal {Z}} \left[ 2 \left(f _ {\left(\mathbf {a} _ {1}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)} (\mathbf {x}) - g ^ {*} - f _ {\left(\mathbf {a} _ {2}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)} (\mathbf {x}) + g ^ {*}\right) \sigma \left(\mathbf {W} ^ {(1) \top} \mathbf {x} - \mathbf {b}\right) \right] \right\| _ {2} (520) \\ \leq \left\| \frac {1}{n} \sum_ {\mathbf {x} \in \mathcal {Z}} \left[ 2 \left(\mathbf {a} _ {1} ^ {\top} \sigma \left(\mathbf {W} ^ {(1) \top} \mathbf {x} - \mathbf {b}\right) - \mathbf {a} _ {2} ^ {\top} \sigma \left(\mathbf {W} ^ {(1) \top} \mathbf {x} - \mathbf {b}\right)\right) \sigma \left(\mathbf {W} ^ {(1) \top} \mathbf {x} - \mathbf {b}\right) \right] \right\| _ {2} (521) \\ \leq \frac {1}{n} \sum_ {\mathbf {x} \in \mathcal {Z}} \left[ 2 \| \mathbf {a} _ {1} - \mathbf {a} _ {2} \| _ {2} \left\| \sigma \left(\mathbf {W} ^ {(1) \top} \mathbf {x} - \mathbf {b}\right) \right\| _ {2} ^ {2} \right]. (522) \\ \end{array}
$$
By the proof of Lemma 25 in [32], we have for  $\forall i \in [4m]$ , with probability at least  $1 - \frac{1}{\mathrm{poly}(m,d)}$ ,  $|\langle \mathbf{w}_i, \mathbf{x} \rangle| \leq 1$ , with some large polynomial  $\mathrm{poly}(m,d)$ . As a result, we have
$$
\frac {1}{n} \sum_ {\mathbf {x} \in \mathcal {Z}} \left\| \mathbf {W} ^ {(1) ^ {\top}} \mathbf {x} \right\| _ {2} ^ {2} \leq m + \frac {1}{\operatorname {p o l y} (m , d)} \leq O (m). \tag {523}
$$
Thus, we have,
$$
\begin{array}{l} L = O \left(\frac {1}{n} \sum_ {\mathbf {x} \in \mathcal {Z}} \left\| \sigma \left(\mathbf {W} ^ {(1) \top} \mathbf {x} - \mathbf {b}\right) \right\| _ {2} ^ {2}\right) (524) \\ \leq O \left(\frac {1}{n} \sum_ {\mathbf {x} \in \mathcal {Z}} \left\| \mathbf {W} ^ {(1) \top} \mathbf {x} - \mathbf {b} \right\| _ {2} ^ {2}\right) (525) \\ \leq O \left(\frac {1}{n} \sum_ {\mathbf {x} \in \mathcal {Z}} \left\| \mathbf {W} ^ {(1) \top} \mathbf {x} \right\| _ {2} ^ {2} + \| \mathbf {b} \| _ {2} ^ {2}\right) (526) \\ \leq O (m). (527) \\ \end{array}
$$
This means that we can let  $\eta = \Theta(m^{-1})$  and we will get our convergence result. We can bound  $\|\mathbf{a}^{(1)}\|_2$  and  $\|\tilde{\mathbf{a}}\|_2$  by  $\|\mathbf{a}^{(1)}\|_2 = O(\sqrt{m})$  and  $\|\tilde{\mathbf{a}}\|_2 = O\left(\frac{r^\tau \kappa^{2\tau}(C_l \log(nmd))^{6\tau}}{m}\right) = O(\epsilon)$ . So, if we choose  $T \geq \Omega\left(\frac{m}{\epsilon\eta}\right)$ , there exists  $t \in [T]$  such that  $\widetilde{\mathcal{L}}_{\mathcal{Z}}\left(f_{(\mathbf{a}^{(t)}, \mathbf{W}^{(1)}, \mathbf{b})}\right) - \widetilde{\mathcal{L}}_{\mathcal{Z}}\left(f_{(\tilde{\mathbf{a}}, \mathbf{W}^{(1)}, \mathbf{b})}\right) \leq O\left(\frac{L\|\mathbf{a}^{(1)} - \tilde{\mathbf{a}}\|_2^2}{T}\right) \leq \epsilon/3$ .
Discussion. We would like to unify [32], which are very closely related to our framework: their analysis for multiple index data follows the same principle and analysis approach as our general framework, although it does not completely fit into our Theorem 3.12 due to some technical differences. We can cover it with our Theorem 3.4.
Our work and [32] share the same principle and analysis approach. [32] shows that the first layer learns good features by one gradient step update, which can approximate the true labels by a low-degree polynomial function. Then, a classifier (the second layer) is trained on top of the learned first layer which leads to the final guarantees. This is consistent with our framework: we first show that the first layer learns good features by one gradient step update, which can approximate the true labels, and then show a good classifier can be learned on the first layer.
Our work and [32] have technical differences. First, in the second stage, [32] fix the first layer and only update the top layer which is a convex optimization. Our framework allows updates in the first layer and uses online convex learning techniques for the analysis. Second, they consider the square loss (this is used to calculate Hermite coefficients explicitly for gradients, which are useful in the low-degree polynomial function approximation). While in our online convex learning analysis, we need boundedness of the derivative of the loss to show that the first layer weights' changes are bounded in the second stage. Given the above two technicalities, we analyze their training algorithm (Algorithm 2) which fixes the first layer weights and fits into our Theorem 3.4.
In this section, we present some Lemmas used frequently.
Lemma F.1 (Lemmas on Gradients).
$$
\nabla_ {\mathbf {W}} \mathcal {L} _ {(\mathbf {x}, y)} (f _ {\Xi}) = \left[ \frac {\partial \mathcal {L} _ {(\mathbf {x} , y)} (f _ {\Xi})}{\partial \mathbf {w} _ {1}}, \dots , \frac {\partial \mathcal {L} _ {(\mathbf {x} , y)} (f _ {\Xi})}{\partial \mathbf {w} _ {i}}, \dots , \frac {\partial \mathcal {L} _ {(\mathbf {x} , y)} (f _ {\Xi})}{\partial \mathbf {w} _ {4 m}} \right], \tag {528}
$$
$$
\frac {\partial \mathcal {L} _ {(\mathbf {x} , y)} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} = \mathbf {a} _ {i} \ell^ {\prime} \left(y f _ {\Xi} (\mathbf {x})\right) y \left[ \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {i}, \mathbf {x} \right\rangle - \mathbf {b} _ {i}\right) \right] \mathbf {x}, \tag {529}
$$
$$
\nabla_ {\mathbf {W}} \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi}\right) = \left[ \frac {\partial \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi}\right)}{\partial \mathbf {w} _ {1}}, \dots , \frac {\partial \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi}\right)}{\partial \mathbf {w} _ {i}}, \dots , \frac {\partial \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi}\right)}{\partial \mathbf {w} _ {4 m}} \right], \tag {530}
$$
$$
\frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} = \mathbf {a} _ {i} \mathbb {E} _ {(\mathbf {x}, y)} \left[ \ell^ {\prime} \left(y f _ {\Xi} (\mathbf {x})\right) y \left[ \sigma^ {\prime} \left(\langle \mathbf {w} _ {i}, \mathbf {x} \rangle - \mathbf {b} _ {i}\right) \right] \mathbf {x} \right], \tag {531}
$$
$$
\frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {a} _ {i}} = \mathbb {E} _ {(\mathbf {x}, y)} \left[ \ell^ {\prime} \left(y f _ {\Xi} (\mathbf {x})\right) y \left[ \sigma \left(\langle \mathbf {w} _ {i}, \mathbf {x} \rangle - \mathbf {b} _ {i}\right) \right] \right]. \tag {532}
$$
Proof. These can be verified by direct calculation.
Lemma F.2 (Property of Symmetric Initialization). For any  $\mathbf{x} \in \mathbb{R}^d$ , we have  $f_{\Xi^{(0)}}(\mathbf{x}) = 0$ . For all  $i \in [2m]$ , we have  $\mathbf{w}_i^{(1)} = -\mathbf{w}_{i+2m}^{(1)}$ . When input data is symmetric, i.e.,  $\mathbb{E}_{(\mathbf{x},y)}[y\mathbf{x}] = \mathbf{0}$ , for all  $i \in [m]$ , we have  $\mathbf{w}_i^{(1)} = \mathbf{w}_{i+m}^{(1)}$ .
Proof of Lemma F.2. By symmetric initialization, we have  $f_{\Xi^{(0)}}(\mathbf{x}) = 0$ . For all  $i \in [2m]$ , we have
$$
\begin{array}{l} \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] (533) \\ = \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i + 2 m} ^ {(0)} \mathbb {E} _ {(\mathbf {x}, y)} \left[ y \sigma^ {\prime} \left[ \left\langle \mathbf {w} _ {i + 2 m} ^ {(0)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i + 2 m} \right] \mathbf {x} \right] (534) \\ = - \mathbf {w} _ {i + 2 m} ^ {(1)}. (535) \\ \end{array}
$$
When  $\mathbb{E}_{(\mathbf{x},y)}[y\mathbf{x}] = \mathbf{0}$ , for all  $i\in [m]$ , we have
$$
\begin{array}{l} \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] (536) \\ = \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i + m} ^ {(0)} \mathbb {E} _ {(\mathbf {x}, y)} \left[ y \sigma^ {\prime} \left[ \left\langle - \mathbf {w} _ {i + m} ^ {(0)}, \mathbf {x} \right\rangle + \mathbf {b} _ {i + m} \right] \mathbf {x} \right] (537) \\ = \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i + m} ^ {(0)} \mathbb {E} _ {(\mathbf {x}, y)} \left[ y \sigma^ {\prime} \left[ \left\langle - \mathbf {w} _ {i + m} ^ {(0)}, \mathbf {x} \right\rangle + \mathbf {b} _ {i + m} \right] \mathbf {x} - y \mathbf {x} \right] (538) \\ = \eta^ {(1)} \ell^ {\prime} (0) \mathbf {a} _ {i + m} ^ {(0)} \mathbb {E} _ {(\mathbf {x}, y)} \left[ - y \sigma^ {\prime} \left[ \left\langle \mathbf {w} _ {i + m} ^ {(0)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i + m} \right] \mathbf {x} \right] (539) \\ = \mathbf {w} _ {i + m} ^ {(1)}. (540) \\ \end{array}
$$
Lemma F.3 (Property of Direction Neighborhood). If  $\mathbf{w} \in \mathcal{C}_{D,\gamma}$ , we have  $\rho \mathbf{w} \in \mathcal{C}_{D,\gamma}$  for any  $\rho \neq 0$ . We also have  $\mathbf{0} \notin \mathcal{C}_{D,\gamma}$ . Also, if  $(D,s) \in S_{p,\gamma,B_G}$ , we have  $(-D,s) \in S_{p,\gamma,B_G}$ .
Lemma F.4 (Maximum Gaussian Tail Bound).  $M_{n}$  is the maximum of  $n$  i.i.d. standard normal Gaussian. Then
$$
\Pr \left(M _ {n} \geq \sqrt {2 \log n} + \frac {z}{\sqrt {2 \log n}}\right) \leq e ^ {- z}. \tag {541}
$$
Lemma F.5 (Chi-squared Tail Bound). If  $X$  is a  $\chi^2(k)$  random variable. Then,  $\forall z \in \mathbb{R}$ , we have
$$
\Pr (X \geq k + 2 \sqrt {k z} + 2 z) \leq e ^ {- z}. \tag {542}
$$
Lemma F.6 (Gaussian Tail Bound). If  $g$  is standard Gaussian and  $z > 0$ , we have
$$
\frac {1}{\sqrt {2 \pi}} \frac {z}{z ^ {2} + 1} e ^ {- z ^ {2} / 2} <   \Pr_ {g \sim \mathcal {N} (0, 1)} [ g > z ] <   \frac {1}{\sqrt {2 \pi}} \frac {1}{z} e ^ {- z ^ {2} / 2}. \tag {543}
$$
Lemma F.7 (Gaussian Tail Expectation Bound). If  $g$  is standard Gaussian and  $z \in \mathbb{R}$ , we have
$$
\left| \mathbb {E} _ {g \sim \mathcal {N} (0, 1)} [ \mathbb {I} [ g > z ] g ] \right| <   2 \Pr_ {g \sim \mathcal {N} (0, 1)} [ g > z ] ^ {0. 9}. \tag {544}
$$
Proof of Lemma F.7. For any  $p \in (0,1)$ , we have
$$
\left| \int_ {- \infty} ^ {\sqrt {2} \operatorname {e r f} ^ {- 1} (2 p - 1)} \frac {e ^ {- \frac {x ^ {2}}{2}} x}{\sqrt {2 \pi}} d x \right| <   2 p ^ {0. 9}, \tag {545}
$$
where  $\sqrt{2}\mathrm{erf}^{-1}(2p - 1)$  is the quantile function of the standard Gaussian. We finish the proof by replacing  $p$  to be  $\operatorname*{Pr}_{g\sim \mathcal{N}(0,1)}[g > z]$ .
Lemma F.8. If a function  $g$  satisfies  $h(n + 2) = 2h(n + 1) - (1 - \rho^2)h(n) + \beta$  for  $n \in \mathbb{N}_+$  where  $\rho, \beta > 0$ , then  $h(n) = -\frac{\beta}{\rho^2} + c_1(1 - \rho)^n + c_2(1 + \rho)^n$ , where  $c_1, c_2$  only depends on  $h(1)$  and  $h(2)$ .
Lemma F.9 (Rademacher Complexity Bounds. Rephrase of Lemma 48 in [32]). For fixed  $\mathbf{W}$ ,  $\mathbf{b}$ , let  $\mathcal{F} = \{f_{(\mathbf{a},\mathbf{W},\mathbf{b})}:\| \mathbf{a}\| \leq B_{a2}\}$ . Then,
$$
\Re (\mathcal {F}) \leq \sqrt {\frac {B _ {a 2} ^ {2} \left(\| \mathbf {W} \| _ {F} ^ {2} B _ {x} ^ {2} + \| \mathbf {b} \| _ {2} ^ {2}\right)}{n}}. \tag {546}
$$