SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1: input: Data matrix  $X \in \{\pm 1\}^{d \times n}$  and label vector  $Y \in \{\pm 1\}^n$
2: initialize: For all  $i \in [N]$ ,  $j \in [d]$ ,  $w_{i,j}^{0} = 0$ ,  $a_{i}^{0} = \kappa$ ,  $b_{i}^{0} = \Delta d + 1$
3:  $X_{1} = \{x^{s}:H(x^{s}) <   d(\frac{1}{2} -\frac{\mu}{4})\}$  and  $Y_{1} = \{y^{s}:x^{s}\in X_{1}\}$   Select sparse inputs
4: for  $t = 0$  to  $T_{1} - 1$ : do
5: Select unused  $x^{s,t} \sim X_1$  and  $y^{s,t} \sim Y_1$ ,  $s \in [B]$ .
6:  $\forall i\in [N],j\in [d]$
7:  $\tilde{w}_{ij}^{t + 1} = w_{ij}^t -\gamma_1\left(\frac{1}{B}\sum_{s\in [B]}\partial_{w_{ij}^t}[L(\theta^t,x^{s,t},y^{s,t})]_A + \mathrm{Unif}[-\tau ,\tau ]\right)$
8:  $w_{ij}^{t + 1} = [\vec{w}_{ij}^{t + 1}]_{\Delta}$ $\triangleright [.\]_{\Delta}$ : proj. on  $[- \Delta, \Delta]$
9:  $b_{i}^{t + 1} = b_{i}^{t},\bar{a}_{i}^{t + 1} = a_{i}^{t}$
0: end for
1: For all  $i \in [N]$ ,  $b_i^{T_1} = -1 + 2(i + 1)\Delta$ ,  $a_i^{T_1} = 0$
2: for  $t = T_1$  to  $T_2 - 1$ : do
3: Select unused  $x^{s,t} \sim X$  and  $y^{s,t} \sim Y$ ,  $s \in [B]$ .
4:  $\forall i\in [N]$
15:  $a_{i}^{t + 1} = a_{i}^{t} - \gamma_{2}\left(\frac{1}{B}\sum_{s\in [B]}\partial_{a_{i}^{t}}[L(\theta^{t},x^{s,t},y^{s,t})]_{A} + \mathrm{Unif}[-\tau ,\tau ]\right)$
16:  $w_{ij}^{t + 1} = w_{ij}^t$ $b_{i}^{t + 1} = b_{i}^{t}$
17: end for
In this section, we report the proof of Theorem 3. The proof follows a similar argument used in [TV19, AGJ21, ABAM23]: we decompose the dynamics into a drift and a martingale contribution. Our work differs from the ones above in considering a mixed input distribution and the covariance loss. For the purposes of this section, we consider a 2-layer neural network:  $\mathrm{NN}(x; \theta) = \sum_{i=1}^{N} a_i \sigma(w_i x + b_i)$ , with  $N = k + 1$  hidden neurons and activation  $\sigma \coloneqq \mathrm{ReLU}$ . We consider the following initialization:
$$
w _ {i j} ^ {0} = 0, \tag {9}
$$
$$
a _ {i} ^ {0} = \kappa \leq \frac {1}{N (2 \Delta d + 2)}, \tag {10}
$$
$$
b _ {i} ^ {0} = d \Delta + 1. \tag {11}
$$
We train the neural network using the layer-wise curriculum-SGD algorithm (see Algorithm 1 for pseudo-code). Without loss of generality, we assume that the target parity is  $\chi_{[k]} \coloneqq \prod_{j \in [k]} x_j$ . The result for other  $k$ -parities follows directly since the initialization and the curriculum-SGD algorithm are invariant to permutations of the input neurons. We show the following:
- With a large enough training set, there are enough sparse inputs, which can be identified by computation of the Hamming weight.
- We train the first layer's weights  $w_{ij}$  on the sparse data, keeping the  $a_i, b_i$  fixed. We first show that with enough sparse samples, the labels' average will concentrate around  $\mathbb{E}_{x \sim \mathcal{D}_{\mu}}[f(x)]$ , with high probability. This term is needed in the computation of the covariance loss. Then, we show that after  $T_1 = \tilde{O}(d)$  steps, with an appropriate learning rate, with high probability, the target parity will belong to the linear span of the hidden units. This is obtained by showing that the weights corresponding to the first  $k$  input coordinates move noticeably, while the ones corresponding to the other  $d - k$  coordinates remain small.
- We train the second layer's weights  $a_{i}$  on the full data, keeping the first layer's weights and biases fixed. Established results on the convergence of SGD on convex losses allow us to conclude.
Bounding the Number of Samples. In the following lemma, we show that if the training set is large enough, the fraction of sparse inputs is large, and their Hamming weights will concentrate around expectation with high probability as  $d$  increases.
Lemma 1. Let  $\{x^s\}_{s\in [m]}$  be  $m$  i.i.d. inputs sampled from  $\mathcal{D}$ , with parameters  $\rho, \mu$ . Let  $X_1 \coloneqq \{x^s : H(x^s) < d\left(\frac{1}{2} - \frac{\mu}{4}\right)\}$  and let  $S_1 \coloneqq \{x^s : x^s \sim \mathcal{D}_\mu\}$ . If  $d \leq m \ll \exp(d)$  and  $m \geq \frac{2m_1}{\rho}$ , for  $C^* > 0$ ,
i) With probability at least  $1 - 2d^{-C^*}$ ,  $|S_1| \geq m_1$ .
ii) With probability at least  $1 - d^{-C^*}$ ,  $S_1 = X_1$ .
Proof. i) For all  $s\in [m],\mathbb{P}(x^s\in S^1) = \rho$  , thus,
$$
\begin{array}{l} \mathbb {P} \left(\sum_ {s \in [ m ]} \mathbb {1} \left(x ^ {s} \in S _ {1}\right) \leq m _ {1}\right) = \mathbb {P} \left(\sum_ {s \in [ m ]} \mathbb {1} \left(x ^ {s} \in S _ {1}\right) - m \rho \leq m _ {1} - m \rho\right) (12) \\ = \mathbb {P} \left(\sum_ {s \in [ m ]} \mathbb {1} \left(x ^ {s} \in S _ {1}\right) - m \rho \leq - \frac {m \rho}{2}\right) (13) \\ \stackrel {(a)} {\leq} \mathbb {P} \left(m \rho - \sum_ {s \in [ m ]} \mathbb {1} \left(x ^ {s} \in S _ {1}\right) \geq \sqrt {C ^ {*} \log (d) m / 2}\right) (14) \\ \stackrel {(b)} {\leq} 2 \exp (- 2 C ^ {*} \log (d)) \leq 2 d ^ {- C ^ {*}}, (15) \\ \end{array}
$$
where in  $(a)$  we used that
$$
\frac {m \rho}{2} \geq \frac {\sqrt {d m} \rho}{2} \geq \sqrt {C ^ {*} \log (d) m / 2}, \tag {16}
$$
for some constant  $C^*$ , that depends only on  $\rho$ , and in  $(b)$  we used Hoeffding's inequality.
ii) It is enough to show that for all  $x^s \in S_1$ ,  $H(x^s) < \bar{h}$ , for  $\bar{h} = d\left(\frac{1}{2} - \frac{\mu}{4}\right)$  and for all  $x^s \notin S_1$ ,  $H(x^s) \geq \bar{h}$ . Thus, for  $x^s \in S_1$ ,
$$
\begin{array}{l} \mathbb {P} \left(H \left(x ^ {s}\right) \geq \bar {h}\right) = \mathbb {P} \left(H \left(x ^ {s}\right) - \frac {1 - \mu}{2} d \geq \bar {h} - \frac {1 - \mu}{2} d\right) (17) \\ \leq \exp \left(- \frac {\mu^ {2} d}{8}\right). (18) \\ \end{array}
$$
Similarly, for  $x^s \notin S_1$ ,  $\mathbb{P}\left(H(x^s) < \bar{h}\right) < \exp(-C'd)$ , for some  $C' > 0$ . The result follows from union bound.
In the rest of the proof, we will assume that the event in Lemma 1 holds.
A.1 First Layer Training
For the first training phase, we only use the sparse samples.
Computing  $\bar{y}_1$ . As a first step, we need to estimate the average of  $y$  over the sparse set. To that purpose, we sample one batch of size  $B_1$  from  $X_1$ . We denote by  $\bar{y}_1$  such estimate
Lemma 2. Let  $\bar{y}_1 = \frac{1}{B_1}\sum_{s\in [B_1]}y^s$ . If  $B_{1}\geq 2C^{*}\log (d) / \zeta^{2}$ , with probability at least  $1 - d^{-C^*}$
$$
\left| \bar {y} _ {1} - \mu^ {k} \right| \leq \zeta . \tag {19}
$$
Proof. By construction, and by Lemma 1,  $\bar{y}_1$  is the mean of  $B_1$  i.i.d random variables sampled from  $\mathcal{D}_{\mu}$ . Thus, by Hoeffding's inequality,
$$
\mathbb {P} \left(\left| \bar {y} _ {1} - \mathbb {E} _ {x \sim \mathcal {D} _ {p}} [ f (x) ] \right| \geq \zeta\right) \leq 2 \exp \left(- \frac {B _ {1} \zeta^ {2}}{2}\right) \leq 2 d ^ {- C ^ {*}}. \tag {20}
$$
We assume that each sample is presented to the network at most once. For instance, this can be done by discarding each sample after it is used, assuming that the number of samples is large enough for covering all training steps. Denoting by  $G_{w_{ij}^t}$  the gradient computed at step  $t < T_1$  for weight  $w_{ij}^t$ , the updates of the first layer weights are given by:
$$
\tilde {w} _ {i j} ^ {t + 1} = w _ {i j} ^ {t} - \gamma G _ {w _ {i j} ^ {t}} \tag {21}
$$
$$
w _ {i j} ^ {t + 1} = \left[ \tilde {w} _ {i j} ^ {t + 1} \right] _ {\Delta} \tag {22}
$$
where by  $[.]_{\Delta}$  we denote the projection of all coordinates that exceed  $\Delta$  (resp.  $-\Delta$  to  $\Delta$  (resp.  $-\Delta$ ). Note that with our assumption on the initialization,  $|\mathrm{NN}(x;\theta^t)| < 1$  for all  $t < T_1$ . Thus, the gradients at each step are given by:
$$
G _ {w _ {i j} ^ {t}} = - \frac {1}{B} \sum_ {s \in [ B ]} \left(y ^ {s, t} - \bar {y} _ {1}\right) \cdot \partial_ {w _ {i j} ^ {t}} \mathrm {N N} \left(x ^ {s, t}; \theta^ {t}\right) + Z _ {w _ {i j} ^ {t}}, \tag {23}
$$
We decompose the dynamic into a drift and a martingale contribution:
$$
G _ {w _ {i j} ^ {t}} = \underbrace {\bar {G} _ {w _ {i j} ^ {t}}} _ {\text {d r i f t}} + \underbrace {\left(G _ {w _ {i j} ^ {t}} - \bar {G} _ {w _ {i j} ^ {t}}\right)} _ {\text {m a r t i n g a l e}}, \tag {24}
$$
where we denoted by
$$
\bar {G} _ {w _ {i j} ^ {t}} = \mathbb {E} _ {x} G _ {w _ {i j} ^ {t}}, \tag {25}
$$
the (deterministic) population gradients.
Drift Contribution. In the following, we denote by  $w_{j}^{t} = w_{ij}^{t}$  for an arbitrary neuron  $i \in [N]$ . We first bound the drift term.
Lemma 3. Let  $L_{\mu} \coloneqq \mu^{k - 1} - \mu^{k + 1}$ . Let  $\zeta > 0$ . If  $B_1 \geq 2C^* \log(d) / \zeta^2$ , with prob.  $1 - d^{-C^*}$  for all  $t \leq T_1$ ,
$$
\forall j \in [ k ]: \quad | \bar {G} _ {w _ {j} ^ {t}} + \kappa L _ {\mu} | \leq \kappa \mu \zeta ; \tag {26}
$$
$$
\forall j \notin [ k ]: \quad | \bar {G} _ {w _ {j} ^ {t}} | \leq \kappa \mu \zeta . \tag {27}
$$
Proof. With our initialization, for all  $t < T_{1}$
$$
\partial_ {w _ {i j} ^ {t}} \mathrm {N N} \left(x ^ {t}; \theta^ {t}\right) = \kappa x _ {j} ^ {t}. \tag {28}
$$
Thus, for all  $j \in [k]$ ,
$$
\bar {G} _ {w _ {j} ^ {t}} = - \kappa \left(\mathbb {E} _ {x} [ f (x) x _ {j} ] - \bar {y} \mathbb {E} _ {x} [ x _ {j} ]\right) \tag {29}
$$
$$
= - \kappa (\mu^ {k - 1} - \bar {y} _ {1} \mu). \tag {30}
$$
Applying Lemma 2, with probability  $1 - \epsilon'$ ,
$$
\left| \bar {G} _ {w _ {j} ^ {t}} + \kappa L _ {\mu} \right| = \left| \kappa \mu \left(\bar {y} - \mu^ {k}\right) \right| \leq \kappa \mu \zeta . \tag {31}
$$
Similarly, for all  $j \notin [k]$ ,
$$
\bar {G} _ {w _ {j} ^ {t}} = - \kappa \mu \left(\mu^ {k} - \bar {y}\right), \tag {32}
$$
thus, applying one more time Lemma 2, we get the result.
Martingale Contribution. We consider  $w_{j}^{t} = w_{ij}^{t}$  for an arbitrary neuron  $i \in [N]$ . For  $t \in \mathbb{N}$  and for all  $j$ , we define the martingale contribution to the dynamics up to time  $t$  by
$$
M _ {j} ^ {t} = \sum_ {s = 1} ^ {t} G _ {w _ {j} ^ {s}} - \bar {G} _ {w _ {j} ^ {s}}. \tag {33}
$$
We show that for appropriate time horizon and learning rate, the martingale contribution stays small.
Lemma 4 (Lemma 4 in [ABAM23]). Fix  $T \leq C_0 d \log(d)^{C_0}$ , for some  $C_0 > 0$ . For all  $C^*$ , there exists  $C$  that depends on  $C_0, C^*$  such that if:
$$
\gamma_ {1} ^ {2} T \leq \frac {1}{C (4 \kappa + \tau) ^ {2} \log (d) ^ {2} d}, \tag {34}
$$
with probability at least  $1 - d^{-C^*}$ :
$$
\max  _ {0 <   t <   T} \max  _ {j \in [ d ]} | \gamma_ {1} M _ {j} ^ {t} | \leq \frac {1}{\sqrt {d \log (d)}}. \tag {35}
$$
Proof. Note that in our setting (Boolean inputs, ReLU activation) we have that for all  $t < T_1$
$$
\left| M _ {j} ^ {t} - M _ {j} ^ {t - 1} \right| \leq \left| G _ {w _ {j} ^ {t}} \right| + \left| \bar {G} _ {w _ {j} ^ {t}} \right| \leq 4 \kappa + \tau . \tag {36}
$$
Thus, by Azuma-Hoeffding inequality:
$$
\mathbb {P} \left(| \gamma_ {1} M _ {j} ^ {t} | \geq \epsilon\right) \leq 2 \exp \left(- \frac {\epsilon^ {2}}{2 T \gamma_ {1} ^ {2} (4 \kappa + \tau) ^ {2}}\right). \tag {37}
$$
The result follows by choosing  $\epsilon = 1 / \sqrt{d\log(d)}$  and by union bound on all  $t\leq C_0d\log (d)^{C_0}$  and  $j\in [d]$ .
Bounding the Contributions to the Dynamics. We define the following stopping times:
$$
\tau_ {j} ^ {\Delta} := \inf  \left\{t \geq 0: \left| w _ {j} ^ {t + 1} \right| \geq \Delta - \gamma_ {1} (4 \kappa + \tau) \right\} \tag {38}
$$
$$
\tau_ {j} ^ {\Delta} = \sup  _ {j \in [ k ]} \tau_ {j} ^ {\Delta}. \tag {39}
$$
We show that for  $t \geq \tau_j^{\Delta}$ ,  $w_j^t$  stays close to  $\Delta$ .
Lemma 5. Fix  $T \leq C_0 d\log(d)^{C_0}$ , for some  $C_0 > 0$ . Let  $\zeta > 0$  and let  $B_1 \geq 2C^*\log(d)/\zeta^2$ . For all  $C^*$ , there exists  $C$  such that if  $\gamma_1$  satisfies (34) and if  $\zeta \leq \frac{1}{\sqrt{d\log(d)^{C_0}}}$ , with probability  $1 - 3d^{-C^*}$  for all  $j \in [d]$ :
$$
\inf  _ {\tau_ {j} ^ {\Delta} <   t <   T} w _ {j} ^ {t} \geq \Delta - \frac {3 + \sqrt {C _ {0}} \mu}{\sqrt {d \log (d)}}. \tag {40}
$$
Proof. For all  $j$ , let  $s = \sup \{t' \leq t : w_j^{t'} \geq \Delta - \gamma_1(4\kappa + \tau)\}$ , then with probability  $1 - 3d^{-C^*}$  for all  $j \in [d]$ :
$$
w _ {j} ^ {t} = w _ {j} ^ {s} - \gamma_ {1} \sum_ {l = s} ^ {t - 1} \bar {G} _ {w _ {j} ^ {l}} + \gamma_ {1} \left(M _ {j} ^ {t - 1} - M _ {j} ^ {s}\right) \tag {41}
$$
$$
\stackrel {(a)} {\geq} \Delta - \gamma_ {1} (4 \kappa + \tau) - \gamma_ {1} t \kappa \mu \zeta - \frac {2}{\sqrt {d \log (d)}} \tag {42}
$$
$$
\stackrel {(b)} {\geq} \Delta - \frac {3}{\sqrt {d \log (d)}} - \frac {\sqrt {t} \zeta \mu}{\sqrt {d \log (d)}}, \tag {43}
$$
$$
\stackrel {(c)} {\geq} \Delta - \frac {3}{\sqrt {d \log (d)}} - \frac {\sqrt {C _ {0}} \mu}{\sqrt {d \log (d)}}, \tag {44}
$$
where in (a) we used that by Lemma 3 with prob.  $1 - 2d^{-C^*}$ ,  $\bar{G}_{w_j^l} < -\kappa L_\mu + \kappa \mu \zeta < \kappa \mu \zeta$  for all  $l \in [t]$  and Lemma 4, in (b) we used that by (34),  $\gamma_1(4\kappa + \tau) \leq \gamma_1(4\kappa + \tau)\sqrt{t} \leq 1 / \sqrt{d\log(d)}$ , and in (c) we used the assumptions on  $T$  and  $\zeta$ . Thus, the result holds.
We will now establish a high-probability bound on the time that it takes for all weights associated with coordinates in the support of the target parity to approach proximity with  $\Delta$ .
Lemma 6. Let  $\zeta > 0$  and let  $B_{1} \geq 2C^{*}\log(d) / \zeta^{2}$ . For all  $C^{*}$ , there exists  $C$  such that if  $\gamma_{1}$  satisfies (34),  $\zeta \leq \frac{1}{\sqrt{d\log(d)^{C_{0}}}}$  and  $d$  is large enough, such that  $d\log(d) \geq \frac{(\sqrt{C_{0}}\mu + 1)^{2}}{\Delta^{2}}$ , with probability  $1 - 3d^{-C^{*}}$ :
$$
\tau^ {\Delta} \leq \frac {2 \Delta}{\gamma_ {1} \kappa L _ {\mu}}. \tag {45}
$$
Proof. Let  $T \leq C_0 d \log(d)^{C_0}$ , for some  $C_0 > 0$ . Let  $\alpha_t = \min \{w_j^t : j \in [k], t \leq \tau_j^\Delta, t \leq T\}$ . Then, with probability  $1 - 3d^{-C^*}$ :
$$
\alpha_ {t} \geq - \gamma_ {1} \sum_ {s = 0} ^ {t - 1} \bar {G} _ {\alpha_ {s}} - \gamma_ {1} \left| M _ {j} ^ {t} \right| \tag {46}
$$
$$
\geq - \gamma_ {1} t (- \kappa L _ {\mu} + \kappa \mu \zeta) + \frac {1}{\sqrt {d \log (d)}} \tag {47}
$$
$$
\stackrel {(a)} {\geq} \gamma_ {1} t \kappa L _ {\mu} - \frac {\sqrt {t} \zeta \mu}{\sqrt {d \log (d)}} - \frac {1}{\sqrt {d \log (d)}} \tag {48}
$$
$$
\geq \gamma_ {1} t \kappa L _ {\mu} - \frac {\sqrt {C _ {0}} \mu + 1}{\sqrt {d \log (d)}}, \tag {49}
$$
where in (a) we used that by (34),  $\sqrt{t}\gamma_{1}\kappa \leq 1 / \sqrt{d\log(d)}$ , and in (b) we used the assumptions on  $T$  and  $\zeta$ . Thus if  $t \geq \frac{2\Delta}{\gamma_1\kappa L_\mu}$  and  $d$  is large enough, such that  $\frac{\sqrt{C_0}\mu + 1}{\sqrt{d\log(d)}} < \Delta$ , we have
$$
\alpha_ {t} \geq 2 \Delta - \Delta > \Delta - \gamma_ {1} (4 \kappa + \tau). \tag {50}
$$
Thus, the result holds.
As our last lemma, we show that during the same time horizon, all coordinates that are not in the support of the parity will stay close to zero.
Lemma 7. Fix  $T \leq C_0 d\log(d)^{C_0}$ , for some  $C_0 > 0$ . Let  $\zeta > 0$  and let  $B_1 \geq 2C^* \log(d) / \zeta^2$ . For all  $C^*$ , there exists  $C$  such that if  $\gamma_1$  satisfies (34),  $\zeta \leq \frac{1}{\sqrt{d\log(d)^{C_0}}}$  and  $d$  is large enough, such that  $d\log(d) \geq \frac{(\sqrt{C_0}\mu + 1)^2}{\Delta^2}$ , with probability  $1 - 3d^{-C^*}$ :
$$
\sup  _ {0 \leq t \leq T} \sup  _ {j \in \{k + 1, \dots , d \}} | w _ {j} ^ {t} | \leq \frac {\sqrt {C _ {0}} \mu + 1}{\sqrt {d \log (d)}}. \tag {51}
$$
Proof. Let  $|\beta_t| = \max \{|w_j^t |:j\in [k + 1,\dots ,d],t\leq T\}$ . Thus, with probability  $1 - 3d^{-C^{*}}$ :
$$
\begin{array}{l} \left| \beta_ {t} \right| = \left| - \gamma_ {1} \sum_ {s = 0} ^ {t - 1} \bar {G} _ {\beta_ {s}} - \gamma_ {1} M _ {j} ^ {t} \right| (52) \\ \leq \gamma_ {1} \sum_ {s = 0} ^ {t - 1} \left| \bar {G} _ {\beta_ {s}} \right| + \gamma_ {1} \left| M _ {j} ^ {t} \right| (53) \\ \leq \gamma_ {1} t \kappa \mu \zeta + \frac {1}{\sqrt {d \log (d)}} (54) \\ \leq \frac {\sqrt {t} \zeta \mu + 1}{\sqrt {d \log (d)}} (55) \\ \leq \frac {\sqrt {C _ {0}} \mu + 1}{\sqrt {d \log (d)}}. (56) \\ \end{array}
$$
Thus, combining Lemmas 5, 6 and 7 and by union bound on  $i \in [N]$ , we get that for all  $C^*$ , there exist  $C_1, C_2$  such that if  $T_1$  and  $\gamma_1$  are such that
$$
\gamma_ {1} ^ {2} T _ {1} \leq \frac {1}{C (4 \kappa + \tau) ^ {2} \log (d) ^ {2} d}, \quad \frac {2 \Delta}{\gamma_ {1} \kappa L _ {\mu}} \leq T _ {1} \leq C _ {0} d \log (d) ^ {C _ {0}}. \tag {57}
$$
for some  $C_0, C$ , with prob.  $1 - 3Nd^{-C^*}$ , for all  $i \in [N]$
$$
\text {F o r a l l} j \in [ k ]: \quad | w _ {i j} ^ {T _ {1}} - \Delta | \leq \frac {C _ {1}}{\sqrt {d \log (d)}}; \tag {58}
$$
$$
\text {F o r a l l} j \notin [ k ]: \quad | w _ {i j} ^ {T _ {1}} | \leq \frac {C _ {2}}{\sqrt {d \log (d)}}. \tag {59}
$$
For  $\tau \leq \kappa$ , we can choose
$$
\gamma_ {1} = \frac {L _ {\mu}}{C ^ {\prime} \Delta \kappa d \log (d) ^ {2}}, \quad T _ {1} = \frac {2 C ^ {\prime} \Delta^ {2} d \log (d) ^ {2}}{L _ {\mu} ^ {2}}, \tag {60}
$$
where  $C' = C / 50$ , and the result follows.
A.2 Second Layer Training
As mentioned, the second phase of training consists of training the second layer's weights  $a_{i}$ , keeping  $w_{ij}$  and  $b_{i}$  fixed. We do not train the biases, and directly assign  $b_{i}^{T_{1}}$  to be well spread in the interval  $[\Delta(2 - k), \Delta(2 + k)]$ , so that the target parity belongs to the linear span of the hidden units at time  $T_{1}$ . We believe that one could extend the argument to cover scenarios where biases are trained, but this would require further technical work. To train the second layer, we use samples drawn from the entire dataset. Again, we take a large enough dataset and we assume that each sample is used at most once. In the following, we assume that the events in (58)-(59) hold. Below, we show that there exist  $a_{i}^{*}$  such that  $\sum_{i=1}^{N} a_{i}^{*} \sigma(w_{i}^{T_{1}} x + b_{i})$  achieves small expected loss. For simplicity, we restrict our attention to  $k$  even, but an analogous Lemma to the below can be written for  $k$  odd. For a function  $f$ , an estimator  $\hat{f}$  and input distribution  $\mathcal{D}$ , let us define
$$
\left. L _ {\operatorname {c o v}} (\hat {f}, f, \mathcal {D}) := \mathbb {E} _ {x \sim \mathcal {D}} \left(1 - f (x) \mathbb {E} _ {x \sim \mathcal {D}} f (x) - \hat {f} (x) (f (x) - \mathbb {E} _ {x \sim \mathcal {D}} f (x))\right) _ {+}. \right. \tag {61}
$$
Lemma 8. Assume  $k$  is an even integer. Let  $N = k + 1$  and for  $i \in \{0, \dots, k\}$  let  $b_{i} = -\Delta k + 2(i + 1)\Delta$ . For all  $\epsilon > 0$ , there exist  $\Delta > 0$ ,  $D > \max \{C_1, C_2\}$  such that if  $\rho < (\frac{\Delta}{2N} - D)\frac{\sqrt{\log(d)}}{\sqrt{d}\mu}$ , and if  $d$  is large enough:
$$
\min  _ {a: \| a \| _ {\infty} \leq 4 / \Delta} L _ {\operatorname {C O V}} \left(f ^ {*}, \chi_ {[ k ]}, \mathcal {D}\right) <   \epsilon / 2. \tag {62}
$$
Proof. We follow a similar proof scheme as in  $\mathrm{[BEG^{+}22]}$ , and choose  $a_{i}^{*}$  such that
$$
a _ {i} ^ {*} = (- 1) ^ {i} \frac {4}{\Delta}, \quad \forall i \in \{0, 1, \dots , k - 2 \}, \tag {63}
$$
$$
a _ {k - 1} ^ {*} = (- 1) ^ {k - 1} \frac {3}{\Delta}, \tag {64}
$$
$$
a _ {k} ^ {*} = (- 1) ^ {k} \frac {1}{\Delta}. \tag {65}
$$
We denote  $f^{*}(x) = \sum_{i=0}^{k} a_{i}^{*} \sigma(w_{i}^{T_{1}} x + b_{i})$ . One can verify that  $2\chi_{[k]}(x) = \sum_{i=0}^{k} a_{i}^{*} \sigma(\Delta \sum_{j=1}^{k} x_{j} + b_{i})$ . For all  $i \in [N]$ , let us denote  $\hat{w}_{ij} = w_{ij}^{T_{1}} - \Delta$  for all  $j \in [k]$  and  $\hat{w}_{ij} = w_{ij}^{T_{1}}$  for all  $j \notin [k]$ . Thus,  $\mathbb{E}_{x \sim \mathcal{D}}[\sum_{j=1}^{d} \hat{w}_{ij} x_{j}] = \rho \mu \sum_{j=1}^{d} \hat{w}_{ij}$  conditioned on the event in (58)-(59),  $\| \hat{w}_{i} \|_{1} \leq \frac{C_{3}\sqrt{d}}{\sqrt{\log(d)}}$  for all
$i\in [N]$  , where  $C_3 = \max \{C_1,C_2\}$  . Let  $\delta = \frac{\Delta}{8N}$  . Then,
$$
\mathbb {P} _ {x \sim \mathcal {D}} \left(\left| \sum_ {j = 1} ^ {d} \hat {w} _ {i j} x _ {j} \right| > \delta\right) \leq \mathbb {P} _ {x \sim \mathcal {D}} \left(\left| \sum_ {j = 1} ^ {d} \hat {w} _ {i j} x _ {j} - \rho \mu \| \hat {w} _ {i} \| _ {1} \right| > \delta - \rho \mu \| \hat {w} _ {i} \| _ {1}\right) \tag {66}
$$
$$
\stackrel {(a)} {\leq} \mathbb {P} _ {x \sim \mathcal {D}} \left(\left| \sum_ {j = 1} ^ {d} \hat {w} _ {i j} x _ {j} - \rho \mu \| \hat {w} _ {i} \| _ {1} \right| > D\right) \tag {67}
$$
$$
\stackrel {(b)} {\leq} 2 \exp \left(- \frac {D ^ {2}}{2 C _ {3} ^ {2}} \log (d)\right) \leq 2 d ^ {- \frac {D ^ {2}}{2 C _ {3} ^ {2}}}, \tag {68}
$$
where in (a) we used the assumption on  $\rho$  and in (b) we used Hoeffding's inequality and the fact that  $\sum_{j=1}^{d}(2\hat{w}_{ij})^2 \leq \frac{4C_3^2}{\log(d)}$ . Thus, with probability  $1 - 2Nd^{-\frac{D^2}{2C_3^2}}$  over  $x$  for all  $i \in \{0, \dots, k\}$ :
$$
\left| \sigma \left(w _ {i} ^ {T _ {1}} x + b _ {i}\right) - \sigma \left(\Delta \sum_ {j = 1} ^ {k} x _ {j} + b _ {i}\right) \right| \leq \left| \sum_ {j = 1} ^ {d} \hat {w} _ {i j} x _ {j} \right| \leq \delta . \tag {69}
$$
$$
\begin{array}{l} \left| 2 \chi_ {[ k ]} (x) - f ^ {*} (x) \right| = \left| \sum_ {i = 0} ^ {k} a _ {i} ^ {*} \left[ \sigma \left(w _ {i} ^ {T _ {1}} x + b _ {i}\right) - \sigma \left(\Delta \sum_ {j = 1} ^ {k} x _ {j} + b _ {i}\right) \right] \right| (70) \\ \leq \sum_ {i = 0} ^ {k} \left| a _ {i} ^ {*} \right| \delta \leq \frac {4 N}{\Delta} \delta = 1 / 2 (71) \\ \end{array}
$$
Thus, for  $d$  large enough, with probability  $1 - 2Nd^{-\frac{D^2}{2C_3^2}}$  we have  $f^{*}(x)\chi_{[k]}(x)\geq 1$ . Moreover, for all  $x$ ,
$$
\begin{array}{l} \left| f ^ {*} (x) \right| \leq N \| a ^ {*} \| _ {\infty} \left(2 \Delta k + \frac {\sqrt {d} C _ {3}}{\sqrt {\log (d)}}\right) (72) \\ \leq \frac {\sqrt {d}}{\sqrt {\log (d)}} \frac {4 N \left(C _ {3} + 2 \Delta k\right)}{\Delta}. (73) \\ \end{array}
$$
$$
\begin{array}{l} L _ {\operatorname {c o v}} \left(f ^ {*}, \chi_ {[ k ]}, \mathcal {D}\right) = \mathbb {E} _ {x \sim \mathcal {D}} \left(\left(1 - \chi_ {[ k ]} (x) \mathbb {E} _ {x} \left[ \chi_ {[ k ]} (x) \right]\right) \left(1 - \chi_ {[ k ]} (x) f ^ {*} (x)\right)\right) _ {+} (74) \\ = \mathbb {E} _ {x \sim \mathcal {D}} \left(1 - \chi_ {[ k ]} (x) \rho \mu^ {k}\right) \cdot \left(1 - \chi_ {[ k ]} (x) f ^ {*} (x)\right) \cdot \mathbb {1} \left(\chi_ {[ k ]} (x) f ^ {*} (x) <   1\right) (75) \\ \leq \frac {9 N ^ {2} \left(C _ {3} + 2 \Delta k\right) \left(1 + \rho \mu^ {k}\right)}{\Delta \sqrt {\log (d)}} \cdot d ^ {- \frac {D ^ {2}}{2 C _ {3} ^ {2}} + \frac {1}{2}} (76) \\ <   \epsilon / 2, (77) \\ \end{array}
$$
for  $d$  large enough.
To conclude, we apply an established result from [SSBD14] on the convergence of SGD on convex losses.
Theorem 6 ([SSBD14]). Let  $\mathcal{L}$  be a convex function and let  $a^* \in \operatorname{argmin}_{\|a\|_2 \leq \mathcal{B}} \mathcal{L}(a)$ , for some  $\mathcal{B} > 0$ . For all  $t$ , let  $\alpha^t$  be such that  $\mathbb{E}[\alpha^t \mid a^t] = -\nabla_{a^t} \mathcal{L}(a^t)$  and assume  $\|\alpha^t\|_2 \leq \xi$  for some  $\xi > 0$ . If  $a^{(0)} = 0$  and for all  $t \in [T]$ $a^{t+1} = a^t + \gamma \alpha^t$ , with  $\gamma = \frac{\mathcal{B}}{\xi \sqrt{T}}$ , then
$$
\frac {1}{T} \sum_ {t = 1} ^ {T} \mathcal {L} \left(a ^ {t}\right) \leq \mathcal {L} \left(a ^ {*}\right) + \frac {\mathcal {B} \xi}{\sqrt {T}}. \tag {78}
$$
We refer to [SSBD14] for a proof. Let us take  $\mathcal{L}(a) := L_{\mathrm{cov}}((a, w^{T_1}, b^{T_1}), \chi_{[k]}, \mathcal{D})$ . Then,  $\mathcal{L}$  is convex in  $a$  and for all  $t \in [T_2]$ ,
$$
\alpha^ {t} = - \frac {1}{B} \sum_ {s \in [ B ]} \nabla_ {a ^ {t}} L _ {\operatorname {c o v}} \left(\left(a ^ {t}, w ^ {T _ {1}}, b ^ {T _ {1}}\right), \chi_ {[ k ]}, x ^ {s, t}\right) + Z _ {a ^ {t}}, \tag {79}
$$
and, recalling  $\sigma \coloneqq \mathrm{ReLU}$ , we have
$$
\left\| \alpha^ {t} \right\| _ {\infty} \leq 2 \sup  _ {x} | \sigma \left(w _ {i} ^ {T _ {1}} x + b _ {i} ^ {T _ {1}}\right) | + \tau \tag {80}
$$
$$
\leq 2 \left(2 \Delta k + \frac {\sqrt {d} C _ {3}}{\sqrt {\log (d)}}\right) + \tau \tag {81}
$$
$$
\leq 2 \left(2 \Delta k + C _ {3} + \tau\right) \sqrt {\frac {d}{\log (d)}}, \tag {82}
$$