SlowGuess/NeurIPS_2023_OCR huggingface dataset:
and by an argument similar to the one used to bound  $A_t'$  (here using Eq. (E.211) from Lemma E.15 instead of Eq. (E.210)), we can thus show that
$$
\left| \left\langle \nabla_ {\hat {u}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\hat {u}} L \left(\rho_ {t}\right), \bar {u} _ {t} - \hat {u} _ {t} \right\rangle \right| \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \gamma_ {2} \| \delta_ {t} \| ^ {2}. \tag {E.110}
$$
Similarly, for the second term we have
$$
\left| \left\langle \nabla_ {\bar {u}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {u}} L \left(\rho_ {t}\right), \bar {u} _ {t} - \hat {u} _ {t} \right\rangle \right| \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \gamma_ {2} \| \delta_ {t} \| ^ {2}. \tag {E.111}
$$
Combining Eq. (E.110), Eq. (E.111) and Eq. (E.102) finishes the proof.
E.2 Upper Bound on  $B_{t}$
$$
B _ {t} (\chi) := - 2 \left\langle \operatorname {g r a d} _ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \operatorname {g r a d} _ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} (\chi) - \bar {u} _ {t} (\chi) \right\rangle . \tag {E.112}
$$
When there is no risk of confusion, we simply use  $B_{t}$ ,  $\hat{u}_{t}$ , and  $\bar{u}_{t}$ , omitting the  $\chi$  dependency. In the next lemma, we first obtain a bound on  $B_{t}(\chi)$ , and after that, we obtain a bound on  $\mathbb{E}_{\chi}[B_t(\chi)]$ . The proof of the following lemma is mostly a straightforward calculation, and follows mostly from Lemma 5.1 and the definition of  $\overline{\Delta}$ .
Lemma E.9. Suppose we are in the setting of Theorem 3.4. Suppose  $T \lesssim \frac{d^C}{\hat{\sigma}_{2,d}^2 + \hat{\sigma}_{4,d}^2}$  for any absolute constant  $C$ . Assume the width  $m$  of  $\bar{\rho}, \hat{\rho}$  is at most  $d^{O(\log d)}$ . Then, with probability at least  $1 - e^{-d^2}$ , we have
$$
\underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\rho_ {t}} (x) - f _ {\hat {\rho} _ {t}} (x)) ^ {2} ] \lesssim \frac {(\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}) d ^ {2} (\log d) ^ {O (1)}}{m} \tag {E.113}
$$
for all times  $t \in [0, T]$ . Additionally, with probability at least  $1 - e^{-d^2}$ , for all  $t \in [0, T]$  and  $\chi \in \hat{\rho}_0$ , we have
$$
B _ {t} (\chi) \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \cdot \left(\frac {d (\log d) ^ {O (1)}}{\sqrt {m}} + \bar {\Delta} _ {t}\right) \| \delta_ {t} \| _ {2}. \tag {E.114}
$$
Proof of Lemma E.9. We first decompose  $B_{t}$  into two terms:
$$
\begin{array}{l} B _ {t} = - 2 \left\langle \operatorname {g r a d} _ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \operatorname {g r a d} _ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle (E.115) \\ = - 2 \left\langle \left(I - \hat {u} _ {t} \hat {u} _ {t} ^ {\top}\right) \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \left(I - \hat {u} _ {t} \hat {u} _ {t} ^ {\top}\right) \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle (E.116) \\ = - 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \left(I - \hat {u} _ {t} \hat {u} _ {t} ^ {\top}\right) \left(\hat {u} _ {t} - \bar {u} _ {t}\right) \right\rangle (E.117) \\ = - 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle + 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} \hat {u} _ {t} ^ {\top} \left(\hat {u} _ {t} - \bar {u} _ {t}\right) \right\rangle (E.118) \\ = - 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle + 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} \right\rangle \cdot \left(1 - \left\langle \hat {u} _ {t}, \bar {u} _ {t} \right\rangle\right) (E.119) \\ = - 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle + 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} \right\rangle \cdot \frac {\left\| \hat {u} _ {t} - \bar {u} _ {t} \right\| _ {2} ^ {2}}{2} (E.120) \\ \end{array}
$$
where the last equality is because  $\frac{\|\hat{u}_t - \bar{u}_t\|_2^2}{2} = 1 - \langle \hat{u}_t, \bar{u}_t \rangle$  since  $\|\hat{u}_t\|_2 = \|\bar{u}_t\|_2 = 1$ . For convenience, define  $\delta_t := \hat{u}_t - \bar{u}_t$ . Then, we have
$$
B _ {t} = - 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \delta_ {t} \right\rangle + \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} \right\rangle \cdot \| \delta_ {t} \| _ {2} ^ {2} \tag {E.121}
$$
We first focus on the first term in Eq. (E.121), which is the main contributor to  $B_{t}$ .
Analysis of First Term in Eq. (E.121). For convenience, we refer to the first term in Eq. (E.121) as  $B_{t}^{\prime}$ . Observe that we can expand  $B_{t}^{\prime}$  as
$$
\begin{array}{l} B _ {t} ^ {\prime} = - 2 \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \delta_ {t} \right\rangle (E.122) \\ = - 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\rho} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \delta_ {t}, x \rangle \right] (ByEq.(3.4)) \\ = - 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\bar {\rho}} (x) - f _ {\rho} (x)\right) \sigma^ {\prime} (\langle \hat {u} _ {t}, x \rangle) \langle \delta_ {t}, x \rangle \right] - 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} (\langle \hat {u} _ {t}, x \rangle) \langle \delta_ {t}, x \rangle \right] (E.123) \\ \end{array}
$$
Let us bound the first term of Eq. (E.123). By Lemma 5.1, with probability  $e^{-d^2}$ , for all  $t \leq T$  we have
$$
\underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho_ {t}} (x) - f _ {\bar {\rho} _ {t}} (x)\right) ^ {2} \right] \lesssim \frac {\left(\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}\right) d ^ {2} (\log d) ^ {O (1)}}{m}. \tag {E.124}
$$
Thus, we can bound the first term of Eq. (E.123) by
$$
\begin{array}{l} \left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\bar {\rho}} (x) - f _ {\rho} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \delta_ {t}, x \rangle \right] \right| \tag {E.125} \\ \leq \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\bar {\rho}} (x) - f _ {\rho} (x)) ^ {2} ] ^ {1 / 2} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ \sigma^ {\prime} (\langle \hat {u} _ {t}, x \rangle) ^ {2} \langle \delta_ {t}, x \rangle^ {2} ] ^ {1 / 2} \\ \end{array}
$$
(By Cauchy-Schwarz inequality)
$$
\lesssim \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\bar {\rho}} (x) - f _ {\rho} (x)) ^ {2} ] ^ {1 / 2} (\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}) ^ {1 / 2} \| \delta_ {t} \| _ {2}
$$
(By Lemma E.15, Eq. (E.211))
$$
\begin{array}{l} \lesssim \frac {\sqrt {\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}} \cdot d (\log d) ^ {O (1)}}{\sqrt {m}} \cdot \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) ^ {1 / 2} \| \delta_ {t} \| _ {2} \quad (ByEq.(E.124)) \\ \lesssim \frac {\left(\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}\right) \cdot d (\log d) ^ {O (1)}}{\sqrt {m}} \| \delta_ {t} \| _ {2}. (E.126) \\ \end{array}
$$
Next, we bound the second term in Eq. (E.123). By applying the Cauchy-Schwarz inequality, we have
$$
\begin{array}{l} \left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} (\langle \hat {u}, x \rangle) \langle \delta_ {t}, x \rangle \right] \right| (E.127) \\ \leq \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} [ (f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)) ^ {2} ] ^ {1 / 2} \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} [ \sigma^ {\prime} (\langle \hat {u}, x \rangle) ^ {2} \langle \delta_ {t}, x \rangle^ {2} ] ^ {1 / 2} (E.128) \\ \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) ^ {1 / 2} \bar {\Delta} _ {t} \cdot \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) ^ {2} \langle \delta_ {t}, x \rangle^ {2} \right] ^ {1 / 2} \quad (\text {B y}) \\ \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) ^ {1 / 2} \bar {\Delta} _ {t} \cdot \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) ^ {1 / 2} \| \delta_ {t} \| _ {2} \quad (\text {B y}) \\ \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \overline {{\Delta}} _ {t} \| \delta_ {t} \| _ {2} (E.129) \\ \end{array}
$$
Combining Eq. (E.126) and Eq. (E.129) with Eq. (E.123), we find that
$$
\left| B _ {t} ^ {\prime} \right| \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \cdot \left(\frac {d (\log d) ^ {O (1)}}{\sqrt {m}} + \overline {{\Delta}} _ {t}\right) \| \delta_ {t} \| _ {2}. \tag {E.130}
$$
Analysis of Second Term in Eq. (E.121). For convenience, we write this term as  $B_{t}^{\prime \prime}$ , i.e. we define  $B_{t}^{\prime \prime} \coloneqq \langle \nabla_{\hat{u}}L(\hat{\rho}_{t}) - \nabla_{\hat{u}}L(\rho_{t}),\hat{u}_{t}\rangle \cdot \| \delta_{t}\|_{2}^{2}$ . It suffices to show that the first dot product in the definition of  $B_{t}^{\prime \prime}$  is at most  $(\hat{\sigma}_{2,d}^{2} + \hat{\sigma}_{4,d}^{2})$  up to constant factors. To start, we have
$$
\begin{array}{l} \left| \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} \right\rangle \right| \lesssim \left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\rho} (x)\right) \sigma^ {\prime} (\langle \hat {u} _ {t}, x \rangle) \langle \hat {u} _ {t}, x \rangle \right] \right| (E.131) \\ \lesssim \left| _ {x \sim \mathbb {S} ^ {d - 1}} \left[ \left(f _ {\rho} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \hat {u} _ {t}, x \rangle \right] \right| (E.132) \\ + \left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\hat {\rho}} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \hat {u} _ {t}, x \rangle \right] \right| (E.133) \\ \end{array}
$$
Both terms in Eq. (E.133) can be bounded as before. For the first term we use Eq. (E.124):
$$
\begin{array}{l} \left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \hat {u} _ {t}, x \rangle \right] \right| \tag {E.134} \\ \leq \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho} (x) - f _ {\bar {\rho}} (x)\right) ^ {2} \right] ^ {1 / 2} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) ^ {2} \langle \hat {u} _ {t}, x \rangle^ {2} \right] ^ {1 / 2} \\ \end{array}
$$
(By Cauchy-Schwarz Inequality)
$$
\lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \cdot \frac {d (\log d) ^ {O (1)}}{\sqrt {m}}. \quad \text {(B y E q . (E . 1 2 4) a n d L e m m a E . 1 5 , E q . (E . 2 1 1))}
$$
For the second term, we use Lemma E.16:
$$
\begin{array}{l} \left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\bar {\rho}} (x) - f _ {\hat {\rho}} (x)) \sigma^ {\prime} (\langle \hat {u} _ {t}, x \rangle) \langle \hat {u} _ {t}, x \rangle ] \right| \tag {E.135} \\ \lesssim \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\bar {\rho}} (x) - f _ {\hat {\rho}} (x)\right) ^ {2} \right] ^ {1 / 2} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) ^ {2} \langle \hat {u} _ {t}, x \rangle^ {2} \right] ^ {1 / 2} \\ \end{array}
$$
$$
\lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \bar {\Delta} _ {t} \quad \text {(B y L e m m a E . 1 6 a n d L e m m a E . 1 5 , E q . (E . 2 1 1))}
$$
Thus, we can conclude that
$$
\left| B _ {t} ^ {\prime \prime} \right| \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \cdot \left(\frac {d (\log d) ^ {O (1)}}{\sqrt {m}} + \overline {{\Delta}} _ {t}\right) \cdot \| \delta_ {t} \| _ {2} ^ {2} \tag {E.136}
$$
Overall Bound on  $B_{t}$ . Combining our bounds on  $B_{t}^{\prime}$  and  $B_{t}^{\prime \prime}$ , and using the fact that  $\| \delta_t\| _2\leq 2$  to absorb the bound for  $B_{t}^{\prime \prime}$  into the bound for  $B_{t}^{\prime}$ , we find that
$$
\left| B _ {t} \right| \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \cdot \left(\frac {d (\log d) ^ {O (1)}}{\sqrt {m}} + \bar {\Delta} _ {t}\right) \| \delta_ {t} \| _ {2} \tag {E.137}
$$
with probability  $1 - e^{-d^2}$  (recall that Eq. (E.124) held with probability  $1 - e^{-d^2}$ ). This completes the proof.
Next, we show a tighter bound when we take the expectation of  $B_{t}$  over the neurons, i.e. over  $\chi$ . Taking the expectation allows us to obtain a tighter bound because of the term
$$
- 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \delta_ {t}, x \rangle \right]. \tag {E.138}
$$
For a single  $\hat{u}_t$ , the best bound we can obtain is by using the Cauchy-Schwarz inequality. However, the key point in the following lemma is that when we take the expectation over  $(\hat{u}_t,\bar{u}_t)\sim \Gamma_t$ , this term becomes
$$
- 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)) ^ {2} ] \leq 0 \tag {E.139}
$$
together with a third-order term which is at most  $(\hat{\sigma}_{2,d}^{2} + \hat{\sigma}_{4,d}^{2})\overline{\Delta}_{t}^{3}$ , which is a significantly better bound than  $(\hat{\sigma}_{2,d}^{2} + \hat{\sigma}_{4,d}^{2})\overline{\Delta}_{t}\| \delta_{t}\|_{2}$ , which we obtained in the previous lemma. This is because to simplify Eq. (E.138), we can use
$$
\sigma \left(\langle \bar {u} _ {t}, x \rangle\right) = \sigma \left(\langle \hat {u} _ {t}, x \rangle\right) + \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \left\langle \bar {u} _ {t} - \hat {u} _ {t}, x \right\rangle + O \left(\sigma^ {\prime \prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \left\langle \delta_ {t}, x \right\rangle^ {2}\right) \tag {E.140}
$$
meaning that we can essentially substitute  $\sigma'(\langle \hat{u}_t, x \rangle) \langle \delta_t, x \rangle \approx \sigma(\langle \hat{u}_t, x \rangle) - \sigma(\langle \bar{u}_t, x \rangle)$  in Eq. (E.138). Note that we need this better bound on the growth rate of  $\overline{\Delta}_t$  since our inductive hypothesis imposes a stricter condition on  $\overline{\Delta}_t$ .
Lemma E.10. Suppose we are in the setting of Lemma E.9, and in particular, suppose  $T$  satisfies the same assumptions as in the statement of Lemma E.9. Then, for all  $t \leq T$ , we have
$$
\underset {\chi \sim \hat {\rho} _ {0}} {\mathbb {E}} \left[ B _ {t} (\chi) \right] \lesssim \frac {\left(\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}\right) d \left(\log d\right) ^ {O (1)}}{\sqrt {m}} \bar {\Delta} _ {t} + \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \bar {\Delta} _ {t} ^ {3}. \tag {E.141}
$$
Proof of Lemma E.10. We again make use of Eq. (E.121), and we recall from the proof of the previous lemma that  $B_{t} = B_{t}^{\prime} + B_{t}^{\prime \prime}$ , where
$$
B _ {t} ^ {\prime} = - 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\bar {\rho}} (x) - f _ {\rho} (x)\right) \sigma^ {\prime} (\langle \hat {u} _ {t}, x \rangle) \langle \delta_ {t}, x \rangle \right] - 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} (\langle \hat {u} _ {t}, x \rangle) \langle \delta_ {t}, x \rangle \right] \tag {E.142}
$$
$$
B _ {t} ^ {\prime \prime} = \left\langle \nabla_ {\hat {u}} L \left(\hat {\rho} _ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} \right\rangle \cdot \left\| \delta_ {t} \right\| _ {2} ^ {2}. \tag {E.143}
$$
The main difference between Lemma E.9 is that we now bound the second term of  $B_t'$  in expectation. By Taylor's theorem, we can write
$$
\sigma \left(\langle \bar {u} _ {t}, x \rangle\right) = \sigma \left(\langle \hat {u} _ {t}, x \rangle\right) + \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \left(\langle \bar {u} _ {t}, x \rangle - \langle \hat {u} _ {t}, x \rangle\right) + \frac {\sigma^ {\prime \prime} (\lambda)}{2} \left(\langle \bar {u} _ {t}, x \rangle - \langle \hat {u} _ {t}, x \rangle\right) ^ {2} \tag {E.144}
$$
for some  $\lambda$  which is a convex combination of  $\langle \hat{u}, x \rangle$  and  $\langle \bar{u}, x \rangle$ . Rearranging gives
$$
\sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \left\langle \delta_ {t}, x \right\rangle = \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \left\langle \hat {u} _ {t} - \bar {u} _ {t}, x \right\rangle = \sigma \left(\left\langle \hat {u} _ {t}, x \right\rangle\right) - \sigma \left(\left\langle \bar {u} _ {t}, x \right\rangle\right) + \frac {\sigma^ {\prime \prime} (\lambda)}{2} \left\langle \delta_ {t}, x \right\rangle^ {2} \tag {E.145}
$$
Thus, taking the expectation of the second term of  $B_{t}^{\prime}$  over  $(\hat{u}_t,\bar{u}_t)\sim \Gamma_t$  gives
$$
\begin{array}{l} - 2 \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \delta_ {t}, x \rangle \right] (E.146) \\ = - 2 \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)) (\sigma (\langle \hat {u} _ {t}, x \rangle) - \sigma (\langle \bar {u} _ {t}, x \rangle)) ] (E.147) \\ - \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime \prime} (\lambda) \left\langle \delta_ {t}, x \right\rangle^ {2} \right] \quad \text {(B y E q . (E . 1 4 5))} \\ = - 2 \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) ^ {2} \right] (E.148) \\ - \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime \prime} (\lambda) \left\langle \delta_ {t}, x \right\rangle^ {2} \right] (ByEq.(3.1)) \\ \leq - \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime \prime} (\lambda) \langle \delta_ {t}, x \rangle^ {2} \right] (E.149) \\ \end{array}
$$
Next, let us bound the inner expectation in Eq. (E.149). By the Cauchy-Schwarz inequality, we have
$$
\begin{array}{l} \left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime \prime} (\lambda) \langle \delta_ {t}, x \rangle^ {2} \right] \right| (E.150) \\ \leq \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} [ (f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)) ^ {2} ] ^ {1 / 2} \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} [ \sigma^ {\prime \prime} (\lambda) ^ {2} \langle \delta_ {t}, x \rangle^ {4} ] ^ {1 / 2} (E.151) \\ \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) ^ {1 / 2} \bar {\Delta} _ {t} \cdot \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \sigma^ {\prime \prime} (\lambda) ^ {2} \langle \delta_ {t}, x \rangle^ {4} \right] ^ {1 / 2} \quad (By Lemma E.16) \\ \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) ^ {1 / 2} \bar {\Delta} _ {t} \cdot \left(\underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ \sigma^ {\prime \prime} (\langle \hat {u} _ {t}, x \rangle) ^ {2} \langle \delta_ {t}, x \rangle^ {4} ] ^ {1 / 2} + \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ \sigma^ {\prime \prime} (\langle \bar {u} _ {t}, x \rangle) ^ {2} \langle \delta_ {t}, x \rangle^ {4} ] ^ {1 / 2}\right) (E.152) \\ \end{array}
$$
Here to obtain the last inequality, we observe that  $\sigma''$  is convex because  $P_{2,d}''$  is a constant function and  $P_{4,d}''$  is a quadratic function by Eq. (C.4) — additionally, we used the inequality  $(a + b)^2 \lesssim a^2 + b^2$  and the fact that  $\lambda$  is a convex combination of  $\langle \hat{u}_t, x \rangle$  and  $\langle \bar{u}_t, x \rangle$ .
By an argument similar to the proof of Lemma E.15, we can show that  $\mathbb{E}_{x\sim \mathbb{S}^{d - 1}}[\sigma^{\prime \prime}(\langle \hat{u}_t,x\rangle)^2\langle \delta_t,x\rangle^4 ]^{1 / 2}\lesssim (\hat{\sigma}_{2,d}^2 +\hat{\sigma}_{4,d}^2)^{1 / 2}\| \delta_t\| _2^2$  and  $\mathbb{E}_{x\sim \mathbb{S}^{d - 1}}[\sigma^{\prime \prime}(\langle \bar{u}_t,x\rangle)^2\langle \delta_t,x\rangle^4 ]^{1 / 2}\lesssim$ $(\hat{\sigma}_{2,d}^2 +\hat{\sigma}_{4,d}^2)^{1 / 2}\| \delta_t\| _2^2.$  Combining this with Eq. (E.149), we have that
$$
\left| \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\hat {\rho}} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \langle \delta_ {t}, x \rangle \right] \right| \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \bar {\Delta} _ {t} \| \delta_ {t} \| _ {2} ^ {2} \tag {E.153}
$$
$$
- 2 \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\hat {\rho}} (x) - f _ {\bar {\rho}} (x)\right) \sigma^ {\prime} \left(\langle \hat {u} _ {t}, x \rangle\right) \left\langle \delta_ {t}, x \right\rangle \right] \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \overline {{\Delta}} _ {t} ^ {3} \tag {E.154}
$$
by the definition of  $\overline{\Delta}_t^3$ . This gives a bound on the second term of  $B_t^\prime$ . Combining this with Eq. (E.126) to bound the first term of  $B_{t}^{\prime}$ , we have
$$
\begin{array}{l} \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \left[ B _ {t} ^ {\prime} \right] \lesssim \frac {\left(\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}\right) d (\log d) ^ {O (1)}}{\sqrt {m}} \underset {(\hat {u}, \bar {u}) \sim \Gamma_ {t}} {\mathbb {E}} \| \delta_ {t} \| _ {2} + \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \overline {{\Delta}} _ {t} ^ {3} \tag {E.155} \\ \lesssim \frac {(\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}) d (\log d) ^ {O (1)}}{\sqrt {m}} \overline {{\Delta}} _ {t} + (\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}) \overline {{\Delta}} _ {t} ^ {3} \quad \text {(B y C a u c h y - S c h w a r z)} \\ \end{array}
$$
Combining this with Eq. (E.136) (for which we also take the expectation over  $(\hat{u}_t,\bar{u}_t)\sim \Gamma_t)$ , we have that