SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Since  $\operatorname{tr}(\pmb{\Sigma}(\pmb{\theta}_t)) \leq \alpha (\mathcal{L}(\pmb{\theta}_t) - \mathcal{L}^*)$ , we have
$$
\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \mid \boldsymbol {\theta} _ {t} \right] \leq \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) - \eta \mu \left(\mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) - \mathcal {L} ^ {*}\right) + \frac {\eta^ {2} \ell \alpha}{2 B} \left(\mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) - \mathcal {L} ^ {*}\right).
$$
Altogether,
$$
\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \right] - \mathcal {L} ^ {*} \leq \left(1 - \eta \mu + \frac {\eta^ {2} \ell \alpha}{2 B}\right) \left(\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right] - \mathcal {L} ^ {*}\right)
$$
Choosing  $\eta = \min \left(\frac{1}{\ell}, \frac{\mu B}{\ell \alpha}\right)$ , we obtain
$$
\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \right] - \mathcal {L} ^ {*} \leq \left(1 - \min  \left(\frac {\mu}{2 \ell}, \frac {\mu^ {2} B}{2 \ell \alpha}\right)\right) \left(\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right] - \mathcal {L} ^ {*}\right).
$$
Therefore we reach a solution with  $\mathbb{E}[\mathcal{L}(\pmb{\theta}_t)] - \mathcal{L}^* \leq \epsilon$  after
$$
t = \max \left(\frac {2 \ell}{\mu}, \frac {2 \ell \alpha}{\mu^ {2} B}\right) \log \left(\frac {\mathcal {L} (\pmb {\theta} _ {0}) - \mathcal {L} ^ {*}}{\epsilon}\right) = O \left(\left(\frac {\ell}{\mu} + \frac {\ell \alpha}{\mu^ {2} B}\right) \log \frac {\mathcal {L} (\pmb {\theta} _ {0}) - \mathcal {L} ^ {*}}{\epsilon}\right)
$$
Proof of Lemma 3. By Corollary 1, ZO-SGD with  $\eta_{\mathrm{ZO}} = \gamma^{-1}\eta_{\mathrm{SGD}}$  yields
$$
\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \mid \boldsymbol {\theta} _ {t} \right] - \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \leq \frac {1}{\gamma} \cdot \left[ - \eta_ {\mathrm {S G D}} \left\| \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right\| ^ {2} + \frac {1}{2} \eta_ {\mathrm {S G D}} ^ {2} \ell \cdot \mathbb {E} \left[ \| \nabla \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \| ^ {2} \right] \right].
$$
As in the proof for SGD, choosing  $\eta_{\mathrm{SGD}} \leq \frac{1}{\ell}$  yields
$$
\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \mid \boldsymbol {\theta} _ {t} \right] - \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \leq \gamma^ {- 1} \cdot \left[ - \frac {\eta_ {\mathrm {S G D}}}{2} \| \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \| ^ {2} + \frac {\eta_ {\mathrm {S G D}} ^ {2} \ell}{2 B} \operatorname {t r} \left(\boldsymbol {\Sigma} \left(\boldsymbol {\theta} _ {t}\right)\right) \right].
$$
Therefore under  $\mu$ -PL and the  $\mathrm{tr}(\Sigma(\pmb{\theta}_t)) \leq \alpha (\mathcal{L}(\pmb{\theta}_t) - \mathcal{L}^*)$  assumption we obtain
$$
\begin{array}{l} \mathbb {E} [ \mathcal {L} (\boldsymbol {\theta} _ {t + 1}) ] - \mathbb {E} [ \mathcal {L} (\boldsymbol {\theta} _ {t}) ] \leq \gamma^ {- 1} \cdot \left[ - \eta_ {\mathrm {S G D}} \mu + \frac {\eta_ {\mathrm {S G D}} ^ {2} \ell \alpha}{2 B} \right] \cdot (\mathbb {E} [ \mathcal {L} (\boldsymbol {\theta} _ {t}) ] - \mathcal {L} ^ {*}) \\ \Longrightarrow \mathbb {E} [ \mathcal {L} (\boldsymbol {\theta} _ {t + 1}) ] - \mathcal {L} ^ {*} \leq \left(1 - \gamma^ {- 1} \left(\eta_ {\mathrm {S G D}} \mu - \frac {\eta_ {\mathrm {S G D}} ^ {2} \ell \alpha}{2 B}\right)\right) (\mathbb {E} [ \mathcal {L} (\boldsymbol {\theta} _ {t}) ] - \mathcal {L} ^ {*}). \\ \end{array}
$$
Choosing  $\eta_{\mathrm{SGD}} = \min \left(\frac{1}{\ell},\frac{\mu B}{\ell\alpha}\right)$  yields
$$
\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \right] - \mathcal {L} ^ {*} \leq \left(1 - \gamma^ {- 1} \cdot \min  \left(\frac {\mu}{2 \ell}, \frac {\mu^ {2} B}{2 \ell \alpha}\right)\right) \left(\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right] - \mathcal {L} ^ {*}\right).
$$
$$
t = \gamma \cdot \max  \left(\frac {2 \ell}{\mu}, \frac {2 \ell \alpha}{\mu^ {2} B}\right) \log \left(\frac {\mathcal {L} (\boldsymbol {\theta} _ {0}) - \mathcal {L} ^ {*}}{\epsilon}\right) = \mathcal {O} \left(\left(\frac {r}{n} + 1\right) \cdot \left(\frac {\ell}{\mu} + \frac {\ell \alpha}{\mu^ {2} B}\right) \log \frac {\mathcal {L} (\boldsymbol {\theta} _ {0}) - \mathcal {L} ^ {*}}{\epsilon}\right)
$$
G.1.1 Verification of assumptions
We show that the  $\mathrm{tr}(\pmb{\Sigma}(\pmb{\theta}_t)) \leq \alpha (\mathcal{L}(\pmb{\theta}_t) - \mathcal{L}^*)$  assumption holds for certain losses.
First, consider optimizing the model  $f(\pmb{x}; \pmb{\theta})$  with square loss, so that
$$
\mathcal {L} (\boldsymbol {\theta}) = \frac {1}{N} \sum_ {i \in [ N ]} (f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) - \boldsymbol {y} _ {i}) ^ {2}.
$$
One then has that
$$
\boldsymbol {\Sigma} (\boldsymbol {\theta}) = \frac {2}{N} \sum_ {i \in [ N ]} (f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) - \boldsymbol {y} _ {i}) ^ {2} \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) ^ {\top} - \nabla \mathcal {L} (\boldsymbol {\theta}) \nabla \mathcal {L} (\boldsymbol {\theta}) ^ {\top}.
$$
$$
\begin{array}{l} \operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta})) \leq \frac {2}{N} \sum_ {i \in [ N ]} (f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) - \boldsymbol {y} _ {i}) ^ {2} \| \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) \| ^ {2} \\ \leq 2 \mathcal {L} (\boldsymbol {\theta}) \sum_ {i \in [ N ]} \| \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) \| ^ {2}. \\ \end{array}
$$
Assume that the data can be interpolated, i.e.,  $\mathcal{L}^* = 0$ . If the function is  $L$ -Lipschitz, i.e.,  $\| \nabla f(\pmb{x};\pmb{\theta}) \| \leq L$ , then the condition holds with  $\alpha = 2NL^2$ . If we are in the kernel regime, i.e.,  $f(\pmb{x}_i; \pmb{\theta}) = \phi(\pmb{x}_i)^\top \pmb{\theta}$  for some feature map  $\phi$ , then
$$
\nabla^ {2} \mathcal {L} (\boldsymbol {\theta}) = \frac {2}{N} \sum_ {i \in [ N ]} f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) ^ {\top}.
$$
$$
\operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta})) \leq N \operatorname {t r} \left(\nabla^ {2} \mathcal {L} (\boldsymbol {\theta})\right) \cdot \mathcal {L} (\boldsymbol {\theta}) \leq N \ell r \cdot \mathcal {L} (\boldsymbol {\theta}).
$$
So the condition holds for  $\alpha = N\ell r$
Next, consider the cross entropy loss function, i.e
$$
\mathcal {L} (\boldsymbol {\theta}) = \frac {1}{N} \sum_ {i \in [ N ]} \exp (- y _ {i} f (\boldsymbol {x} _ {i}; \boldsymbol {\theta})).
$$
$$
\boldsymbol {\Sigma} (\boldsymbol {\theta}) = \frac {1}{N} \sum_ {i \in [ N ]} \exp (- 2 y _ {i} f (\boldsymbol {x} _ {i}; \boldsymbol {\theta})) y _ {i} ^ {2} \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) ^ {\top} - \mathcal {L} (\boldsymbol {\theta}) \mathcal {L} (\boldsymbol {\theta}) ^ {\top},
$$
Assume that the targets  $y_{i}$  are bounded in  $[-1, 1]$  (which is true for binary classification tasks), and that  $\mathcal{L}^{*} = 0$  (which can be achieved if  $|f(\pmb{x};\pmb{\theta})|$  can be sent to  $\infty$ ) we have that
$$
\operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta})) \leq \frac {1}{N} \sum_ {i \in [ N ]} \exp (- 2 y _ {i} f (\boldsymbol {x} _ {i}; \boldsymbol {\theta})) \| \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) \| ^ {2}.
$$
In the kernel regime,  $f(\pmb{x}_i;\pmb{\theta}) = \phi(\pmb{x}_i)^\top \pmb{\theta}$ , and thus
$$
\nabla^ {2} \mathcal {L} (\boldsymbol {\theta}) = \frac {1}{N} \sum_ {i \in [ N ]} \exp (- y _ {i} f (\boldsymbol {x} _ {i}; \boldsymbol {\theta})) \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) \nabla f (\boldsymbol {x} _ {i}; \boldsymbol {\theta}) ^ {\top}.
$$
Therefore the condition holds with  $\alpha = N\ell r$  as well.
G.2 Proofs for Gaussian perturbations
The first lemma computes the second moment of the covariance estimate  $\widehat{\nabla}\mathcal{L}(\pmb {\theta};\mathcal{B})$  when  $z$  is drawn  $\mathcal{N}(0,\pmb {I})$
Lemma 5. Let  $z_{i}\sim \mathcal{N}(0,I)$  i.i.d. Then
$$
\begin{array}{l} \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \right] = \left(1 + \frac {1}{n}\right) \cdot \left(\nabla \mathcal {L} (\boldsymbol {\theta}) \nabla \mathcal {L} (\boldsymbol {\theta}) ^ {\top} + \frac {1}{B} \boldsymbol {\Sigma} _ {M B} (\boldsymbol {\theta})\right) \tag {8} \\ + \frac {1}{n} \boldsymbol {I} \cdot \left(\| \nabla \mathcal {L} (\boldsymbol {\theta}) \| ^ {2} + \frac {1}{B} \operatorname {t r} \left(\boldsymbol {\Sigma} _ {M B} (\boldsymbol {\theta})\right)\right). \\ \end{array}
$$
Proof. As in the proof of Lemma 2, we have that in the  $\epsilon \to 0$  limit
$$
\begin{array}{l} \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \right] \\ = \frac {1}{B ^ {2} n ^ {2}} \sum_ {\left(\boldsymbol {x} _ {1}, \boldsymbol {y} _ {1}\right), \left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right) \in \mathcal {B}} \sum_ {i, j \in [ n ]} \mathbb {E} \left[ \right.\left(\boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {1}, \boldsymbol {y} _ {1}\right)\right\}\right)\left. \right)\left(\boldsymbol {z} _ {j} \boldsymbol {z} _ {j} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right)\right\})\right) ^ {\top} \left. \right] \\ \end{array}
$$
For vectors  $\mathbf{u}, \mathbf{v}$ , we have that
$$
\mathbb {E} _ {\boldsymbol {z} _ {i}, \boldsymbol {z} _ {j}} [ \boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} \boldsymbol {u v} ^ {\top} \boldsymbol {z} _ {j} \boldsymbol {z} _ {j} ^ {\top} ] = \boldsymbol {u v} ^ {\top}
$$
$$
\mathbb {E} _ {\boldsymbol {z} _ {i}} [ \boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} \boldsymbol {u} \boldsymbol {v} ^ {\top} \boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} ] = \mathbb {E} _ {\boldsymbol {z}} [ \boldsymbol {z} ^ {\otimes 4} ] (\boldsymbol {u}, \boldsymbol {v}) = 3 \mathrm {S y m} (\boldsymbol {I} ^ {\otimes 2}) (\boldsymbol {u}, \boldsymbol {v}) = \boldsymbol {u} ^ {\top} \boldsymbol {v} \cdot \boldsymbol {I} + 2 \boldsymbol {u} \boldsymbol {v} ^ {\top}.
$$
$$
\begin{array}{l} \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \right] \\ = \frac {1}{B ^ {2}} \sum_ {\left(\boldsymbol {x} _ {1}, \boldsymbol {y} _ {1}\right), \left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right) \in \mathcal {B}} \left(\frac {n - 1}{n} + \frac {2}{n}\right) \mathbb {E} \left[ \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {1}, \boldsymbol {y} _ {1}\right) \right\}) \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right) \right\}) ^ {\top} \right] \\ + \frac {1}{n} \cdot \mathbb {E} \left[ \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {1}, \boldsymbol {y} _ {1}\right) \right\}) ^ {\top} \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right) \right\}) \right] \boldsymbol {I}. \\ \end{array}
$$
In the proof of Lemma 2 we showed that
$$
\frac {1}{B ^ {2}} \sum_ {\left(\boldsymbol {x} _ {1}, \boldsymbol {y} _ {1}\right), \left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right) \in \mathcal {B}} \mathbb {E} \left[ \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {1}, \boldsymbol {y} _ {1}\right) \right\}) \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right) \right\}) ^ {\top} \right] = \nabla \mathcal {L} (\boldsymbol {\theta}) \nabla \mathcal {L} (\boldsymbol {\theta}) ^ {\top} + \frac {1}{B} \boldsymbol {\Sigma} (\boldsymbol {\theta}).
$$
Plugging this yields
$$
\begin{array}{l} \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \right] = \left(\frac {n + 1}{n}\right) \cdot \left(\nabla \mathcal {L} (\boldsymbol {\theta}) \nabla \mathcal {L} (\boldsymbol {\theta}) ^ {\top} + \frac {1}{B} \boldsymbol {\Sigma} (\boldsymbol {\theta})\right) \tag {9} \\ + \frac {1}{n} \boldsymbol {I} \cdot \left(\| \nabla \mathcal {L} (\boldsymbol {\theta}) \| ^ {2} + \frac {1}{B} \operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta}))\right). \\ \end{array}
$$
We can prove an analog to Theorem 1 in the case where the  $z_{i}$  are Gaussian. One challenge is that  $\| \pmb{\theta}_{t + 1} - \pmb{\theta}_t\|$  is no longer bounded; instead we the  $r$ -local effective rank assumption only holds with high probability, and thus to bound the expected loss decrease we must control the probability of the  $\| \pmb{\theta}_{t + 1} - \pmb{\theta}_t\|$  being large.
Consider the following modified version of the local  $r$ -effective rank assumption, where the upper bound on the Hessian is measured over a ball of radius twice as large as the one in Assumption 1.
Assumption 2 (Local  $r$ -effective rank, Gaussian). Let  $G(\pmb{\theta}_t) = \max_{(\pmb{x},\pmb{y})\in \mathcal{D}}\|\nabla\mathcal{L}(\pmb{\theta}_t;\{(\pmb{x},\pmb{y})\})\|$ . There exists a matrix  $\pmb{H}(\pmb{\theta}_t)$  such that:
1. For all  $\pmb{\theta}$  such that  $\| \pmb {\theta} - \pmb {\theta}_t\| \leq 2\eta dG(\pmb {\theta}_t)$ , we have  $\nabla^2\mathcal{L}(\pmb {\theta})\preceq \pmb {H}(\pmb {\theta}_t)$ .
2. The effective rank of  $\mathbf{H}(\pmb{\theta}_t)$ , i.e.,  $\mathrm{tr}(\pmb{H}(\pmb{\theta}_t)) / \| \pmb{H}(\pmb{\theta}_t)\|_{\mathrm{op}}$ , is at most  $r$ .
Theorem 2 (Dimension-Free Rate, Gaussian  $\mathbf{z}$ ). Assume the loss exhibits local  $r$ -effective rank (Assumption 2). If  $\pmb{\theta}_{t+1} = \pmb{\theta}_t - \eta_{ZO}\widehat{\nabla}\mathcal{L}(\pmb{\theta}_t; \mathcal{B})$  is a single step of ZO-SGD using the  $n$ -SPSA estimate with a minibatch of size  $B$ , then there exists a  $\gamma = \Theta(r/n)$  such that the expected loss decrease can be bounded as
$$
\begin{array}{l} \mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \mid \boldsymbol {\theta} _ {t} \right] - \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \\ \leq - \eta_ {Z O} \| \nabla \mathcal {L} (\pmb {\theta} _ {t}) \| ^ {2} + \frac {1}{2} \eta_ {Z O} ^ {2} \ell \cdot \gamma \cdot \mathbb {E} [ \| \nabla \mathcal {L} (\pmb {\theta} _ {t}; \mathcal {B}) \| ^ {2} ] + \eta_ {Z O} ^ {2} \ell G (\pmb {\theta} _ {t}) ^ {2} \exp (- \Omega (n d)). \\ \end{array}
$$
Proof of Theorem 2. Let  $\mathcal{A}$  be the event that  $\| \pmb{\theta}_{t + 1} - \pmb{\theta}_t\| \leq 2\eta dG(\pmb{\theta}_t)$ . On  $\mathcal{A}$ , we have that
$$
\mathcal {L} (\boldsymbol {\theta} _ {t + 1}) \leq \mathcal {L} (\boldsymbol {\theta} _ {t}) - \eta \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) ^ {\top} \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) + \frac {1}{2} \eta^ {2} \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) ^ {\top} \boldsymbol {H} (\boldsymbol {\theta}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}).
$$
Likewise, since  $\mathcal{L}$  is  $\ell$ -smooth, we have that
$$
\mathcal {L} (\boldsymbol {\theta} _ {t + 1}) \leq \mathcal {L} (\boldsymbol {\theta} _ {t}) - \eta \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) ^ {\top} \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) + \frac {1}{2} \eta^ {2} \ell \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) \right\| ^ {2}.
$$
$$
\begin{array}{l} \mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \mid \boldsymbol {\theta} _ {t} \right] \leq \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) - \eta \left\| \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right\| ^ {2} + \frac {1}{2} \eta^ {2} \left\langle \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \cdot \mathbf {1} (\mathcal {A}) \right], \boldsymbol {H} \left(\boldsymbol {\theta} _ {t}\right) \right\rangle \\ + \frac {1}{2} \eta^ {2} \ell \mathbb {E} \left[ \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) \right\| ^ {2} \cdot \mathbf {1} (\neg \mathcal {A}) \right] \\ = \mathcal {L} (\boldsymbol {\theta} _ {t + 1}) - \eta \| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) \| ^ {2} + \frac {1}{2} \eta^ {2} \left\langle \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \right], \boldsymbol {H} (\boldsymbol {\theta} _ {t}) \right\rangle \\ \frac {1}{2} \eta^ {2} \left\langle \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \cdot \mathbf {1} (\neg \mathcal {A}) \right], \ell I - \boldsymbol {H} (\boldsymbol {\theta} _ {t}) \right\rangle . \\ \end{array}
$$
The latter term can be bounded as follows
$$
\begin{array}{l} \frac {1}{2} \eta^ {2} \left\langle \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \cdot \mathbf {1} (\neg \mathcal {A}) \right], \ell I - \boldsymbol {H} (\boldsymbol {\theta} _ {t}) \right\rangle \leq \eta^ {2} \ell \mathbb {E} \left[ \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right\| ^ {2} \cdot \mathbf {1} (\neg \mathcal {A}) \right] \\ \leq \eta^ {2} \ell \mathbb {E} \left[ \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right\| ^ {4} \right] ^ {\frac {1}{2}} \Pr [ \neg \mathcal {A} ] ^ {1 / 2}. \\ \end{array}
$$
The gradient estimate  $\widehat{\nabla}\mathcal{L}(\pmb {\theta};\mathcal{B})$  satisfies
$$
\left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right\| \leq \frac {1}{n} \sum_ {i \in [ n ]} \left| \boldsymbol {z} _ {i} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right| \cdot \| \boldsymbol {z} _ {i} \|
$$
The expectation term is upper bounded as
$$
\begin{array}{l} \mathbb {E} \left[ \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right\| ^ {4} \right] \leq \frac {1}{n} \sum_ {i \in [ n ]} \mathbb {E} \left[ \left| \boldsymbol {z} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right| ^ {4} \cdot \| \boldsymbol {z} \| ^ {4} \right] \\ \leq \mathbb {E} \left[ \left| \boldsymbol {z} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right| ^ {8} \right] ^ {1 / 2} \mathbb {E} \left[ \| \boldsymbol {z} \| ^ {8} \right] ^ {1 / 2} \\ \leq \sqrt {1 0 5} (d + 6) ^ {2} G (\boldsymbol {\theta} _ {t}) ^ {4}, \\ \end{array}
$$
where we have plugged in explicit formulas for moments of Gaussian and  $\chi^2$  random variables.
Next, note that on the event  $\neg A$ , we have
$$
2 \eta d G (\boldsymbol {\theta} _ {t}) \leq \| \boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t} \| = \eta \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) \right\| \leq \eta \cdot \frac {1}{n} \sum_ {i \in [ n ]} \| z _ {i} \| ^ {2} G (\boldsymbol {\theta} _ {t}).
$$
$$
\Pr [ \neg \mathcal {A} ] \leq \Pr \left[ \sum_ {i \in [ n ]} \| z _ {i} \| ^ {2} \geq 2 n d \right]
$$
Lemma 6 (Standard  $\chi^2$ -tail bound). Let  $Z$  be a  $\chi^2$  random variable with  $k$  degrees of freedom. Then
$$
\Pr [ Z \geq k + u ] \leq \exp \left(- \min  \left(\frac {u ^ {2}}{1 6 k}, \frac {u}{1 6}\right)\right)
$$
Since  $\sum_{i\in [n]}\| \pmb {z}_i\|^2$  is a  $\chi^2$  random variable with nd degrees of freedom, we thus have that
$$
\Pr [ \neg \mathcal {A} ] \leq \exp \left(- \frac {n d}{1 6}\right).
$$
$$
\begin{array}{l} \frac {1}{2} \eta^ {2} \left\langle \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \cdot \mathbf {1} (\neg \mathcal {A}) \right], \ell I - \boldsymbol {H} (\boldsymbol {\theta} _ {t}) \right\rangle \leq \eta^ {2} \ell 1 0 5 ^ {1 / 4} (d + 6) G (\boldsymbol {\theta} _ {t}) ^ {2} \exp (- \frac {n d}{3 2}) \\ = \eta^ {2} \ell G (\boldsymbol {\theta} _ {t}) ^ {2} \exp (- \Omega (n d)). \\ \end{array}
$$
Finally, plugging in (8), along with the fact that  $\| H(\pmb{\theta}_t)\|_{op} \leq \ell$  and  $\mathrm{tr}(H(\pmb{\theta}_t)) \leq \ell r$ ,
$$
\begin{array}{l} \left\langle \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) ^ {\top} \right], \boldsymbol {H} (\boldsymbol {\theta} _ {t}) \right\rangle = \frac {r + n + 1}{n} \cdot \ell \left(\| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) \| ^ {2} + \frac {1}{B} \operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta} _ {t}))\right) \\ = \frac {r + n + 1}{n} \cdot \mathbb {E} \left[ \| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) \| ^ {2} \right] \\ \end{array}
$$
Thus letting  $\gamma = \frac{r + n + 1}{n}$  yields
$$
\begin{array}{l} \mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \mid \boldsymbol {\theta} _ {t} \right] - \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \\ \leq - \eta \| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) \| ^ {2} + \frac {1}{2} \eta^ {2} \ell \cdot \gamma \cdot \mathbb {E} [ \| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) \| ^ {2} ] + \eta^ {2} \ell G (\boldsymbol {\theta} _ {t}) ^ {2} \exp (- \Omega (n d)), \\ \end{array}
$$