SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sadhika Malladi*
Tianyu Gao*
Eshaan Nichani
Alex Damian
Jason D. Lee
Danqi Chen
Sanjeev Arora
{smalladi, tianyug, eshnich, ad27, jasonlee, danqic, arora}@princeton.edu
Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a prohibitively large amount of memory. Zeroth-order (ZO) methods can in principle estimate gradients using only two forward passes but are theorized to be catastrophically slow for optimizing large models. In this work, we propose a memory-efficient zeroth-order optimizer (MeZO), adapting the classical ZO-SGD method to operate in-place, thereby fine-tuning LMs with the same memory footprint as inference. For example, with a single A100 80GB GPU, MeZO can train a 30-billion parameter model, whereas fine-tuning with backpropagation can train only a 2.7B LM with the same budget. We conduct comprehensive experiments across model types (masked and autoregressive LMs), model scales (up to 66B), and downstream tasks (classification, multiple-choice, and generation). Our results demonstrate that (1) MeZO significantly outperforms in-context learning and linear probing; (2) MeZO achieves comparable performance to fine-tuning with backpropagation across multiple tasks, with up to  $12 \times$  memory reduction and up to  $2 \times$  GPU-hour reduction in our implementation; (3) MeZO is compatible with both full-parameter and parameter-efficient tuning techniques such as LoRA and prefix tuning; (4) MeZO can effectively optimize non-differentiable objectives (e.g., maximizing accuracy or F1). We support our empirical findings with theoretical insights, highlighting how adequate pre-training and task prompts enable MeZO to fine-tune huge models, despite classical ZO analyses suggesting otherwise. $^{1}$
Fine-tuning pre-trained language models (LMs) has been the dominant methodology for solving many language tasks [28], adapting to specialized domains [42], or incorporating human instructions and preferences [73]. However, as LMs are scaled up [13, 72], computing gradients for backpropagation requires a prohibitive amount of memory – in our test, up to  $12 \times$  the memory required for inference – because it needs to cache activations during the forward pass, gradients during the backward pass, and, in the case of Adam [52], also store gradient history (see Section 3.4 for a detailed analysis).
As a result, while it is possible to run inference with a 30-billion (30B) parameter LM on a single Nvidia A100 GPU (with 80GB memory), backpropagation with Adam is feasible only for a 2.7B LM. Parameter-efficient fine-tuning methods (PEFT [46, 57, 54]) update just a fraction of the network
*Equal contribution and corresponding authors.
<sup>1</sup>Our code is available at https://github.com/princeton-nlp/MeZ0.
Figure 1: OPT-13B results with zero-shot, in-context learning (ICL), MeZO (we report the best among MeZO/MeZO (LoRA)/MeZO (prefix)), and fine-tuning with Adam (FT). MeZO demonstrates superior results over zero-shot and ICL and performs on par with FT (within  $1\%$  ) on 7 out of 11 tasks, despite using only 1/12 memory. See Table 1 for detailed numbers and Figure 3 for memory profiling.
parameters, but still need to cache many activations, because the tuned parameters are scattered throughout the model. In our tests, fine-tuning an OPT-13B model with full parameter tuning or PEFT requires  $12 \times$  and  $6 \times$  more memory than inference respectively.
In-context learning (ICL [13]) has allowed solving many tasks with a single inference pass, during which the model processes labeled examples (demonstrations) in its context and then outputs a prediction on a test example. While this allows for quick adaptation of the model to specific use cases, current models allow a limited context size (and thus, limited demonstrations) and the performance is sensitive to the formatting and choice of demonstrations [60, 66]. ICL can slow with the number of demonstrations, and it often performs worse than fine-tuning of medium-sized models [13].
Backpropagation also cannot optimize non-differentiable criteria, which have gained popularity in fine-tuning LMs according to human preference scores or set safety standards [89, 73]. Typically, these adaptations involve expensive reinforcement learning from human feedback (RLHF [20]).
A classical zeroth-order optimization method, ZO-SGD [88], uses only differences of loss values to estimate the gradients. Thus, in principle, the method can update neural networks with just forward passes, though naive implementation still doubles the memory overhead and classical lower bounds [69, 32] suggest that convergence slows linearly with model size. As such, ZO methods have been applied in deep learning settings to find adversarial examples or tune input embeddings [91, 90] but not to directly optimize large-scale models (see Liu et al. [61] for a survey).
In this work, we propose a memory-efficient zeroth-order optimizer (MeZO), which adapts the classical ZO-SGD algorithm and reduces its memory consumption to the same as inference. We apply MeZO to fine-tune large LMs and show that, both empirically and theoretically, MeZO can successfully optimize LMs with billions of parameters. Specifically, our contributions are:
1. In MeZO, we adapt the ZO-SGD algorithm [88] and a number of variants to operate in-place on arbitrarily large models with almost no memory overhead (see Algorithm 1 and Section 2).
2. We conduct comprehensive experiments across model types (masked LM and autoregressive LM), model scales (from 350M to 66B), and downstream tasks (classification, multiple-choice, and generation). MeZO consistently outperforms zero-shot, ICL, and linear probing. Moreover, with RoBERTa-large, MeZO achieves performance close to standard fine-tuning within  $5\%$  gap; with OPT-13B, MeZO outperforms or performs comparably to fine-tuning on 7 out of 11 tasks, despite requiring roughly  $12 \times$  less memory (Figure 1 and Section 3). In our implementation, MeZO requires only half as many GPU-hours as Adam fine-tuning for a 30B model (see Appendix F.6).
3. We demonstrate MeZO's compatibility with full-parameter tuning and PEFT (e.g., LoRA [46] and prefix-tuning [57]) in Section 3.
4. Further exploration showcases that MeZO can optimize non-differentiable objectives such as accuracy or F1 score, while still requiring only the same memory as inference (Section 3.3).
5. Our theory suggests that adequate pre-training ensures the per-step optimization rate (Theorem 1) and global convergence rate (Lemma 3) of MeZO depend on a certain condition number of the landscape (i.e., the local effective rank, see Assumption 1) instead of numbers of parameters. This
Algorithm 1: MeZO
Require: parameters  $\theta \in \mathbb{R}^d$  , loss  $\mathcal{L}:\mathbb{R}^d\to \mathbb{R}$  , step budget  $T$  , perturbation scale  $\epsilon$  , batch size  $B$  , learning rate schedule  $\{\eta_t\}$    
for  $t = 1,\dots,T$  do Sample batch  $\mathcal{B}\subset \mathcal{D}$  and random seed s  $\theta \leftarrow$  PerturbParameters(0,ε,s)  $\ell_{+}\gets \mathcal{L}(\theta ;\mathcal{B})$ $\theta \leftarrow$  PerturbParameters(0,-2ε,s)  $\ell_{-}\leftarrow \mathcal{L}(\theta ;\mathcal{B})$ $\theta \leftarrow$  PerturbParameters(0,ε,s) Reset parameters before descent projected_grad  $\leftarrow (\ell_{+} - \ell_{-}) / (2\epsilon)$  Reset random number generator with seed s For sampling z for  $\theta_{i}\in \theta$  do  $z\sim \mathcal{N}(0,1)$ $\theta_{i}\gets \theta_{i} - \eta_{t}*$  projected_grad  $^* z$    
end   
Subroutine PerturbParameters(0,ε,s) Reset random number generator with seed s For sampling z for  $\theta_{i}\in \theta$  do  $z\sim \mathcal{N}(0,1)$ $\theta_{i}\gets \theta_{i} + \epsilon z$  Modify parameters in place   
end return  $\theta$
result is in sharp contrast to existing ZO lower bounds [69, 32] suggesting that the convergence rate can slow proportionally to the number of parameters (Section 4).
2 Zeroth-order optimization
Zeroth-order (ZO) optimizers have long been studied in the context of convex and strongly convex objectives. In the following, we first introduce a classical ZO gradient estimator, SPSA (Definition 1 [88]) and the corresponding SGD algorithm, ZO-SGD (Definition 2). Then we describe MeZO, our in-place implementation that requires the same memory as inference in Section 2.1 and Algorithm 1. We highlight that SPSA can also be used in more complex optimizers, such as Adam, and we provide memory-efficient implementations for those algorithms too (Section 2.2).
Consider a labelled dataset  $\mathcal{D} = \{(x_i, y_i)\}_{i \in [|\mathcal{D}|]}$  and a minibatch  $\mathcal{B} \subset \mathcal{D}$  of size  $B$ , we let  $\mathcal{L}(\boldsymbol{\theta}; \mathcal{B})$  denote the loss on the minibatch. We introduce a classical ZO gradient estimate in this setting.
Definition 1 (Simultaneous Perturbation Stochastic Approximation or SPSA [88]). Given a model with parameters  $\pmb{\theta} \in \mathbb{R}^d$  and a loss function  $\mathcal{L}$ , SPSA estimates the gradient on a minibatch  $\mathcal{B}$  as
$$
\widehat {\nabla} \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) = \frac {\mathcal {L} (\boldsymbol {\theta} + \epsilon z ; \mathcal {B}) - \mathcal {L} (\boldsymbol {\theta} - \epsilon z ; \mathcal {B})}{2 \epsilon} z \approx z z ^ {\top} \nabla \mathcal {L} (\boldsymbol {\theta}; \mathcal {B}) \tag {1}
$$
where  $\mathbf{z} \in \mathbb{R}^d$  with  $\mathbf{z} \sim \mathcal{N}(0, I_d)$  and  $\epsilon$  is the perturbation scale. The  $n$ -SPSA gradient estimate averages  $\widehat{\nabla}\mathcal{L}(\boldsymbol{\theta};\mathcal{B})$  over  $n$  randomly sampled  $\mathbf{z}$ .
SPSA requires only two forward passes through the model to compute the gradient estimate (for  $n$ -SPSA, each estimate requires  $2n$  forward passes). As  $\epsilon \to 0$ , the SPSA estimate can be understood as a rank-1 reconstruction of the gradient. During training,  $n$  can be treated as a hyperparameter and follow a schedule [11, 15], though in cursory experiments (Appendix A),  $n = 1$  is the most efficient.
The original SPSA algorithm [88] perturbs the model by  $1 / z$  and thus requires that  $z$  has finite inverse moments, precluding the choice of  $z$  as Gaussian.  $1 / z$  is very large with high probability for a zero-mean Gaussian  $z$ , so we adopt the standard in many theoretical [70, 32] and empirical [64] works and perturb the parameters by  $z$  with  $z$  as a Gaussian random variable.
We use  $n = 1$  as the default. It is widely known that the SPSA estimate can be used to replace the backpropagation gradient in any optimizer such as SGD.
Definition 2 (ZO-SGD). ZO-SGD is an optimizer with learning rate  $\eta$  that updates parameters as  $\pmb{\theta}_{t + 1} = \pmb{\theta}_t - \eta \widehat{\nabla}\mathcal{L}(\pmb {\theta};\mathcal{B}_t)$  where  $\mathcal{B}_t$  is the minibatch at time  $t$  and  $\widehat{\nabla}\mathcal{L}$  is the SPSA gradient estimate.
2.1 Memory-efficient ZO-SGD (MeZO)
The vanilla ZO-SGD algorithm costs twice the memory of inference, as it needs to store  $z \in \mathbb{R}^d$ . We propose a memory-efficient implementation of ZO-SGD called MeZO, as illustrated in Algorithm 1. At each step, we first sample a random seed  $s$ , and then for each of  $z$ 's four uses in Algorithm 1, we reset the random number generator by  $s$  and resample the relevant entry of  $z$ . Using this in-place implementation, MeZO has a memory footprint equivalent to the inference memory cost.
We note that Algorithm 1 describes perturbing each parameter separately, which may be time-consuming for large models. In practice, we can save time by perturbing an entire weight matrix instead of each scalar independently. This incurs an additional memory cost as large as the largest weight matrix; usually, this is the word embedding matrix (e.g., 0.86GB for OPT-66B).
Storage Efficiency of MeZO. Parameter-efficient fine-tuning (PEFT) techniques fine-tune just a fraction of the network parameters and have thus been proposed as a way to reduce the storage costs of fine-tuned model checkpoints. Fine-tuning with MeZO reduces the storage cost of the resulting checkpoint far more than popular PEFT techniques (e.g., LoRA [46] and prefix tuning [57]). We reconstruct the MeZO trajectory using a single seed, which spawns step-wise seeds to sample  $z$ , and the projected_grad at each step. As such, for fine-tuning a 66B model, MeZO requires saving the seed plus 20,000 (steps) × 2 bytes, which is less than 0.1MB. LoRA fine-tunes 19M parameters and requires 38MB storage, and prefix tuning fine-tunes 6M parameters and requires 12MB storage.
2.2 MeZO extensions
We note that SPSA is a popular ZO gradient estimator but not the only one. Many one-point gradient estimators have been proposed in past works [34, 87, 95], and using such estimators in place of SPSA would halve the training time. However, cursory experiments with one such promising estimator [113] reveal that these are not as efficient as SPSA when fixing the number of forward passes (Appendix B.5). As such, we implement MeZO with the SPSA estimator.
MeZO can also be combined with other gradient-based optimizers, including SGD with momentum or Adam. Though naive implementation would require additional memory to store the gradient moment estimates, MeZO-momentum and MeZO-Adam alleviate such overhead by recomputing the moving average of the gradients using saved past losses and  $z$  (see Appendix B for a full discussion).
We also note that all of the coordinates of the SPSA gradient estimate have the same scale, but deep Transformers can have gradients of different scales for each layer [59, 61]. As such, we draw inspiration from layerwise adaptive optimizers [109, 110] to design several MeZO variants. Cursory experiments showed that these algorithms are not more efficient (in terms of forward passes), but we nevertheless present them as potential optimizers for more complex objectives. See Appendix B.
Forward Auto-Differentiation Note that  $z^\top \nabla \mathcal{L}(\theta; \mathcal{B})$  is a Jacobian-vector product (JVP), which can be computed in parallel with an inference pass with excess memory consumption equivalent to that of the largest activation in the network [40]. In this case,  $z$  must be stored on the GPU in order to construct the gradient estimate, so this procedure requires slightly more than two times the memory needed for inference. We analyze this algorithm in detail in Appendix D. Note that using a non-zero  $\epsilon$  in SPSA, which is not possible through the JVP method, may boost generalization by promoting a sharpness-minimizing term. Past works (e.g., Baydin et al. [9]) have also studied JVP-based training but achieved limited empirical success.
Note that this reconstruction requires no additional forward passes through the model and no access to the data used during fine-tuning, since projected_grad implicitly encodes this information.
Figure 2: Experiments on RoBERTa-large. We report zero-shot, linear probing (LP), and MeZO and fine-tuning (FT) with full parameter, LoRA, and prefix-tuning. MeZO outperforms zero-shot and LP and approaches FT (within  $5\%$  for  $k = 512$ ) with much less memory. Detailed numbers in Table 18.
Preliminary experiments (Appendix A) show that MeZO only works when using prompts [13, 84, 35]. Past works [83, 67] have demonstrated how the inclusion of a suitable prompt ensures the fine-tuning objective is closely related to the pre-training one. In Section 4, we extend these ideas to show how using a simple prompt simplifies the fine-tuning optimization procedure, thereby enabling zeroth order methods to work efficiently. All experiments below use prompts detailed in Appendix E.2. All fine-tuning with backpropagation (FT) experiments follow convention and use Adam, though we also report results when performing FT with SGD in Appendix F.
We conduct comprehensive experiments on both medium-sized masked LMs (RoBERTa-large, 350M [65]) and large autoregressive LMs (OPT-13B, 30B, 66B [112]) in few-shot and many-shot settings with prompts. We also explore both full-parameter tuning and PEFT including LoRA [46] and prefix-tuning [57] (see Appendix E.5 for details). We compare MeZO with zero-shot, in-context learning (ICL), linear-probing (LP), and fine-tuning with Adam (FT). MeZO uses substantially less memory than FT but requires significantly more training steps.
We first show that MeZO improves substantially over zero-shot, ICL, and LP across model types, sizes, and task types. Moreover, MeZO performs comparably to FT over a number of tasks, while drastically reducing the memory cost by, for example,  $12 \times$  on OPT-13B. Further experiments demonstrate that MeZO can optimize non-differentiable objectives, such as accuracy and F1 score (Section 3.3). We compare the memory consumption of ICL, FT, LP, and MeZO in Figures 3 and 4.
3.1 Medium-sized masked language models
We conduct experiments with RoBERTa-large on sentiment classification, natural language inference, and topic classification tasks. We follow past works [35, 67] in studying the few-shot and many-shot settings, sampling  $k$  examples per class for  $k = 16$  and  $k = 512$  (details in Appendix E). We run MeZO for 100K steps and fine-tuning for 1000 steps, noting that one MeZO step is substantially faster than one fine-tuning step (see Appendix F.6 for a comparison). We summarize the results from Figure 2 and Table 18 below.
MeZO works significantly better than zero-shot, linear probing, and other memory-equivalent methods. On all six diverse tasks, MeZO can optimize the pre-trained model and consistently perform better than zero-shot and linear probing. We also show for several tasks that MeZO can outperform another ZO algorithm, BBTv2 [90], by up to  $11\%$  absolute (Appendix F.4).<sup>4</sup>
With enough data, MeZO achieves comparable performance (up to  $5\%$  gap) to FT. MeZO achieves close-to-fine-tuning performance on  $k = 16$ , with some tasks only having  $2\%$  gaps. When using  $k = 512$  data, the gap between MeZO and FT further reduced to within  $5\%$  across all tasks.
MeZO works well on both full-parameter tuning and PEFT. Full-parameter tuning (MeZO) and PEFT (MeZO with LoRA and prefix-tuning) achieve comparable performance, while MeZO (prefix)
$^{4}$ BBTv2 can only train low-dimensional projected prefixes instead of the full model.
Task
Task type	SST-2	RTE	CB	BoolQ classification	WSC	WIC	MultiRC	COPA - multiple choice	ReCoRD	SQuAD generation
Zero-shot	58.8	59.6	46.4	59.0	38.5	55.0	46.9	80.0	81.2	46.2 14.6
ICL	87.0	62.1	57.1	66.9	39.4	50.5	53.1	87.0	82.5	75.9 29.6
LP	93.4	68.6	67.9	59.3	63.5	60.2	63.5	55.0	27.1	3.7 11.1
MeZO	91.4	66.1	67.9	67.6	63.5	61.1	60.1	88.0	81.7	84.7 30.9
MeZO (LoRA)	89.6	67.9	66.1	73.8	64.4	59.7	61.5	87.0	81.4	83.8 31.4
MeZO (prefix)	90.7	70.8	69.6	73.1	57.7	59.9	63.7	84.0	81.2	84.2 28.9
FT (12x memory)	92.0	70.8	83.9	77.1	63.5	70.1	71.1	79.0	74.1	84.9 31.3
Table 1: Experiments on OPT-13B (with 1000 examples). ICL: in-context learning; LP: linear probing; FT: full fine-tuning with Adam. MeZO outperforms zero-shot, ICL, and LP across the board, and achieves comparable (within  $1\%$  ) or better performance than FT on 7 out of 11 tasks.
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/MeZO (prefix)	90.6	72.6	73.5	63.5	59.1	85.2
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/MeZO (prefix)	93.6	66.4	73.7	63.5	58.9	85.0
Table 2: Experiments on OPT-30B and OPT-66B (with 1000 examples). We report the best of MeZO and MeZO (prefix). See Appendix F.2 for more results. We see that on most tasks MeZO effectively optimizes up to 66B models and outperforms zero-shot and ICL.
sometimes outperforms MeZO. We also show in Appendix F.3 that the three variants converge at similar rates, agreeing with our theory in Section 4, which shows that MeZO converges at a rate independent of the number of parameters being optimized.
We show additional results with more FT and MeZO variants in Appendix F.1. We see that (1) ZO-Adam sometimes outperforms ZO-SGD but is not consistent across tasks; (2) LP and then MeZO, as suggested for fine-tuning [53], can sometimes improve the performance.
3.2 Large autoregressive language models
With the promising results from RoBERTa-large, we extend MeZO to the OPT family [112], on a scale of 13B (Table 1), 30B, and 66B (Table 2). We select both SuperGLUE [98] tasks $^5$  (including classification and multiple-choice) and generation tasks. We randomly sample 1000, 500, and 1000 examples for training, validation, and test, respectively, for each dataset. We run MeZO for 20K steps and fine-tuning for 5 epochs, or 625 steps, noting that each step of MeZO is substantially faster than fine-tuning (see Appendix F.6 for a comparison). Please refer to Appendix E for details. Table 1 yields the following observations.
MeZO outperforms memory-equivalent methods and closely approaches fine-tuning results. We see that on a 13B-parameter scale, MeZO and its PEFT variants outperform zero-shot, ICL, and LP across almost all tasks. When comparing to FT, which costs  $12 \times$  more memory (Section 3.4), MeZO achieves comparable (within  $1\%$  ) or better performance on 7 out of the 11 tasks.
MeZO exhibits strong performance across classification, multiple-choice, and generation tasks. We investigate MeZO on generation tasks, which are regarded as more intricate than classification or multiple-choice tasks. We evaluate on two question answering datasets, SQuAD [80] and DROP [31]. We use teacher forcing for training and greedy decoding for inference (details in Appendix E).
Table 1 shows that, on all generation tasks, MeZO outperforms zero-shot, ICL, and LP, and achieves comparable performance to FT. Considering that many applications of fine-tuning LMs – including instruction tuning or domain adaptation – target generation tasks, our results underscore the potential of MeZO as a memory-efficient technique to optimize large LMs for realistic and exciting applications.
We also include SST-2, which is a simple sentiment classification task that we use for development.
Model	RoBERTa-large (350M)	OPT-13B
Task	SST-2	SST-5	SNLI	TREC	SQuAD
Zero-shot	79.0	35.5	50.2	32.0	46.2
Cross entropy (FT)	93.9	55.9	88.7	97.3	84.2
Cross entropy (MeZO)	93.3	53.2	83.0	94.3	84.7
Accuracy/F1 (MeZO)	92.7	48.9	82.7	68.6	78.5
Figure 3: GPU memory consumption with different OPT models and tuning methods on MultiRC (400 tokens per example on average).
Figure 4: Largest OPT models that one can tune with specific hardwares and algorithms.  $\dagger$  : projected results without actual testing.
Table 3: MeZO with non-differentiable objectives. For classification  $(k = 512)$ , we use MeZO with full-parameter and optimize accuracy; for SQuAD (1,000 examples), we use MeZO (prefix) and F1.
Hardware	Largest OPT that can fit
FT	FT-prefix	Inference
1×A100 (80GB)	2.7B	6.7B	30B
2×A100 (160GB)	6.7B	13B	66B
4×A100 (320GB)	13B	30B	66B
8×A100 (640GB)	30B	66B	\( 175\mathrm{B}^{\dagger} \)
MeZO scales up to 66 billion parameter models. We demonstrate the efficacy of MeZO on even larger models, up to 66B, in Table 2. While directly fine-tuning models at such scales is extremely costly (Section 3.4), MeZO can effectively optimize these models and outperform zero-shot and ICL.
3.3 Training with non-differentiable objectives
We demonstrate the efficacy of MeZO for optimizing non-differentiable objectives through initial experiments. Accuracy and F1 are used as the respective objectives (details in Appendix E.6). Table 3 reveals that MeZO with accuracy/F1 successfully optimizes LMs with superior performance to zero-shot. Although minimizing cross entropy results in stronger performance, these preliminary findings highlight the promising potential of applying MeZO to optimize non-differentiable objectives without clear differentiable surrogates, such as human preferences [73].
3.4 Memory usage and wall-clock time analysis