SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lemma E.1 (Downstream classification accuracy). Suppose we draw labeled data points  $(u, y) \in \mathbb{R}^{K \times d} \times \{+1, 1\}$ , where as before,  $u_k \sim \mathcal{N}(0, I_d)$  for  $k \in [K]$ , and the label is given by  $\mathrm{sign}(u_k^T \mu_k)$ . Then the best linear classifier  $\mathbf{a} \in \mathbb{R}^K$  on the representations  $f_\Theta(u) \in \mathbb{R}^K$  achieves an test error of  $\frac{1}{\pi} \arccos \left( \frac{|\mu_k^T \theta_k|}{\|\theta_k\|_2} \right)$ . That is
$$
\min  _ {\boldsymbol {a} \in \mathbb {R} ^ {K}} \Pr_ {u} [ \operatorname {s i g n} (\boldsymbol {a} ^ {T} f _ {\Theta} (u)) \neq \operatorname {s i g n} (\mu_ {k} ^ {T} u _ {k}) ] = \frac {\operatorname {a r c c o s} \left(\frac {\left| \mu_ {k} ^ {T} \theta_ {k} \right|}{\left\| \theta_ {k} \right\| _ {2}}\right)}{\pi}. \tag {1}
$$
Figure 10: Alignment of features with verses without added noise. From left to right:  $\alpha_{1} = 0.125, 0.25, 0.375, 0.5$ . The top plots show the alignment of Features 1 (weak) and 2 (dominant) to the ground truth; the bottom plots show the probability of predicting the correct augmentation pair from the batch. Standard deviation bars are shown for the mean alignment over 200 runs. We used dimension  $d = 5$ , and a batch size of  $m = 25$ .
Thus if  $\theta_{k}$  and  $\mu_{k}$  are orthogonal, then the test error is  $50\%$ . If the angle between  $\theta_{k}$  and the  $\pm \mu_{k}$  is zero, then we achieve perfect classification accuracy.
Proof. It is easy to see that the best linear classifier  $\pmb{a}$  will (up to scaling) be equal to the vector  $\mathrm{sign}(\mu_k^T\theta_k)e_k$ . Such a classifier predicts the correct sign whenever  $\mathrm{sign}(\pmb{a}^T f_\Theta(u)) = \mathrm{sign}(\mu_k^T\theta_k)\mathrm{sign}(\theta_k^T u_k)$  equals  $\mathrm{sign}(\mu_k^T u_k)$ , which occurs exactly a 1 -  $\frac{\arccos\left(\frac{|\mu_k^T\theta_k|}{\|\theta_k\|_2}\right)}{\pi}$  fraction of the time.
In the rest of this section, we prove our main theoretical result, Theorem 5.1, which shows that  $\arccos \left(\frac{|\mu_k^T\theta_k|}{\|\theta_k\|_2}\right)$  decreases faster in expectation during gradient descent if we add noise to the  $k^{\prime}$  feature.
E.1 Notation.
We let  $\delta_{ij}$  denote the  $\delta$ -function which equals 1 if  $i = j$  and 0 otherwise. For a parameter  $\Theta = \{\theta_k\}_{k \in [K]}$ , we let  $\theta_k^{\parallel} := \mu_k \mu_k^T \theta_k$  be the projection of  $\theta_k$  in the  $\mu_k$  direction. We let  $\theta_k^{\perp} = \theta_k - \theta_k^{\parallel}$  be the projection of  $\theta_k$  orthogonal to the feature  $\mu_k$ .
Throughout this section, we consider the ground truth directions to be fixed, and we fix some initial correlation vector  $\alpha$ . We let  $\mathbb{P}_{\alpha}$  denote the distribution from which the pair  $(u,v)$  is drawn from the Gaussian distribution described in Section 5 with correlation coefficients  $\alpha$ . When unspecified, the variables  $U,V$  are drawn from the distribution  $\mathbb{P}_{\alpha}^{m}$ . Since we study what happens when we vary  $\alpha_{k'}$ , for  $x\in [0,1]$ , we use the shorthand  $\mathbb{P}_x$  to denote the distribution  $\mathbb{P}_{\alpha (x)}^{m}$ , where  $\alpha (x)_{k'} = x$ , and  $\alpha (x)_k = \alpha_k$  for all other  $k$ .
We denote  $\mathcal{L}_i(\Theta; U, V) = \mathrm{CE}(\{p_{ij}\}_{j \in [m]}, e_i) = -\log(p_{ii})$ , which we abbreviate by  $\mathcal{L}_i$ . When it is clear that we are considering  $\mathcal{L}_i$  for some fixed  $i$ , we omit the superscripts on the  $i$ th data point or its pair. That is, we denote  $u_k := u_k^{(i)}$  and  $v_k := v_k^{(i)}$ .
E.2 Preliminaries
The following facts about of the derivative of the cross entropy loss are easy derived.
Lemma E.2.
$$
\frac {\partial \mathcal {L} _ {i}}{\partial \Theta} = \sum_ {j} \left(p _ {i j} - \delta_ {i j}\right) \frac {\partial z _ {i j}}{\partial \Theta} = \sum_ {i} \sum_ {j \neq i} p _ {i j} \left(\frac {\partial z _ {i j}}{\partial \Theta} - \frac {\partial z _ {i i}}{\partial \Theta}\right), \tag {2}
$$
$$
\frac {\partial z _ {i j}}{\partial \theta_ {k}} = \left(u _ {k} ^ {(i)} v _ {k} ^ {(j) ^ {T}} + v _ {k} ^ {(j)} u _ {k} ^ {(i) ^ {T}}\right) \theta_ {k}. \tag {3}
$$
We will also need the following facts on Gaussian random variables. The first, Stein's Lemma, is well known.
Lemma E.3 (Stein's Lemma).
$$
\mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} [ X f (X) ] = \sigma^ {2} \mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} [ f ^ {\prime} (X) ]. \tag {4}
$$
The next two lemmas are proved in Section E.4.
Lemma E.4. There exists some constant  $C$  such that following holds. If  $\sigma \leq \frac{1}{C}$ , and  $0 \leq t \leq \frac{1}{\sigma}$  then for any  $c \in \{0, 1, 2, 3\}$ , and  $X \sim \mathcal{N}(0, \sigma^2)$  we have
$$
\mathbb {E} _ {X} \left[ | X | ^ {c} \exp (t | X |) \exp (t X ^ {2}) \right] \leq C \sigma^ {c}. \tag {5}
$$
If additionally  $d\in \{0,1,2,3\}$ $\rho \leq \frac{1}{C}$  and  $Y\sim \mathcal{N}(0,\rho^2)$  , then
$$
\mathbb {E} _ {X} \left[ | X | ^ {c} | Y | ^ {d} \exp (t | X |) \exp (| X Y |) \right] \leq C \sigma^ {c} \rho^ {d}. \tag {6}
$$
Lemma E.5. For some universal constant  $C$ , for any  $\sigma \in [0,1]$ ,  $t \geq 0$ ,  $c \in \{0,1,2,3,4\}$ , we have
$$
\mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} \left[ \left(\exp (t | X |) - 1\right) | X | ^ {c} \right] \leq C t \sigma^ {c}.
$$
E.3 Approach and Lemmas
Intuition for proof of Theorem 5.1. Our proof involves comparing the gradient of the loss in the  $\theta_{k}$  direction,  $\nabla_{k} := \frac{\partial}{\partial\theta_{k}}\mathcal{L}$  in the setting with noise to the setting without noise. Loosely, our goal is to show that for any  $k$ , the projection of the this gradient onto the ground truth direction,  $\mu_k^T\nabla_k\mathrm{sign}(\mu_k^T\theta_k)$ , increases when when increase the noise. The main intuition comes from an expansion of this gradient in Lemma E.7, which shows that  $\mathbb{E}\mu_k^T\nabla_k\mathrm{sign}(\mu_k^T\theta_k)$  approximately scales with  $\sum_{i}(1 - p_{ii})$ . Now observe that  $p_{ii}$ , the probability of correctly matching the  $i$ th view to its pair, decreases when we add noise to feature  $k'$ . Thus adding noise will increase  $\mu_k^T\nabla_k\mathrm{sign}(\mu_k^T\theta_k)$ , thereby improving the alignment.
In the remainder of this section, we outline our proof of Theorem 5.1 in this section. We prove all the lemmas below in Section E.4.
To understand  $\mathbb{E}_{U,V}\left[\arccos \left(\frac{|\mu_k^T\theta_k^{(t + 1)}|}{||\theta_k^{(t + 1)}||_2}\right)\right]$  for a small enough step size, we first claim that it suffices to understand the expected projection of the gradient with respect to  $\theta_{k}$  in the  $\mu_{k}$  direction and in the  $\theta_{k}$  direction. We use the notation  $\nabla_{k} = \frac{\partial\mathcal{L}(\Theta;U,V)}{\partial\theta_{k}}$
Lemma E.6. Let  $\theta_{k}^{+} = \theta_{k} - \eta (\nabla_{k} + \lambda \theta_{k})$ . Then
$$
\lim  _ {\eta \rightarrow 0} \frac {1}{\eta} \left(\mathbb {E} _ {U, V} \left[ \operatorname {a r c c o s} \left(\frac {| \mu_ {k} ^ {T} \theta_ {k} ^ {+} |}{\| \theta_ {k} ^ {+} \| _ {2}}\right)\right] - \operatorname {a r c c o s} \left(\frac {| \mu_ {k} ^ {T} \theta_ {k} |}{\| \theta_ {k} \| _ {2}}\right)\right) = N \mathbb {E} _ {U, V} \left[ - (\mu_ {k} ^ {T} \theta_ {k}) (\mu_ {k} ^ {T} \nabla_ {k}) + \frac {\theta_ {k} ^ {T} \nabla_ {k} (\mu_ {k} ^ {T} \theta_ {k}) ^ {2}}{\| \theta_ {k} \| _ {2} ^ {2}} \right],
$$
where  $N$  is some negative value that depends only on  $\theta_{k}$
Now, since we care about the quantity  $\mathbb{E}_{U,V}\left[\arccos \left(\frac{|\mu_k^T\theta_k^{(t + 1)}|}{\|\theta_k^{(t + 1)}\|_2}\right)\right] - \mathbb{E}_{U,\tilde{V}}\left[\arccos \left(\frac{|\mu_k^T\tilde{\theta}_k^{(t + 1)}|}{\|\tilde{\theta}_k^{(t + 1)}\|_2}\right)\right]$  being positive, it suffices to show that derivative
$$
\frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ - (\mu_ {k} ^ {T} \theta_ {k}) (\mu_ {k} ^ {T} \nabla_ {k}) + \frac {\theta_ {k} ^ {T} \nabla_ {k} (\mu_ {k} ^ {T} \theta_ {k}) ^ {2}}{\| \theta_ {k} \| _ {2} ^ {2}} \right],
$$
is negative for all  $x \in [\tilde{\alpha}_{k'}, \alpha_{k'}]$ . Indeed, from Lemma E.6, we have that
$$
\begin{array}{l} \lim  _ {\eta \rightarrow 0} \frac {1}{\eta} \left(\mathbb {E} _ {U, V \sim \mathbb {P} _ {\alpha_ {k ^ {\prime}}}} \left[ \operatorname {a r c c o s} \left(\frac {\left| \mu_ {k} ^ {T} \theta_ {k} ^ {+} \right|}{\left\| \theta_ {k} ^ {+} \right\| _ {2}}\right)\right] - \mathbb {E} _ {U, V \sim \mathbb {P} _ {\tilde {\alpha} _ {k ^ {\prime}}}} \left[ \operatorname {a r c c o s} \left(\frac {\left| \mu_ {k} ^ {T} \theta_ {k} \right|}{\left\| \theta_ {k} \right\| _ {2}}\right)\right]\right) (8) \\ = N \int_ {\tilde {\alpha} _ {k ^ {\prime}}} ^ {\alpha_ {k ^ {\prime}}} \frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ - \left(\mu_ {k} ^ {T} \theta_ {k}\right) \left(\mu_ {k} ^ {T} \nabla_ {k}\right) + \frac {\theta_ {k} ^ {T} \nabla_ {k} \left(\mu_ {k} ^ {T} \theta_ {k}\right) ^ {2}}{\| \theta_ {k} \| _ {2} ^ {2}} \right] d x, (9) \\ \end{array}
$$
so if the derivative is negative for the full range, then the difference in arccosines is positive.
In the following lemma we compute the derivative of  $\mathbb{E}[\nabla_k]$  with respect to  $x$ .
Lemma E.7.
$$
\begin{array}{l} \frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ \nabla_ {k} \right] = m \frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ \frac {\partial \mathcal {L} _ {i}}{\partial \theta_ {k}} \right] \\ = \frac {- m}{1 - x ^ {2}} \theta_ {k ^ {\prime}} ^ {T} \mu_ {k ^ {\prime}} \sum_ {j \neq i} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ p _ {i j} p _ {i i} \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}}\right) \left(\mu_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} ^ {(i)} - x \mu_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(i)}\right) \left(\frac {\partial (z _ {i j} - z _ {i i})}{\partial \theta_ {k}}\right) \right]. \\ \end{array}
$$
We will analyze this quantity by explicitly taking the expectation with respect to some set of random variables. Let  $S = \{U_k, V_k, U_{k'}, V_{k'}\}$  consist of the random variables  $u_{k'}^{(i)}$ ,  $u_{k'}^{(i)}$ , and  $v_{k'}^{(i)}$ ,  $v_{k'}^{(i)}$  for all  $i \in [m]$ . Define  $q_{ij}$  to be the logits when all variables in  $S$  are set to 0 (Thus explicitly,  $q_{ij} = \frac{\exp\left(\sum_{\tilde{k} \neq k, k'} \theta_{\tilde{k}}^{T} u_{\tilde{k}}^{(i)} \theta_{\tilde{k}}^{T} v_{\tilde{k}}^{(j)}\right)}{\sum_{j'} \exp\left(\sum_{\tilde{k} \neq k, k'} \theta_{\tilde{k}}^{T} u_{\tilde{k}}^{(i)} \theta_{\tilde{k}}^{T} v_{\tilde{k}}^{(j')}\right)}$ ). We will use the notation  $j \sim q$  to denote the distribution on  $[m]$  with mass  $q_{ij}$  on  $j$ .
$$
h (S) := \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}}\right) \left(\mu_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} ^ {(i)} - x \mu_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(i)}\right) \left(\frac {\partial \left(z _ {i j} - z _ {i i}\right)}{\partial \theta_ {k}}\right), \tag {10}
$$
$$
h _ {1} (S) = \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}}\right) \left((1 - x ^ {2}) \mu_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} ^ {(i)}\right) 2 \alpha_ {k} \left(\left(\mu_ {k} ^ {T} u _ {k}\right) \left(\theta_ {k} ^ {\parallel} u _ {k}\right) \mu_ {k} ^ {T}\right), \tag {11}
$$
which are the terms that appear in the right hand side of Lemma E.7 after  $p_{ii}p_{ij}$ . Observe that
$$
\mathbb {E} _ {S} [ h (S) - h _ {1} (S) ] = 0.
$$
The following four lemmas serve to bound  $\frac{d}{dx}\mathbb{E}_S\left[\mu_k^T\nabla_k\right]$  and  $\frac{d}{dx}\mathbb{E}_S\left[\theta_k^T\nabla_k\right]$ . We call the terms of the form  $\mathbb{E}p_{ii}p_{ij}(h(S) - h_1(S))$  "junk" terms, and our goal will be to show that these terms are small. We will control more closely the terms of the form  $\mathbb{E}p_{ii}p_{ij}(h_1(S))$ .
Lemma E.8 (Junk Terms for  $\mu_{k}$  term.). If  $\| \theta_k\| \leq 1$  and  $\| \theta_{k^{\prime}}\| \leq 1$ , then for some universal constant  $C$
$$
\left| \mathbb {E} _ {S} \left[ p _ {i i} p _ {i j} \mu_ {k} ^ {T} (h (S) - h _ {1} (S)) \right] \right| \leq C q _ {i i} q _ {i j} \left(\| \theta_ {k ^ {\prime}} \| ^ {3} \| \theta_ {k} \| ^ {3} + \| \theta_ {k ^ {\prime}} ^ {\parallel} \| \| \theta_ {k} \| ^ {3} + \alpha_ {k} \left(\| \theta_ {k ^ {\prime}} \| ^ {3} \| \theta_ {k} ^ {\parallel} \|\right)\right).
$$
Lemma E.9 (Good Term for  $\mu_{k}$  term.). If  $\| \theta_k\| \leq 1$  and  $\| \theta_{k^{\prime}}\| \leq 1$ , then for some universal constant  $C$
$$
\left| \mathbb {E} _ {S} \left[ p _ {i i} p _ {i j} \mu_ {k} ^ {T} h _ {1} (S) \right] \right| \geq 2 \alpha_ {k} (1 - x ^ {2}) q _ {i i} q _ {i j} \left(\| \theta_ {k ^ {\prime}} ^ {\parallel} \| \| \theta_ {k} ^ {\parallel} \|\right) \left(1 - C (\| \theta_ {k ^ {\prime}} \| ^ {2} + \| \theta_ {k} \| ^ {2})\right).
$$
Plugging these two lemmas into Lemma E.7 yields the following corollary.
Corollary E.9.1 (Total  $\mu_{k}$  term.). If for a sufficiently large constant  $C$ ,  $|\theta_k^T\mu_k| \leq \frac{1 - \alpha_{k'}^2}{C} \| \theta_k\|$ ,  $||\theta_{k'}||^3 \leq |\theta_{k'}'^T\mu_k|$ , and  $||\theta_k||^2 \leq \frac{\alpha_k(1 - \alpha_{k'}^2)}{C}$ , then
$$
\left(\mu_ {k} ^ {T} \theta_ {k}\right) \frac {d}{d x} \mathbb {E} _ {\mathbb {P} _ {x}} \left[ \mu_ {k} ^ {T} \nabla_ {k} \right] \geq \frac {m}{2} \mathbb {E} _ {U, V \backslash S} \left[ \sum_ {i, j} q _ {i i} q _ {i j} 2 \alpha_ {k} \| \theta_ {k ^ {\prime}} ^ {\parallel} \| ^ {2} \| \theta_ {k} ^ {\parallel} \| ^ {2} \right].
$$
Lemma E.10 (Junk Terms for  $\theta_{k}$  term.). If  $\| \theta_k\| \leq 1$  and  $\| \theta_{k^{\prime}}\| \leq 1$ , then for some universal constant  $C$
$$
\left| \mathbb {E} _ {S} \left[ p _ {i i} p _ {i j} \theta_ {k} ^ {T} (h (S) - h _ {1} (S)) \right] \right| \leq C q _ {i i} q _ {i j} \left(\| \theta_ {k ^ {\prime}} \| ^ {3} \| \theta_ {k} \| ^ {4} + \| \theta_ {k ^ {\prime}} ^ {\parallel} \| \| \theta_ {k} \| ^ {4} + \alpha_ {k} \left(\| \theta_ {k ^ {\prime}} \| ^ {3} \| \theta_ {k} \| \| \theta_ {k} ^ {\parallel} \| + \| \theta_ {k ^ {\prime}} ^ {\parallel} \| \| \theta_ {k} \| ^ {3} \| \theta_ {k} ^ {\parallel} \|\right)\right).
$$
Lemma E.11 (Good Term for  $\theta_{k}$  term.). If  $\| \theta_k\| \leq 1$  and  $\| \theta_{k^{\prime}}\| \leq 1$ , then for some universal constant  $C$
$$
\left| \mathbb {E} _ {S} \left[ p _ {i i} p _ {i j} \theta_ {k} ^ {T} h _ {1} (S) \right] \right| \leq (1 - x ^ {2}) 2 \alpha_ {k} q _ {i i} q _ {i j} \left(\| \theta_ {k ^ {\prime}} ^ {\parallel} \| \| \theta_ {k} ^ {\parallel} \| ^ {2}\right) \left(1 + C (\| \theta_ {k ^ {\prime}} \| ^ {2} + \| \theta_ {k} \| ^ {2})\right).
$$
Corollary E.11.1 (Total  $\theta_{k}$  term.). If for a sufficiently large constant  $C$ ,  $\| \theta_k^{\parallel}\| \leq \frac{1 - x^2}{C}\| \theta_k\|$ ,  $\| \theta_{k'}\|^3 \leq \| \theta_{k'}^{\parallel}\|$ ,  $\| \theta_k\|^2 \leq \frac{\alpha_k(1 - x^2)}{C}$ , then
$$
\frac {(\mu_ {k} ^ {T} \theta_ {k}) ^ {2}}{\| \theta_ {k} \| ^ {2}} \left| \frac {d}{d x} \mathbb {E} _ {\mathbb {P} _ {x}} [ \theta_ {k} ^ {T} \nabla_ {k} ] \right| \leq \frac {m}{2} \mathbb {E} _ {U, V \setminus S} \left[ \sum_ {i, j} q _ {i i} q _ {i j} \alpha_ {k} \| \theta_ {k ^ {\prime}} ^ {\parallel} \| ^ {2} \| \theta_ {k} ^ {\parallel} \| ^ {2} \right].
$$
Combining Corollaries E.9.1 and E.11.1, we obtain the following lemma.
Lemma E.12. If for a sufficiently large constant  $C$ ,  $\| \theta_k^{\parallel}\| \leq \frac{1 - x^2}{C}\| \theta_k\|$ ,  $\| \theta_{k'}\| ^3\leq \| \theta_{k'}^{\parallel}\|$ ,  $\| \theta_k\|^2\leq \frac{\alpha_k(1 - x^2)}{C}$ , then
$$
\mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ - \left(\mu_ {k} ^ {T} \theta_ {k}\right) \left(\mu_ {k} ^ {T} \nabla_ {k}\right) + \frac {\theta_ {k} ^ {T} \nabla_ {k} \left(\mu_ {k} ^ {T} \theta_ {k}\right) ^ {2}}{\| \theta_ {k} \| _ {2} ^ {2}} \right] <   0. \tag {12}
$$
Theorem 5.1 now follows.
E.4 Proofs of Lemmas
To prove the Lemmas E.4 and E.5, we will use the following well-known formula for the moment generating function (MGF) of the half-normal distribution.
Lemma E.13 (MGF of half-normal distribution). The MGF of the half-normal distribution is
$$
\mathbb {E} _ {X \sim \mathcal {N} (0, 1) | X > 0} \left[ e ^ {t | X |} \right] = 2 e ^ {t ^ {2} / 2} \Phi (t),
$$
where  $\Phi(t)$  is the cumulative distribution of a normal random variable.
Proof of Lemma E.4.
$$
\begin{array}{l} \mathbb {E} _ {X} \left[ | X | ^ {c} \exp (t | X |) \exp (t X ^ {2}) \right] = \frac {1}{\sigma \sqrt {2 \pi}} \int_ {- \infty} ^ {\infty} | x | ^ {c} \exp (t | x |) \exp (t x ^ {2}) \exp \left(- \frac {x ^ {2}}{2 \sigma^ {2}}\right) d x \\ = \frac {\sqrt {1 - 2 \sigma^ {2} t}}{\left(\frac {\sigma}{\sqrt {1 - 2 \sigma^ {2} t}}\right) \sqrt {2 \pi}} \int_ {- \infty} ^ {\infty} | x | ^ {c} \exp (t | x |) \exp \left(- \frac {x ^ {2}}{2 \left(\frac {\sigma}{\sqrt {1 - 2 \sigma^ {2} t}}\right) ^ {2}}\right) d x \\ = \sqrt {1 - 2 \sigma^ {2} t} \mathbb {E} _ {Z \sim \mathcal {N} (0, r) | Z \geq 0} [ Z ^ {c} \exp (t Z) ], \\ \end{array}
$$
where  $r = \frac{\sigma}{\sqrt{1 - 2\sigma^2t}}$ . To evaluate this, we use the MGF of the half-normal distribution in Lemma E.13. Thus for some constant  $C$ , for all  $c \in \{1, 2, 3, 4\}$ ,
$$
\begin{array}{l} \mathbb {E} _ {X \sim \mathcal {N} (0, 1) | X > 0} \left[ c! | X | ^ {c} e ^ {t | X |} \right] \leq \mathbb {E} _ {X \sim \mathcal {N} (0, 1) | X > 0} \left[ \frac {d ^ {c}}{d t ^ {c}} e ^ {t | X |} \right] \\ \leq C \left(1 + t ^ {c}\right) e ^ {t ^ {2} / 2}. \\ \end{array}
$$
So for some constant  $C$  (whose value changes throughout this equation), so long as  $\sigma \leq \frac{1}{C}$
$$
\begin{array}{l} \sqrt {1 - 2 \sigma^ {2} t} \mathbb {E} _ {Z \sim \mathcal {N} (0, r) | Z \geq 0} [ Z ^ {c} \exp (t Z) ] = \sqrt {1 - 2 \sigma^ {2} t} \mathbb {E} _ {X \sim \mathcal {N} (0, 1) | Z \geq 0} [ r ^ {c} Z ^ {c} \exp (r t Z) ] \\ \leq \sqrt {1 - 2 \sigma^ {2} t} C r ^ {c} (1 + (t r) ^ {c}) e ^ {(t r) ^ {2} / 2} \\ \leq C \sigma^ {c}. \\ \end{array}
$$
This proves the first statement in the lemma. To prove the second, we first take the expectation over  $X$ , and using the half-Gaussian MGF as before, we obtain
$$
\mathbb {E} _ {X} \mathbb {E} _ {Y} \left[ | X | ^ {c} | Y | ^ {d} \exp (t | X |) \exp (| X Y |) \right] \leq C \mathbb {E} _ {Y} \left[ | Y | ^ {d} \sigma^ {c} (1 + (t + | Y |) ^ {c}) e ^ {(t + | Y |) ^ {2} / 2} \right]
$$
Now applying the first statement to take the expectation over  $Y$ , we obtain
$$
\mathbb {E} _ {Y} \left[ | Y | ^ {d} (1 + (t + | Y |) ^ {c}) e ^ {(t + | Y |) ^ {2} / 2} \right] \leq C \sigma^ {c} \rho^ {d}.
$$
Proof of Lemma E.5. We prove the lemma by induction on  $c$ . Suppose  $c = 0$ . Then by plugging in the MGF for the half-normal distribution from Lemma E.13, for some constant  $C$ , we have
$$
\begin{array}{l} \mathbb {E} _ {X \sim \mathcal {N} (0, 1) | X > 0} [ \left(e ^ {t | X |} - 1\right) ] = 2 e ^ {t ^ {2} / 2} \Phi (t) - 1 (13) \\ \leq 2 e ^ {t ^ {2} / 2} \left(\frac {1 + t}{2}\right) - 1 (14) \\ \leq \left(e ^ {t ^ {2} / 2} - 1\right) + t e ^ {t ^ {2} / 2} (15) \\ \leq C t, (16) \\ \end{array}
$$
$$
\mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} [ (e ^ {t | X |} - 1) ] = \mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2}) | X > 0} [ (e ^ {\sigma t | X |} - 1) ] \leq C t \sigma .
$$
Now for  $c \geq 1$ , by Stein's Lemma, we have (for a new constant  $C$ ),
$$
\begin{array}{l} \mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} [ | X | ^ {c} (e ^ {t | X |} - 1) ] = \mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} [ X | X | ^ {c - 1} \operatorname {s i g n} (X) (e ^ {t | X |} - 1) ] (17) \\ = \sigma^ {2} \mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} \left[ \frac {d}{d X} \left(| X | ^ {c - 1} \operatorname {s i g n} (X) \left(e ^ {t | X |} - 1\right)\right) \right] (18) \\ = \sigma^ {2} \mathbb {E} _ {X \sim \mathcal {N} (0, \sigma^ {2})} \left[ (c - 2) \left(| X | ^ {c - 2} \left(e ^ {t | X |} - 1\right)\right) + \left(| X | ^ {c - 1} \left(t e ^ {t | X |}\right)\right) \right] (19) \\ \leq C t \sigma^ {c + 1}. (20) \\ \end{array}
$$
where in the last step we used the inductive hypothesis and Lemma E.4.
Proof of Lemma E.6. First observe that
$$
\begin{array}{l} \lim  _ {\eta \rightarrow 0} \frac {1}{\eta} \left(\mathbb {E} _ {U, V} \left[ \arccos  \left(\frac {| \mu_ {k} ^ {T} \theta_ {k} ^ {+} |}{\| \theta_ {k} ^ {+} \| _ {2}}\right)\right] - \arccos  \left(\frac {| \mu_ {k} ^ {T} \theta_ {k} |}{\| \theta_ {k} \| _ {2}}\right)\right) \\ = \lim  _ {\eta \rightarrow 0} \frac {1}{\eta} \left(\mathbb {E} _ {U, V} \left[ \arccos  \left(\frac {| \mu_ {k} ^ {T} (\theta_ {k} (1 - \eta \lambda) - \eta \nabla_ {k}) |}{\| \theta_ {k} (1 - \eta \lambda) - \eta \nabla_ {k} \| _ {2}}\right)\right] - \arccos  \left(\frac {| \mu_ {k} ^ {T} \theta_ {k} |}{\| \theta_ {k} \| _ {2}}\right)\right) \\ = \lim  _ {\eta \rightarrow 0} \frac {1}{\eta} \left(\mathbb {E} _ {U, V} \left[ \operatorname {a r c c o s} \left(\frac {\left| \mu_ {k} ^ {T} \left(\theta_ {k} - \frac {\eta}{1 - \eta \lambda} \nabla_ {k}\right)\right|}{\left\| \theta_ {k} - \frac {\eta}{1 - \eta \lambda} \nabla_ {k} \right\| _ {2}}\right)\right] - \operatorname {a r c c o s} \left(\frac {\left| \mu_ {k} ^ {T} \theta_ {k} \right|}{\left\| \theta_ {k} \right\| _ {2}}\right)\right) \\ = \mathbb {E} _ {U, V} \left[ \frac {d}{d \eta} \operatorname {a r c c o s} \left(\frac {\left| \mu_ {k} ^ {T} \left(\theta_ {k} - \eta \nabla_ {k}\right) \right|}{\left\| \theta_ {k} - \eta \nabla_ {k} \right\| _ {2}}\right) (0) \right], \\ \end{array}
$$
since  $\lim_{\eta \to 0}\frac{\eta}{1 - \eta\lambda} = 0$  Now
$$
\begin{array}{l} \frac {d}{d \eta} \operatorname {a r c c o s} \left(\frac {\left| \mu_ {k} ^ {T} \left(\theta_ {k} - \eta \nabla_ {k}\right) \right|}{\left\| \theta_ {k} - \eta \nabla_ {k} \right\| _ {2}}\right) (0) = \operatorname {a r c c o s} ^ {\prime} \left(\frac {\left| \mu_ {k} ^ {T} \theta_ {k} \right|}{\left\| \theta_ {k} \right\| _ {2}}\right) \frac {d}{d \eta} \left(\frac {\left| \mu_ {k} ^ {T} \left(\theta_ {k} - \eta \nabla_ {k}\right) \right|}{\left\| \theta_ {k} - \eta \nabla_ {k} \right\| _ {2}}\right) (0) \\ = \operatorname {a r c c o s} ^ {\prime} \left(\frac {| \mu_ {k} ^ {T} \theta_ {k} |}{\| \theta_ {k} \| _ {2}}\right) \left(\frac {- \operatorname {s i g n} (\mu_ {k} ^ {T} \theta_ {k}) \mu_ {k} ^ {T} \nabla_ {k} \| \theta_ {k} \| + | \mu_ {k} ^ {T} \theta_ {k} | \frac {\theta_ {k} ^ {T} \nabla_ {k}}{\| \theta_ {k} \|}}{\| \theta_ {k} \| _ {2} ^ {2}}\right) \\ = N \left(- \mu_ {k} ^ {T} \theta_ {k} \mu_ {k} ^ {T} \nabla_ {k} + (\mu_ {k} ^ {T} \theta_ {k}) ^ {2} \frac {\theta_ {k} ^ {T} \nabla_ {k}}{\| \theta_ {k} \| ^ {2}}\right), \\ \end{array}
$$
where  $N = \arccos' \left( \frac{|\mu_k^T \theta_k|}{\|\theta_k\|_2} \right) \frac{1}{\|\theta_k\| |\mu_k^T \theta_k|}$ . The lemma follows by taking the expectation over  $U, V$ , and observing derivative of  $\arccos(x)$  is negative whenever  $x$  is positive.
Proof of Lemma E.7. First observe that by symmetry, we have
$$
\frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} [ \nabla_ {k} ] = m \frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ \frac {\partial \mathcal {L} _ {i}}{\partial \theta_ {k}} \right].
$$
To make this expectation easier to analyze, we express the random variable  $(U(x), V(x)) \sim \mathbb{P}_x$  as an interpolation of Gaussians in the coordinate  $\mu_{k'}^T v_{k'}^{(i)}$ . Let  $\xi \sim \mathcal{N}(0,1)$ , and define  $(U,V) \sim \mathbb{P}_1$ , such that  $\mu_{k'}^T v_{k'}^{(i)} = \mu_{k'}^T u_{k'}^{(i)}$ . For  $x \in [0,1)$ , define  $(U(x), V(x))$  to have
$$
\mu_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(i)} (x) = x \mu_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} ^ {(i)} + \sqrt {1 - x ^ {2}} \xi , \tag {21}
$$
and otherwise be the same as  $(U,V)$ . It is easy to check that  $(U(x),V(x))\sim \mathbb{P}_x$