SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ \frac {\partial \mathcal {L} _ {i} (\Theta ; U , V)}{\partial \theta_ {k}} \right] = \mathbb {E} _ {U, V \sim \mathbb {P} _ {1}, \xi} \left[ \frac {d}{d x} \frac {\partial \mathcal {L} _ {i} (\Theta ; U (x) , V (x))}{\partial \theta_ {k}} \right].
$$
Taking the derivative of the cross-entropy loss, we have
$$
\begin{array}{l} \frac {d}{d x} \frac {\partial \mathcal {L} _ {i} (\Theta ; U (x) , V (x))}{\partial \theta_ {k}} = \frac {d}{d x} \left(\sum_ {j \neq i} p _ {i j} \left(\frac {\partial \left(z _ {i j} - z _ {i i}\right)}{\partial \theta_ {k}}\right)\right) \\ = \sum_ {j \neq i} \frac {d p _ {i j}}{d \mu_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(i)} (x)} \frac {d \mu_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(i)} (x)}{d x} \frac {\partial (z _ {i j} - z _ {i i})}{\partial \theta_ {k}} \\ = \sum_ {j \neq i} - p _ {i j} p _ {i i} \frac {d z _ {i i}}{d \mu_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(i)} (x)} \left(\mu_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} ^ {(i)} - \frac {x}{\sqrt {1 - x ^ {2}}} \xi\right) \left(\frac {\partial (z _ {i j} - z _ {i i})}{\partial \theta_ {k}}\right) \\ \end{array}
$$
where the variables  $z_{ij}$  and  $p_{ij}$  are the similarity scores and the softmaxes from the data  $(U(x), V(x))$ . Here the first line is by Lemma E.2, and the second line holds by chain rule since  $\frac{\partial z_{ij}}{\partial \theta_k} - \frac{\partial z_{ii}}{\partial \theta_k}$  does not depend on  $v_{k'}^{(i)}$ . The third line uses the proof of Claim E.14 to take the derivative of  $p_{ij}$ , and Equation 21 to take the derivative of  $\mu_{k'}^T v_{k'}^{(i)}(x)$ .
Now we reparameterize  $\mu_{k'}^T u_{k'}^{(i)} - \frac{x}{\sqrt{1 - x^2}}\xi$  as follows:
$$
\mu_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} ^ {(i)} - \frac {x}{\sqrt {1 - x ^ {2}}} \xi = \left(\frac {1}{1 - x ^ {2}}\right) \mu_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} ^ {(i)} - \frac {x}{1 - x ^ {2}} \mu_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(i)} (x).
$$
Plugging in this reparameterization and  $\frac{dz_{ii}}{d\mu_{k'}^T v_{k'}^{(i)}(x)} = \theta_{k'}^T \mu_{k'} \theta_{k'}^T u_{k'}$ , we obtain
$$
\frac {d}{d x} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ \frac {\partial \mathcal {L} _ {i} (\Theta ; U , V)}{\partial \theta_ {k}} \right] = \frac {- 1}{1 - x ^ {2}} \sum_ {j \neq i} \mathbb {E} _ {U, V \sim \mathbb {P} _ {x}} \left[ p _ {i j} p _ {i i} \left(\theta_ {k ^ {\prime}} ^ {T} \mu_ {k ^ {\prime}} \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].
$$
We now prove Lemmas E.8, E.9, E.10, and E.11.
Notation. Since  $i$  is fixed throughout, we drop the  $(i)$  superscripts and let  $u_{k} = u_{k}^{(i)}$  and  $v_{k} = v_{k}^{(i)}$ . We will introduce the following random variables, which are all independent, to simplify the exposition:
-  $\xi_{j} \coloneqq \theta_{k}^{T} v_{k}^{(j)}$  for  $j \neq i$ . Thus  $\xi_{j} \sim \mathcal{N}(0, \| \theta_{k} \|^{2})$ .
-  $\xi_j' \coloneqq \theta_{k'}^T v_{k'}^{(j)}$  for  $j \neq i$ . Thus  $\xi_j' \sim \mathcal{N}(0, \| \theta_{k'} \|^2)$ .
-  $\xi_{i}:= (\theta_{k}^{\perp})^{T}v_{k} + (\theta_{k}^{\parallel})^{T}(v_{k} - \alpha_{k}u_{k})$ . Thus  $\xi_{i}\sim \mathcal{N}(0,\| \theta_{k}^{\perp}\|^{2} + (1 - \alpha_{k}^{2})\| \theta_{k}^{\parallel}\|^{2})$
-  $\xi_{i}^{\prime} := (\theta_{k^{\prime}}^{\perp})^{T}v_{k^{\prime}}$ . Thus  $\xi_{i}^{\prime} \sim \mathcal{N}(0, \| \theta_{k^{\prime}}^{\perp} \|^{2} \| \theta_{k^{\prime}}^{\parallel} \|^{2})$ .
$\zeta_i^{\prime} \coloneqq (\theta_{k^{\prime}}^{\parallel})^T (v_{k^{\prime}} - \alpha_{k^{\prime}} u_{k^{\prime}})$ . Thus  $\zeta_i^{\prime} \sim \mathcal{N}(0, (1 - \alpha_{k^{\prime}}^2) \| \theta_{k^{\prime}}^{\parallel} \|^2)$ .
-  $y = (\theta_k^{\parallel})^T u_k$ . Thus  $y \sim \mathcal{N}(0, \| \theta_k^{\parallel} \|^2)$ .
-  $y' = (\theta_{k'}^{\parallel})^T u_{k'}$ . Thus  $y' \sim \mathcal{N}(0, \| \theta_{k'}^{\parallel} \|^2)$ .
-  $\eta_{i} \coloneqq (\theta_{k}^{\perp})^{T} u_{k}$ . Thus  $\eta_{i} \sim \mathcal{N}(0, \| \theta_{k}^{\perp} \|^{2})$ .
-  $\eta_i' := (\theta_{k'}^\perp)^T u_{k'}$ . Thus  $\eta_i' \sim \mathcal{N}(0, \| \theta_{k'}^\perp \|^2)$ .
For any such random variable  $X$ , we use  $\sigma_X^2$  to denote its variance. Observe that
$$
\frac {p _ {i i} p _ {i j}}{q _ {i i} q _ {i j}} = \frac {\exp \left(\theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k}\right) \exp \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}}\right)}{\mathbb {E} _ {j ^ {\prime} \sim q} \exp \left(\theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})}\right) \exp \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j ^ {\prime})}\right)} \frac {\exp \left(\theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j)}\right) \exp \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j)}\right)}{\mathbb {E} _ {j ^ {\prime} \sim q} \exp \left(\theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})}\right) \exp \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {\mathrm {T}} v _ {k ^ {\prime}} ^ {(j ^ {\prime})}\right)}.
$$
We will use the following two claims in the proofs of all four lemmas.
Claim E.14. For  $\beta \in \{\xi_j,\xi_j',\xi_i,\xi_i',\zeta_i',\eta_i,\eta_i',x,x'\}$ , let  $\bar{\beta}_{j'} \coloneqq \frac{\partial}{\partial\beta}\left(\theta_k^T u_k\theta_k^T v_k^{(j')} + \theta_{k'}^T u_{k'}\theta_{k'}^T v_{k'}^{(j')}\right)$ .
$$
\left| \frac {\partial p _ {i i} p _ {i j}}{\partial \beta} \right| \leq p _ {i i} p _ {i j} \left(| \bar {\beta} _ {j} | + | \bar {\beta} _ {i} | + 2 \mathbb {E} _ {j ^ {\prime} \sim q} | \bar {\beta} _ {j ^ {\prime}} |\right).
$$
If additionally  $\gamma \in \{\xi_j,\xi_j',\xi_i,\xi_i',\zeta_i',\eta_i,\eta_i'\}$  and  $\gamma \perp \{\bar{\beta}_{j'}\}_{j'\in [m]}$ , then
$$
\begin{array}{l} \left| \frac {\partial}{\partial \gamma} \frac {\partial p _ {i i} p _ {i j}}{\partial \beta} \right| \leq p _ {i i} p _ {i j} \left(\left(| \bar {\beta} _ {j} | + | \bar {\beta} _ {i} | + 2 \mathbb {E} _ {j ^ {\prime} \sim q} | \bar {\beta} _ {j ^ {\prime}} |\right) \left(| \bar {\gamma} _ {j} | + | \bar {\gamma} _ {i} | + 2 \mathbb {E} _ {j ^ {\prime} \sim q} | \bar {\gamma} _ {j ^ {\prime}} |\right)\right) \\ + p _ {i i} p _ {i j} \left(2 \mathbb {E} _ {j ^ {\prime} \sim q} | \bar {\beta} _ {j ^ {\prime}} \bar {\gamma} _ {j ^ {\prime}} | + 2 (\mathbb {E} _ {j ^ {\prime} \sim q} | \bar {\beta} _ {j ^ {\prime}} |) (\mathbb {E} _ {j ^ {\prime} \sim q} | \bar {\gamma} _ {j ^ {\prime}} |\right). \\ \end{array}
$$
Proof. By a straightforward quotient-rule computation of the derivative of  $\frac{p_{ij}}{q_{ij}}$ , recalling that  $q_{ij}$  is independent of  $S$ , we obtain
$$
\frac {\partial p _ {i j}}{\partial \beta} = p _ {i j} \left(\bar {\beta} _ {j} - \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} p _ {i j ^ {\prime}}\right).
$$
By applying product to the expression above, we obtain
$$
\frac {\partial p _ {i i} p _ {i j}}{\partial \beta} = p _ {i i} p _ {i j} \left(\bar {\beta} _ {j} + \bar {\beta} _ {i} - 2 \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} p _ {i j ^ {\prime}}\right).
$$
Taking absolute values and using the fact that  $p_{ij'} \leq 1$ , we obtain the first result.
Next we take the derivative of  $p_{ij}$  with respect to both  $\beta$  and  $\gamma$ . Using the expression above for  $\frac{\partial p_{ij}}{\partial\beta}$ , we obtain
$$
\frac {\partial}{\partial \gamma} \frac {\partial p _ {i j}}{\partial \beta} = p _ {i j} \left(\left(\bar {\beta} _ {j} - \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} p _ {i j ^ {\prime}}\right) (\bar {\gamma} _ {j} - \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\gamma} _ {j ^ {\prime}} p _ {i j ^ {\prime}}) - \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} \bar {\gamma} _ {j ^ {\prime}} p _ {i j ^ {\prime}} + (\mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} p _ {i j ^ {\prime}}) (\mathbb {E} _ {j ^ {\prime} \sim q} \bar {\gamma} _ {j ^ {\prime}} p _ {i j ^ {\prime}})\right),
$$
$$
\begin{array}{l} \frac {\partial}{\partial \gamma} \frac {\partial p _ {i i} p _ {i j}}{\partial \beta} = p _ {i i} p _ {i j} \left(\left(\bar {\beta} _ {j} + \bar {\beta} _ {i} - 2 \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} p _ {i j ^ {\prime}}\right) \left(\bar {\gamma} _ {j} + \bar {\gamma} _ {i} - 2 \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\gamma} _ {j ^ {\prime}} p _ {i j ^ {\prime}}\right)\right) \\ + p _ {i i} p _ {i j} \left(- 2 \mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} \bar {\gamma} _ {j ^ {\prime}} p _ {i j ^ {\prime}} + 2 \big (\mathbb {E} _ {j ^ {\prime} \sim q} \bar {\beta} _ {j ^ {\prime}} p _ {i j ^ {\prime}} \big) \big (\mathbb {E} _ {j ^ {\prime} \sim q} \bar {\gamma} _ {j ^ {\prime}} p _ {i j ^ {\prime}} \big)\right). \\ \end{array}
$$
The second result follows by taking absolute values and the fact that  $p_{ij'} \leq 1$ .
Claim E.15.
$$
\frac {p _ {i j}}{q _ {i j}} \leq \exp \left(| \theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j)} |\right) \exp \left(| \theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j)} |\right) \mathbb {E} _ {j ^ {\prime} \sim q} \left[ \exp \left(| \theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})} |\right) \exp \left(| \theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j ^ {\prime})} |\right) \right].
$$
Proof. This follows directly from using Jenson's inequality on the distribution  $j' \sim q$  to show that
$$
\begin{array}{l} \frac {1}{\mathbb {E} _ {j ^ {\prime} \sim q} \left[ \exp \left(\theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})}\right) \exp \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j ^ {\prime})}\right) \right]} \leq \mathbb {E} _ {j ^ {\prime} \sim q} \left[ \exp \left(- \theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})}\right) \exp \left(- \theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j ^ {\prime})}\right) \right] \\ \leq \mathbb {E} _ {j ^ {\prime} \sim q} \left[ \exp \left(| \theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})} |\right) \exp \left(| \theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j ^ {\prime})} |\right) \right]. \\ \end{array}
$$
Claim E.16.
$$
\left| 1 - \frac {p _ {i j}}{q _ {i j}} \right| \leq Z _ {j} - 1,
$$
where  $Z_{j}\coloneqq \exp \left(|\theta_{k}^{T}u_{k}\theta_{k}^{T}v_{k}^{(j)}|\right)\exp \left(|\theta_{k^{\prime}}^{T}u_{k^{\prime}}\theta_{k^{\prime}}^{T}v_{k^{\prime}}^{(j)}|\right)\mathbb{E}_{j^{\prime}\sim q}\left[\exp \left(|\theta_{k}^{T}u_{k}\theta_{k}^{T}v_{k}^{(j^{\prime})}|\right)\exp \left(|\theta_{k^{\prime}}^{T}u_{k^{\prime}}\theta_{k^{\prime}}^{T}v_{k^{\prime}}^{(j^{\prime})}|\right)\right]$
Proof. Note that for any  $x \geq 0$ , we have  $|1 - x| \leq \max \left(x - 1, \frac{1}{x} - 1\right)$ . By Claim E.15,  $\frac{p_{ij}}{q_{ij}} - 1$  is at most the desired value given in this claim.
$$
\begin{array}{l} \frac {q _ {i j}}{p _ {i j}} = \frac {\mathbb {E} _ {j ^ {\prime} \sim q} \left[ \exp \left(\theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})}\right) \exp \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j ^ {\prime})}\right) \right]}{\exp \left(\theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j)}\right) \exp \left(\theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j)}\right)} \\ \leq \exp \left(\left| \theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j)} \right|\right) \exp \left(\left| \theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j)} \right|\right) \mathbb {E} _ {j ^ {\prime} \sim q} \left[ \exp \left(\left| \theta_ {k} ^ {T} u _ {k} \theta_ {k} ^ {T} v _ {k} ^ {(j ^ {\prime})} \right|\right) \exp \left(\left| \theta_ {k ^ {\prime}} ^ {T} u _ {k ^ {\prime}} \theta_ {k ^ {\prime}} ^ {T} v _ {k ^ {\prime}} ^ {(j ^ {\prime})} \right|\right) \right]. \\ \end{array}
$$
This yields the claim.
Proof of Lemma E.8. Expanding  $h(S) - h_1(S)$ , we see that we need to control the following terms:
1. (a)  $\left| \mathbb{E}_S \left[ p_{ii} p_{ij} \left( \eta_i' \left( \mu_{k'}^T u_{k'} - x \mu_{k'}^T v_{k'} \right) \right) \left( \mu_k^T u_k \xi_j \right) \right] \right|,$  (b)  $\left| \mathbb{E}_S \left[ p_{ii} p_{ij} \left( y' \left( \mu_{k'}^T u_{k'} - x \mu_{k'}^T v_{k'} \right) \right) \left( \mu_k^T u_k \xi_j \right) \right] \right|$
2. (a)  $\left| \mathbb{E}_S \left[ p_{ii} p_{ij} \left( \eta_i' \left( \mu_{k'}^T u_{k'} - x \mu_{k'}^T v_{k'} \right) \right) \left( \mu_k^T u_k \xi_i \right) \right] \right|,$  (b)  $\left| \mathbb{E}_S \left[ p_{ii} p_{ij} \left( \left( y' \left( \mu_{k'}^T u_{k'} - x \mu_{k'}^T v_{k'} \right) \right) \right) \left( \mu_k^T u_k \xi_i \right) \right] \right|$
3. (a)  $\left| \alpha_{k} \mathbb{E}_{S} \left[ p_{ii} p_{ij} \left( \eta_{i}' \left( \mu_{k'}^{T} u_{k'} - x \mu_{k'}^{T} v_{k'} \right) \right) \left( \mu_{k}^{T} u_{k} y \right) \right] \right|, (\mathsf{b}) \left| \alpha_{k} \mathbb{E}_{S} \left[ p_{ii} p_{ij} \left( y'(-x \xi_{i}') \right) \left( \mu_{k}^{T} u_{k} y \right) \right] \right|$
4. (a)  $\left|\mathbb{E}_S\left[p_{ii}p_{ij}\left(\eta_i'\left(\mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'}\right)\right)\left(\xi_i(v_k - v_k^{(j)})^T\mu_k\right)\right]\right|$
(b)  $\left|\mathbb{E}_S\left[p_{ii}p_{ij}\left(y'\left(\mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'}\right)\right)\left(\xi_i(v_k - v_k^{(j)})^T\mu_k\right)\right]\right|$
5. (a)  $\left|\alpha_{k}\mathbb{E}_{S}\left[p_{ii}p_{ij}\left(\eta_{i}^{\prime}\left(\mu_{k^{\prime}}^{T}u_{k^{\prime}} - x\mu_{k^{\prime}}^{T}v_{k^{\prime}}\right)\right)\left(y(v_{k} - v_{k}^{(j)})^{T}\mu_{k}\right)\right]\right|$
(b)  $\left|\mathbb{E}_S\left[p_{ii}p_{ij}\left(y'\left(\mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'}\right)\right)\left(y(v_k - \alpha_ku_k - v_k^{(j)})^T\mu_k\right)\right]\right|$
We begin by bounding the terms where the expression after  $p_{ii}p_{ij}$  has two independent mean-0 terms, mainly (1a), (2a), (4a). The first step is to apply Stein's Lemma (Lemma E.3) twice to these two terms, which we will call  $\beta$  and  $\gamma$ . Let  $\beta \gamma g(S \setminus \{\beta, \gamma\})$  be the terms after  $p_{ii}p_{ij}$ . Then we have
$$
\left. \right.\left| \mathbb {E} _ {S} \left[ p _ {i i} p _ {i j} \beta \gamma g (S \setminus \{\beta , \gamma \}) \right]\right| \leq \sigma_ {\beta} ^ {2} \sigma_ {\gamma} ^ {2} \left| \mathbb {E} _ {S} \left[\left| \frac {\partial}{\partial \gamma} \frac {\partial p _ {i i} p _ {i j}}{\partial \beta} \right| | g (S \setminus \{\beta , \gamma \}) | \right]\right|.
$$
Next we apply the final result in Claim E.14 to bound the absolute value of  $\left|\frac{\partial}{\partial\gamma}\frac{\partial p_{ii}p_{ij}}{\partial\beta}\right|$ . Once we do this, we achieve
$$
| \mathbb {E} _ {S} [ p _ {i i} p _ {i j} \beta \gamma g (S \setminus \{\beta , \gamma \}) ] | \leq \sigma_ {\beta} ^ {2} \sigma_ {\gamma} ^ {2} q _ {i i} q _ {i j} \mathbb {E} _ {S} \left[ Z | g (S \setminus \{\beta , \gamma \}) | \sum_ {j ^ {\prime}, \ell \in [ m ]} c _ {j ^ {\prime}, \ell} | \bar {\beta} _ {j ^ {\prime}} | | \bar {\gamma} _ {\ell} | \right],
$$
where  $\sum_{j',\ell \in [m]}c_{j',\ell}\leq C$  for some constant  $C$ , and  $Z\coloneqq \frac{p_{ii}p_{ij}}{q_{ii}q_{ij}}$ . Finally, we use the bound on  $Z$  from Claim E.15, and then Lemma E.4 to take the expectation over  $S$ , iteratively applying Lemma E.4 to each variable in  $S$ . Thus we have, for some (different) constant  $C$ ,
1.  $\left| {\mathbb{E}}_{S}\left\lbrack  {{p}_{ii}{p}_{ij}\left( {{\eta }_{i}^{\prime }\left( {{\mu }_{{k}^{\prime }}^{T}{u}_{{k}^{\prime }} - x{\mu }_{{k}^{\prime }}^{T}{v}_{{k}^{\prime }}}\right) }\right) \left( {{\mu }_{k}^{T}{u}_{k}{\xi }_{j}}\right) }\right\rbrack  \right| \; \leq  \;C{q}_{ii}{q}_{ij}{\sigma }_{{\eta }_{i}^{\prime }}^{2}{\sigma }_{{\xi }_{j}}^{2}\begin{Vmatrix}{\theta }_{{k}^{\prime }}\end{Vmatrix}\begin{Vmatrix}{\theta }_{k}\end{Vmatrix}$
$$
C q _ {i i} q _ {i j} \| \theta_ {k ^ {\prime}} ^ {\perp} \| ^ {2} \| \theta_ {k ^ {\prime}} \| \| \theta_ {k} \| ^ {3} \leq C q _ {i i} q _ {i j} \| \theta_ {k ^ {\prime}} \| ^ {3} \| \theta_ {k} \| ^ {3}.
$$
2.  $\left| \mathbb{E}_S \left[ p_{ii} p_{ij} \left( \eta_i' \left( \mu_{k'}^T u_{k'} - x \mu_{k'}^T v_{k'} \right) \right) \left( \mu_k^T u_k \xi_i \right) \right] \right| \leq C q_{ii} q_{ij} \sigma_{\eta_i'}^2 \sigma_{\xi_i}^2 \| \theta_{k'} \| \| \theta_k \| \leq$
3.  $\left|\mathbb{E}_S\left[p_{ii}p_{ij}\left(\eta_i'\left(\mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'}\right)\right)\left(\xi_i(v_k - v_k^{(j)})^T\mu_k\right)\right]\right|\leq Cq_{ii}q_{ij}\sigma_{\eta_i'}^2\sigma_{\xi_i}^2\| \theta_{k'}\| \| \theta_k\| \leq$ $Cq_{ii}q_{ij}\| \theta_{k'}^{\perp}\| ^2\| \theta_{k'}\| \| \theta_k\| ^3\leq Cq_{ii}q_{ij}\| \theta_{k'}\| ^3\| \theta_k\| ^3.$
Now we consider the remaining 7 terms. Here we decompose the expression inside the expectation as  $p_{ii}p_{ij}\beta g(S\setminus \beta)$ , where  $\beta \in S$ . We proceed as before, but we only apply Stein's Lemma once, to  $\beta$ . Applying Steins, the expression for  $\frac{\partial p_{ii}p_{ij}}{\partial\beta}$  given in the first result of Claim E.14, we obtain
$$
\left. \right.\left| \mathbb {E} _ {S} \left[ p _ {i i} p _ {i j} \beta g (S \setminus \beta) \right]\right| \leq \sigma_ {\beta} ^ {2} \left| \mathbb {E} _ {S} \left[\left| \frac {\partial p _ {i i} p _ {i j}}{\partial \beta} \right| | g (S \setminus \beta) | \right]\right| \leq \sigma_ {\beta} ^ {2} q _ {i i} q _ {i j} \mathbb {E} _ {S} \left[ Z | g (S \setminus \beta) | \sum_ {j ^ {\prime} \in [ m ]} c _ {j ^ {\prime}} | \bar {\beta_ {j ^ {\prime}}} | \right], \tag {22}
$$
where  $\sum_{j' \in [m]} c_{j'} \leq C$  for some constant  $C$ , and  $Z := \frac{p_{ii} p_{ij}}{q_{ii} q_{ij}}$ . Finally, we plug in a bound for  $Z$  in Claim E.15, an use Lemma E.4 to take the expectation over  $S$ , again iteratively over each variable.
Thus we have, for some (different) constant  $C$
1.  $\left| {\mathbb{E}}_{S}\left\lbrack  {{p}_{ii}{p}_{ij}\left( {{y}^{\prime }\left( {{\mu }_{{k}^{\prime }}^{T}{u}_{{k}^{\prime }} - x{\mu }_{{k}^{\prime }}^{T}{v}_{{k}^{\prime }}}\right) }\right) \left( {{\mu }_{k}^{T}{u}_{k}{\xi }_{j}}\right) }\right\rbrack  \right|$
2.  $\left| {\mathbb{E}}_{S}\left\lbrack  {{p}_{ii}{p}_{ij}\left( {{y}^{\prime }\left( {{\mu }_{{k}^{\prime }}^{T}{u}_{{k}^{\prime }} - x{\mu }_{{k}^{\prime }}^{T}{v}_{{k}^{\prime }}}\right) }\right) \left( {{\mu }_{k}^{T}{u}_{k}{\xi }_{i}}\right) }\right\rbrack  \right|$
3.  $\left| \alpha_{k} \mathbb{E}_{S} \left[ p_{ii} p_{ij} \left( \eta_{i}' \left( \mu_{k'}^{T} u_{k'} - x \mu_{k'}^{T} v_{k'} \right) \right) \left( \mu_{k}^{T} u_{k} y \right) \right] \right| \leq C \alpha_{k} q_{ii} q_{ij} \sigma_{\eta_{i}'}^{2} \| \theta_{k'} \| \| \theta_{k}^{\parallel} \| = C \alpha_{k} q_{ii} q_{ij} \| \theta_{k'} \|^{3} \| \theta_{k}^{\parallel} \|$ .
4.  $\left|\alpha_{k}\mathbb{E}_{S}\left[p_{ii}p_{ij}\left(y^{\prime}(-x\zeta_{i}^{\prime})\right)\left(\mu_{k}^{T}u_{k}y\right)\right]\right|$ $\leq$ $C\alpha_{k}q_{ii}q_{ij}\sigma_{\zeta_{i}^{\prime}}^{2}\| \theta_{k^{\prime}}\| \| \theta_{k}^{\parallel}\|$ $= C\alpha_{k}q_{ii}q_{ij}\| \theta_{k^{\prime}}^{\parallel}\|^{2}\| \theta_{k^{\prime}}\| \| \theta_{k}^{\parallel}\| .$
5.  $\left|\mathbb{E}_S\left[p_{ii}p_{ij}\left(y'\left(\mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'}\right)\right)\left(\xi_i(v_k - v_k^{(j)})^T\mu_k\right)\right]\right|\leq Cq_{ii}q_{ij}\sigma_{\xi_i}^2\| \theta_k\| \| \theta_{k'}^{\parallel}\| \leq$ $C q_{i i}q_{i j}\| \theta_{k}\|^{3}\| \theta_{k^{\prime}}^{\parallel}\| .$
6.  $\left|\alpha_{k}\mathbb{E}_{S}\left[p_{ii}p_{ij}\left(\eta_{i}^{\prime}\left(\mu_{k^{\prime}}^{T}u_{k^{\prime}} - x\mu_{k^{\prime}}^{T}v_{k^{\prime}}\right)\right)\left(x(v_{k} - v_{k}^{(j)})^{T}\mu_{k}\right)\right]\right|$ $C\alpha_{k}q_{ii}q_{ij}\sigma_{\eta_{i}^{\prime}}^{2}\| \theta_{k^{\prime}}\| \| \theta_{k}^{\parallel}\| = C\alpha_{k}q_{ii}q_{ij}\| \theta_{k^{\prime}}^{\perp}\|^{2}\| \theta_{k^{\prime}}\| \| \theta_{k}^{\parallel}\| \leq C\alpha_{k}q_{ii}q_{ij}\| \theta_{k^{\prime}}\|^{3}\| \theta_{k}^{\parallel}\| .$
7.  $\left|\mathbb{E}_S\left[p_{ii}p_{ij}\left(y'\left(\mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'}\right)\right)\left(x(v_k - \alpha_ku_k - v_k^{(j)})^T\mu_k\right)\right]\right|$ $Cq_{ii}q_{ij}\sigma_x^2\| \theta_k\| \| \theta_{k'}^{\parallel}\| = Cq_{ii}q_{ij}\| \theta_k^{\parallel}\| ^2\| \theta_k\| \| \theta_{k'}^{\parallel}\| .$
Combining the bounds on these 10 terms proves the lemma:
$$
\begin{array}{r l} & {\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).} \\ & {\square} \end{array}
$$
Proof of Lemma E.10. The proof of Lemma E.10 is nearly identical, besides some differences in the terms we need to bound. We list them below:
1. (a)  $\left| \mathbb{E}_S \left[ p_{ii} p_{ij} \left( \eta_i' \left( \mu_{k'}^T u_{k'} - x \mu_{k'}^T v_{k'} \right) \right) \left( \theta_k^T u_k \xi_j \right) \right] \right|$  (b)  $\left| \mathbb{E}_S \left[ p_{ii} p_{ij} \left( y' \left( \mu_{k'}^T u_{k'} - x \mu_{k'}^T v_{k'} \right) \right) \left( \theta_k^T u_k \xi_j \right) \right] \right|$
2. (a)  $\left| \mathbb{E}_S \left[ p_{ii}p_{ij} \left( \eta_i' \left( \mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'} \right) \right) \left( \theta_k^T u_k \xi_i \right) \right] \right|$  (b)  $\left| \mathbb{E}_S \left[ p_{ii}p_{ij} \left( \left( y' \left( \mu_{k'}^T u_{k'} - x\mu_{k'}^T v_{k'} \right) \right) \left( \theta_k^T u_k \xi_i \right) \right] \right|$
3. (a)  $\left| \alpha_{k} \mathbb{E}_{S} \left[ p_{ii} p_{ij} \left( \eta_{i}' \left( \mu_{k'}^{T} u_{k'} - x \mu_{k'}^{T} v_{k'} \right) \right) \left( \theta_{k}^{T} u_{k} y \right) \right] \right|$  (b)  $\left| \alpha_{k} \mathbb{E}_{S} \left[ p_{ii} p_{ij} \left( y' \left( \mu_{k'}^{T} u_{k'} - x \mu_{k'}^{T} v_{k'} \right) \right) (\eta_{i} y) \right] \right|$
4.  $\left|\alpha_{k}\mathbb{E}_{S}\left[p_{ii}p_{ij}\left(y^{\prime}\left(-x\zeta_{i}^{\prime}\right)\right)\left(\theta_{k}^{T}u_{k}y\right)\right]\right|$
We use the same approach as before. For the terms (1a) and (2a) we apply Stein's Lemma to  $(\eta_i', \xi_j)$  and  $(\eta_i', \xi_i)$  respectively. For (1b), (2b), (3a) and (3b) and (4), we apply Stein's Lemma to  $\xi_j, \xi_i, \eta_i'$ ,  $\eta_i$ , and  $\xi_i'$  respectively. Using Claim E.15 and then Lemma E.4 as before, we obtain the following result: