SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Corollary 1. For any  $u \in S^{d-1}$  with  $u \perp w$ ,
$$
\mathbb {E} _ {x \sim N (0, I _ {d})} \left[ (u \cdot \nabla_ {w} \mathcal {L} _ {\lambda} (\sigma (w \cdot x))) ^ {2} \right] \lesssim \frac {\min \left(1 + \lambda^ {2} , \sqrt {d}\right) ^ {- (k ^ {\star} - 1)}}{1 + \lambda^ {2}}.
$$
$$
\mathcal {L} _ {\lambda} (\sigma (w \cdot x)) = \sum_ {k} \frac {c _ {k}}{k !} \langle H e _ {k} (x), T _ {k} (w) \rangle
$$
$$
u \cdot \nabla_ {w} \mathcal {L} _ {\lambda} (\sigma (w \cdot x)) = \sum_ {k} \frac {c _ {k}}{k !} \left\langle H e _ {k} (x), \nabla_ {w} T _ {k} (w) [ u ] \right\rangle .
$$
By Lemma 4, these terms are orthogonal so by Lemma 12,
$$
\begin{array}{l} \mathbb {E} _ {x \sim N (0, I _ {d})} \left[ (u \cdot \nabla_ {w} \mathcal {L} _ {\lambda} (\sigma (w \cdot x))) ^ {2} \right] \\ \lesssim \sum_ {k} \frac {c _ {k} ^ {2}}{k !} \left(\frac {k ^ {2}}{1 + \lambda^ {2}} \mathcal {L} _ {\lambda} \left(\alpha^ {k - 1}\right) + \frac {\lambda^ {4} k ^ {4}}{(1 + \lambda^ {2}) ^ {2} d ^ {2}} \mathcal {L} _ {\lambda} ^ {(d + 2)} \left(\alpha^ {k - 2}\right)\right) \Bigg | _ {\alpha = \frac {1}{\sqrt {1 + \lambda^ {2}}}} \\ \lesssim \frac {1}{1 + \lambda^ {2}} \mathcal {L} _ {\lambda} \left(\frac {\alpha^ {k ^ {\star} - 1}}{(1 - \alpha) ^ {3}}\right) + \frac {\lambda^ {4}}{(1 + \lambda^ {2}) ^ {2} d ^ {2}} \mathcal {L} _ {\lambda} ^ {(d + 2)} \left(\frac {\alpha^ {k ^ {\star} - 2}}{(1 - \alpha) ^ {5}}\right) \\ \lesssim \frac {1}{1 + \lambda^ {2}} s _ {k ^ {\star} - 1} \left(\frac {1}{\sqrt {1 + \lambda^ {2}}}; \lambda\right) + \frac {\lambda^ {4}}{(1 + \lambda^ {2}) ^ {2} d} s _ {k ^ {\star} - 2} \left(\frac {1}{\sqrt {1 + \lambda^ {2}}}; \lambda\right) \\ \lesssim \frac {1}{1 + \lambda^ {2}} s _ {k ^ {*} - 1} (\lambda^ {- 1}; \lambda) + \frac {\lambda^ {4}}{(1 + \lambda^ {2}) ^ {2} d} s _ {k ^ {*} - 2} (\lambda^ {- 1}; \lambda). \\ \end{array}
$$
Now plugging in the formula for  $s_k(\alpha ;\lambda)$  gives:
$$
\begin{array}{l} \mathbb{E}_{x\sim N(0,I_{d})}\left[(u\cdot \nabla_{w}\mathcal{L}_{\lambda}(\sigma (w\cdot x)))^{2}\right]\lesssim (1 + \lambda^{2})^{-\frac{k^{\star} + 1}{2}}\left\{ \begin{array}{ll}(1 + \lambda^{2})^{-\frac{k^{\star} - 1}{2}} & 1 + \lambda^{2}\leq \sqrt{d}\\ \left(\frac{1 + \lambda^{2}}{d}\right)^{\frac{k^{\star} - 1}{2}} & 1 + \lambda^{2}\geq \sqrt{d} \end{array} \right. \\ \lesssim \frac {\operatorname* {m i n} \left(1 + \lambda^ {2} , \sqrt {d}\right) ^ {- (k ^ {\star} - 1)}}{1 + \lambda^ {2}}. \\ \end{array}
$$
The following lemma shows that  $\nabla_w\mathcal{L}_\lambda (\sigma (w\cdot x))$  inherits polynomial tails from  $\sigma^{\prime}$
Lemma 13. There exists an absolute constant  $C$  such that for any  $u \in S^{d-1}$  with  $u \perp w$  and any  $p \in [0, d/C]$ ,
$$
\mathbb {E} _ {x \sim N (0, I _ {d})} \left[ \left(u \cdot \nabla_ {w} \mathcal {L} _ {\lambda} (\sigma (w \cdot x))\right) ^ {p} \right] ^ {1 / p} \lesssim \frac {p ^ {C}}{\sqrt {1 + \lambda^ {2}}}.
$$
Proof. Following the proof of Lemma 9 we have
$$
\begin{array}{l} u \cdot \nabla_ {w} \mathcal {L} _ {\lambda} (\sigma (w \cdot x)) \\ = (u \cdot x) \mathbb {E} _ {z _ {1} \sim \operatorname {U n i f} (S ^ {d - 2})} \left[ \sigma^ {\prime} \left(\frac {w \cdot x + \lambda z _ {1} \| P _ {w} ^ {\perp} x \|}{\sqrt {1 + \lambda^ {2}}}\right) \left(\frac {1}{\sqrt {1 + \lambda^ {2}}} - \frac {\lambda z _ {1} (w \cdot x)}{\| P _ {w} ^ {\perp} x \| \sqrt {1 + \lambda^ {2}}}\right) \right]. \\ \end{array}
$$
First, we consider the first term. Its  $p$  norm is bounded by
$$
\frac {1}{\sqrt {1 + \lambda^ {2}}} \mathbb {E} _ {x} \left[ (u \cdot x) ^ {2 p} \right] \mathbb {E} _ {x} \left(\mathbb {E} _ {z _ {1} \sim \operatorname {U n i f} \left(S ^ {d - 2}\right)} \left[ \sigma^ {\prime} \left(\frac {w \cdot x + \lambda z _ {1} \| P _ {w} ^ {\perp} x \|}{\sqrt {1 + \lambda^ {2}}}\right) \right] ^ {2 p}\right).
$$
By Jensen we can pull out the expectation over  $z_{1}$  and use Assumption 1 to get
$$
\begin{array}{l} \frac {1}{\sqrt {1 + \lambda^ {2}}} \mathbb {E} _ {x} \left[ (u \cdot x) ^ {2 p} \right] \mathbb {E} _ {x \sim N \left(0, I _ {d}\right), z _ {1} \sim \operatorname {U n i f} \left(S ^ {d - 2}\right)} \left[ \sigma^ {\prime} \left(\frac {w \cdot x + \lambda z _ {1} \| P _ {w} ^ {\perp} x \|}{\sqrt {1 + \lambda^ {2}}}\right) ^ {2 p} \right] \\ = \frac {1}{\sqrt {1 + \lambda^ {2}}} \mathbb {E} _ {x \sim N (0, 1)} \left[ x ^ {2 p} \right] \mathbb {E} _ {x \sim N (0, 1)} \left[ \sigma^ {\prime} (x) ^ {p} \right] \\ \lesssim \frac {\operatorname {p o l y} (p)}{\sqrt {1 + \lambda^ {2}}}. \\ \end{array}
$$
Similarly, the  $p$  norm of the second term is bounded by
$$
\frac {\lambda}{\sqrt {1 + \lambda^ {2}}} \cdot \mathrm {p o l y} (p) \cdot \mathbb {E} _ {x, z _ {1}} \left[ \left(\frac {z _ {1} (x \cdot w)}{\| P _ {w} ^ {\perp} \|}\right) ^ {2 p} \right] ^ {\frac {1}{2 p}} \lesssim \frac {\lambda}{d \sqrt {1 + \lambda^ {2}}} \mathrm {p o l y} (p) \ll \frac {\mathrm {p o l y} (p)}{\sqrt {1 + \lambda^ {2}}}.
$$
Finally, we can use Corollary 1 and Lemma 13 to bound the  $p$  norms of the gradient:
Lemma 14. Let  $(x,y)$  be a fresh sample and let  $v = -\nabla L_{\lambda}(w;x;y)$ . Then there exists a constant  $C$  such that for any  $u\in S^{d - 1}$  with  $u\perp w$ , any  $\lambda \leq d^{1 / 4}$  and all  $2\le p\le d / C$
$$
\mathbb {E} _ {x, y} \left[ (u \cdot v) ^ {p} \right] ^ {1 / p} \lesssim \operatorname {p o l y} (p) \cdot \tilde {O} \Big ((1 + \lambda^ {2}) ^ {- \frac {1}{2} - \frac {k ^ {\star} - 1}{p}} \Big).
$$
$$
\begin{array}{l} \mathbb {E} _ {x, y} [ (u \cdot v) ^ {p} ] ^ {1 / p} = \mathbb {E} _ {x, y} [ (u \cdot \nabla L _ {\lambda} (w; \mathcal {B})) ^ {p} ] ^ {1 / p} \\ = \mathbb {E} _ {x, y} [ y ^ {p} (u \cdot \nabla_ {w} \mathcal {L} _ {\lambda} (\sigma (w \cdot x)) ^ {p} ] ^ {1 / p}. \\ \end{array}
$$
Applying Lemma 23 with  $X = (u\cdot \nabla_w\mathcal{L}_\lambda (\sigma (w\cdot x)))^2$  and  $Y = y^{p}(u\cdot \nabla_{w}\mathcal{L}_{\lambda}(\sigma (w\cdot x)))^{p - 2}$  gives:
$$
\begin{array}{l} \mathbb {E} _ {x, y} [ (u \cdot v) ^ {p} ] ^ {1 / p} \lesssim \operatorname {p o l y} (p) \tilde {O} \left(\frac {\min  \left(1 + \lambda^ {2} , \sqrt {d}\right) ^ {- \frac {k ^ {*} - 1}{p}}}{\sqrt {1 + \lambda^ {2}}}\right) \\ \lesssim \operatorname {p o l y} (p) \cdot \tilde {O} \left(\left(1 + \lambda^ {2}\right) ^ {- \frac {1}{2} - \frac {k ^ {*} - 1}{p}}\right) \\ \end{array}
$$
Corollary 2. Let  $v, \epsilon$  be as in Lemma 14. Then for all  $2 \leq p \leq d / C$ ,
$$
\mathbb {E} _ {x, y} [ \| v \| ^ {2 p} ] ^ {1 / p} \lesssim \operatorname {p o l y} (p) \cdot d \cdot \tilde {O} \left((1 + \lambda^ {2}) ^ {- 1 - \frac {k ^ {*} - 1}{p}}\right).
$$
$$
\| v \| ^ {2 p} = \mathbb {E} \left[ \left(\sum_ {i = 1} ^ {d} (v \cdot e _ {i}) ^ {2}\right) ^ {p} \right] \lesssim d ^ {p - 1} \mathbb {E} \left[ \sum_ {i = 1} ^ {d} (v \cdot e _ {i}) ^ {2 p} \right] \lesssim d ^ {p} \max  _ {i} \mathbb {E} [ (z \cdot e _ {i}) ^ {2 p} ].
$$
Taking  $p$ th roots and using Lemma 14 finishes the proof.
B.4 Analyzing the Dynamics
Throughout this section we will assume  $1 \leq \lambda \leq d^{1/4}$ . The proof of the dynamics is split into three stages.
In the first stage, we analyze the regime  $\alpha \in [\alpha_0,\lambda d^{-1 / 2}]$ . In this regime, the signal is dominated by the smoothing.
In the second stage, we analyze the regime  $\alpha \in [\lambda d^{-1/2}, 1 - o_d(1)]$ . This analysis is similar to the analysis in Ben Arous et al. [1] and could be equivalently carried out with  $\lambda = 0$ .
Finally in the third stage, we decay the learning rate linearly to achieve the optimal rate
$$
n \gtrsim d ^ {\frac {k ^ {*}}{2}} + \frac {d}{\epsilon}.
$$
All three stages will use the following progress lemma:
Lemma 15. Let  $w \in S^{d - 1}$  and let  $\alpha \coloneqq w \cdot w^{\star}$ . Let  $(x,y)$  be a fresh batch and define
$$
v := - \nabla_ {w} L _ {\lambda} (w; x; y), \quad z := v - \mathbb {E} _ {x, y} [ v ], \quad w ^ {\prime} = \frac {w + \lambda v}{\| w + \lambda v \|} \quad a n d \quad \alpha^ {\prime} := w ^ {\prime} \cdot w ^ {\star}.
$$
Then if  $\eta \lesssim \alpha \sqrt{1 + \lambda^2}$
$$
\alpha^ {\prime} = \alpha + \eta (1 - \alpha^ {2}) c _ {\lambda} (\alpha) + Z + \tilde {O} \bigg (\frac {\eta^ {2} d \alpha}{(1 + \lambda^ {2}) ^ {k ^ {\star}}} \bigg).
$$
where  $\mathbb{E}_{x,y}[Z] = 0$  and for all  $2\leq p\leq d / C$
$$
\mathbb {E} _ {x, y} [ Z ^ {p} ] ^ {1 / p} \leq \tilde {O} (\mathrm {p o l y} (p)) \left[ \eta (1 + \lambda^ {2}) ^ {- \frac {1}{2} - \frac {(k ^ {\star} - 1)}{p}} \right] \left[ \sqrt {1 - \alpha^ {2}} + \frac {\eta d \alpha}{\sqrt {1 + \lambda^ {2}}} \right].
$$
Furthermore, if  $\lambda = O(1)$  the  $\tilde{O} (\cdot)$  can be replaced with  $O(\cdot)$
Proof. Because  $v \perp w$  and  $1 \geq \frac{1}{\sqrt{1 + x^2}} \geq 1 - \frac{x^2}{2}$ ,
$$
\alpha^ {\prime} = \frac {\alpha + \eta (v \cdot w ^ {\star})}{\sqrt {1 + \eta^ {2} \| v \| ^ {2}}} = \alpha + \eta (v \cdot w ^ {\star}) + r
$$
where  $|r| \leq \frac{\eta^2}{2} \| v \|^2 [\alpha + \eta |v \cdot w^\star|]$ . Note that by Lemma 14,  $\eta(v \cdot w^\star)$  has moments bounded by  $\frac{\eta}{\lambda} \mathrm{poly}(p) \lesssim \alpha \mathrm{poly}(p)$ . Therefore by Lemma 23 with  $X = \| v \|^2$  and  $Y = \alpha + \eta |v \cdot w^\star|$ ,
$$
\mathbb {E} _ {x, y} [ r ] \leq \tilde {O} \Big (\eta^ {2} \mathbb {E} [ \| v \| ^ {2} ] \alpha \Big).
$$
Plugging in the bound on  $\mathbb{E}[\| v\|^2]$  from Corollary 2 gives
$$
\mathbb {E} _ {x, y} [ \alpha^ {\prime} ] = \alpha + \eta (1 - \alpha^ {2}) c _ {\lambda} (\alpha) + \tilde {O} \Big (\eta^ {2} d \alpha (1 + \lambda^ {2}) ^ {- k ^ {\star}} \Big).
$$
In addition, by Lemma 14,
$$
\begin{array}{l} \mathbb {E} _ {x, y} \left[ | \eta (v \cdot w ^ {\star}) - \mathbb {E} _ {x, y} [ \eta (v \cdot w ^ {\star}) ] | ^ {p} \right] ^ {1 / p} \lesssim \operatorname {p o l y} (p) \cdot \eta \cdot \left\| P _ {w} ^ {\perp} w ^ {\star} \right\| \cdot \tilde {O} \left(\left(1 + \lambda^ {2}\right) ^ {- \frac {1}{2} - \frac {k ^ {\star} - 1}{p}}\right) \\ = \operatorname {p o l y} (p) \cdot \eta \cdot \sqrt {1 - \alpha^ {2}} \cdot \tilde {O} \Big ((1 + \lambda^ {2}) ^ {- \frac {1}{2} - \frac {k ^ {\star} - 1}{p}} \Big). \\ \end{array}
$$
Similarly, by Lemma 23 with  $X = \| v\| ^2$  and  $Y = \| v\|^{2(p - 1)}[\alpha +\eta |v\cdot w|]^{p}$ , Lemma 14, and Corollary 2,
$$
\begin{array}{l} \mathbb {E} _ {x, y} \left[ | r _ {t} - \mathbb {E} _ {x, y} [ r _ {t} ] | ^ {p} \right] ^ {1 / p} \lesssim \mathbb {E} _ {x, y} \left[ | r _ {t} | ^ {p} \right] ^ {1 / p} \\ \lesssim \eta^ {2} \alpha \operatorname {p o l y} (p) \tilde {O} \left(\left(\frac {d}{(1 + \lambda^ {2}) ^ {k \star}}\right) ^ {1 / p} \left(\frac {d}{1 + \lambda^ {2}}\right) ^ {\frac {p - 1}{p}}\right) \\ = \operatorname {p o l y} (p) \cdot \eta^ {2} d \alpha \cdot \tilde {O} \left((1 + \lambda^ {2}) ^ {- 1 - \frac {k ^ {\star} - 1}{p}}\right). \\ \end{array}
$$
We can now analyze the first stage in which  $\alpha \in [d^{-1/2}, \lambda \cdot d^{-1/2}]$ . This stage is dominated by the signal from the smoothing.
Lemma 16 (Stage 1). Assume that  $\lambda \geq 1$  and  $\alpha_0 \geq \frac{1}{C} d^{-1/2}$ . Set
$$
\eta = \frac {d ^ {- \frac {k ^ {\star}}{2}} (1 + \lambda^ {2}) ^ {k ^ {\star} - 1}}{\log (d) ^ {C}} \quad a n d \quad T _ {1} = \frac {C (1 + \lambda^ {2}) d ^ {\frac {k ^ {\star} - 2}{2}} \log (d)}{\eta} = \tilde {O} \left(d ^ {k ^ {\star} - 1} \lambda^ {- 2 k ^ {\star} + 4}\right)
$$
for a sufficiently large constant  $C$ . Then with high probability, there exists  $t \leq T_1$  such that  $\alpha_t \geq \lambda d^{-1/2}$ .
Proof. Let  $\tau$  be the hitting time for  $\alpha_{\tau} \geq \lambda d^{-1/2}$ . For  $t \leq T_1$ , let  $E_t$  be the event that
$$
\alpha_ {t} \geq \frac {1}{2} \left[ \alpha_ {0} + \eta \sum_ {j = 0} ^ {t - 1} c _ {\lambda} (\alpha_ {j}) \right].
$$
We will prove by induction that for any  $t \leq T_1$ , the event:  $\{E_t \text{ or } t \geq \tau\}$  happens with high probability. The base case of  $t = 0$  is trivial so let  $t \geq 0$  and assume the result for all  $s < t$ . Note that  $\eta / \lambda \ll \frac{d^{-1/2}}{C} \leq \alpha_j$  so by Lemma 15 and the fact that  $\lambda \geq 1$ ,
$$
\alpha_ {t} = \alpha_ {0} + \sum_ {j = 0} ^ {t - 1} \left[ \eta \left(1 - \alpha_ {j} ^ {2}\right) c _ {\lambda} (\alpha_ {j}) + Z _ {j} + \tilde {O} \left(\eta^ {2} d \alpha_ {j} \lambda^ {- 2 k ^ {\star}}\right) \right].
$$
Now note that  $\mathbb{P}[E_t$  or  $t\geq \tau ] = 1 - \mathbb{P}\left[!E_t\text{and} t <   \tau \right]$  so let us condition on the event  $t <   \tau$  . Then by the induction hypothesis, with high probability we have  $\alpha_{s}\in [\frac{\alpha_{0}}{2},\lambda d^{-1 / 2}]$  for all  $s <   t$  . Plugging in the value of  $\eta$  gives:
$$
\begin{array}{l} \eta \left(1 - \alpha_ {j} ^ {2}\right) c _ {\lambda} (\alpha_ {j}) + \tilde {O} \left(\eta^ {2} d \alpha_ {j} \lambda^ {- 2 k ^ {\star}}\right) \\ \geq \eta (1 - \alpha_ {j} ^ {2}) c _ {\lambda} (\alpha_ {j}) - \frac {\eta d ^ {- \frac {k ^ {\star} - 2}{2}} \lambda^ {- 2} \alpha_ {j}}{C} \\ \geq \frac {\eta c _ {\lambda} (\alpha_ {j})}{2}. \\ \end{array}
$$
Similarly, because  $\sum_{j=0}^{t-1} Z_j$  is a martingale we have by Lemma 22 and Lemma 24 that with high probability,
$$
\begin{array}{l} \sum_ {j = 0} ^ {t - 1} Z _ {j} \lesssim \tilde {O} \left(\left[ \sqrt {T _ {1}} \cdot \eta \lambda^ {- k ^ {\star}} + \eta \lambda^ {- 1} \right] \left[ 1 + \max  _ {j <   t} \frac {\eta d \alpha_ {j}}{\lambda} \right]\right) \\ \lesssim \tilde {O} \left(\sqrt {T _ {1}} \cdot \eta \lambda^ {- k ^ {\star}} + \eta \lambda^ {- 1}\right) \\ \leq \frac {d ^ {- 1 / 2}}{C}. \\ \end{array}
$$
where we used that  $\eta d\alpha_{j} / \lambda \leq \eta \sqrt{d}\ll 1$ . Therefore conditioned on  $t\leq \tau$  we have with high probability that for all  $s\leq t$ :
Now we split into two cases depending on the parity of  $k^{\star}$ . First, if  $k^{\star}$  is odd we have that with high probability, for all  $t \leq T_1$ :
$$
\alpha_ {t} \gtrsim \alpha_ {0} + \eta t \lambda^ {- 1} d ^ {- \frac {k ^ {\star} - 1}{2}} \quad \text {o r} \quad t \geq \tau .
$$
Now let  $t = T_{1}$ . Then we have that with high probability,
$$
\alpha_ {t} \geq \lambda d ^ {- 1 / 2} \quad \text {o r} \quad \tau \leq T _ {1}
$$
which implies that  $\tau \leq T_1$  with high probability. Next, if  $k^{\star}$  is even we have that with high probability
$$
\alpha_ {t} \gtrsim \alpha_ {0} + \frac {\eta \cdot d ^ {- \frac {k ^ {*} - 2}{2}}}{\lambda^ {2}} \sum_ {s = 0} ^ {t - 1} \alpha_ {s} \quad \text {o r} \quad t \geq \tau .
$$
As above, by Lemma 27 the first event implies that  $\alpha_{T_1} \geq \lambda d^{-1/2}$  so we must have  $\tau \leq T_1$  with high probability.
Next, we consider what happens when  $\alpha \geq \lambda d^{-1/2}$ . The analysis in this stage is similar to the online SGD analysis in [1].
Lemma 17 (Stage 2). Assume that  $\alpha_0 \geq \lambda d^{-1/2}$ . Set  $\eta, T_1$  as in Lemma 16. Then with high probability,  $\alpha_{T_1} \geq 1 - d^{-1/4}$ .
Proof. The proof is almost identical to Lemma 16. We again have from Lemma 15
$$
\alpha_ {t} \geq \alpha_ {0} + \sum_ {j = 0} ^ {t - 1} \left[ \eta (1 - \alpha_ {j} ^ {2}) c _ {\lambda} (\alpha_ {j}) + Z _ {j} - \tilde {O} \left(\eta^ {2} d \alpha_ {j} \lambda^ {- 2 k ^ {*}}\right) \right].
$$
First, from martingale concentration we have that
$$
\begin{array}{l} \sum_ {j = 0} ^ {t - 1} Z _ {j} \lesssim \tilde {O} \left(\left[ \sqrt {T _ {1}} \cdot \eta \lambda^ {- k ^ {\star}} + \eta \lambda^ {- 1} \right] \left[ 1 + \frac {\eta d}{\lambda} \right]\right) \\ \lesssim \tilde {O} \left(\left[ \sqrt {T _ {1}} \cdot \eta \lambda^ {- k ^ {\star}} + \eta \lambda^ {- 1} \right] \cdot \lambda\right) \\ \lesssim \frac {\lambda d ^ {- 1 / 2}}{C} \\ \end{array}
$$
where we used that  $\eta \ll \frac{\lambda^2}{d}$ . Therefore with high probability,
$$
\begin{array}{l} \alpha_ {t} \geq \frac {\alpha_ {0}}{2} + \sum_ {j = 0} ^ {t - 1} \left[ \eta \left(1 - \alpha_ {j} ^ {2}\right) c _ {\lambda} (\alpha_ {j}) - \tilde {O} \left(\eta^ {2} d \alpha_ {j} \lambda^ {- 2 k ^ {*}}\right) \right] \\ \geq \frac {\alpha_ {0}}{2} + \eta \sum_ {j = 0} ^ {t - 1} \left[ (1 - \alpha_ {j} ^ {2}) c _ {\lambda} (\alpha_ {j}) - \frac {\alpha_ {j} d ^ {- \frac {k ^ {*} - 2}{2}}}{C \lambda^ {2}} \right]. \\ \end{array}
$$
Therefore while  $\alpha_{t} \leq 1 - \frac{1}{k^{s}}$ , for sufficiently large  $C$  we have
$$
\alpha_ {t} \geq \frac {\alpha_ {0}}{2} + \frac {\eta}{C ^ {1 / 2} \lambda^ {k ^ {\star}}} \sum_ {j = 0} ^ {t - 1} \alpha_ {j} ^ {k ^ {\star} - 1}.
$$
Therefore by Lemma 27, we have that there exists  $t \leq T_1 / 2$  such that  $\alpha_t \geq 1 - \frac{1}{k^\star}$ . Next, let  $p_t = 1 - \alpha_t$ . Then applying Lemma 15 to  $p_t$  and using  $(1 - \frac{1}{k^\star})^{k^\star} \gtrsim 1 / e$  gives that if
$$
r := \frac {d ^ {- \frac {k ^ {\star} - 2}{2}}}{C \lambda^ {2}} \quad \text {a n d} \quad c := \frac {1}{C ^ {1 / 2} \lambda^ {k ^ {\star}}}
$$
$$
\begin{array}{l} p _ {t + 1} \leq p _ {t} - \eta c p _ {t} + \eta r + Z _ {t} \\ = (1 - \eta c) p _ {t} + \eta r + Z _ {t}. \\ \end{array}
$$
$$
p _ {t + s} \leq (1 - \eta c) ^ {s} p _ {t} + r / c + \sum_ {i = 0} ^ {s - 1} (1 - \eta c) ^ {i} Z _ {t + s - 1 - i}.
$$
With high probability, the martingale term is bounded by  $\lambda d^{-1/2} / C$  as before as long as  $s \leq T_1$ , so for  $s \in [\frac{C\log(d)}{\eta c}, T_1]$  we have that  $p_{t+s} \lesssim C^{-1/2}\left(\left(\frac{\lambda^2}{d}\right)^{\frac{k^\star - 2}{2}} + \lambda d^{-1/2}\right) \lesssim C^{-1/2}d^{-1/4}$ . Setting  $s = T_1 - t$  and choosing  $C$  appropriately yields  $p_{T_1} \leq d^{-1/4}$ , which completes the proof.
Finally, the third stage guarantees not only a hitting time but a last iterate guarantee. It also achieves the optimal sample complexity in terms of the target accuracy  $\epsilon$ :
Lemma 18 (Stage 3). Assume that  $\alpha_0 \geq 1 - d^{-1/4}$ . Set  $\lambda = 0$  and
$$
\eta_ {t} = \frac {C}{C ^ {4} d + t}.
$$
for a sufficiently large constant  $C$ . Then for any  $t \leq \exp(d^{1/C})$ , we have that with high probability,
$$
\alpha_ {t} \geq 1 - O \left(\frac {d}{d + t}\right).
$$
Proof. Let  $p_t = 1 - \alpha_t$ . By Lemma 15, while  $p_t \leq 1 / k^\star$ :
$$
\begin{array}{l} p _ {t + 1} \leq p _ {t} - \frac {\eta_ {t} p _ {t}}{2 C} + C \eta_ {t} ^ {2} d + \eta_ {t} \sqrt {p _ {t}} \cdot W _ {t} + \eta_ {t} ^ {2} d \cdot Z _ {t} \\ = \left(\frac {C ^ {4} d + t - 2}{C ^ {4} d + t}\right) p _ {t} + C \eta_ {t} ^ {2} d + \eta_ {t} \sqrt {p _ {t}} \cdot W _ {t} + \eta_ {t} ^ {2} d \cdot Z _ {t}. \\ \end{array}
$$
where the moments of  $W_{t}, Z_{t}$  are each bounded by  $\mathrm{poly}(p)$ . We will prove by induction that with probability at least  $1 - t\exp \left(-Cd^{1 / C} / e\right)$ , we have for all  $s\leq t$ :
$$
p _ {s} \leq \frac {2 C ^ {3} d}{C ^ {4} d + s} \leq \frac {2}{C} \leq \frac {1}{k ^ {\star}}.
$$
The base case is clear so assume the result for all  $s \leq t$ . Then from the recurrence above,
$$
p _ {t + 1} \leq p _ {0} \frac {C ^ {8} d ^ {2}}{(C ^ {4} d + t) ^ {2}} + \frac {1}{(C ^ {4} d + t) ^ {2}} \sum_ {j = 0} ^ {t} (C ^ {4} d + j) ^ {2} \big [ C \eta_ {j} ^ {2} d + \eta_ {t} \sqrt {p _ {t}} \cdot W _ {t} + \eta_ {t} ^ {2} d \cdot Z _ {t} \big ].
$$
First, because  $p_0 \leq d^{-1/4}$ ,
$$
\frac {C ^ {8} p _ {0} d ^ {2}}{(C ^ {4} d + t) ^ {2}} \leq \frac {C ^ {4} p _ {0} d}{C ^ {4} d + t} \ll \frac {d}{C ^ {4} d + t}.
$$
$$
\begin{array}{l} \frac {1}{(C ^ {4} d + t) ^ {2}} \sum_ {j = 0} ^ {t} (C ^ {4} d + j) ^ {2} C \eta_ {j} ^ {2} d = \frac {1}{(C ^ {4} d + t) ^ {2}} \sum_ {j = 0} ^ {t} C ^ {3} d \\ = \frac {C ^ {3} d t}{(C ^ {4} d + t) ^ {2}} \\ \leq \frac {C ^ {3} d}{C ^ {4} d + t}. \\ \end{array}
$$
The next error term is:
$$
\frac {1}{(C ^ {4} d + t) ^ {2}} \sum_ {j = 0} ^ {t} (C ^ {4} d + j) ^ {2} \eta_ {t} \sqrt {p _ {t}} \cdot W _ {t}.
$$
Fix  $p = \frac{d^{1 / C}}{e}$ . Then we will bound the  $p$ th moment of  $p_t$ :
$$
\begin{array}{l} \mathbb {E} \left[ p _ {t} ^ {p} \right] \leq \frac {2 C ^ {3} d}{C ^ {4} d + t} + \mathbb {E} \left[ p _ {t} ^ {p} \mathbf {1} _ {p _ {t} \geq \frac {2 C ^ {3} d}{C ^ {4} d + s}} \right] \\ \leq \left(\frac {2 C ^ {3} d}{C ^ {4} d + t}\right) ^ {p} + 2 ^ {p} \mathbb {P} \left[ p _ {t} \geq \frac {2 C ^ {3} d}{C ^ {4} d + t} \right] \\ \leq \left(\frac {2 C ^ {3} d}{C ^ {4} d + t}\right) ^ {p} + 2 ^ {p} t \exp \left(- C d ^ {1 / C}\right). \\ \end{array}
$$
Now note that because  $t \leq \exp \left(d^{1 / C}\right)$ ,
$$
\log \left(t \exp \left(- C d ^ {1 / C}\right)\right) = \log (t) - C d ^ {1 / C} \leq - \log (t) (C - 1) d ^ {1 / C} \leq - p \log (t).
$$
Therefore  $\mathbb{E}[p_t^p ]^{1 / p}\leq \frac{4C^3d}{C^4d + t}$ . Therefore the  $p$  norm of the predictable quadratic variation of the next error term is bounded by:
$$
\begin{array}{l} \operatorname {p o l y} (p) \sum_ {j = 0} ^ {t} (C ^ {4} d + j) ^ {4} \eta_ {t} ^ {2} \mathbb {E} [ p _ {t} ^ {p} ] ^ {1 / p} \leq \operatorname {p o l y} (p) \sum_ {j = 0} ^ {t} C ^ {5} d (C ^ {4} d + j) \\ \lesssim \operatorname {p o l y} (p) C ^ {5} d t \left(C ^ {4} d + t\right). \\ \end{array}
$$
In addition, the  $p$  norm of the largest term in this sum is bounded by
$$
\operatorname {p o l y} (p) \sqrt {C ^ {5} d \left(C ^ {4} d + t\right)}.
$$
Therefore by Lemma 22 and Lemma 24, we have with probability at least  $1 - \exp \left(Cd^{-1 / C} / e\right)$ , this term is bounded by
$$
\frac {\sqrt {C ^ {5} d t}}{(C ^ {4} d + t) ^ {3 / 2}} \cdot d ^ {1 / 2} \leq \frac {C ^ {3} d}{C ^ {4} d + t}.
$$
Finally, the last term is similarly bounded with probability at least  $1 - \exp \left(-Cd^{-1 / C} / e\right)$  by
$$
\frac {C ^ {2} d \sqrt {t}}{(C ^ {4} d + t) ^ {2}} \cdot d ^ {1 / 2} \ll \frac {C ^ {3} d}{C ^ {4} d + t}
$$
which completes the induction.
We can now combine the above lemmas to prove Theorem 1:
Proof of Theorem 1. By Lemmas 16 to 18, if  $T = T_{1} + T_{2}$  we have with high probability for all  $T_{2} \leq \exp \left(d^{1 / C}\right)$ :
$$
\alpha_ {T} \geq 1 - O \left(\frac {d}{d + T _ {2}}\right).
$$
Next, note that by Bernoulli's inequality  $((1 + x)^n \geq 1 + nx)$ , we have that  $1 - \alpha^k \leq k(1 - \alpha)$ . Therefore,
$$
\begin{array}{l} L (w _ {T}) = \sum_ {k \geq 0} \frac {c _ {k} ^ {2}}{k !} [ 1 - \alpha_ {T} ^ {k} ] \\ \leq \left(1 - \alpha_ {T}\right) \sum_ {k \geq 0} \frac {c _ {k} ^ {2}}{(k - 1) !} \\ = \left(1 - \alpha_ {T}\right) \mathbb {E} _ {x \sim N (0, 1)} \left[ \sigma^ {\prime} (x) ^ {2} \right] \\ \lesssim \frac {d}{d + T _ {2}} \\ \end{array}
$$
which completes the proof of Theorem 1.
B.5 Proof of Theorem 2
We directly follow the proof of Theorem 2 in Damian et al. [6] which is reproduced here for completeness. We begin with the following general CSQ lemma which can be found in Szörenyi [43], Damian et al. [6]:
Lemma 19. Let  $\mathcal{F}$  be a class of functions and  $\mathcal{D}$  be a data distribution such that
$$
\mathbb {E} _ {x \sim \mathcal {D}} [ f (x) ^ {2} ] = 1 \quad a n d \quad | \mathbb {E} _ {x \sim \mathcal {D}} [ f (x) g (x) ] | \leq \epsilon \quad \forall f \neq g \in \mathcal {F}.
$$
Then any correlational statistical query learner requires at least  $\frac{|\mathcal{F}|(\tau^2 - \epsilon)}{2}$  queries of tolerance  $\tau$  to output a function in  $\mathcal{F}$  with  $L^2 (\mathcal{D})$  loss at most  $2 - 2\epsilon$ .
First, we will construct a function class from a subset of  $\mathcal{F} \coloneqq \left\{\sigma(w \cdot x) : w \in S^{d-1}\right\}$ . By [6, Lemma 3], for any  $\epsilon$  there exist  $\frac{1}{2} e^{c\epsilon^2 d}$  unit vectors  $w_1, \ldots, w_s$  such that their pairwise inner products are all bounded by  $\epsilon$ . Let  $\widehat{\mathcal{F}} \coloneqq \{\sigma(w_i \cdot x) : i \in [s]\}$ . Then for  $i \neq j$ ,
$$
\left| \mathbb {E} _ {x \sim N (0, I _ {d})} [ \sigma (w _ {i} \cdot x) \sigma (w _ {j} \cdot x) ] \right| = \left| \sum_ {k \geq 0} \frac {c _ {k} ^ {2}}{k !} (w _ {i} \cdot w _ {j}) ^ {k} \right| \leq | w _ {i} \cdot w _ {j} | ^ {k ^ {\star}} \leq \epsilon^ {k ^ {\star}}.
$$
Therefore by Lemma 19,
$$
4 q \geq e ^ {c \epsilon^ {2} d} (\tau^ {2} - \epsilon^ {k ^ {\star}}).
$$
Now set
$$
\epsilon = \sqrt {\frac {\log \left(4 q (c d) ^ {k ^ {\star} / 2}\right)}{c d}}
$$