SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \frac {d}{d t} \bar {\Delta} _ {t} = \frac {1}{\bar {\Delta} _ {t}} \frac {d}{d t} \bar {\Delta} _ {t} ^ {2} (E.52) \\ = \frac {1}{\bar {\Delta} _ {t}} \underset {x} {\mathbb {E}} [ \langle \hat {u} _ {t} - \bar {u} _ {t}, \hat {u} _ {t} ^ {\prime} - \bar {u} _ {t} ^ {\prime} \rangle ] (E.53) \\ \leq \frac {1}{\Delta_ {t}} \mathbb {E} _ {\chi} [ \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} \| \hat {u} _ {t} ^ {\prime} - \bar {u} _ {t} ^ {\prime} \| _ {2} ] \quad (\text {B y C a u c h y - S h w a r z I n e q u a l i t y}) \\ \leq \frac {\left(\hat {\sigma} _ {2 , d} ^ {2} + \hat {\sigma} _ {4 , d} ^ {2}\right) d ^ {4}}{\bar {\Delta} _ {t}} \underset {\chi} {\mathbb {E}} \left[ \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} \right] (By Lemma E.14) \\ \leq \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) d ^ {4}. (E.54) \\ \end{array}
$$
Thus, for  $\mu_{2} = \left(\frac{1}{d^{\phi}} -\frac{1}{d^{\phi}(\log d)^{O(1)}}\right)\cdot \frac{1}{(\hat{\sigma}_{2,d}^{2} + \hat{\sigma}_{4,d}^{2})d^{4}}$  we still have the inductive hypothesis for  $\overline{\Delta}$  until time  $T + \mu_{2}$ , i.e.  $\overline{\Delta}_{T + \mu_2}\leq \frac{1}{d^\phi}$ . As a result, assuming the inductive hypothesis holds up to time  $T$ , it holds up to time  $T + \min (\mu_1,\mu_2)$ .
Our above proof strategy applies to all times  $T \leq T_1$ . For the induction when  $T \in [T_1, T_{*,\epsilon}]$  we follow a similar argument but instead use the bound on  $A_t$  obtained from Lemma E.8. Since we know from Lemma D.36 and the assumption that  $\epsilon = O\left(\frac{1}{\log\log d}\right)$  that
$$
\int_ {T _ {1}} ^ {T _ {* , \epsilon}} \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \gamma_ {2} d t \leq \left(T _ {* , \epsilon} - T _ {1}\right) \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \gamma_ {2} \leq \frac {\operatorname {p o l y} (\log \log d)}{\gamma_ {2} ^ {O (1)}}, \tag {E.55}
$$
we know that this additional phase will at most add another  $\exp (\mathrm{poly}(\log \log d))$  factor to the growth of  $\Delta_{\max ,t}$  which is smaller than any polynomial in  $d$ . Thus, Eq. (E.47) and Eq. (E.51) still hold for this period, which finishes the proof.
Finally, we prove Lemma 5.2:
Proof of Lemma 5.2. The proof directly follows Lemma E.3 with  $\phi = \frac{\mu - 1}{4}$  and  $\psi = \frac{\mu - 3}{8}$ .
E.1 Upper bound on  $A_{t}$
The goal of this section is to obtain an upper bound on
$$
A _ {t} (\chi) := - 2 \left\langle \operatorname {g r a d} _ {\hat {u}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} (\chi) - \bar {u} _ {t} (\chi) \right\rangle . \tag {E.56}
$$
When there is no risk of confusion, we simply use  $A_{t}$ ,  $\hat{u}_t$ , and  $\bar{u}_t$ , omitting the dependence on  $\chi$ .
We first obtain a strong upper bound on this inner product during Phase 1 of the population projected gradient flow. To obtain an upper bound on this inner product, we separately consider the contributions from the first coordinates of the vectors, and the remaining  $(d - 1)$  coordinates — in other words, the gradients with respect to  $w$  and the gradients with respect to  $z$ . For convenience, let  $\operatorname{grad}_w L(\rho)$  denote the gradient with respect to  $w$  and  $\operatorname{grad}_z L(\rho)$  denote the gradient with respect to  $z$ . The following lemma is about the contribution of the first coordinates of the particles to  $A_t$ , and is useful during Phase 1.
Lemma E.4 (Contribution of  $w$ 's). Suppose we are in the setting of Theorem 3.4, and let  $\rho_t$  be the projected gradient flow on population loss, initialized with the uniform distribution on  $\mathbb{S}^{d-1}$ , defined
in Eq. (3.7). Suppose  $D_{2,t} < 0$  and  $D_{4,t} < 0$ , where  $D_{2,t}$  and  $D_{4,t}$  are as defined in Section 4.1. Then, the contribution of the first coordinates to  $A_{t}$  can be bounded by
$$
\begin{array}{l} - 2 \left(\operatorname {g r a d} _ {\hat {w}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {w}} L \left(\rho_ {t}\right)\right) \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right) (E.57) \\ \leq 4 \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2} + O \left(\hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | w _ {\max } \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2}\right) (E.58) \\ + O \left(\hat {\sigma} _ {4, d} ^ {2} \left| D _ {4, t} \right| \Delta_ {\max } \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2}\right) (E.59) \\ + O \left(\frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d} \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2}\right). (E.60) \\ \end{array}
$$
where  $w_{\mathrm{max}}$  is as defined in Definition 4.5.
The constant 4 at the beginning of the first term is particularly important since it determines the growth rate of  $\| \delta_t\| _2^2$ , and in turn affects the number of samples we will need. At least during Phase 1, the contribution of the first coordinates will be the dominant term in  $A_{t}$ , as we will see in the next few lemmas.
Proof of Lemma E.4. The contribution of the first coordinate to the dot product defining  $A_{t}$  is
$$
\begin{array}{l} - 2 \left(\operatorname {g r a d} _ {\bar {w}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {w}} L \left(\rho_ {t}\right)\right) \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right) (E.61) \\ = 2 \left(v \left(\hat {w} _ {t}\right) - v \left(\bar {w} _ {t}\right)\right) \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right) (E.62) \\ \end{array}
$$
where  $v(w)$  is as defined in Lemma 4.4. We can further simplify  $v(\hat{w}_t) - v(\bar{w}_t)$  using Lemma 4.4. Let  $P_t$  and  $Q_t$  be defined as in Lemma 4.4. Then,
$$
v \left(\hat {w} _ {t}\right) - v \left(\bar {w} _ {t}\right) = - \left(1 - \hat {w} _ {t} ^ {2}\right) \left(P _ {t} \left(\hat {w} _ {t}\right) + Q _ {t} \left(\hat {w} _ {t}\right)\right) + \left(1 - \bar {w} _ {t} ^ {2}\right) \left(P _ {t} \left(\bar {w} _ {t}\right) + Q _ {t} \left(\bar {w} _ {t}\right)\right). \tag {E.63}
$$
Let us upper bound the term corresponding to  $Q_{t}$  first. Since each coefficient of  $Q_{t}$  is bounded above by  $O\left(\frac{\hat{\sigma}_{2,d}^{2}|D_{2,t}| + \hat{\sigma}_{4,d}^{2}|D_{4,t}|}{d}\right)$  (by the conclusion of Lemma 4.4) and  $Q_{t}$  only has linear and cubic terms, we have
$$
\left| - \left(1 - \hat {w} _ {t} ^ {2}\right) Q _ {t} \left(\hat {w} _ {t}\right) + \left(1 - \bar {w} _ {t} ^ {2}\right) Q _ {t} \left(\bar {w} _ {t}\right) \right| \cdot | \hat {w} _ {t} - \bar {w} _ {t} | \lesssim \frac {\hat {\sigma} _ {2 , d} ^ {2} \left| D _ {2 , t} \right| + \hat {\sigma} _ {4 , d} ^ {2} \left| D _ {4 , t} \right|}{d} | \hat {w} _ {t} - \bar {w} _ {t} | ^ {2} \tag {E.64}
$$
by Lemma I.16. The terms corresponding to  $P_{t}(\bar{w}_{t})$  in  $(v(\hat{w}_t) - v(\bar{w}_t))(\hat{w}_t - \bar{w}_t)$  can be bounded as follows:
$$
\left.\left.\left.\left(- \left(1 - \hat {w} _ {t} ^ {2}\right) P _ {t} \left(\hat {w} _ {t}\right) + \left(1 - \bar {w} _ {t} ^ {2}\right) P _ {t} \left(\bar {w} _ {t}\right)\right) \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right)\right.\right.\right.\left. \right. \tag {E.65}
$$
$$
\begin{array}{l} = 2 \hat {\sigma} _ {2, d} ^ {2} \left| D _ {2, t} \right| \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2} + 4 \hat {\sigma} _ {4, d} ^ {2} \left| D _ {4, t} \right| \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2} \left(\hat {w} _ {t} ^ {2} + \hat {w} _ {t} \bar {w} _ {t} + \bar {w} _ {t} ^ {2}\right) (E.66) \\ - 2 \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2} \left(\hat {w} _ {t} ^ {2} + \hat {w} _ {t} \bar {w} _ {t} + \bar {w} _ {t} ^ {2}\right) - 4 \hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2} \left(\hat {w} _ {t} ^ {2} + \hat {w} _ {t} \bar {w} _ {t} + \bar {w} _ {t} ^ {2}\right) (E.67) \\ \end{array}
$$
$$
\leq 2 \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2} + 4 \hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2} (\hat {w} _ {t} ^ {2} + \hat {w} _ {t} \bar {w} _ {t} + \bar {w} _ {t} ^ {2})
$$
(Since last two terms are nonnegative)
$$
\leq 2 \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2} + O \left(\hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | | \hat {w} _ {t} + \bar {w} _ {t} | \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2}\right) \quad (\text {S i n c e} \bar {w} _ {t}, \hat {w} _ {t} \leq 1)
$$
Thus, the overall bound on the contribution of the first coordinates to  $A_{t}$  is
$$
\begin{array}{l} 2 \left(v \left(\hat {w} _ {t}\right) - v \left(\bar {w} _ {t}\right)\right) \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right) \leq 2 \cdot \left(- \left(1 - \hat {w} _ {t} ^ {2}\right) P _ {t} (\hat {w} _ {t}) + \left(1 - \bar {w} _ {t} ^ {2}\right) P _ {t} (\bar {w} _ {t})\right) \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right) (E.68) \\ + \left| - \left(1 - \hat {w} _ {t} ^ {2}\right) Q _ {t} \left(\hat {w} _ {t}\right) + \left(1 - \bar {w} _ {t} ^ {2}\right) Q _ {t} \left(\bar {w} _ {t}\right) \right| \cdot | \hat {w} _ {t} - \bar {w} _ {t} | (E.69) \\ \leq 4 \hat {\sigma} _ {2, d} ^ {2} \left| D _ {2, t} \right| \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2} + O \left(\hat {\sigma} _ {4, d} ^ {2} \left| D _ {4, t} \right| \left| \hat {w} _ {t} + \bar {w} _ {t} \right| \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2}\right) (E.70) \\ + O \left(\frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d} (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2}\right) (E.71) \\ \end{array}
$$
We can simplify this bound further in terms of  $w_{\mathrm{max}}$  and  $\Delta_{\mathrm{max},t}$ . First observe that
$$
\left| \hat {w} _ {t} - \bar {w} _ {t} \right| \leq \left\| \delta_ {t} \right\| _ {2} \leq \Delta_ {\max , t} \tag {E.72}
$$
Thus, by the triangle inequality, we have  $|\bar{w}_t| \leq w_{\max}$  and  $|\hat{w}_t| \leq w_{\max} + \Delta_{\max,t}$  meaning
$$
\begin{array}{l} 2 \left(v \left(\hat {w} _ {t}\right) - v \left(\bar {w} _ {t}\right)\right) \cdot \left(\hat {w} _ {t} - \bar {w} _ {t}\right) \leq 4 \hat {\sigma} _ {2, d} ^ {2} \left| D _ {2, t} \right| \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2} + O \left(\hat {\sigma} _ {4, d} ^ {2} \left| D _ {4, t} \right| \left| \hat {w} _ {t} + \bar {w} _ {t} \right| \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2}\right) (E.73) \\ + O \left(\frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d} (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2}\right) (E.74) \\ \leq 4 \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2} + O \left(\hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | w _ {\max } (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2}\right) (E.75) \\ + O \left(\hat {\sigma} _ {4, d} ^ {2} \left| D _ {4, t} \right| \Delta_ {\max , t} \left(\hat {w} _ {t} - \bar {w} _ {t}\right) ^ {2}\right) (E.76) \\ + O \left(\frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d} (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2}\right) (E.78) \\ \end{array}
$$
Next, we study the contribution of the last  $(d - 1)$  coordinates of the particles to  $A_{t}(\chi)$  during Phase 1. We first introduce the following formula for the projected gradient of the population loss with respect to the last  $(d - 1)$  coordinates of a particle.
Lemma E.5. Let  $\rho$  be a distribution which is rotationally invariant, as in Definition 4.1. Let  $\operatorname{grad}_u L(\rho)$  be defined as in Eq. (3.4). In addition, let  $\operatorname{grad}_z L(\rho)$  denote the last  $(d - 1)$  coordinates of  $\operatorname{grad}_u L(\rho)$ , for any particle  $u = (w, z)$ . Then, for any particle  $u = (w, z)$ , we have  $\operatorname{grad}_z L(\rho) = -\left(\frac{w}{1 - w^2} \cdot \operatorname{grad}_w L(\rho)\right)z$ .
Proof of Lemma E.5. Notice that  $w^2 + z^2 = 1$  for any  $u = (w, z) \in \mathbb{S}^{d-1}$ , and this equality is maintained by projected gradient descent on the population loss, as in Eq. (3.7). Also, by Lemma I.15, we have  $\mathrm{grad}_z L(\rho) = cz$  for some scalar  $c$ , since  $\rho$  is rotationally invariant. Thus,
$$
\begin{array}{l} 0 = \frac {d}{d t} \left(w ^ {2} + z ^ {2}\right) (E.79) \\ = - 2 w \operatorname {g r a d} _ {w} L (\rho) - 2 \langle z, \operatorname {g r a d} _ {z} L (\rho) \rangle (E.80) \\ = - 2 w \operatorname {g r a d} _ {w} L (\rho) - 2 c \| z \| ^ {2} (E.81) \\ = - 2 w \operatorname {g r a d} _ {w} L (\rho) - 2 c \left(1 - w ^ {2}\right). (E.82) \\ \end{array}
$$
Rearranging gives us
$$
c = - \frac {w}{1 - w ^ {2}} \operatorname {g r a d} _ {w} L (\rho). \tag {E.83}
$$
By directly applying this formula for the last  $(d - 1)$  coordinates of the projected gradient, we have the following upper bound on the contribution of the last  $(d - 1)$  coordinates to  $A_{t}$  during Phase 1.
Lemma E.6 (Contribution of  $z$ 's During Phase 1.). Suppose we are in the setting of Theorem 3.4. Let  $t \leq T_1$  (where  $T_1$  is as defined in Definition 4.5). Let  $\rho_t$  be defined according to the population projected gradient flow as in Eq. (3.7). Let  $\operatorname{grad}_u L(\rho)$  be as in Eq. (3.4), and let  $\operatorname{grad}_z L(\rho)$  denote the last  $(d - 1)$  coordinates of  $\operatorname{grad}_u L(\rho)$  for any particle  $u = (w, z)$ . Then, we have the following upper bound:
$$
\begin{array}{l} - 2 \left\langle \operatorname {g r a d} _ {\bar {z}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {z}} L \left(\rho_ {t}\right), \hat {z} _ {t} - \bar {z} _ {t} \right\rangle (E.84) \\ \lesssim \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | \left(| \bar {w} _ {t} | + | \hat {w} _ {t} |\right) \| \bar {u} _ {t} - \hat {u} _ {t} \| ^ {2} + \hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | \left(| \bar {w} _ {t} ^ {3} | + | \hat {w} _ {t} ^ {3} |\right) \| \bar {u} _ {t} - \hat {u} _ {t} \| ^ {2} (E.85) \\ + \frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d} \| \bar {u} _ {t} - \hat {u} _ {t} \| _ {2} ^ {2}. (E.86) \\ \end{array}
$$
Proof of Lemma E.6. Using the formula for projected gradient in Lemma E.5, we have that
$$
\begin{array}{l} - \left\langle \operatorname {g r a d} _ {\bar {z}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {z}} L \left(\rho_ {t}\right), \hat {z} _ {t} - \bar {z} _ {t} \right\rangle (E.87) \\ = \left\langle \frac {\hat {w} _ {t}}{1 - \hat {w} _ {t} ^ {2}} \cdot \operatorname {g r a d} _ {\hat {w}} L \left(\rho_ {t}\right) \cdot \hat {z} _ {t} - \frac {\bar {w} _ {t}}{1 - \bar {w} _ {t} ^ {2}} \cdot \operatorname {g r a d} _ {\bar {w}} L \left(\rho_ {t}\right) \cdot \bar {z} _ {t}, \hat {z} _ {t} - \bar {z} _ {t} \right\rangle (E.88) \\ = \left(\frac {\hat {w} _ {t} \operatorname {g r a d} _ {\hat {w}} L (\rho_ {t})}{1 - \hat {w} _ {t} ^ {2}} - \frac {\bar {w} _ {t} \operatorname {g r a d} _ {\bar {w}} L (\rho_ {t})}{1 - \bar {w} _ {t} ^ {2}}\right) \hat {z} ^ {\top} \left(\hat {z} _ {t} - \bar {z} _ {t}\right) + \frac {\bar {w} _ {t} \operatorname {g r a d} _ {\bar {w}} L (\rho_ {t})}{1 - \bar {w} _ {t} ^ {2}} \| \hat {z} _ {t} - \bar {z} _ {t} \| _ {2} ^ {2} (E.89) \\ \end{array}
$$
The second term is always negative for  $t \leq T_1$ , because Lemma D.9 implies that  $v(\bar{w}_t)$  has the same sign as  $\bar{w}_t$ , meaning that  $\mathrm{grad}_{\bar{w}} L(\rho_t)$  has the opposite sign. We bound the first term as
$$
\begin{array}{l} \left(\frac {\hat {w} _ {t} \operatorname {g r a d} _ {\hat {w}} L (\rho_ {t})}{1 - \hat {w} _ {t} ^ {2}} - \frac {\bar {w} _ {t} \operatorname {g r a d} _ {\bar {w}} L (\rho_ {t})}{1 - \bar {w} _ {t} ^ {2}}\right) \hat {z} _ {t} ^ {\top} \left(\bar {z} _ {t} - \hat {z} _ {t}\right) (E.90) \\ \leq \left| \frac {\hat {w} _ {t} \operatorname {g r a d} _ {\hat {w}} L (\rho_ {t})}{1 - \hat {w} _ {t} ^ {2}} - \frac {\bar {w} _ {t} \operatorname {g r a d} _ {\bar {w}} L (\rho_ {t})}{1 - \bar {w} _ {t} ^ {2}} \right| \cdot \| \bar {z} _ {t} - \hat {z} _ {t} \| \quad \text {(B . c .} | \hat {z} _ {t} | _ {2} \leq \| \hat {u} _ {t} \| _ {2} = 1) \\ = \left| \hat {w} _ {t} \left(P _ {t} \left(\hat {w} _ {t}\right) + Q _ {t} \left(\hat {w} _ {t}\right)\right) - \bar {w} _ {t} \left(P _ {t} \left(\bar {w} _ {t}\right) + Q _ {t} \left(\bar {w} _ {t}\right)\right) \right| \cdot \| \bar {z} _ {t} - \hat {z} _ {t} \| _ {2} (By Lemma 4.4) \\ \leq 2 \hat {\sigma} _ {2, d} ^ {2} \left| D _ {2, t} \right| \left| \hat {w} _ {t} ^ {2} - \bar {w} _ {t} ^ {2} \right| \| \bar {z} _ {t} - \hat {z} _ {t} \| _ {2} + 4 \hat {\sigma} _ {4, d} ^ {2} \left| D _ {4, t} \right| \left| \hat {w} _ {t} ^ {4} - \bar {w} _ {t} ^ {4} \right| \| \bar {z} _ {t} - \hat {z} _ {t} \| _ {2} (E.91) \\ + \left| Q _ {t} \left(\bar {w} _ {t}\right) - Q _ {t} \left(\hat {w} _ {t}\right) \right| \cdot \| \bar {z} _ {t} - \hat {z} _ {t} \| _ {2} \quad (\text {B y}) \\ \leq 2 \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | | \hat {w} _ {t} ^ {2} - \bar {w} _ {t} ^ {2} | | \bar {z} _ {t} - \hat {z} _ {t} | | _ {2} + 4 \hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | | \hat {w} _ {t} ^ {4} - \bar {w} _ {t} ^ {4} | | \bar {z} _ {t} - \hat {z} _ {t} | | _ {2} (E.92) \\ + O \Big (\frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d} \Big) | \bar {w} _ {t} - \hat {w} _ {t} | \| \bar {z} _ {t} - \hat {z} _ {t} \| _ {2} \\ \end{array}
$$
(By Lemma I.16 and Def. of  $Q_{t}$  from Lemma 4.4)
$$
\begin{array}{l} \lesssim \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | \left(| \bar {w} _ {t} | + | \hat {w} _ {t} |\right) \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} ^ {2} + \hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} | \left(\left| \bar {w} _ {t} \right| ^ {3} + \left| \hat {w} _ {t} \right| ^ {3}\right) \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} ^ {2} \tag {E.93} \\ + \frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d} \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} ^ {2} \\ (\mathrm {B . c .} | \bar {w} _ {t} - \hat {w} _ {t} | \leq \| \bar {u} _ {t} - \hat {u} _ {t} \| _ {2} \text {a n d} | \bar {z} _ {t} - \hat {z} _ {t} \| _ {2} \leq \| \bar {u} _ {t} - \hat {u} _ {t} \| _ {2}) \\ \end{array}
$$
We can combine Lemma E.4 and Lemma E.6 and get the following bound on  $A_{t}$  during Phase 1.
Lemma E.7. Suppose we are in the setting of Theorem 3.4. Let  $t \leq T_1$  (where  $T_1$  is as defined in Definition 4.5). Then,
$$
\begin{array}{l} A _ {t} \leq 4 \hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | (\hat {w} _ {t} - \bar {w} _ {t}) ^ {2} + O \left(\left(\hat {\sigma} _ {2, d} ^ {2} | D _ {2, t} | + \hat {\sigma} _ {4, d} ^ {2} | D _ {4, t} |\right) w _ {\max }\right) \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} ^ {2} (E.94) \\ + O \left(\left(\hat {\sigma} _ {2, d} ^ {2} \left| D _ {2, t} \right| + \hat {\sigma} _ {4, d} ^ {2} \left| D _ {4, t} \right|\right) \Delta_ {\max }\right) \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} ^ {2} (E.95) \\ + O \left(\frac {\hat {\sigma} _ {2 , d} ^ {2} | D _ {2 , t} | + \hat {\sigma} _ {4 , d} ^ {2} | D _ {4 , t} |}{d}\right) \| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} ^ {2}. (E.96) \\ \end{array}
$$
Proof of Lemma E.7. This directly follows from Lemma E.4 and Lemma E.6.
Next, we give a potentially looser bound for  $A_{t}$  which holds even after Phase 1. The running time after Phase 1 is at most  $(\log \log d)^{O(1)}$  (aside from the dependency on  $\hat{\sigma}_{2,d}, \hat{\sigma}_{4,d}$ , etc.), meaning that even the following growth rate is enough to ensure that the coupling error grows by only a sub-polynomial factor after Phase 1.
Lemma E.8. Suppose we are in the setting of Theorem 3.4. Then, for all  $t \geq 0$ , we have
$$
\left| A _ {t} \right| \lesssim \left(\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}\right) \gamma_ {2} \| \bar {u} _ {t} - \hat {u} _ {t} \| ^ {2}. \tag {E.97}
$$
Proof of Lemma E.8. We first consider the part of  $A_{t}$  that originates from the non-projected gradients, which is defined as
$$
\begin{array}{l} A _ {t} ^ {\prime} := - 2 \left\langle \nabla_ {\bar {u}} L \left(\rho_ {t}\right) - \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \bar {u} _ {t} - \hat {u} _ {t} \right\rangle (E.98) \\ = - 2 \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} \left[ \left(f _ {\rho} (x) - y (x)\right) \left(\sigma^ {\prime} \left(\bar {u} _ {t} ^ {\top} x\right) - \sigma^ {\prime} \left(\hat {u} _ {t} ^ {\top} x\right)\right) \left(\bar {u} _ {t} ^ {\top} x - \hat {u} _ {t} ^ {\top} x\right) \right]. (E.99) \\ \end{array}
$$
By the Cauchy-Schwarz inequality, we have
$$
\left| A _ {t} ^ {\prime} \right| \lesssim \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} \left[ \left(f _ {\rho} (x) - y (x)\right) ^ {2} \right] ^ {\frac {1}{2}} \cdot \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} \left[ \left(\sigma^ {\prime} \left(\bar {u} _ {t} ^ {\top} x\right) - \sigma^ {\prime} \left(\hat {u} _ {t} ^ {\top} x\right)\right) ^ {2} \left(\bar {u} _ {t} ^ {\top} x - \hat {u} _ {t} ^ {\top} x\right) ^ {2} \right] ^ {\frac {1}{2}}. \tag {E.100}
$$
We bound the first factor in the equation above by the loss at initialization, which is
$$
\frac {1}{2} \mathbb {E} _ {x \sim \mathbb {S} ^ {d - 1}} [ (f _ {\rho_ {0}} (x) - y) ^ {2} ] \lesssim (\hat {\sigma} _ {2, d} ^ {2} + \hat {\sigma} _ {4, d} ^ {2}) \gamma_ {2} ^ {2}. \tag {E.101}
$$
by Lemma D.8. Thus, plugging Eq. (E.101) and Eq. (E.210) (from the statement of Lemma E.15) into Eq. (E.100) gives us
$$
\left| A _ {t} ^ {\prime} \right| \lesssim \left(\hat {\sigma} _ {4, d} ^ {2} + \hat {\sigma} _ {2, d} ^ {2}\right) \gamma_ {2} \| \delta_ {t} \| ^ {2}. \tag {E.102}
$$
Next, we bound the difference between  $A_{t}$  and  $A_{t}^{\prime}$ . We have
$$
\begin{array}{l} A _ {t} - A _ {t} ^ {\prime} = - \left\langle \operatorname {g r a d} _ {\bar {u}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle + \left\langle \nabla_ {\hat {u}} L \left(\rho_ {t}\right) - \nabla_ {\bar {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle (E.103) \\ = \left\langle \nabla_ {\hat {u}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle - \left\langle \nabla_ {\bar {u}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\bar {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle (E.104) \\ \end{array}
$$
For the first term of Eq. (E.104), we have
$$
\begin{array}{l} \left\langle \nabla_ {\hat {u}} L \left(\rho_ {t}\right) - \operatorname {g r a d} _ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle = \left\langle \hat {u} _ {t} \hat {u} _ {t} ^ {\top} \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \hat {u} _ {t} - \bar {u} _ {t} \right\rangle (E.105) \\ = \left\langle \nabla_ {\hat {u}} L \left(\rho_ {t}\right), \left(\hat {u} _ {t} \hat {u} _ {t} ^ {\top}\right) \hat {u} _ {t} - \left(\hat {u} _ {t} \hat {u} _ {t} ^ {\top}\right) \bar {u} _ {t} \right\rangle (E.106) \\ = \left\langle \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.107) \\ = \left\langle \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} \quad (\mathrm {B . c .} \| \hat {u} _ {t} \| _ {2} = \| \bar {u} _ {t} \| _ {2} = 1) \\ = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} [ (f _ {\rho_ {t}} (x) - y (x)) \sigma^ {\prime} (\hat {u} _ {t} ^ {\top} x) \hat {u} _ {t} ^ {\top} x ] \cdot \frac {\| \hat {u} _ {t} - \bar {u} _ {t} \| _ {2} ^ {2}}{2} (E.108) \\ = \underset {x \sim \mathbb {S} ^ {d - 1}} {\mathbb {E}} \left[ \left(f _ {\rho_ {t}} (x) - y (x)\right) \sigma^ {\prime} \left(\hat {u} _ {t} ^ {\top} x\right) \hat {u} _ {t} ^ {\top} x \right] \cdot \frac {\left\| \delta_ {t} \right\| _ {2} ^ {2}}{2} (E.109) \\ \end{array}
$$