SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Prefix-tuning [57] adds a prefix of  $m$  tunable representations at each layer and freezes the rest of the model. The representations are added as new keys and values and treated as additional context during the attention operation. We initialize these tunable representations by randomly sampling tokens from the vocabulary and passing them through the LLM to get their keys and values at different attention layers. We found this crucial to make prefix tuning stable with MeZO, and this trick additionally boosts the performance of prefix tuning with backpropagation, as shown in Table 17. We also tried the reparameterization trick in [57], which does not help MeZO training. In our experiments, we find  $m = 5$  to be sufficient to achieve good performance on most tasks.
We also show that MeZO is compatible with parameter-efficient fine-tuning methods, such as prefix tuning and LoRA. Surprisingly, the performance of MeZO does not improve substantially when tuning much fewer parameters, as one might expect from classical analyses (see Section 4). Accordingly, our theoretical analysis in Section 4 suggests that the convergence rate of ZO-SGD does not depend on the parameter dimension during fine-tuning.
Table 16: The hyperparameter grids used for OPT experiments. All weight decay is set to 0. FT uses 5 epochs and linear scheduled learning rates and MeZO uses 20K steps and constant learning rates. We check validation performance and save the best checkpoint every 1/5 total training steps.
Task	SST-2	SST-5	SNLI	MNLI	RTE	TREC
Type	— sentiment —	— natural language inference —	— topic —
FT (prefix, random init)	90.7 (1.7)	47.2 (2.0)	70.7 (2.8)	62.6 (3.3)	63.5 (4.4)	83.4 (4.7)
FT (prefix, real act init)	91.9 (1.0)	47.7 (1.1)	77.2 (1.3)	66.5 (2.5)	66.6 (2.0)	85.7 (1.3)
Table 17: Prefix-tuning ablations. We compare randomly-initialized prefixes and real word activation prefixes. Using real word activations significantly outperforms random initialization.
E.6 Training with non-differentiable objectives
The experiments maximizing the accuracy of a RoBERTa-large model were all conducted using the same grid as MeZO in Table 15.
For OPT experiments on SQuAD with F1 as objective, we use a batch size of 16. For MeZO, we use a learning rate of  $\{1\mathrm{e} - 6,5\mathrm{e} - 6,1\mathrm{e} - 5\}$  and  $\epsilon = 1\mathrm{e} - 3$ . For MeZO (prefix), we use a learning rate of  $\{1\mathrm{e} - 1,5\mathrm{e} - 2,1\mathrm{e} - 2\}$  and  $\epsilon = 1\mathrm{e} - 1$ .
E.7 Memory profiling
In memory profiling, we use standard implementation with Huggingface's transformers [104] package. We did not turn on any advance memory-saving options, e.g., gradient checkpointing. We set the per-device batch size as 1 to test the minimum hardware requirement to run the model with specific optimization algorithms. For multi-GPU backpropagation, we use fully sharded data parallel (FSDP) [33] provided by PyTorch [76]. For multi-GPU MeZO, we use transformers multi-GPU inference of large models. We use Nvidia's nvidia-smi command to monitor the GPU memory usage. We call a run "successful" if there is no out of memory error from GPUs for at least 100 steps. We also profile fine-tuning with LoRA, but find its memory usage similar to that of fine-tuning with prefix-tuning. Hence here we only show the analysis with prefix-tuning.
F More experiment results
F.1 RoBERTa-large experiments
Table 18 contains the detailed numbers corresponding to Figure 2 and also reports the performance of MeZO-Adam.
Task Type	SST-2 —— sentiment ——	SNLI —— natural language inference ——	RTE —— topic ——
Zero-shot	79.0	35.5	50.2	48.8	51.4	32.0
Gradient-free methods: k = 16
LP	76.0 (2.8)	40.3 (1.9)	66.0 (2.7)	56.5 (2.5)	59.4 (5.3)	51.3 (5.5)
MeZO	90.5 (1.2)	45.5 (2.0)	68.5 (3.9)	58.7 (2.5)	64.0 (3.3)	76.9 (2.7)
MeZO (LoRA)	91.4 (0.9)	43.0 (1.6)	69.7 (6.0)	64.0 (2.5)	64.9 (3.6)	73.1 (6.5)
MeZO (prefix)	90.8 (1.7)	45.8 (2.0)	71.6 (2.5)	63.4 (1.8)	65.4 (3.9)	80.3 (3.6)
MeZO-Adam	90.4 (1.4)	45.4 (1.5)	74.1 (2.7)	64.3 (0.8)†	59.2 (11.1)†	78.3 (1.4)
Gradient-based methods: k = 16
FT	91.9 (1.8)	47.5 (1.9)	77.5 (2.6)	70.0 (2.3)	66.4 (7.2)	85.0 (2.5)
FT (LoRA)	91.4 (1.7)	46.7 (1.1)	74.9 (4.3)	67.7 (1.4)	66.1 (3.5)	82.7 (4.1)
FT (prefix)	91.9 (1.0)	47.7 (1.1)	77.2 (1.3)	66.5 (2.5)	66.6 (2.0)	85.7 (1.3)
Gradient-free methods: k = 512
LP	91.3 (0.5)	51.7 (0.5)	80.9 (1.0)	71.5 (1.1)	73.1 (1.5)	89.4 (0.5)
MeZO	93.3 (0.7)	53.2 (1.4)	83.0 (1.0)	78.3 (0.5)	78.6 (2.0)	94.3 (1.3)
MeZO (LoRA)	93.4 (0.4)	52.4 (0.8)	84.0 (0.8)	77.9 (0.6)	77.6 (1.3)	95.0 (0.7)
MeZO (prefix)	93.3 (0.1)	53.6 (0.5)	84.8 (1.1)	79.8 (1.2)	77.2 (0.8)	94.4 (0.7)
MeZO-Adam	93.3 (0.6)	53.9 (0.8)	85.3 (0.8)	79.6 (0.4)	79.2 (1.2)	95.1 (0.3)
Gradient-based methods: k = 512
FT	93.9 (0.7)	55.9 (0.9)	88.7 (0.8)	84.4 (0.8)	82.7 (1.4)	97.3 (0.2)
FT (LoRA)	94.2 (0.2)	55.3 (0.7)	88.3 (0.5)	83.9 (0.6)	83.2 (1.3)	97.0 (0.3)
FT (prefix)	93.7 (0.3)	54.6 (0.7)	88.3 (0.7)	83.3 (0.5)	82.5 (0.8)	97.4 (0.2)
Table 18: Experiments on RoBERTa-large (350M parameters). LP: Linear probing; ZO, ZO (LoRA), and ZO (prefix): our memory-efficient ZO-SGD (Section 2.1) with full-parameter tuning, LoRA, and prefix-tuning respectively; FT: fine-tuning with Adam. All reported numbers are averaged accuracy (standard deviation). All experiments use prompts (Appendix E.2). ZO outperforms zero-shot and LP by a large margin and approaches FT performance with much less memory cost.
LP-MeZO We also compare MeZO to performing linear probing and then subsequently performing fine-tuning via MeZO, following the analogous suggestion for fine-tuning in Kumar et al. [53]. We use the MeZO grid described in Table 15. Note that the linear probing checkpoints used here have early stopping, unlike the ones reported in Table 18. We heuristically implement early stopping by limiting the number of iterations (from 5000 to 1000) and increasing the convergence tolerance (from  $1\mathrm{e} - 4$  to 0.01) in the scipy solver. Experiments on a few settings show that LP-MeZO can sometimes improve performance without increasing the memory consumption (see Table 19). However, sometimes, linear probing first can severely hurt performance.
Task	SST-2	SST-5	SNLI	TREC
Zero-shot	79.0	35.5	50.2	32.0
FT	91.9 (1.8)	47.5 (1.9)	77.5 (2.6)	85.0 (2.5)
MeZO	90.5 (1.2)	45.5 (2.0)	68.5 (3.9)	76.9 (2.7)
LP-MeZO	91.4 (1.4)	41.9 (3.3)	70.7 (3.4)	54.0 (4.5)
Table 19: Performing linear probing before fine-tuning with MeZO, as suggested previously [53], can sometimes improve performance without increasing the memory overhead. We use  $k = 16$  for these experiments.
F.2 OPT experiments
Table 20 present the full results of OPT-30B and OPT-66B, with detailed MeZO numbers.
Task	SST-2	RTE	BoolQ	WSC	WIC	SQuAD
30B zero-shot	56.7	52.0	39.1	38.5	50.2	46.5
30B ICL	81.9	66.8	66.2	56.7	51.3	78.0
30B MeZO	90.6	66.4	67.2	63.5	56.3	85.2
30B MeZO (prefix)	87.5	72.6	73.5	55.8	59.1	83.9
66B zero-shot	57.5	67.2	66.8	43.3	50.6	48.1
66B ICL	89.3	65.3	62.8	52.9	54.9	81.3
66B MeZO	91.2	65.7	72.7	63.5	58.9	*
66B MeZO (prefix)	93.6	66.4	73.7	57.7	58.6	85.0
Table 20: Experiments on OPT-30B and OPT-66B (with 1,000 examples).*: MeZO requires further tuning to successfully optimize.
F.3 Convergence of MeZO with full-parameter and PEFT
We demonstrate the convergence rate of MeZO, MeZO (LoRA) and MeZO (prefix) on SST-2 and SNLI for the first 5,000 steps in Figures 5. We see that despite the different number of parameters they optimize, MeZO demonstrates similar training speed on full parameter and PEFT. This agrees with our theory in Section 4, which shows that MeZO's optimization speed is independent of the number of parameters.
Figure 5: MeZO does not optimize significantly faster when tuning fewer parameters, agreeing with our theory in Section 4.
F.4 ZO vs BBTv2
We compare ZO with BBTv2 [90] on mutually assessed tasks in Table 21. ZO significantly outperforms BBTv2. Furthermore, BBTv2 is limited to optimize in low-dimensional space and requires prefix-tuning and a down-projection to reduce the number of optimized parameters. BBTv2 also employs an iterative scheme which only optimizes one layer at a time. In contrast, ZO works with both full-parameter tuning and PEFT, as shown in our experiments (Section 3) and theory (Section 4).
Task	SST-2	SNLI	RTE
Task type	sentiment	- natural language inference
Zero-shot	79.0	50.2	51.4
BBTv2	90.3 (1.7)	57.3 (2.3)	56.7 (3.3)
MeZO	90.5 (1.2)	68.5 (3.9)	64.0 (3.3)
MeZO (LoRA)	91.4 (0.9)	69.7 (6.0)	64.9 (3.6)
MeZO (prefix)	90.8 (1.7)	71.6 (2.5)	65.4 (3.9)
F.5 Memory profiling
We show the detailed numbers of memory profiling results Table 22, which also corresponds to Figure 3. For how we profile the memory usage, please refer to Appendix E.7.
Table 21: ZO vs BBTv2 with RoBERTa-large. BBTv2 performance is from Sun et al. [90].
Method	Zero-shot / MeZO	ICL	Prefix FT	Full-parameter FT
1.3B	1xA100 (4GB)	1xA100 (6GB)	1xA100 (19GB)	1xA100 (27GB)
2.7B	1xA100 (7GB)	1xA100 (8GB)	1xA100 (29GB)	1xA100 (55GB)
6.7B	1xA100 (14GB)	1xA100 (16GB)	1xA100 (46GB)	2xA100 (156GB)
13B	1xA100 (26GB)	1xA100 (29GB)	2xA100 (158GB)	4xA100 (316GB)
30B	1xA100 (58GB)	1xA100 (62GB)	4xA100 (315GB)	8xA100 (633GB)
66B	2xA100 (128GB)	2xA100 (134GB)	8xA100	16xA100
Table 22: Memory usage on the MultiRC (avg #tokens=400) dataset.
F.6 Wallclock time efficiency
In this section, we measure the wallclock time efficiency of MeZO compared to full-parameter FT, with respect to different model sizes. We conduct our experiments with 80GB A100s connected by NVLink and InfiniteBand, which are state-of-the-art solutions for distributed training. As shown in Table 23, on the MultiRC datasets, training with MeZO brings  $7.74 \times$  speedup per step compared to full-parameter FT on a 30B model. This is due to (1) MeZO does not require costly backpropagation and (2) MeZO requires fewer GPUs and reduces the multi-GPU communication overhead. We can see that MeZO has a bigger advantage when training larger models—the multi-GPU overhead for fine-tuning is larger.
Note that even though MeZO has better per-step wallclock efficiency, it requires significantly more steps than standard FT. Taking our OPT-30B experiments as an example: MeZO takes  $32 \times$  more steps than standard FT, while FT takes  $8 \times$  more GPUs and  $7.74 \times$  more time per step. Overall, MeZO requires only half as many GPU-hours as FT for a 30B model.
1.3B	2.7B	13B	30B	66B
MeZO (bsz=16)	0.815s (1)	1.400s (1)	2.702s (1)	5.896s (1)	12.438s (4)
MeZO (bsz=8)	0.450s (1)	0.788s (1)	1.927s (1)	4.267s (1)	7.580s (2)
FT (bsz=8)	0.784s (1)	1.326s (1)	13.638s (4)	45.608s (8)	84.098s (20)
bspd=2, ga=4	bspd=2, ga=4	bspd=1, ga=2	bspd=1, ga=1	bspd=1, ga=1
Table 23: Wallclock time per step of different training methods. Numbers in brackets are numbers of GPUs required. It is measured on 80GB A100s with NVLink and InfiniteBand connections. The wallclock time is averaged over 100 steps. It is measured on the MultiRC task with the OPT family. We use a batch size ("bsz") of 8 for FT and 16 for MeZO (consistent with our main experiment setting). For comparison we also add MeZO with a batch size of 8. For FT (FSDP), we show the following additional information. "bspd": batch size per device. "ga": gradient accumulation steps. The effective batch size is  $\text{bspd} \times \text{ga} \times \# \text{GPUs}$ . Note that for FT 66B, the effective batch size 20.
Proof of Lemma 2. We first note that in the  $\epsilon \to 0$  limit, we have
$$
\widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) = \frac {1}{B n} \sum_ {(\boldsymbol {x}, \boldsymbol {y}) \in \mathcal {B}} \sum_ {i \in [ n ]} \boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \{(\boldsymbol {x}, \boldsymbol {y}) \}).
$$
Taking expectation over the batch  $\mathcal{B}$  and the  $z_{i}$ , we have  $\mathbb{E}[\widehat{\nabla}\mathcal{L}(\pmb {\theta};\mathcal{B})] = \nabla \mathcal{L}(\pmb {\theta})$ , so  $\widehat{\nabla}\mathcal{L}(\pmb {\theta};\mathcal{B})$  is an unbiased estimator of the gradient.
Computing the second moment, we get
$$
\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[ \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(\boldsymbol {z} _ {j} \boldsymbol {z} _ {j} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \left\{\left(\boldsymbol {x} _ {2}, \boldsymbol {y} _ {2}\right) \right\})\right) ^ {\top} \right] \\ \end{array}
$$
Let  $\pmb{u},\pmb{v}$  be two arbitrary vectors. We have that
$$
\mathbb {E} _ {\boldsymbol {z} _ {i}, \boldsymbol {z} _ {j}} \left[ \boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} \boldsymbol {u v} ^ {\top} \boldsymbol {z} _ {j} \boldsymbol {z} _ {j} ^ {\top} \right] = \boldsymbol {u v} ^ {\top}
$$
when  $i\neq j$  , and
$$
\begin{array}{l} \mathbb {E} _ {\boldsymbol {z} _ {i}} \left[ \boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} \boldsymbol {u v} ^ {\top} \boldsymbol {z} _ {i} \boldsymbol {z} _ {i} ^ {\top} \right] = \mathbb {E} _ {\boldsymbol {z}} \left[ \boldsymbol {z} ^ {\otimes 4} \right] (\boldsymbol {u}, \boldsymbol {v}) \\ = \frac {3 d}{d + 2} \operatorname {S y m} \left(\boldsymbol {I} ^ {\otimes 2}\right) (\boldsymbol {u}, \boldsymbol {v}) \\ = \frac {d}{d + 2} \cdot \boldsymbol {u} ^ {\top} \boldsymbol {v} \cdot \boldsymbol {I} + \frac {2 d}{d + 2} \cdot \boldsymbol {u v} ^ {\top}. \\ \end{array}
$$
$$
\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 d}{n (d + 2)}\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 {d}{n (d + 2)} \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}
$$
Next, note that when  $(\pmb{x}_1,\pmb{y}_1)\neq (\pmb{x}_2,\pmb{y}_2)$ , we have
$$
\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},
$$
and when  $(\pmb{x}_1, \pmb{y}_1) = (\pmb{x}_2, \pmb{y}_2)$  we have
$$
\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} + \boldsymbol {\Sigma} _ {M B} (\boldsymbol {\theta}).
$$
$$
\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}),
$$
and 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(1 + \frac {d - 2}{n (d + 2)}\right) \cdot \left(\nabla \mathcal {L} (\boldsymbol {\theta}) \nabla \mathcal {L} (\boldsymbol {\theta}) ^ {\top} + \frac {1}{B} \boldsymbol {\Sigma} (\boldsymbol {\theta})\right) \tag {7} \\ + \frac {d}{n (d + 2)} \boldsymbol {I} \cdot \left(\| \nabla \mathcal {L} (\boldsymbol {\theta}) \| ^ {2} + \frac {1}{B} \operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta}))\right). \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} \left[ \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right\| ^ {2} \right] = \left(1 + \frac {d ^ {2} + d - 2}{n (d + 2)}\right) \cdot \left(\| \nabla \mathcal {L} (\boldsymbol {\theta}) \| ^ {2} + \frac {1}{B} \operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta}))\right) \\ = \frac {d + n - 1}{n} \cdot \mathbb {E} \left[ \| \nabla \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \| ^ {2} \right]. \\ \end{array}
$$
Proof of Theorem 1. By Taylor's theorem with remainder, we have that
$$
\begin{array}{l} \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) = \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) + \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) ^ {\top} \left(\boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t}\right) \\ + \int_ {0} ^ {1} \lambda (\boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t}) ^ {\top} \nabla^ {2} \mathcal {L} (\lambda \boldsymbol {\theta} _ {t + 1} + (1 - \lambda) \boldsymbol {\theta} _ {t}) (\boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t}) ^ {\top} d \lambda \\ \end{array}
$$
Next, note that
$$
\left\| \boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t} \right\| = \eta \left\| \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \right\| \leq \eta \sqrt {d} \cdot \frac {1}{B n} \sum \left| \boldsymbol {z} _ {i} ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \{(\boldsymbol {x}, \boldsymbol {y}) \}) \right| \leq \eta d G (\boldsymbol {\theta} _ {t}).
$$
Therefore  $\| \lambda \pmb{\theta}_{t + 1} + (1 - \lambda) \pmb{\theta}_t - \pmb{\theta}_t \| \leq \eta dG(\pmb{\theta}_t)$ . By the assumption we have the upper bound  $\nabla^2 \mathcal{L}(\lambda \pmb{\theta}_{t + 1} + (1 - \lambda) \pmb{\theta}_t) \preceq \pmb{H}(\pmb{\theta}_t)$ , and thus
$$
\begin{array}{l} \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \leq \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) + \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) ^ {\top} \left(\boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t}\right) + \left(\boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t}\right) ^ {\top} \boldsymbol {H} \left(\boldsymbol {\theta} _ {t}\right) \left(\boldsymbol {\theta} _ {t + 1} - \boldsymbol {\theta} _ {t}\right) \\ = \mathcal {L} (\boldsymbol {\theta} _ {t}) - \eta \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) ^ {\top} \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) + \frac {1}{2} \eta^ {2} \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) ^ {\top} \boldsymbol {H} (\boldsymbol {\theta} _ {t}) \widehat {\nabla} \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}). \\ \end{array}
$$
Taking the conditional expectation with respect to  $\pmb{\theta}_t$  and plugging in (9), the formula for the covariance of our ZO estimate  $\widehat{\nabla}\mathcal{L}(\pmb{\theta}_t;\mathcal{B})$ , yields
$$
\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}\right) - \eta \left\| \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right\| ^ {2} + \frac {\eta^ {2}}{2} \left\langle \boldsymbol {H} \left(\boldsymbol {\theta} _ {t}\right), \mathbb {E} \left[ \widehat {\nabla} \mathcal {L} \left(\boldsymbol {\theta}; \mathcal {B}\right) \widehat {\nabla} \mathcal {L} \left(\boldsymbol {\theta}; \mathcal {B}\right) ^ {\top} \right] \right\rangle \\ = \mathcal {L} (\boldsymbol {\theta} _ {t}) - \eta \| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) \| ^ {2} + \frac {\eta^ {2}}{2} \cdot \frac {d}{n (d + 2)} \left(\| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) \| ^ {2} + \frac {1}{B} \operatorname {t r} (\boldsymbol {\Sigma} (\boldsymbol {\theta} _ {t}))\right) \operatorname {t r} (\boldsymbol {H} (\boldsymbol {\theta} _ {t})) \\ + \frac {\eta^ {2}}{2} \left(1 + \frac {d - 2}{n (d + 2)}\right) \left(\nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) ^ {\top} \boldsymbol {H} \left(\boldsymbol {\theta} _ {t}\right) \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) + \frac {1}{B} \langle \boldsymbol {\Sigma} \left(\boldsymbol {\theta} _ {t}\right), \boldsymbol {H} \left(\boldsymbol {\theta} _ {t}\right) \rangle\right) \\ \end{array}
$$
By assumption, the Hessian upper bound  $\mathbf{H}(\pmb{\theta}_t)$  satisfies  $\| \mathbf{H}(\pmb{\theta}_t) \|_{op} \leq \ell$  and  $\mathrm{tr}(\mathbf{H}(\pmb{\theta}_t)) \leq \ell r$ . Thus
$$
\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}\right) - \eta \left\| \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right\| ^ {2} + \frac {\eta^ {2} \ell}{2} \cdot \left(\frac {d r + d - 2}{n (d + 2)} + 1\right) \cdot \left(\left\| \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right\| ^ {2} + \frac {1}{B} \operatorname {t r} \left(\boldsymbol {\Sigma} \left(\boldsymbol {\theta} _ {t}\right)\right)\right) \\ = \mathcal {L} (\boldsymbol {\theta} _ {t}) - \eta \left\| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}) \right\| ^ {2} + \frac {\eta^ {2} \ell}{2} \cdot \left(\frac {d r + d - 2}{n (d + 2)} + 1\right) \cdot \mathbb {E} \left[ \left\| \nabla \mathcal {L} (\boldsymbol {\theta} _ {t}; \mathcal {B}) \right\| ^ {2} \right], \\ \end{array}
$$
G.1 Proofs of Global Convergence
Lemma 4. Let  $\mathcal{L}(\pmb{\theta})$  be  $\mu$ -PL and let there exist  $\alpha$  such that  $\mathrm{tr}(\pmb{\Sigma}(\pmb{\theta})) \leq \alpha (\mathcal{L}(\pmb{\theta}) - \mathcal{L}^{*})$  for all  $\pmb{\theta}$ . Then after
$$
t = O \left(\left(\frac {\ell}{\mu} + \frac {\ell \alpha}{\mu^ {2} B}\right) \log \frac {\mathcal {L} (\boldsymbol {\theta} _ {0}) - \mathcal {L} ^ {*}}{\epsilon}\right)
$$
iterations of SGD we have  $\mathbb{E}[\mathcal{L}(\pmb{\theta}_t)] \leq \mathcal{L}^* + \epsilon$ .
Proof of Lemma 4. The descent lemma for SGD yields
$$
\mathbb {E} [ \mathcal {L} (\pmb {\theta} _ {t + 1}) \mid \pmb {\theta} _ {t} ] - \mathcal {L} (\pmb {\theta} _ {t}) \leq - \eta \left\| \nabla \mathcal {L} (\pmb {\theta} _ {t}) \right\| ^ {2} + \frac {1}{2} \eta^ {2} \ell \cdot \mathbb {E} [ \| \nabla \mathcal {L} (\pmb {\theta} _ {t}; \mathcal {B}) \| ^ {2} ].
$$
Plugging in  $\mathbb{E}[\| \nabla \mathcal{L}(\pmb{\theta}_t; \mathcal{B})\|^2] = \| \nabla \mathcal{L}(\pmb{\theta}_t)\|^2 + \frac{1}{B}\operatorname{tr}(\pmb{\Sigma}(\pmb{\theta}_t))$  and selecting a learning rate  $\eta \leq \frac{1}{\ell}$  yields
$$
\mathbb {E} \left[ \mathcal {L} \left(\boldsymbol {\theta} _ {t + 1}\right) \mid \boldsymbol {\theta} _ {t} \right] \leq \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) - \frac {\eta}{2} \left\| \nabla \mathcal {L} \left(\boldsymbol {\theta} _ {t}\right) \right\| ^ {2} + \frac {\eta^ {2} \ell}{2 B} \operatorname {t r} \left(\boldsymbol {\Sigma} \left(\boldsymbol {\theta} _ {t}\right)\right)
$$
Since  $\mathcal{L}$  is  $\mu$ -PL, we get
$$
\mathbb {E} [ \mathcal {L} (\pmb {\theta} _ {t + 1}) \mid \pmb {\theta} _ {t} ] \leq \mathcal {L} (\pmb {\theta} _ {t}) - \eta \mu (\mathcal {L} (\pmb {\theta} _ {t}) - \mathcal {L} ^ {*}) + \frac {\eta^ {2} \ell}{2 B} \operatorname {t r} (\pmb {\Sigma} (\pmb {\theta} _ {t})).
$$