SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left| \mathbf {a} _ {i} ^ {(t)} \right| \leq \left| \mathbf {a} _ {i} ^ {(t - 1)} \right| + \eta Z _ {i} + \left(\eta B _ {x 1}\right) ^ {2} \sum_ {l = 1} ^ {t - 2} \left| \mathbf {a} _ {i} ^ {(l)} \right| (78) \\ \Leftrightarrow \sum_ {l = 1} ^ {t} \left| \mathbf {a} _ {i} ^ {(l)} \right| \leq 2 \left(\sum_ {l = 1} ^ {t - 1} \left| \mathbf {a} _ {i} ^ {(l)} \right|\right) - (1 - (\eta B _ {x 1}) ^ {2}) \left(\sum_ {l = 1} ^ {t - 2} \left| \mathbf {a} _ {i} ^ {(l)} \right|\right) + \eta Z _ {i}. (79) \\ \end{array}
$$
Let  $h(1) = |\mathbf{a}_i^{(1)}|$ ,  $h(2) = 2|\mathbf{a}_i^{(1)}| + \eta Z_i$  and  $h(t + 2) = 2h(t + 1) - (1 - (\eta B_{x1})^2)h(t) + \eta Z_i$  for  $n \in \mathbb{N}_+$ , by Lemma F.8, we have
$$
h (t) = - \frac {Z _ {i}}{\eta B _ {x 1} ^ {2}} + c _ {1} \left(1 - \eta B _ {x 1}\right) ^ {(t - 1)} + c _ {2} \left(1 + \eta B _ {x 1}\right) ^ {(t - 1)} \tag {80}
$$
$$
c _ {1} = \frac {1}{2} \left(\left| \mathbf {a} _ {i} ^ {(1)} \right| + \frac {Z _ {i}}{\eta B _ {x 1} ^ {2}} - \frac {\left| \mathbf {a} _ {i} ^ {(1)} \right| + \eta Z _ {i}}{\eta B _ {x 1}}\right) \tag {81}
$$
$$
c _ {2} = \frac {1}{2} \left(\left| \mathbf {a} _ {i} ^ {(1)} \right| + \frac {Z _ {i}}{\eta B _ {x 1} ^ {2}} + \frac {\left| \mathbf {a} _ {i} ^ {(1)} \right| + \eta Z _ {i}}{\eta B _ {x 1}}\right). \tag {82}
$$
Thus, by  $|c_{1}| \leq c_{2}$ , and  $0 < T\eta B_{x1} \leq o(1)$ , we have
$$
\begin{array}{l} \left| \mathbf {a} _ {i} ^ {(t)} \right| \leq h (t) - h (t - 1) (83) \\ = - \eta B _ {x 1} c _ {1} \left(1 - \eta B _ {x 1}\right) ^ {(t - 2)} + \eta B _ {x 1} c _ {2} \left(1 + \eta B _ {x 1}\right) ^ {(t - 2)} (84) \\ \leq 2 \eta B _ {x 1} c _ {2} \left(1 + \eta B _ {x 1}\right) ^ {t} (85) \\ \leq O \left(2 \eta B _ {x 1} c _ {2}\right). (86) \\ \end{array}
$$
Similarly, by binomial approximation, we also have
$$
\begin{array}{l} \left\| \mathbf {w} _ {i} ^ {(t)} - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} \leq \eta B _ {x 1} h (t - 1) (87) \\ = \eta B _ {x 1} \left(- \frac {Z _ {i}}{\eta B _ {x 1} ^ {2}} + c _ {1} (1 - \eta B _ {x 1}) ^ {(t - 2)} + c _ {2} (1 + \eta B _ {x 1}) ^ {(t - 2)}\right) (88) \\ \leq \eta B _ {x 1} O \left(- \frac {Z _ {i}}{\eta B _ {x 1} ^ {2}} + c _ {1} (1 - (t - 2) \eta B _ {x 1}) + c _ {2} (1 + (t - 2) \eta B _ {x 1})\right) (89) \\ \leq \eta B _ {x 1} O \left(- \frac {Z _ {i}}{\eta B _ {x 1} ^ {2}} + c _ {1} + c _ {2} + \left(c _ {2} - c _ {1}\right) t \eta B _ {x 1}\right) (90) \\ \leq \eta B _ {x 1} O \left(\left| \mathbf {a} _ {i} ^ {(1)} \right| + \frac {\left| \mathbf {a} _ {i} ^ {(1)} \right| + \eta Z _ {i}}{\eta B _ {x 1}} t \eta B _ {x 1}\right) (91) \\ \leq O \left(\left(\eta \left| \mathbf {a} _ {i} ^ {(1)} \right| + \eta^ {2} Z _ {i}\right) t B _ {x 1}\right). (92) \\ \end{array}
$$
We finish the proof by plugging  $Z_{i}, c_{2}$  into the bound.
Lemma D.8 (Bound of Loss Gap and Gradient). Assume the same conditions as in Lemma D.7, for all  $t \in [T]$ , we have
$$
\left| \mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(t)}, \mathbf {b})}\right) - \mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) \right| \leq B _ {x 1} \| \tilde {\mathbf {a}} \| _ {2} \sqrt {\| \tilde {\mathbf {a}} \| _ {0}} \max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} ^ {(t)} - \mathbf {w} _ {i} ^ {(1)} \| _ {2} \tag {93}
$$
and for all  $t\in [T]$  ,for all  $i\in [4m]$  ,we have
$$
\left| \frac {\partial \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(t)}}\right)}{\partial \mathbf {a} _ {i} ^ {(t)}} \right| \leq B _ {x 1} \left(\left\| \mathbf {w} _ {i} ^ {(t)} - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} + \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}\right) + \tilde {b}. \tag {94}
$$
Proof of Lemma D.8. It follows from that
$$
\begin{array}{l} \left| \mathcal {L} _ {\mathcal {D}} \left(f _ {(\hat {\mathbf {a}}, \mathbf {W} ^ {(t)}, \mathbf {b})}\right) - \mathcal {L} _ {\mathcal {D}} \left(f _ {(\hat {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) \right| (95) \\ \leq \mathbb {E} _ {(\mathbf {x}, y)} | f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(t)}, \mathbf {b})} (\mathbf {x}) - f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})} (\mathbf {x}) | (96) \\ \leq \mathbb {E} _ {(\mathbf {x}, y)} \left[ \| \tilde {\mathbf {a}} \| _ {2} \sqrt {\| \tilde {\mathbf {a}} \| _ {0}} \max  _ {i \in [ 4 m ]} \left| \sigma \left[ \left\langle \mathbf {w} _ {i} ^ {(t)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i} \right] - \sigma \left[ \left\langle \mathbf {w} _ {i} ^ {(1)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i} \right] \right| \right] (97) \\ \leq B _ {x 1} \| \tilde {\mathbf {a}} \| _ {2} \sqrt {\| \tilde {\mathbf {a}} \| _ {0}} \max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} ^ {(t)} - \mathbf {w} _ {i} ^ {(1)} \| _ {2}. (98) \\ \end{array}
$$
$$
\begin{array}{l} \left. \left| \frac {\partial \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(t)}}\right)}{\partial \mathbf {a} _ {i} ^ {(t)}} \right| = \left| \mathbb {E} _ {(\mathbf {x}, y)} \left[ \ell^ {\prime} \left(y f _ {\Xi^ {(t)}} (\mathbf {x})\right) y \left[ \sigma \left(\left\langle \mathbf {w} _ {i} ^ {(t)}, \mathbf {x} \right\rangle - \mathbf {b} _ {i}\right) \right] \right] \right| \right. (99) \\ \leq B _ {x 1} \left\| \mathbf {w} _ {i} ^ {(t)} \right\| _ {2} + \tilde {b} (100) \\ \leq B _ {x 1} \left(\left\| \mathbf {w} _ {i} ^ {(t)} - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} + \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2}\right) + \tilde {b}. (101) \\ \end{array}
$$
We are now ready to prove the main theorem.
Theorem D.9 (Online Convex Optimization. Full Statement of Theorem D.1). Consider training by Algorithm 3, and any  $\delta \in (0,1)$ . Assume  $d\geq \log m$ . Set
$$
\sigma_ {w} > 0, \quad \tilde {b} > 0, \quad \eta^ {(t)} = \eta , \lambda^ {(t)} = 0 f o r a l l t \in \{2, 3, \dots , T \},
$$
$$
\eta^ {(1)} = \Theta \left(\frac {\min \{O (\eta) , O (\eta \tilde {b}) \}}{- \ell^ {\prime} (0) (B _ {x 1} \sigma_ {w} \sqrt {\tilde {d}} + \tilde {b})}\right), \lambda^ {(1)} = \frac {1}{\eta^ {(1)}}, \sigma_ {a} = \Theta \left(\frac {\tilde {b} (m p) ^ {\frac {1}{4}}}{- \ell^ {\prime} (0) \eta^ {(1)} B _ {x 1} \sqrt {B _ {G} B _ {b}}}\right).
$$
Let  $0 < T\eta B_{x1} \leq o(1)$ ,  $m = \Omega \left(\frac{1}{\sqrt{\delta}} + \frac{1}{p}\left(\log \left(\frac{r}{\delta}\right)\right)^2\right)$ . With probability at least  $1 - \delta$  over the initialization, there exists  $t \in [T]$  such that
$$
\begin{array}{l} \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} \left(\frac {2 B _ {x 1}}{(m p) ^ {\frac {1}{4}}} \sqrt {\frac {B _ {b}}{B _ {G}}} + B _ {x 1} \sqrt {2 \gamma}\right) \tag {102} \\ + \eta \left(\sqrt {r} B _ {a 2} B _ {b} T \eta B _ {x 1} ^ {2} + m \tilde {b}\right) O \left(\frac {\sqrt {\log m} B _ {x 1} (m p) ^ {\frac {1}{4}}}{\sqrt {B _ {b} B _ {G}}} + 1\right) + O \left(\frac {B _ {a 2} ^ {2} B _ {b} ^ {2}}{\eta T \tilde {b} ^ {2} (m p) ^ {\frac {1}{2}}}\right). \\ \end{array}
$$
Furthermore, for any  $\epsilon \in (0,1)$ , set
$$
\tilde {b} = \Theta \left(\frac {B _ {G} ^ {\frac {1}{4}} B _ {a 2} B _ {b} ^ {\frac {3}{4}}}{\sqrt {r B _ {a 1}}}\right), \quad m = \Omega \left(\frac {1}{p \epsilon^ {4}} \left(r B _ {a 1} B _ {x 1} \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 {103}
$$
$$
\eta = \Theta \left(\frac {\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} B _ {x 1} (m p) ^ {\frac {1}{4}}}{\sqrt {B _ {b} B _ {G}}} + 1\right)}\right), \quad T = \Theta \left(\frac {1}{\eta B _ {x 1} (m p) ^ {\frac {1}{4}}}\right), \tag {104}
$$
we have there exists  $t\in [T]$  with
$$
\Pr \left[ \operatorname {s i g n} \left(f _ {\Xi^ {(t)}}\right) (\mathbf {x}) \neq y \right] \leq \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(t)}}\right) \leq \operatorname {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}} + r B _ {a 1} B _ {x 1} \sqrt {2 \gamma} + \epsilon . \tag {105}
$$
Proof of Theorem D.9. By  $m = \Omega\left(\frac{1}{\sqrt{\delta}} + \frac{1}{p}\left(\log\left(\frac{r}{\delta}\right)\right)^2\right)$  we have  $2re^{-\sqrt{mp}} + \frac{1}{m^2} \leq \delta$ . For any  $B_{\epsilon} \in (0, B_b)$ , when  $\sigma_a = \Theta\left(\frac{\tilde{b}}{-\ell'(0)\eta^{(1)}B_GB_e}\right)$ , by Theorem D.5, Lemma D.4, Lemma D.8, with probability at least  $1 - \delta$  over the initialization, we have
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 1} ^ {T} \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(t)}}\right) (106) \\ \leq \frac {1}{T} \sum_ {t = 1} ^ {T} \left| \left(\mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(t)}, \mathbf {b})}\right) - \mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right) \right| + \mathcal {L} _ {\mathcal {D}} \left(f _ {(\tilde {\mathbf {a}}, \mathbf {W} ^ {(1)}, \mathbf {b})}\right)\right) (107) \\ + \frac {\left\| \tilde {\mathbf {a}} \right\| _ {2} ^ {2}}{2 \eta T} + \left(2 \left\| \mathbf {a} ^ {(1)} \right\| _ {2} \sqrt {m} + 4 \eta m\right) \max  _ {i \in [ 4 m ]} \left| \frac {\partial \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(T)}}\right)}{\partial \mathbf {a} _ {i} ^ {(T)}} \right| (108) \\ \leq \mathrm {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}} + 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) (109) \\ + B _ {x 1} \| \tilde {\mathbf {a}} \| _ {2} \sqrt {\| \tilde {\mathbf {a}} \| _ {0}} \max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} ^ {(T)} - \mathbf {w} _ {i} ^ {(1)} \| _ {2} (110) \\ + \frac {\left\| \tilde {\mathbf {a}} \right\| _ {2} ^ {2}}{2 \eta T} + 4 m B _ {x 1} \left(\left\| \mathbf {a} ^ {(1)} \right\| _ {\infty} + \eta\right) \left(\max  _ {i \in [ 4 m ]} \left\| \mathbf {w} _ {i} ^ {(T)} - \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} + \max  _ {i \in [ 4 m ]} \left\| \mathbf {w} _ {i} ^ {(1)} \right\| _ {2} + \frac {\tilde {b}}{B _ {x 1}}\right). (111) \\ \end{array}
$$
By Lemma D.4, Lemma D.6, Lemma D.7, when  $\eta^{(1)} = \Theta \left(\frac{\min\{O(\eta), O(\eta\tilde{b})\}}{-\ell'(0)(B_{x1}\sigma_w\sqrt{\tilde{d} + \tilde{b}})}\right)$ , we have
$$
\left\| \tilde {\mathbf {a}} \right\| _ {0} = O \left(r (m p) ^ {\frac {1}{2}}\right), \quad \left\| \tilde {\mathbf {a}} \right\| _ {2} = O \left(\frac {B _ {a 2} B _ {b}}{\tilde {b} (m p) ^ {\frac {1}{4}}}\right) \tag {112}
$$
$$
\begin{array}{l} \left\| \mathbf {a} ^ {(1)} \right\| _ {\infty} = O \left(- \eta^ {(1)} \ell^ {\prime} (0) \left(B _ {x 1} \sigma_ {w} \sqrt {d} + \tilde {b}\right)\right) (113) \\ = \min  \left\{O (\eta), O (\eta \tilde {b}) \right\} (114) \\ \end{array}
$$
$$
\begin{array}{l} \max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} ^ {(1)} \| _ {2} = O \left(\frac {\tilde {b} \sqrt {\log m} B _ {x 1}}{B _ {G} B _ {\epsilon}}\right) (115) \\ \max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} ^ {(T)} - \mathbf {w} _ {i} ^ {(1)} \| _ {2} = O \left(T \eta B _ {x 1} \| \mathbf {a} ^ {(1)} \| _ {\infty} + T \eta^ {2} B _ {x 1} ^ {2} \max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} ^ {(1)} \| _ {2} + T \eta^ {2} B _ {x 1} \tilde {b}\right) (116) \\ = O \left(T \eta^ {2} B _ {x 1} ^ {2} \left(\max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} ^ {(1)} \| _ {2} + \frac {\tilde {b}}{B _ {x 1}}\right)\right). (117) \\ \end{array}
$$
Set  $B_{\epsilon} = \frac{B_{x1}}{(mp)^{\frac{1}{4}}}\sqrt{\frac{B_{b}}{B_{G}}}$ , we have  $\sigma_{a} = \Theta \left(\frac{\tilde{b}(mp)^{\frac{1}{4}}}{-\ell'(0)\eta^{(1)}B_{x1}\sqrt{B_{G}B_{b}}}\right)$  which satisfy the requirements. Then,
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 1} ^ {T} \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(t)}}\right) (118) \\ \leq \mathrm {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}} + r B _ {a 1} \left(\frac {2 B _ {x 1}}{(m p) ^ {\frac {1}{4}}} \sqrt {\frac {B _ {b}}{B _ {G}}} + B _ {x 1} \sqrt {2 \gamma}\right) (119) \\ + \left(\sqrt {r} B _ {a 2} B _ {b} T \eta^ {2} B _ {x 1} ^ {2} \frac {B _ {x 1}}{\tilde {b}} + m \eta B _ {x 1}\right) O \left(\frac {\tilde {b} \sqrt {\log m} B _ {x 1}}{B _ {G} B _ {\epsilon}} + \frac {\tilde {b}}{B _ {x 1}}\right) + O \left(\frac {B _ {a 2} ^ {2} B _ {b} ^ {2}}{\eta T \tilde {b} ^ {2} (m p) ^ {\frac {1}{2}}}\right) \\ \leq \mathrm {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}} + r B _ {a 1} \left(\frac {2 B _ {x 1}}{(m p) ^ {\frac {1}{4}}} \sqrt {\frac {B _ {b}}{B _ {G}}} + B _ {x 1} \sqrt {2 \gamma}\right) (120) \\ + \eta \left(\sqrt {r} B _ {a 2} B _ {b} T \eta B _ {x 1} ^ {2} + m \tilde {b}\right) O \left(\frac {\sqrt {\log m} B _ {x 1} (m p) ^ {\frac {1}{4}}}{\sqrt {B _ {b} B _ {G}}} + 1\right) + O \left(\frac {B _ {a 2} ^ {2} B _ {b} ^ {2}}{\eta T \tilde {b} ^ {2} (m p) ^ {\frac {1}{2}}}\right). (121) \\ \end{array}
$$
$$
\tilde {b} = \Theta \left(\frac {B _ {G} ^ {\frac {1}{4}} B _ {a 2} B _ {b} ^ {\frac {3}{4}}}{\sqrt {r B _ {a 1}}}\right), \quad m = \Omega \left(\frac {1}{p \epsilon^ {4}} \left(r B _ {a 1} B _ {x 1} \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 {122}
$$
$$
\eta = \Theta \left(\frac {\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} B _ {x 1} (m p) ^ {\frac {1}{4}}}{\sqrt {B _ {b} B _ {G}}} + 1\right)}\right), \quad T = \Theta \left(\frac {1}{\eta B _ {x 1} (m p) ^ {\frac {1}{4}}}\right), \tag {123}
$$
$$
\begin{array}{l} \frac {1}{T} \sum_ {t = 1} ^ {T} \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi^ {(t)}}\right) \leq \operatorname {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}} + r B _ {a 1} \left(\frac {2 B _ {x 1}}{(m p) ^ {\frac {1}{4}}} \sqrt {\frac {B _ {b}}{B _ {G}}} + B _ {x 1} \sqrt {2 \gamma}\right) + \frac {\epsilon}{2} (124) \\ + O \left(\frac {B _ {x 1} B _ {a 2} ^ {2} B _ {b} ^ {2}}{\tilde {b} ^ {2} (m p) ^ {\frac {1}{4}}}\right) (125) \\ \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 . (126) \\ \end{array}
$$
We finish the proof as the 0-1 classification error is bounded by the loss function, e.g.,  $\mathbb{I}[\mathrm{sign}(f(\mathbf{x}))\neq y]\leq \frac{\ell(yf(\mathbf{x}))}{\ell(0)}$ , where  $\ell (0) = 1$
D.3 More Discussion abut Setting
Range of  $\sigma_w$ . In practice, the value of  $\sigma_w$  cannot be arbitrary, because its choice will have an effect on the Gradient Feature set  $S_{p,\gamma,B_G}$ . On the other hand,  $d \geq \log m$  is a natural assumption, otherwise, the two-layer neural networks may fall in the NTK regime.
Parameter Choice. We use  $\lambda = 1 / \eta$  in the first step so that the neural network will totally forget its initialization, leading to the feature emergence here. This is a common setting for analysis convenience in previous work, e.g., [32, 33, 105]. We can extend this to other choices (e.g., small initialization and large step size for the first few steps), as long as after the gradient update, the gradient dominates the neuron weights. We use  $\lambda = 0$  afterward as the regularization effect is weak in our analysis. We can extend our analysis to  $\lambda$  being a small value.
Early Stopping. Our analysis divides network learning into two stages: the feature learning stage, and then classifier learning over the good features. The feature learning stage is simplified to one gradient step for the convenience of analysis, while in practice feature learning can happen in multiple steps. The current framework focuses on the gradient features in the early gradient steps, while feature learning can also happen in later steps, in particular for more complicated data. It is an interesting direction to extend the analysis to a longer training horizon.
Role of  $s$ . The  $s$  encodes the sign of the bias term, which is important. Recall that we do not update the bias term for simplicity. Let's consider a simple toy example. Assume we have  $f_{1}(x) = a_{1}\sigma (w_{1}^{\top}x + 1)$ ,  $f_{2}(x) = a_{2}\sigma (w_{2}^{\top}x - 1)$  and  $f_{3}(x) = a_{3}\sigma (w_{3}^{\top}x + 2)$ , where  $\sigma$  is ReLU activation function which is a homogeneous function.
1. The sign of the bias term is important. We can see that we always have  $a_1 \sigma(w_1^\top x + 1) \neq a_2 \sigma(w_2^\top x - 1)$  for any  $a_1, w_1, a_2, w_2$ . This means that  $f_1(x)$  and  $f_2(x)$  are intrinsically different and have different active patterns. Thus, we need to handle the sign of the bias term carefully.
2. The scaling of the bias is absorbed. On the other hand, we can see that  $a_1 \sigma(w_1^\top x + 1) = a_3 \sigma(w_3^\top x + 2)$  when  $a_1 = 2a_3$ ,  $2w_1 = w_3$ . It means that the scale of the bias term is less important, which can be absorbed into other terms.
Thus, we only need to handle bias with different signs carefully.
Gradient Feature Distribution. We may define a gradient feature distribution rather than a gradient feature set. However, we find that the technical tools used in this continuous setting are pretty different from the discrete version.
Activation Functions. We can change the ReLU activation function to a sublinear activation function, e.g. leaky ReLU, sigmoid, to get a similar conclusion. First, we need to introduce a corresponding gradient feature set, and then we can make it by following the same analysis pipeline. For simplicity, we present ReLU only.
D.4 Gradient Feature Learning Framework under Empirical Risk with Sample Complexity
In this section, we consider training with empirical risk. Intuitively, the proof is straightforward from the proof for population loss. We can simply replace the population loss with the empirical loss, which will introduce an error term in the gradient analysis. We use concentration inequality to control the error term and show that the error term depends inverse-polynomially on the sample size  $n$ .
Definition D.10 (Empirical Simplified Gradient Vector). Recall  $\mathcal{Z} = \{(\mathbf{x}^{(l)},y^{(l)})\}_{l\in [n]}$ , for any  $\mathbf{w}\in \mathbb{R}^d$ ,  $b\in \mathbb{R}$ , an Empirical Simplified Gradient Vector is defined as
$$
\widetilde {G} (\mathbf {w}, b) := \frac {1}{n} \sum_ {l \in [ n ]} \left[ y ^ {(l)} \mathbf {x} ^ {(l)} \mathbb {I} \left[ \mathbf {w} ^ {\top} \mathbf {x} ^ {(l)} > b \right] \right]. \tag {127}
$$
Definition D.11 (Empirical Gradient Feature). Recall  $\mathcal{Z} = \{(\mathbf{x}^{(l)},y^{(l)})\}_{l\in [n]}$ , let  $\mathbf{w}\in \mathbb{R}^d$ ,  $b\in \mathbb{R}$  be random variables drawn from some distribution  $\mathcal{W},\mathcal{B}$ . An Empirical Gradient Feature set with parameters  $p,\gamma, B_G$  is defined as:
$$
\widetilde{S}_{p,\gamma ,B_{G}}(\mathcal{W},\mathcal{B}):= \Big\{(D,s)\bigg|  \operatorname *{Pr}_{\mathbf{w},b}\left[\widetilde{G} (\mathbf{w},b)\in \mathcal{C}_{D,\gamma}  and  \| \widetilde{G} (\mathbf{w},b)\|_{2}\geq B_{G}  and  s = \frac{b}{|b|}\right]\geq p\Bigg\}.
$$
When clear from context, write it as  $\widetilde{S}_{p,\gamma ,B_G}$
Considering training by Algorithm 1, we have the following results.
See the full statement and proof in Theorem D.17. Below, we show some lemmas used in the analysis under empirical loss.
Lemma D.12 (Empirical Gradient Concentration Bound). When  $\frac{n}{\log n} > \frac{B_x^2}{B_{x2}}$ , with probability at least  $1 - O\left(\frac{1}{n}\right)$  over training samples, for all  $i \in [4m]$ , we have
$$
\left\| \frac {\partial \widetilde {\mathcal {L}} _ {\mathcal {Z}} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} - \frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} \right\| _ {2} \leq O \left(\frac {| \mathbf {a} _ {i} | \sqrt {B _ {x 2} \log n}}{n ^ {\frac {1}{2}}}\right), \tag {128}
$$
$$
\left| \frac {\partial \widetilde {\mathcal {L}} _ {\mathcal {Z}} (f _ {\Xi})}{\partial \mathbf {a} _ {i}} - \frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {a} _ {i}} \right| \leq O \left(\frac {\left\| \mathbf {w} _ {i} \right\| _ {2} \sqrt {B _ {x 2} \log n}}{n ^ {\frac {1}{2}}}\right), \tag {129}
$$
$$
\left| \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\Xi}\right) - \mathcal {L} _ {\mathcal {D}} \left(f _ {\Xi}\right) \right| \leq O \left(\frac {\left(\| \mathbf {a} \| _ {0} \| \mathbf {a} \| _ {\infty} \left(\max  _ {i \in [ 4 m ]} \| \mathbf {w} _ {i} \| _ {2} B _ {x} + \tilde {b}\right) + 1\right) \sqrt {\log n}}{n ^ {\frac {1}{2}}}\right). \tag {130}
$$
Proof of Lemma D.12. First, we define,
$$
\begin{array}{l} \mathbf {z} ^ {(l)} = \ell^ {\prime} \left(y ^ {(l)} f _ {\Xi} \left(\mathbf {x} ^ {(l)}\right)\right) y ^ {(l)} \left[ \sigma^ {\prime} \left(\left\langle \mathbf {w} _ {i}, \mathbf {x} ^ {(l)} \right\rangle - \mathbf {b} _ {i}\right) \mathbf {x} ^ {(l)} \right] (131) \\ - \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]. (132) \\ \end{array}
$$
As  $|\ell'(z)| \leq 1, |y| \leq 1, |\sigma'(z)| \leq 1$ , we have  $\mathbf{z}^{(l)}$  is zero-mean random vector with  $\left\| \mathbf{z}^{(l)} \right\|_2 \leq 2B_x$  as well as  $\mathbb{E}\left[\left\| \mathbf{z}^{(l)} \right\|_2^2\right] \leq B_{x2}$ . Then by Vector Bernstein Inequality, Lemma 18 in [66], for  $0 < z < \frac{B_{x2}}{B_x}$  we have
$$
\begin{array}{l} \Pr \left(\left\| \frac {\partial \widetilde {\mathcal {L}} _ {\mathcal {Z}} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} - \frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} \right\| _ {2} \geq | \mathbf {a} _ {i} | z\right) = \Pr \left(\left\| \frac {1}{n} \sum_ {l \in [ n ]} \mathbf {z} ^ {(l)} \right\| _ {2} \geq z\right) (133) \\ \leq \exp \left(- n \cdot \frac {z ^ {2}}{8 B _ {x 2}} + \frac {1}{4}\right). (134) \\ \end{array}
$$
Thus, let  $z = n^{-\frac{1}{2}} \sqrt{B_{x2} \log n}$ , with probability at least  $1 - O\left(\frac{1}{n}\right)$ , we have
$$
\left\| \frac {\partial \widetilde {\mathcal {L}} _ {\mathcal {Z}} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} - \frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {w} _ {i}} \right\| _ {2} \leq O \left(\frac {| \mathbf {a} _ {i} | \sqrt {B _ {x 2} \log n}}{n ^ {\frac {1}{2}}}\right). \tag {135}
$$
On the other hand, by Bernstein Inequality, for  $z > 0$  we have
$$
\begin{array}{l} \Pr \left(\left| \frac {\partial \widetilde {\mathcal {L}} _ {\mathcal {Z}} (f _ {\Xi})}{\partial \mathbf {a} _ {i}} - \frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {a} _ {i}} \right| > z \| \mathbf {w} _ {i} \| _ {2}\right) (136) \\ = \Pr \left( \right.\left| \frac {1}{n} \sum_ {l \in [ n ]} \left(\ell^ {\prime} \left(y ^ {(l)} f _ {\Xi} \left(\mathbf {x} ^ {(l)}\right)\right) y ^ {(l)} \left[ \sigma \left(\left\langle \mathbf {w} _ {i}, \mathbf {x} ^ {(l)} \right\rangle - \mathbf {b} _ {i}\right)\right]\right.\right. (137) \\ \left. - \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]\right) \Bigg | > z \| \mathbf {w} _ {i} \| _ {2} \Bigg) (138) \\ \leq 2 \exp \left(- \frac {\frac {1}{2} n z ^ {2}}{B _ {x 2} + \frac {1}{3} B _ {x} z}\right). (139) \\ \end{array}
$$
Thus, when  $\frac{n}{\log n} > \frac{B_x^2}{B_{x2}}$ , let  $z = n^{-\frac{1}{2}} \sqrt{B_{x2} \log n}$ , with probability at least  $1 - O\left(\frac{1}{n}\right)$ , we have
$$
\left| \frac {\partial \widetilde {\mathcal {L}} _ {\mathcal {Z}} (f _ {\Xi})}{\partial \mathbf {a} _ {i}} - \frac {\partial \mathcal {L} _ {\mathcal {D}} (f _ {\Xi})}{\partial \mathbf {a} _ {i}} \right| \leq O \left(\frac {\left\| \mathbf {w} _ {i} \right\| _ {2} \sqrt {B _ {x 2} \log n}}{n ^ {\frac {1}{2}}}\right). \tag {140}
$$