Fine-Tuning Language Models with Just Forward Passes
Authors: Sadhika Malladi
Tianyu Gao
Eshaan Nichani
Alex Damian
Jason D. Lee
Danqi Chen
Sanjeev Arora
Venue: NeurIPS 2023 oral (2023)
Keywords: language models
fine-tuning
zeroth order optimization
memory efficiency

Abstract: 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 zerothorder 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× memory reduction and up to 2× 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.

Fine-Tuning Language Models with Just Forward Passes
Sadhika Malladì"
Tianyu Gao
Eshaan Nichani
Alex Damian
Jason D. Lee
Danqi Chen
Sanjeev Arora
Princeton University
{smalladi, tianyug, eshnich, ad27, jasonlee, danqic, arora)@princeton.ede
Abstract
Fine-tuning language models (LMs) has yielded success diverse downstream tasks, but as LMs grow in size, backpropagation requires a prohibitively large amount of memory. Zeroth-order (ZO) methods 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 memory memory-efficient efficienn 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 L.ss), model scales (up to 66B), and downstream tasks (classification, multiple-choice, and generation). Our results demonstrate that (1) Me70 significantly outperforms in-context learning and linear probing; (2) MeZO achieves comparable performance to fine-tuning with backpropagation across multiple tasks, with up to 12x memory reduction and up to 2x GPU-hour reduction in our implementation; (3) MeZO is compatible with both full-parameter and parameter-efficien tuning techniques such as LoRA and prefix tuning: (4) MeZO can effectively optimize non- con-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."
Introduction
Fine-tuning pre-traine 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 [13, 72], computing gradients for backpropagation requires a prohibitive amount of memory
our test, up to 12x the memory required for inference because needs to cache activations during the forward pass, gradients during the backward pass, and, in the case Adam [52], also store gradient history (see Section 3.4 for a detailed analysis).
As result, while it is possible to run inference with 30-billion (30B) parameter LM on a single Nvidia A.00 GPU (with 80GB memory), fackpropagation with Adam is feasible only for a 2.7B LM. crrameter-efficien fine-tuning methods (PEFT [46, 57, 54]) update just fraction of the network
Equual contribution and corresponding authors.
Our code is available https: //githubb hrtttuub.oom/prrcccttooonn/////e
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Accuracy/F1 (%%
70
30
SST-2 RTE
BoolQ
MultiRC
Copa
ReCoRD
SQUAD
shot
[12x memory)
Figure 1: OРT-13B results with zero-shot, in-context learning (ICL), MeZO (we report the best among MeZO/MeZO (LoRAYMeZO (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 for detailed numbers and Figure for memory profiling.
parameters, but still need to cache many activations, because the tuned parameters are scattered throughout the model. our tests, fine-tuning an OPT-13B model with full parameter tuning or PEFT requires 12x and 6x more memory than inference respectively.
In-context learning (ICL [13]) has allowed solving many tasks with single inference pass, during which the model processes labeled examples (demonstrations) in its context and then outputs prediction on 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 sensitive 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]).
classical zeroth-order optimization method, ZO-SGD [88], uses only differences of loss values to estimate the gradients. Thus, in principle, the method 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 directly optimize large-scale models (see Liu al. [61] for survey).
update neural networks with just
In this work, we propose a memory-efficient zeroth-order optimizer (Me70), which adapts the classical ZO-SGD algorithm and reduces its memory consumption to the same 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 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). Me70 consistently outperforms zero-shot, ICI. and linear probing. Moreover, with -araae, 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 12x less memory (Figure and Section 3). In our implemen- tation, MeZO requires only half many GPU-hours Adam fine-tuning for a 30B model (see Appendix F.6).
3. We demonstrate MeZO's compatibility with full parammerret 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 Fl 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 Me70 depend on certain condition number of the landscape (i.c., the local effective rank, see Assumption 1) instead of numbers of parameters. This

---

Algorithm 1: MeZO
Require:
parameters A E Rd, loss IRd + R, step budget T, perturbation scale e, batch size
B, learning rate schedule {ne}
for i 1,.... do
Sample batch B C D and random seed 8
< ParturoParameters(O....
__ + L(0: B)
< PerturoParrmeters((.
8- <- L(0:B)
<- PerturbPPrrmeeerss((.....1
Reset parameters before descent
pro projected_gre _ereed <- (4+ - (-)
Reset random number generator with seed for 0g E f do
For sampling
2 ~~~ (0,
0, "tt * projected_grad
end
end
Subroutine PerturbParameters! (0, s)
Reset random number generator with seed for 0g 0 do 2 ~ (0,
For sampling z
8
Modify parameters in place
end return θ
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 (7O) optimizers have long been studied the context of convex and strongly convex objectives. In the following, we first introduce classical 70 gradient estimator, SPSA (Defini- tion [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 I. We highlight that SPSA can also be used in more complex optimizers, such as Adam, and we provide memory- implementations for those algorithms too (Section 2.2).
Consider labelled dataset {( (e)tae)pp)) and a minibatch C. D of size B, we let L(O: B) denote the loss on the minibatch. We introduce classical ZO gradient estimate in this setting.
Definition (Simultaneous Perturbation Stochastic Approximation or SPSA [88]). Given a model with parameters f E IRd and a loss function C, SPSA estimates the gradient minibatch B as
CC B) C(O + ez; B) L(0 <2; B) R= zz (L(;
(1)
2e
where z Rdd with z NO Ia) and e the perturbation scale. The n-SPSA gradient estimate averages VL(; B) over n randomly sampled z.
SPSA requires only two forward passes through the model compute the gradient estimate (for n-SPSA, each estimate requires 2n forward passes). As - 0, the SPSA estimate can be understood as rank-1 reconstruction of the gradient. During training, n can treated hyperparameter and follow schedule [11, 15), though in cursory experiments (Appendix A),
is the most efficient.
Thhe original SPSA algorithm [88] perturbs the model by 1/x thus requires that has finite inverse moments, precluding the choice of z as Gaussian. 1/z is very large with high probability for a zero-me zero-mean mean Gaussian so we adopt the standard in many theoretical [70, 32) and empirical [64] works and perturb the parameters by with Gaussian random variable.

---

We use n = 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 that updates parameters as 01+1 0, 102-77 (1 B.) where B, is the minibatch at time t and V 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 2 E R We propose a memory-efficien implementation of ZO-SGD called MeZO, as illustrated Algorithm 1. At each step, we first sample a random seed and then for each of z's four uses in Algorithm I, we reset the random number generator by and resample the relevant entry of z. Using this in-place implementation, Me70 has memory footprint equivalent to the inference memory cost.
We note that Algorithm 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 Me7O. Parameter-efficien fine-tuning (PEFT) techniques fine-tune just 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 single seed, which spawns step-wind wise seeds to sample and the projected _rraddddddd at each step. As such, for tin-ttuning& a 66B model, MeZO requires saving the seed plus 20,000 (steps) x 2 bytes, which is less than I.I... 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 a popular 70 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 baasee optimizers, including SGD with momentum or Adam. Though naive implementation would require additional memory to store the gradient moment estimates, MeZO-momentee 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- into-Differentiation Note that = T V ((;; B3) is Jacobian-vector product (JVP), which can be computed 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 non-zero e in SPSA, which is not possible through the JVP method, may boost generalization by promoting a sharpness miniimiingn 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-tum fine-tuning, since projected_gree implicitly encodes this information.

---

RoBERTa-large
k=512 RoBERTa-large
100
100
90
90
Accuracy (%)
70
70
60
60
40
40
30
30
SST-2
SST-5
SNLI
MNLI
RTE
TREC
SST-2
SST-5
SNLI
MNLI
RTE
TREC
Zero-shot
LP
MezO
Mezo (LoRA)
MezO (prefix)
FT
FT (LoRA)
FT (prefix)
Figure 2: Experiments on RoBERTallarge. 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.
3
Experiments
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 suitable prompt ensures the fine-tuning objective is closely related to the pre-training one. Section 4, we extend these ideas show how using simple prompt simplifies the fine-tuning optimization procedure. thereby enabling zeroth order methods to work efficiently. All experiments below use prompts detailed in Appendix F.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
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
[57] (see Appendix E.5 for details). We compare MeZO with zero shot, in-context learning (ICL). linear peearr-roobig (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, ICI., and IP across model types, sizes, and task types. Moreover, MeZ0 performs comparably FT over number of tasks, while drastically reducing the memory cost by, for example, 12x on ОPT-13B. Further experiments demonstrate that MeZO can optimize non-differentiable objectives, such as accuracy and Fl score (Section 3.3). We compare the memory consumption of ICL, FT, LP, and MeZO in Figures 3 and 4.
3.1 Medium-sizer masked language models
We conduct experiments with RoBERTa-large 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 & 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 comparison). We summarize the results from Figure and Table 18 below.
MeZO works significantly better than zero-shot, linear probing, and other memory-equivalee methods. On all six diverse tasks, McZO 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).
With enough data, MeZO achieves comparable performance (up to 5% gap) to FT. MeZO achieves close-to- coosse-t-fine-tumiiin finn-tuniing performance on 16, with some tasks only having 2% gaps. When using k = 512 data, the gap between McZO 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)
'BBTv2 can only train low -iwdeeensional projected prefixes instead of the full model.
5

---

Task
SST-2
RTE
CB
BoolQ
WSC
WIC
MultiRC
COPA
ReCoRD -multiple choice
SQuAD DROP
Task type
classification
generation
sero-shot
58.8
59.6
46.4
59.0
38.5
55.0
46.9
80.0
81.2
46.2
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
MeZ0
91.4
66.1
67.9
67.6
63.5
61.1
60.1
88.0
81.7
84.7
30,9
McZO (LoRA)
89.6
67.9
66.1
73.8
64.4
59.7
61.5
87.0
81.4
83.8
31.4
McZ0 (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 I: 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 McZO/Mc20 (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/MeO (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 McZO (prefix), See Appendix F.2 for more results. We see that on most tasks McZO 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 McZO 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 scale of 13B (Table 1), 30B, and 66B (Table 2). We select both SuperGLUE [98] tasks (including classification and multiple-choice) and generation tasks. We randomly sample 1000, 500, and 1000 examples for training, validation, and test, respectively, for each datset. We run MeZO for 20K steps and fine-tuming for 5 epochs, or 625 steps, noting that each step of McZO is substantially faster than fine-tuning (see Appendix F.6 for 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 13B-parameter scale, MeZO and its PEFT variants outperform zero-shot, ICL, and LP across almost all tasks. When comparing to FT, which costs 12x 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, 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 nemory-efficient technique to optimize large LMs for realistic and exciting applications.
We also include SST-2, which is simple sentiment classification task that for development.

---

Model
RoBERTa-Iang (350M)
OPT-13B
Task
SST-2
SST-5
SNLI
TREC
SQuAD
Zero-shot
79.0
35.5
32,0
46.2
Cross entropy (FT)
93.9
97.3
84.2
Cross entropy (Me20)
93.3
83.0
94.3
84.7
Accuracy/F1 (Mc20)
92.7
48.9
82.7
68.6
78.5
Table 3: McZO with con-differentiable objectives. For classification (k 512), McZO with full-paree parammeeere and optimize accuracy; for SQuAD (1,000 examples), we use McZO (prefix) and FI.
1000
Zara-shat
ICL
GPU Memory (GB)
100
FT (pref(i)
Largest OPT that can fit
Hardware
FT
FT-prrfix
Inference
10
1XA100 (80GB)
2.7B
6.7B
30B
Ax (160GB)
6.7B
13B
66B
4xA100 (320GB)
138
30B
66B
8xA100 (640GB)
30B
66B
175B
2.78
6.78
Parrmmeters
Figure 3: GPU memory consumption with different Figure 4: Largest OPT models that one can OPT models and tuning methods on MultiRC (400 tune with specific hardwares and algorithms. kens per example on average).
t projected results without actual testing.
MeZO scales up to 66 billion parameter models. We demonstrate the efficacy of Me70 on even larger models, up to 66B, Table 2. While directly fine-tuning models at such scales extremely costly (Section 3.4), MeZO can effectively optimize these models and outperform zero-shot and ICL.
3.3 Training with non-diffeeretiiable objectives
We demonstrate the efficacy of MeZO for optimizing cinfeifferentiabee objectives through initial experiments. Accuracy and F1 are used as the respective objectives (details in Appendix E.6). Table reveals that Me70 with accuracy/F1 successfully optimizes L.Ms 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
In this section profile the memory usage of zero-shot, ICL, FT, FT (prefix), and MeZO. We test OPT models of various sizes with Nvidia A A100 GPUs (80GB memory) on MultiRC (average #tokens=400), and report the peak GPU memory consumption (details in Appendix E.7).
As shown in Figure (refer to Appendix F.5 for detailed numbers), MeZO exhibits the same memory consumption as zero-shot while offering memory savings of up to 12 times compared to standard FT and 6 times compared to FT (prefix). This advantage enables training larger models within a fixed hardware budget, illustrated in Figure 4. Specifically, using single A100 GPU, MeZO allows for tuning a model that is 11 times larger than what is feasible with FT.
In Appendix F6,, we compare the wall-clock time efficiencies of our implementations of MeZ0 and Adam fine-tuning. McZO achieves 7.74x per-step speedup and requires 8x fewer GPUs with 30B model, but takes more steps to converge. Overall, McZO only requires half as many GPU-hours to fine-tune a 30B model compared to full-parameter fine-tuning. The wall-clock benefit MeZO is not inherent to the algorithm and is highly dependent the implementation. We primarily provide this information as a demonstration that MeZO does not take a prohibitively long time to run.
The above measurements are dependent on the computing infrastructure. In Appendix C, we compare the theoretical time-memory tradeoff MeZO and backpropagation and find that MeZO always more memory-efficien than backpropagation and is often more time-efficient. The above analyses also do not consider recent advances (e.g., gradient checkpointing [18], FlashAttention [23], and quantization [27]). We leave investigating the how MeZO works with these methods to future work.

---

Thus, in the usual case where n < d, McZO has much larger gradient norm than SGD. The descent lemma also shows that to guarantee loss decrease, one needs to choose the learning rate as 4 Theory
Our theoretical analysis highlights why MeZO can optimize large LMs, although a number of classical results [69, 47, 79, 3, 70] suggest that optimization should be catastrophically slow when training so many parameters. The inclusion of simple prompt is crucial for MeZO to succeed (Appendix A). Past works [83, 67] have suggested that including such prompt ensures that the fine-tuning objective closely related to the pre-training one. As such, here, we make the assumption that the model has already been trained for many steps on the fine-tuning objective, which implies that the loss landscape exhibits favorable conditions (Assumption 1). Then, we derive a convergence rate independent of the number of parameters. We show that the loss decreases per step at a rate independent of the parameter dimension d (Theorem 1), and that, under stronger conditions, the algorithm converges in time independent of d (Lemma 3). Together, these results imply that McZO is not catastrophically slower than SGD when fine-tuning." For ease of illustration, we assume that z is sampled from sphere with radius vd, and in Appendix G.2, we derive the rate for a general Gaussian x, which was used in the experiments.
We follow classical analyses of SGD and replace the minibatch gradient estimate with SPSA (Defi- nition 1). Consider the minibatch SGD update oc+1 t 0, L(V B.) where B, is minibatch drawn uniformly from DB, Crucially, the SGD minibatch gradient estimate unbiased.
Definition 3 (Unbiased Gradient Estimate). Any minibatch gradient estimate g B) is said to be
unbiased if (lg((, B)] = VZ(())
4.1 Per-step analysis
The classical descent lemma uses Taylor expansion to study how SGD reduces the loss at each optimization step. It highlights that when the gradient covariance is large, the maximum possible decrease in loss at each optimization step is small, thereby resulting in slower optimization.
Lemma (Descent Lemma). Let ((0) be e-smooth. ? For any unbiased gradient estimate g(0, B)),
EE(((0+1) 0.] C(O.) S -IVCC(0) 1 znde (0,,)(,
(2)
The descent lemma highlights the importance of the gradient norm, which we derive for MeZO below. Lemma 2. Let B be a random minibatch of size B. Then, the gradient norm of MeZO is
Ez ||
d + n 1 E [IV
n
where n is the number of sampled in n-SPSA (Definition 1) and d is the number of parameters.
[(.,))
Lemma 2
n S L [[llg 0. 131
TZO
NSGD
d +
n
where 7200 and 7/SD are the maximum permissible learning rates for MeZO and SGD respectively. Thus see that without any further assumptions, MeZO can slow optimization by decreasing the largest permissible learning rate by a factor d. Moreover, MeZO reduces the loss decrease that can be obtained at each step and, consequence, slows convergence by factor of d as well.
Surprisingly, our experiments show that MeZO can quickly optimize pre-trained models with billions of parameters, and reducing the number of tuned parameters via PEFT techniques does not substan- tially accelerate optimization (Appendix F.3). We attribute these phenomena to the Hessian of the loss exhibiting small local effective rank. It is prohibitively expensive to directly measure the effective rank of the Hessian of large IM on a reasonably sized dataset. However, many previous works have shown that the Hessian of the loss for deep neural networks trained by SGD has remarkably low
SSection 3 uses the standard choice of Adam for FT; provide SGD experiments in Appendix F.I. "This satisfied for the standard cross-entropy objective.
Alll of our experiments n 1.
8

---

effective rank [74, 75, 36, 107, 105, 82]. In particular, the bulk of the spectrum concentrates around with only a small number of outliers, and the number of these outliers is an upper bound on the effective rank. In addition, prior works [4, 56] have demonstrated that LM fine-tuning can occur in very low dimensional subspace (< 200 parameters), which further supports the below assumption. We formalize the assumption the effective rank below. In particular, we require an upper bound on the Hessian a neighborhood around the current iterate to have effective rank at most 7.
Assumption 1 (Local r-effective rank). Let G(O,) (D VC (ar,y))) There exists a matrix H (0:) 3 E La such that:
1. For all 0 such that ||0 Oell S ndG(())) we have V2 3 (
2. The effective rank of H (O),, ie tr( tr(H H (0.) llep is at most r.
Under this assumption, we show that the convergence rate of ZO-SGD does not depend on the number of parameters. Instead, the slowdown factor only depends on the effective rank of the Hessian.
Theorem 1 (Dimension- Free Rate). Assume the loss exhibits local r-effective rank (Assumption 1). Pct11 0. - Tizo v (e,;; B) a single step of ZO-SGD using the n-SPSA estimate with a minibatch of size B, then there exists (r (r) such that the expected loss decrease can be bounded
B(C(0(+1) 0,] ((e)) v nzo /C())) + niol (
(4)
By applying Equation (3), we can directly compare the SGD descent lemma.
Corollary 1. Choosing the learning rate Yzzo = 7- VIGG, ZO-SGD obtains a loss decrease of
1
B[C(0(41) 0.1 - ((O,) - TSGD |VC(0) + 2 nant E[I
(5)
Here that comparing SGD, the slowdown factor of ZO-SGD scales with the local effective rank r, which we argue is much smaller than the number of parameters d. The above analysis focuses on how much ZO-SGD and SGD decrease the loss at each step. Below, we show that under stronger assumptions about the loss landscape, we can obtain rates for how quickly the ZO-SGD algorithm converges an optimal value.
4.2 Global convergence analysis
We show that the global convergence rate also slows by factor proportional to the local effective rank under stronger assumptions about the loss landscape. We assume that the landscape obeys the classical PL inequality: the gradient norm grows quadratically with the suboptimality the iterate.
Definition (PL Inequality). Let C = ming L(O). The loss μ-PL if. for all 0, || 2 ((((0)~L))
The PL. inequality is not strong assuming that optimization exhibits kernel-like dynamics, but it ensures that the landscape is amenable to analysis [50]. In addition to the PL inequality, we assume the trace of the gradient covariance bounded, so noise does disrupt the trajectory drastically, Definition 5 (Gradient Covariance). The SGD gradient estimate on a minibatch of size has covariance Σ(0) (E [V VT(00B)VLLC () B) v (L()))((()) T).
As show Appendix G.I, this assumption holds for common loss functions such as square loss or binary cross entropy for several settings (e.g., kernel behavior [67]). With these two assumptions, we show that ZO-SGD has slowdown proportional to the effective rank r, not the parameter dimension. Lemma 3 (Global Convergence of ZO-SGD). Let L(0) be μ-PL and let there exist a such that (r(E()) < a(C(O) L') for all Then after
(" +1) 11
la log C(Og) L
t t= 0
p +
12B
SGD rate (Lemma 4)
iterations of ZO-SGD we have [[[((,)) < L' +e.

---

5
Related work
Zeroth-order optimization Many classical lower bounds have been derived for ZO-SGD in the
strongly convex and convex settings [47, 3, 79, 32, 85, 69] as well as non-convex [101]. These bounds generally depended on the number of parameters d. More recently, [100, 6, 15] showed that if the gradient has low-dimensional structure, then the query complexity scales linearly with the intrinsic dimension and logarithmically with the number of parameters, though the estimation has at least (sdd log d) memory cost. Additional tricks such as sampling schedules [11] and other variance reduction methods [48, 62] can be added to ZO-SGD. ZO has inspired distributed methods [93, 43] and box adversarial example generation [14, 63, 17, 64] in deep learning. Ye et al. [108], Balasubramanian and Ghadimi [7] estimate the Hessian to perform ZO optimization along important directions. There are also ZO methods that optimize without estimating the gradient [38, 68, 44].
Memory- backpropagation Several algorithms have been proposed to efficiently approx- imate backpropagation by sparsifying gradients [92, 102], approximating Jacobians [1, 19]. and subsampling the computational graph [71, 2]. However, these methods may accrue large approxima- tion errors for deep networks. Gradient checkpointing [18] reduces memory cost of backpropagation at the cost of recomputing some activations. FlashAttention [23] also reduces memory cost by recomputing attention matrices. Dettmers et al. [26, 27] explore quantization of large LMs' weights and optimizer states, which leads to memory reduction in both training and inference.
Gradient-free adaptation of large language models BBT and BBTv2 [91. 90] use evolutionary algorithms to achieve gradient-free optimization; however, due to its sensitivity to high dimensionality, BBT is limited only optimize low-dimension projection of prefixes and they focus on RoBERTa- large size models and few-shot settings. Other works in "black-box tuning" of LMs focus on optimizing discrete prompts without updating the model, either via reinforcement learning [16, 25, 29]. ensemble [45], or iterative search [78]. Concurrent work in [106] uses iterative forward passes to improve in-context learning performance.
Conclusion
We have shown that MeZO can effectively optimize large LMs across many tasks and scales. Further experiments suggest that MeZO can optimize non-differentiable objectives, which backpropagation usually cannot do. Our theory illustrates why MeZO is not catastrophically slow when tuning billions of parameters. As limitation, MeZO takes many steps order to achieve strong performance, though we show that the per-step speedup in MeZO can often make fine-tuning with MeZ0 run faster than a standard implementation of mue-tuning with backpropagation. We did not explore combining McZO with other memory- efficiinnt methods, such as FlashAttention [23] and quantization [26], though we hope to investigate this in the future.
We are excited to explore the applicability of MeZO to number of areas, including but not limited to: pruning, distillation, saliency, interpretability, and dataset selection for fine-tuning. Non-differentiable objectives are particularly exciting area, given recent advances in tuning large LMs to adapt to human feedback. Conducting theoretical analyses for how these efficient gradient estimates impact the performance of different applications is also of interest.
Acknowledgements
We thank Xinyi Chen, Yin Tat Lee, Kaifeng Lyu, Tengyu Ma, Abhishek Panigrahi, Nikunj Saunshi, and Mengzhou Xia for their helpful feedback. SA and SM are funded by NSR, ONR, SRC, and Simons Foundation. JDL., AD, and EN acknowledge the support of the ARO under MURI Award W911NFF-111-0304 the Sloan Research Fellowship. NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994. TG is supported by an IBM PhD Fellowship. This work is also partially funded by the National Science Foundation (IS7-211777).
10

---

References
[1] Hany Abdel-Khalik, Paul D Hovland, Andrew Lyons, Tracy Stover, and Jean Utke. A low rank approach to automatic differentiation. In Advances in Automatic Differentiation, pages 55-65, 2008.
[2] Menachem Adelman, Kfir Levy, Ido Hakimi, and Mark Silberstein. Faster neural network training with approximate tensor operations. Advances in Neural Information Processing Systems, 34:27877-27889, 2021.
[3] Alekh Agarwal, Peter L. Bartleu, Pradeep Ravikumar, and Martin J. Wainwright. Information- theoretic lower bounds on the oracle complexity of stochastic convex optimization. IEEE Transactions on Information Theory, 58(5):3235-3244 May 2012.
[4] Armen Aghajanyan, Sonal Gupta, and Luke Zeuteemyer. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers), pages 7319-7328, 2021.
[5] Stephen H Bach, Victor Sanh, Zheng-Xing Min Yong, Albert Webson, Colin Raffel, Nihal V Nayak, Abheesht Sharma, Taewoon Kim, M Saiful Bari, Thibault Fevry, et al. Promptsource: An integrated development environment and repository for natural language prompts. arXiv preprint arxiv:2202.01279, 2022.
[6] Krishnakumar Balasubramanian and Saeed Ghadimi. Zeroth-order (non)-conve stochastic optimization via conditional gradient and gradient updates. In Advances in Neural Information Processing Systems, volume 31, 2018.
[7] Krishnakumar Balasubramanian and Saeed Ghadimi. Zeroth-order nonconvex stochastic optimization: Handling constraints, high dimensionality, and saddle points. Foundations of Computational Mathematics, pages 1-42, 2022.
[8] Roy Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Szpektor. The second PASCAL recognising textual entailment challenge. Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment, 2006.
[9] Atilmm Güneş Baydin, Barak A. Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradi- without backpropagation, 2022.
[10] Luisa Bentivogli, Peter Clark, Ido Dagan, and Danilo Giampiccolo. The fifth PASCAL recognizing textual entailment challenge. In TAC, 2009.
[11] Raghu Bollapragada, Richard Byrd, and Jorge Nocedal. Adaptive sampling strategies for stochastic optimization. SIAM Journal on Optimization, 28(4):3312-334 2018.
[12] Samuel R. Bowman, Gabor Angeli, Christopher Potts, and Christopher D. Manning. A large annotated corpus for learning natural language inference. In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing, pages 632-642, 2015.
[13] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in neural information processing systems, vol- ume 33, pages 1877-1901, 2020.
[14] HanQin Cai, Yuchen Lou, Daniel McKenzie, and Wotao Yin. A zeroth-order block coordinate descent algorithm for huge- black-box box optimization. In International Conference on Muchine Learning, pages 193-1203, 2021.
[15] HanQin Cai, Daniel McKenzie, Wotao Yin, and Zhenliang Zhang. Zeroth-order regularized optimization (zoro): Approximately sparse gradients and adaptive sampling. SIAM Journal on Optimization, (3(2):687-714 2022.
[16] Yekun Chai, Shuohuan Wang, Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. Clip- tuning: Towards derivative-free prompt learning with a mixture of rewards. arXiv preprint arXiv:2210.12550 2022.

---

[17] Pin- Yu Chen, Huan Zhang. Yash Sharma, Jinfeng Yi, and Cho-Jui Hsich. Zoo: Zeroth order optimization based black-box attacks to deep neural networks without training substitute models. In Proceedings of the 10th ACM workshop artificial intelligence and security, pages 15-26, 2017.
[18] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost, arXiv preprint arXiv:1604.065 176600006744 2016.
[19] Krzysztof M Choromanski and Vikas Sindhwani. On blackbox backpropagation and jacobian sensing, In Advances Neural Information Processing Systems, volume 30, 2017.
[20] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Advances in neural information processing systems, volume 30, 2017.
[21] Christopher Clark, Kenton Lee, Ming-Wer Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long and Short Papers), pages 2924-2936, 2019.
[22] Ido Dagan, Oren Glickman, and Bernardo Magnini. The PASCAL recognising textual en- tailment challenge. In the First International Conference on Machine Learning Challenges: Evaluating Predictive Uncertainty Visual Object Classification, and Recognizing Textual Entailment, 2005.
[23] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Advances in Neural Information Processing Systems, volume 35, pages 16344-16359, 2022.
[24] Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. The commitmenthank: Investigating projection in naturally occurring discourse. In Sinn und Bedeutung, volume 23, pages 107-124, 2019.
[25] Mingkai Deng. Jianyu Wang, Cheng-Ping Hsich, Yihan Wang, Han Guo, Tianmin Shu, Meng Song, Eric Xing, and Zhiting Hu. RLPrompt: Optimizing discrete text prompts with reinforcement learning. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 3369-3391, 2022.
[26] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. GPT3.int80) 8-bit matrix multiplication for transformers at scale. In Advances in Neural Information Processing Systems, 2022.
[27] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block- wise quantization. In International Conference on Learning Representations, 2022.
[28] Jacob Devlin, Ming-Wer Beig-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171 1171-4186, 2019.
[29] Shizhe Diao, Xuechun Li, Yong Lin, Zhichao Huang, and Tong Zhang. Black-box prompt learning for pre-trained language models. arXiv preprint urXiv:2201.0853/ 2022.
[30] Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et al. Delta tuning: A comprehensive study of parameter efficient methods for pre-trained language models. arXiv preprint arXiv:2203.06904 2022.
[31] Dheeru Dua, Yizhong Wang, Pradeep Dasigi. Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proceedings of the 2019 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 2368-2378, 2019.
12

---

[32] John C. Duchi, Michael I. Jordan, Martin J. Wainwright, and Andre Wibisono. Optimal rates for zero-order convex optimization: The power of two function evaluations. IEEE Transactions Information Thheory 61 (5):2788-2888 2015.
[33] FairScale authors. Fairscale: A general purpose modular pytorch library for high performance and large scale training, 2021.
[34] Abraham D. Flaxman, Adam Tauman Kalai, and H. Brendan McMahan. Online convex optimization in the bandit setting: Gradient descent without a gradient. In Proceedings of the Sixteenth Annual ACM-SIAM Symposium on Discrete Algorithms, SODA '05, page 385-394, USA, 2005. Society for Industrial and Applied Mathematics. ISBN 08987158577
[35] Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better few- shot learners. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Pupers), pages 3816-3830, 2021.
[36] Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pages 2232-22442 2221, 2019.
[37] Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The third PASCAL recognizing textual entailment challenge. In the ACL-PASCAL Workshop on Textual Entailment and Paraphrasing, 2007.
[38] Daniel Golovin, John Karro, Greg Kochanski, Chansoo Lee, Xingyou Song, and Qiuyi Zhang. Gradientless descent: High-dimensiona zeroth-order optimization. In International Conference on Learning Representations, 2020.
[39] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint /rXiv:/706.0020 2017.
[40] Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.
[41] José Grimm, Loic Pottier, and Nicole time product for reversing a certain class of programs. PhD thesis, INRIA, 1996.
Optimal time and minimum space-
[42] Suchin Gururangan, Ana Marasovié, Swabha Swayamdipta, Kyle Lo, Beltagy, Doug Downey, and Noah A. Smith. Don't stop pretraining: Adapt language models to domains and tasks. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 0342-8360, 2020.
[43] Davood Hajinezhad and Michael M 7avlanos. Gradient-free multi-agent nonconvex nonsmooth optimization. In 2018 IEEE Conference on Decision and Control (CDC), pages 4944, 2018.
[44] Geoffrey Hinton. The forward-forwan algorithm: Some preliminary investigations. arXiv preprint arXiv:2212.13345 2022.
[45] Bairu Hou, Joe O'Connor, Jacob Andreas, Shiyu Chang, and Yang Zhang. Promptboosting: Black-box text classification with ten forward passes. arXiv preprint arX1v:22/2.09257, 2022.
[46] Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
[47] Kevin G Jamieson, Robert Nowak, and Ben Recht. Query complexity of derivative-free optimization. In Advances in Neural Information Processing Systems, volume 25, 2012.
[48] Kaiyi Ji, Zhe Wang, Yi Zhou, and Yingbin Liang. Improved zeroth-order variance reduced algorithms and analysis for nonconvex optimization. In International conference on machine learning, pages 3100-3109, 2019.
13

---

[49] Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive vari- ance reduction. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26, 2013.
[50] Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-lojasiewiez condition, 2020.
[51]
Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Look- ing beyond the surface: challenge set for reading comprehension over multiple sentences. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Papers), pages 252-262, 2018
[52] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[53] Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine- tuning can distort pretrained features and underperform out- out-of-distribution In International Conference on Learning Representations, 2022.
[54] Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 202/ Conference on Empirical Methods in Natural Language Processing, pages 3045-3059, 2021.
[55] Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012.
[56] Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations, 2018.
[57] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 4582 4597, 2021.
[58] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). In A. Beygelzimer. Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[59] Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar. Robust training of neural networks using scale invariant architectures. In International Conference on Muchine Learning, pages 12656-12684, 2022.
[60] Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for GPT-3? In Proceedings of Deep Learning Inside Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pages 140-114,70 2022.
[61] Liyuan Liu, Xiaodong Liu, Jianfeng Gao. Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 5747-5763, 2020.
[62] Sijia Liu, Bhavya Kailkhura, Pin- YinYu Chen, Paishun Ting. Shiyu Chang, and Lisa Amini. Zeroth-order stochastic variance reduction for nonconvex optimization. In Advances in Neural Information Processing Systems, volume 31, 2018
[63] Sijia Liu, Pin-Yu Chen, Xiangyi Chen, and Mingyi Hong. signSGD via zeroth-order oracle. In International Conference on Learning Representations, 2019.
[64] Sijia Liu, Pin- Yu Chen, Bhavya Kailkhura, Gaoyuan Zhang, Alfred Hero III, and Pramod Varshney. A primer on zeroth-order optimization in signal processing and machine learning: Principals, recent advances, and applications. IEEE Signal Processing Magazine, 33355443 2020.
14

---

[65] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy. Mike Lewis, Luke Zeytlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint arXiv: 1907. 1692, 2019.
[66] Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeling of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8086-8098, 2022.
[67] Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. kernel-based view of language model fine-tuning, arXiv preprint arXiv:2210.0566 2022.
[68] Horia Mania, Aurelia Guy, and Benjamin Recht. Simple random search of static linear policies is competitive for reinforcement learning. In Advances in Neural Information Processing Systems, volume 31, 2018.
[69] Arkadij Semenović Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
[70] Yurii Nesterov and Vladimir Spokoiny. Random gradient-free minimization of convex func- tions. Foundations of Computational Mathematicc, 17:527-566, 2017.
[71] Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P Adams. Randomized automatic differentiation. arXiv preprint arXiv:2007. 2020.
[72] OpenAL. Gpt-4 technical report. arXiv preprint arXiv:2303.0877/ 2023.
[73] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730-27744, 2022.
[74] Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint / 2018.. 2018.
[75] Vardan Papyan. Traces of class/cross-class structure pervade deep learning spectra. Journal of Machine Learning Research, 21(25)):1-64, 2020.
[76] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-8035. 2019.
[77] Mohammad Taher Pilehvar and Jose Camacho-Collados. WiC: the word-in-context dataset for evaluating context-sensitive meaning representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 1267-1273, 2019.
[78] Archiki Prasad, Peter Hase, Xiang Zhou, and Mohit Bansal. Grips: Gradient-free edit-based instruction search for prompting large language models. arXiv preprint w.XV2:2203.07281, 2022.
[79] Maxim Raginsky and Alexander Rakhlin. Information-based complexity, feedback and dynamics in convex programming. IEEE Transactions on Information Theory, 57(10):7036- 7056, 2011.
[80] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2383-2392, 2016.
[81] Melissa Roemmele, Cosmin Adrian Bejan, and Andrew Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. 2011.
15

---

[82] Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis the hessian of over- wer-parametrized neural networks. arXiv preprint orXiv:1706.0444 04454, 2017.
[83] Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora. A mathematical exploration of why language models help solve downstream tasks. In International Conference on Learning Representations, 2021.
[84] Timo Schick and Hinrich Schütze. Exploiting cloze-questions for few-shot text classification and natural language inference. In Proceedings of the 16th Conference of the European Chapter of the Association for Computational Linguistics: Main Volume, pages 255-269, 2021.
[85] Ohad Shamir. An optimal algorithm for bandit and zero-order convex optimization with two-point feedback. The Journal of Machine Learning Research, 18(1):1703-177 1713, 2017.
[86] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, 2013.
[87] James C Spall. A one measurement form of simultancous perturbation stochastic approxima- tion. Automatica, (3(1):109-112, 1997.
[88] J.C. Spall. Multivariate stochastic approximation using a simultancous perturbation gradient approximation. IEEE Transactions on Automatic Control, 3733):332-341, 1992.
[89] Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano. Learning to summarize with human feedback. In Advances in Neural Information Processing Systems, volume 33, pages 3008-3021, 2020.
[90] Tianxiang Sun, Zhengfu He, Hong Qian, Yunhua Zhou, Xuanjing Huang, and Xipeng Qiu. BBTv2: Towards gradient-free future with large language models. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 39 9916-3930, 2022.
[91] Tianxiang Sun, Yunfan Shao, Hong Qian, Xuanjing Huang, and Xipeng Qiu. Black-box tuning for megguage-moeel. sell-assa--evviie a-service. In International Conference on Machine Learning, pages 20841-200 20855, 2022.
[92] Xu Sun, Xuancheng Ren, Shuming Ma, and Houfeng Wang. meProp: Sparsified back propagation for accelerated deep learning with reduced overfitting. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pages 3299-3308 2017.
[93] Yujie Tang and Na Li. Distributed zero-order algorithms for nonconvex multi-agent optimiza- tion. In 2019 57th Annual Allerton Conference on Communication, Control, and Computing (Allerton), pages 781 781-786, 2019.
[94] Zhiwei Tang, Dmitry Rybin, and Esung-Hui Chang. Zeroth-order optimization meets human feedback: Provable learning via ranking oracles, 2023.
[95] Alexander Timurovich Vakhitov, Oleg Nikolaevich Granichin, and LS Gurevich. Algorithm for stochastic approximation with trial input perturbation in the nonstationary problem of optimization. Automation and Remote Control, 70: 70:1827-1835 2009.
[96] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, volume 30, 2017.
[97] Ellen M Voorhees and Dawn M Tice. Building a question answering test collection. In the 23rd annual international ACM SIGIR conference on Research and development information retrieval, 2000.
16

---

[98] Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: stickier benchmark for general-purpose language understanding systems. In Advances in neural information processing systems, volume 32, 2019.
[99] Chong Wang, Xi Chen, Alexander J Smola, and Eric P Xing. Variance reduction for stochastic gradient optimization. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
[100] Yining Wang. Simon Du, Sivaraman Balakrishnan, and Aarti Singh. Stochastic zeroth-order optimization in high dimensions. In Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, volume 84, pages 1356-1365. 2018
[101] Zhongruo Wang. Krishnakumar Balasubramanian. Shiqian Ma, and Meisam Razaviyayn. Zeroth-order algorithms for nonconvex minimax problems with improved complexities. arXiv preprint arXiv:2001. 078819 2020.
[102] Bingzhen Wei, Xu Sun, Xuancheng Ren, and Jingjing Xu. Minimal effort back propagation for convolutional neural networks. arXiv preprint 17X.vv170.0055000 2017.
[103] Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computationnal Linguistics: Human Language Technologies, Volume (Long Papers), 2018
[104] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi. Pierric Cistac. Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer. Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: Sfate-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38 38-45, 2020.
[105] Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Under- standing common structure of hessian in neural networks. arXiv preprint arXiv:2010.04261, 2020.
[106] Jiaxi Yang, Binyuan Hui, Min Yang. Binhua Li, Fei Huang, and Yongbin Li. Iterative forward tuning boosts in-context learning in language models, 2023.
[107] Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pages 581-590, 2020.
[108] Haishan Ye, Zhichao Huang, Cong Fang, Chris Junchi Li, and Tong Zhang. Hessian-aware zeroth-order optimization for black-box adversarial attack. arXiv preprint arXiv: 18 11377, 2018.
[109] Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 03888, 2017.
[110] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 0ttt::/904.0009 2019.
[111] Sheng Zhang, Xiaodong Liu, Jingjing Liu, Jianfeng Gao, Kevin Duh, and Benjamin Van Durme. Record: Bridging the gap between human and machine commonsense reading comprehension. arXiv preprint arXiv: AxXiv:///00 18/0. 12885, 2018.
[112] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, al. Opt: Open pre-trained transformer language models. arXiv preprint 27X0v22205.01068, 2022.
[113] Yan Zhang, Yi Zhou, Kaiyi Ji, and Michael M Zavlanos. A new one-point residual-feedhack oracle for black-box learning and control. Automatica, 136: 1 10006, 2022.
17

---

A Algorithmic Ablations
We perform number of ablations to select the best algorithm. As is standard in ZO literature, we consider the main computational cost to be the number of forward passes. In our case, the number of forward passes can be affected by the number gradient steps taken, any usage of gradient accumulation, and using more noise samples to reduce the variance of the gradient estimate.
We observed that the performance of MeZO improves monotonically with the number of steps, and there does not appear to be any overfitting. Hence, when performing algorithmic ablations, we can focus on the efficiency of different algorithms without considering implicit bias. This is also reflected in our theoretical analysis. To ease the computational load, we fix the number of forward passes to 10, 000 and compare many different algorithms for RoBERTa-large on a smaller set of tasks that span sentiment analysis, entailment, and topic classification: SST-2, SNLI, and TREC. We emphasize that 10, 000 a small budget and is only used setting compare these ZO algorithms to each other. We find that using linearly decreasing learning rate schedule during training, as was done for fine-tuning with backpropagation in [65]. does not help or hurt MeZO. Similarly, using a learning rate warmup leads to identical results on these three tasks. For simplicity, we use constant learning rate schedule with no warmup for all the below experiments. We perform few-shot experiments
with & 16 and average the results across 5 seeds.
Experiment
Hyperparameters
Values
MeZO
Batch size Learning rate
{16, 64) X
{le-5, 5, le 1c-6, 1e-7) 7} x {le-3, le-5] (0,0.1}
Weight Decay
Table 4: The hyperparameter grid used in our ablation experiments. For simplicity, we use constant learning rate schedule.
A.1 Prompting
We study if adding prompt is crucial the ability of McZO to optimize the network. We prompts from Gao et al. [35]. Malladi et al. [67] claimed the prompt makes the optimization trajectory behaved, though models while the previous work only studied RoBERTa-base. We note the similarity between kernel behavior and theoretical setting in Section 4. MeZO succeeds tasks that reported to not exhibit kernel behavior in Malladi et al. [67], so we investigate whether or not the prompt is necessary.
note that the current paper considers RoBERTa-large and large autoregressive
SST-2
SNLI
TREC
Prompt
89.6 (1.2)
65.1 (6.2)
66.7 (6.2)
No Prompt 51.9 (2.9)
34.8 (2.1)
19.5 (9.0)
Table Experiments using MeZO to fine-tune models with and without a prompt.
Both experiments followed the grid in Table but we also expanded the grid include learning rate of le - 4 for the no prompt case. As result of these experiments, we fix the setting to prompt-based fine-tuning for all of the below experiments.
A.2 Sample schedules
One can sample noise vectors at the ith step and use SPSS to compute the gradient estimate. Similar ideas were proposed in Bollapragada et al. [11], Cai et al. [15]. We study the effect of linearly increasing and constant sampling schedules in the ablation setting. The intuition for the linearly increasing schedule that the optimizer may need a higher fidelity gradient as it approaches the minimum. Increasing the number of =s can speed up optimization by reducing the gradient variance, but doing so also increases the number of forward passes required for each optimization step, so there is trade-off to study. We note that increasing the number of 2s should be accompanied by
18

---

proportional scaling of the learning rate, analogous to the linear scaling rule proposed in [39] (theoretical justification can follow the SDE technique [58]). Table shows no consistent benefit in one schedule over the other, and it demonstrates that increasing the in n-SPSA while fixing the number of forward passes allowed results in only marginal gains best.
Y
Schedule
SST-2
SNLI
TREC
Y1
Constant
89.6 (1.2)
65.1 (6.2)
66.7 (6.2)
T1
Constant
89.5 (1.1)
68.6 (3.2)
62.3 (5.6)
T1 4
Linear
89.6 (1.4)
65.3 (6.4)
66.1 (5.5)
N2 16
Constant
90.4 (0.7)
67.0 (3.4)
62.8 (6.3)
16
Linear
88.9 (1.2)
62.8 (5.9)
64.2 (5.3)
Table 6: Experiments using MeZ0 with different schedules for n. proportionally to the number of z's sampled.
We scale the learning rate
B
MeZO Variants
There is rich history of transferring ideas from first order optimization enhance ZO algorithms. Below, we highlight several variants of MeZO that did not achieve as high performance as the algorithm presented in Algorithm
B.1 Memory-eflicient n-SPSA
We highlight how MeZO can perform n-SPSA (Definition 1) efficiently for n > 1 in Algorithm In particular, if sampling vectors and averaging the projected gradients, we require storing 2n additional scalars: the random seeds and the projected gradients. The same caveat about perturbing individual weights versus entire weight matrices still applies here (see Section 2).
B.2 Augmenting Me70 with Gradient History
The n-SPSA algorithm merely provides gradient estimate that subsequently be used in place of the gradient any muddent--bssdd optimizer. Many popular optimizers, such as Adam and SGD with momentum, require storing some historical information about gradients (e.g., moving average). This requirement causes such algorithms to require 2x or 3x the memory that is needed for SGD.
However, one advantage of MeZO is that the gradient history can be recomputed at each step without requiring much additional memory. reference to Algorithm I, note that the gradient only needs projected grad and the random seed s used to compute the perturbation
so we need to only store scalars per step to reproduce the gradient history (i.c., up to 2T' scalars during training). This is substantial reduction in added memory overhead that is usually needed for using Adam or momentum instead of vanilla SGD.
Table 18 illustrates that MeZO-Adam can sometimes improve the performance of MeZO, though each gradient step requires additional computation (but no additional forward passes). We leave it future work to investigate when MeZO -dzm-Adammm may be more useful than MeZO.
Experiment
Hyperparameters
Values
MeZO-Adam
Batch size
64
Learning rate
{le-6, 6, le-5,
le-4, 5e - le-3] le-3
Weight Decay
Table The hyperparameter grid used for MeZO-Adam. For simplicity, constant learning
rate schedule.
19

---

Algorithm 2: MeZO with n >>>>
Require: parameters A E R2, loss L IRd + R, step budget T, perturbation scale c, batch size B, learning rate schedule {ne), for n-SPSA estimate (Definition
for 1 1,.., do
seeds, projected_grads
Will each contain scalars
for j = L.....7 12 do
Sample batch B D and random seed s
0 PerturdParameters((..
4 f- L():B)
0 < PerturbParameters
L + [(0;
0 < PerturbParameterss projected_gret f- (f, - (2e)
D Reset parameters
projected_grade [j] < projected_grad
seeds [j] +
end
for
j 1, .... do
Reset random number generator with seed seeds [j]
for 0, E 0 do
2 ~~ (0,
0 + 0, (nt/n) projected_grads [j] * z
>> Avg grad for 2,, Zn
end
end
end
Subroutine PerturbParameters. (0, s)
Reset random number generator with seed
For sampling
for 0: E 0 do
10) 0, + 0, -6
b Modify parameters in place
end
return A
B.3 Modifying the Variance of MeZO
Our theory in Section 4 sketches the well-known fact that the variance of the stochastic gradient estimate can impact the rate of optimization. ZO methods can be combined with standard variance reduction techniques to possibly improve optimization speed. For example, Liu et al. [62] designed variance reduced ZO algorithm, analogous SVRG [49], to improve the speed of convergence. Below, we show that several variance reduction methods (e.g., using the gradient norm) can be implemented memory-efficien manner. However, when controlling for the total budget of forward passes (i.e., function queries), these methods are not as performant as MeZO. We nevertheless present them to demonstrate the ease with which MeZO can be adapted, and we suggest these methods may be useful for optimizing more complex objectives.
First, we define general SPSA estimate that has the same expectation (i.e., the true gradient) but has a scaled variance.
Definition (Variance-Modiffed SPSA), Given a matrix = (diag(d), the variance modified SPSA computes
B)), C( + e(d a 2); C(O <(d <(d e 2); (d 2)
2e
where d E Rdd has nonzero entries and d-1 denotes the coordinatewise reciprocal.
The above SPSA variant is an unbiased estimator of the gradient, because L B)] E|D-Izz DVL B)] = E [(( B)). We will draw inspiration from classical methods (i.e., "control variates") and choose d to be a block vector with gradient norms parameter norms [99].
20

---

To select the parameter groups, we split the model by layer, keeping the embedding and the head separate (i.... RoBERTa-large has 24 + 2 = 26 parameter groups). straightforward measure the parameter norms without consuming additional memory. We can measure the gradient norms without performing backpropagation, as shown below.
Proposition (ZO Estimate of Gradient Norm of êth Layer). Define zt to have 2 NN 1) in each coordinate corresponding to parameters in the Eth layer and O everywhere else. Then, we can estimate the norm of the gradient of the loss w.r.t. the Eth layer Vo, as
C( ez B)
L( E B)
IVo. ((0:B11) ~
2c
As is true for SPSA, increasing the number of zę's sampled for each value of and averaging the result reduces the variance of the estimate. The rationale for this estimate is that for any vector v, Eyl((v, 2 = || for Gaussian z. It is clear that this estimate can be computed in a memory efficient way, although it requires 2T. forward passes to compute gradient norms for I, parameter groups.
We show the experimental results for modifying the variance below. We follow the ablation setting and use fixed budget of 10, 000 steps (Appendix A). Generally, using the gradient norm reduce the variance substantially hurts performance (Table 8). If we "cheat" and allow one backpropagation through the network estimate the gradient norm, then we see that reducing the variance using the gradient norm does not substantially hurt or help performance. Modifying the variance using the parameter norm, analogous layerwise adaptive rate methods, does not substantially impact the performance of McZO (Table 9).
Our observation is that decreasing the variance by setting d as the gradient norm does not improve optimization. This empirical result agrees with the exposition in Section 4 that the straightforward variance analysis (which yields dependence the number of parameters d) is not the best lens to study the rate of optimization when fine-tuning with MeZO. Our effective rank view in Theorem I and Lemma 3 is likely a better characterization of fine-tuning dynamics. We leave future work explore if these methods can be useful for other more complex objectives.
Recompute d ZO estimate of d
SST-2
SNLI
TREC
Baseline MeZO (Algorithm 1)
89.6 (1.2)
65.1 (6.2)
66.7 (6.2)
X
89.7 (0.8)
65.2 (5.2)
64.3 (6.4)
87.0 (2.5)
49.6 (9.2)
32.6 (7.7)
79.0 (10.3)
48.9 (2.2)
38.7 (7.5)
Table 8: Experiments modifying the variance of McZO using as the gradient norm (see Definition 6). We sometimes recompute d at the start of each epoch or Proposition to estimate d without requiring backpropagation.
Recompute
SST-2
SNLI
TREC
Baseline MeZ0 (Algorithm 1)
89.6 (1.2)
65.1 (6.2)
66.7 (6.2)
89.2 (2.1)
65.4 (4.2)
64.8 (5.6)
88.2 (4.7)
65.2 (4.0)
64.7 (5.5)
Table 9: Experiments modifying the variance of MeZO using d as the parameter norm (see Defini- tion 6). We sometimes recompute d at the start of each epoch.
B.4 Modifying the Expectation of MeZO
The above experiments show that modifying the variance of MeZO cannot consistently accelerate its convergence. However, a simple modification of Definition 6 allows us to change the expectation of Me7. as well. This can be used to efficiently estimate coordinate- wise normalized gradient- based optimizer updates (e.g., Adam).
21

---

Definition 7 (Expectation-Modiine SPSA). Given matrix D diag(d, (d), the variance modified
SPSA computes
() +( 2 B): [(O =(d 2):
2c
where d E Rd
Now, we see that estimator for VZ If we choose d to be the gradient norm, for example, then SPSA can estimate the normalized gradient. Concurrent work in Tang et al. [94] gives another ZO estimate of the normalized gradient while assuming access to only rankings of inputs (instead of the noisy function evaluations available in setting). We find that estimating the normalized gradient does perform as well as directly estimating the gradient (Table 10). Regardless, we present this algorithm as a way to highlight that any coordinate-wing operation to the gradient can be applied in memory-efficien manner.
so the SPSA estimate is no longer an unbiased
Method
SST-2
SNLI
TREC
Bascline McZO (Algorithm 1)
89.6 (1.2)
65.1 (6.2)
66.7 (6.2)
Estimate normalized gradient (Definition 7)
88.0 (1.2)
60.0 (2.4)
44.0 (14.0)
Table 10: Experiments modifying the expectation of MeZO using d the gradient (see Definition 7). We use the 70 estimate of the gradient norm (Proposition 1).
B.5 One-point estimate
Here, we investigate the efficacy of one-point gradient estimators in place of the two-point SPSA method. Using one-point estimator instead of SPSA can reduce the MeZO running time by half. Many one-point estimators have been proposed in the past [34, 87, 95]. For simplicity, we focus one estimator [113] that has form reminiscent of SPSA but requires only one forward pass to estimate the gradient at each step.
Definition 8 Point Gradient Estimate, Zhang et al. [113]). For loss function L evaluated batch B, with parameters θ₂ at step t, we can draw random noise 2; ~ N(0 Ia) and compute the gradient estimate using hyperparameter
written below.
(.);;; [(O, + ezg; B:) L( + 12
Notably, this one-point gradient estimate uses the loss at the previous iterate instead of evaluating the loss again at the current iterate. As such, this estimator requires only forward pass at each iterate to estimate the gradient. For well-behaved loss functions and slow-moving optimization, these two formulas intuitively similar. However, Table 11 finds this estimator be much less efficient than SPSA when fixing the number of forward passes.
Steps
SST-2 sentiment
SST
SNLI
MNLI
RTE
TREC topic
natural language inference
SPSA [88]
20K
92.8 (0.5)
(0.9)
82.9 (1.0)
74.4 (0.8)
76.7 (1.7)
92.7 (0.6)
One-point estimate [113]
20K
90.0 (0.4)
44.6 (2.0)
70.1 (1.5)
57.2 (0.9)
64.1 (1.0)
57.3 (5.7)
One-point estimate [113]
40K
91.8 (0.5)
45.9 (1.7)
(0.8)
61.0
68.7 (1.2)
(3.1)
Table 11:
Comparison between SPSA and one-point estimate Zhang et al. [113]. One-point estimate only does one forward pass per step. thus is twice as fast as two-point estimate per step. As such, the number of forward passes after 40K steps with the one-point estimate is the same as the number of forward passes with SPSA after 20K steps. The results show that two-point estimate is much more effective than one-point estimate.
C Memory Analysis
The compute-memory tradeoff of backpropagation is complex analyze. Griewank and Walther [40] provides rigorous theoretical treatment of the problem. We empirically measure the memory
22

---

consumption of different methods for commonly used large language models, but here we hope to provide a more rigorous comparison of different gradient estimation algorithms, independent of the software used to implement them. Below, we summarize some key points that may help readers to understand how the MeZ0 compute-memory tradeoff compares to backpropagation.
Given a network, the first step to perform backpropagation is to decompose the model into easily differentiable blocks. We note that this decomposition is not unique. For each block, one can choose to cache the resulting output during the forward pass (thereby consuming memory) or instead recompute the output when it is needed (thereby consuming compute). The below proposition. adapted from Rule 21 in Griewank and Walther [40], captures this tradeoff.
Tradeoff for Backpropagation. Griewank and Walther [40]). Consider network containing N bits. For any time-memory tradeoff hyperparameter c (1), there exists backpropagation algorithm that runs in time (c:N) and consumes memory proportional to O(N1/o).
Proposition
Grimm al. [41] also gave sharp bounds for the memory-tine product. Note that the popular gradient
checkpointing [18] method allows one to tune e with limited precision (i.e., one cannot always further split a differentiable block and observe savings). Experiments in Chen al. [18] choose c to achieve O( DVVN) memory while consuming O(2N) computation. In the extreme case, gradient checkpointing allows one to use O(N log N) computation and Ollog N) memory.
MeZO always consumes 2NV compute and (1) memory, so it is more compute--fficient at the same memory cost as gradient checkpointing. Our exposition in Section 2 discusses that we can perturb groups of parameters together to save time while consuming additional memory. However, not consider that variant here because it is somewhere in the middle of the compute-memor pareto curve, where we cannot reason about what backpropagation will do. In particular, MeZO can split groups differently than backpropagation can, since MeZO does not require that each parameter group easily differentiable, so it is hard compare the two algorithms along the entire pareto curve.
do
We also compare backpropagation for the e = case (i.e., storing everything during the forward pass). When storing everything, backpropagation consumes O(N) time and (NN) memory. Hence, SPSA consumes slightly more time and substantially less memory than backpropagation at this end of the tradeoff.
Unlike gradient checkpointing, MeZO computes only an approximation of the gradient. This approximation is only useful for fine-tuning with a prompt, making it less broadly useful than gradient checkpointing. There are other methods that approximate the gradient with less memory consumption than gradient checkpointing (see the Related Work section), though it is unclear how the memory consumption those algorithms compare to MeZO.
D Forward Auto-Differentiation
We discuss the merits of using forward duto-iifferentiation instead of two forward passes to construct a gradient estimate for mie-ttning. As e - O, the SPSA gradient estimate (Definition 1) can be written as == T V ((e; B). The term x (L(: B3) is a Jacobian- vector product (JVP), and it is well-known that this can be computed in parallel with single forward pass while consuming additional memory equivalent to that of the largest activation the model. To fully compute the gradient estimate, one must store z on the GPU while performing inference, so we observe that this algorithm requires more memory than MeZO.
We note that implementation of the forward auto-differentiation algorithm is not well-supported in PyTorch at the time of writing. The autograd JVP function computes the JVP in a memory-inefficient way, as noted in the documentation, and the other available methods to compute a JVP are not sophisticated enough to easily scale to a complex LLM. Computing the JVP is straightforward when using JAX, so we profile the memory consumption of inference and the JVP for CoBERTa-large when using JAX. We use batch size 16 with the MultiRC task. Note that JAX may automatically use rematerialization to avoid out of memory errors so we focus on settings in which the memory utilization remains below 50%. The resulting memory usage during inference, backpropagation, and forward auto-differentiaaion are reported in Table 12.
23

---

We see that forward auto-diffein cittoddifferentittoon is substantially more memory efficient than backpropagation but less memory efficient than inference. Furthermore, forward uuto-differentiation selects which removes potentially beneficial third-and-higher order Taylor expansion terms from the estimate.
Task
Inference (and MeZO)
Backpropagation
Forward Auto-Differentiation
Excess Memory (MB)
327.50
24166.
830.66
Memory consumption of RoBERTa-large when using batch size 16 with the MultiRC task. The reported memory does not include the cost storing the model on the GPU, which is required for all three cases.
Table 12:
E Experiment setup
E.1 Datasets
For RoBERTa--arge, consider classification datasets: SST-2 [86), SST-5 [86], TREC [97], MNLI [103], SNLI [12], and RTE [22, 8, 37, 10]. We follow Malladi et al. [67] in limiting the test to 000 examples for fast iteration. For training and validation, we have two settings: k 16
and k = 512, which mean that we have 16 or 512 examples per class for both training and validation.
For OPT experiments, consider the SuperGLUE dataset collection [98], including: BoolQ [21], CB [24]. CÓPA [81], MultiRC [51], ReCoRD [1 RTE [22, 8, 37, 10), WiC [77]. and WSC [55].
We also include SST-2 [86] and two question answering (QA) datasets, SQuAD [80] and DROP [31].
We randomly sample 1,000 examples for training, 500 examples for validation, and 1,000 examples
for testing.
E.2 Prompts
Table 13 shows the set of downstream tasks and prompts with which we fine-tune RoBERTa-large, which are adapted from [35].
Dataset
C
Type
Prompt
Label words
SST-2
2
sentiment cls.
<S1> It was [MASK]
(great, terrible]
SST-5
5
sentiment cls.
<S1> It was [MASK]
(great, good, okay, bad, terrible) (Description, Expression, Entity,
TREC
topic cls.
[MASK] <S
Human, Location, Number) (Yes, Maybe, No Yes, Maybe, No} Yes, No)
MNLI
3
NLI
<S,> ? [MASK]
<S2>
SNLI
3
NLI
<S,> ? [MASK]
<S2>
RTE
NLI
<S,> [MASK]
<S2>
Table 13: The prompts of the datasets we used in our RoBERTa- experiments (Table 18 and
Figure 2). The prompts are adapted from [35] and include a template and a set of label words that can fill in the [MASK] token. <Si> and <SS>> refer to the first and the second (if any) input sentence.
Table 14 demonstrates the prompts we use for OPT. Note that in OPT experiments we have three types of tasks: classification, multiple-choice, and question answering. Prompts are adopted from GPT-3 [13] and PromptSource with minor changes [5].
E.3 Hyperparameters
We use the hyperparameters in Table 15 for MeZO experiments on RoBERTa-large (Table 18 and Figure 2). Experiments in Appendix A informed the grid; in particular, the choice of seemed to not significantly impact performance, and using a larger batch size consistently yielded faster optimization. We use the hyperparameters in Table 16 for MeZO experiments OPT.
Regarding learning rate scheduling and early stopping, we use linear learning scheduling for all fine-tuning with backpropagation experiments and constant learning rate for all MeZO experiments. For RoBERTa experiments, we evaluate the model on validation sets every 1/10 of total training steps and save the best validation checkpoint. All FT experiments use IK steps and MeZO experiments use
24

---

Dataset
Type
Prompt
SST-2
cls.
<text> It was terrible/great
RTE
cls.
<premise>
Does this mean that "<hypothesis>" is truc? Yes or No?
Yes/No
CB
cls.
Suppose <premise> Can we infer that "<hypothesis>"? Yes, No, Maybe?
Yes/No//Maybe
BoolQ
cls.
<passage> <question>?
Yes/No
WSC
cls.
<text>
In the previous sentence, does the pronoun "<span2>" refer to <spani>? Yes or No?
Yes/No
WIC
cls.
Does the word "<uord>" have the same meaning in these two sentences? Yes, No?
<sent1>
<sent2>
Yes/No
MultiRC
cls.
<paragraph>
Question: <question>
found this answer "<answer". Is that correct? Yes or No?
Yes/No
COРA
meh.
<premise> so/because <candidate> <passage>
ReCoRD
meh.
<query>.replace ("@ppplccccceellde"""" <candidate>)
SQuAD
QA
Title: <title>
Context: <context>
Question: <queation>
Answer:
DROP
QA
Passage: <context>
Question: <question>
Answer:
Table 14: The prompts of the datasets we used our OPT experiments. There are three types of tasks: classification (cls.), multiple-choice (mch.), and question answering (QA). Prompts adopted from GPT-3 [13] and PromptSource [5] with minor changes. <text> represents input from the dataset and Yes represents label words. For inference multiple choice tasks, put in different candidates the prompt and calculate the average log-likelihood for each candidate, and choose the candidate with the highest score. For inference on QA tasks, we use greedy decoding to generate the answer.
100K steps. For OPT experiments, we evaluate the model on validation sets every 1/5 of the total training steps and save the best validation checkpoint. All FT experiments train for epochs and all MeZO experiments use 20K steps. Note that FT experiments mostly converge within 5 epochs but we observe that MeZO performance can still improve with more training steps.
E.4 Modeling and implementation
For RoBERTa experiments, we follow [35] for the prompt-based fine-tum Ane-tuning paradigm for masked language models. Please refer to the original paper for more details.
In OPT experiments, for classification tasks, train the model similarly [35], i.e., take the logits corresponding to the label words and apply cross entropy loss on them; for multiple choice tasks and generation tasks (QA), we only keep the correct candidate and use teacher forcing train on the correct examples. We only keep the loss on tokens in the candidate part and exclude the prompt part.
For OPT inference on classification and multiple-choice tasks, we use the model to get the average log likelihood (by tokens) of all the candidates/label words, and predict the one with the highest average log-likelihnn likeelooo.odd. For generation tasks, we use greedy decoding to generate the answer.
For in-context learning, we use 32 examples in the context. We also try filling in as many examples as possible in the context this does not improve performance and sometimes leads to unstable results. Thus we keep the 32-example results.
25

---

Experiment
Hyperparameters
Values
McZ0
Batch size
64
Learning rate
{le- 1e-6, 6, le-5)
1e-3
Weight Decay
McZ0 (prefix)
Batch size
64
Learning rate
{le-2 2, 5e-3, le-3)
le-1
Weight Decay
prefix tokens
MeZO (LoRA)
Batch size Learning rate
(le-5, 5, 5e-5, 1e-4)
le-3
Weight Decay (r,a)
0.1
(8, 16)
FT with Adam
Batch size (k = 16)
[2,4, 8}
Batch size (k 512)
[8, 16,32)
Learning Rates
{le-b 3,3e-5, 5e-5)
Weight Decay
with SGD
Batch size (k = 16)
{2,4,
Batch size (k = 512)
{8, 16, 32)
Learning Rates
(1e-4, 5e 1,1e,,, 5e-3, 1e-2)
Weight Decay
FT (prefix)
Batch size
{8, 16, 16,
Learning Rates
{le-2 2, 2e-3e-2 5e-22
Weight Decay
prefix tokens
FT (LoRA)
Batch size
{4,8, 16) {le- 4, 3e-,, 5e-4) 4]
Learning Rates
(r, a)
(8, 16)
Table 15: The hyperparameter grids used for RoBERTa- large experiments. MeZO uses a constant learning rate schedule, and FT uses linear scheduling. All FT experiments use 1K steps and McZO experiments use 100K steps. We check validation performance every 1/10 total training steps.
For linear probing of classification tasks, take the output feature and scipy package train a linear classifier. For multiple-choice tasks and generation tasks, we found that this leads to poor results since the output space is the whole vocabulary: instead, do head-tuning, where the whole model is fixed except for the LM projection head. We use batch size of and learning rate of (le-4 4 4e-4), and train the head for epochs.
For experiments on 30B and 66B OPT models, we largely follow the OPT hyperparameters except that we do not evaluate the intermediate validation performance and directly use the last checkpoint for evaluation, due to the high storage cost of intermediate checkpoints of large models.
E.5
Parameter-efficient fine-tuning
Fine-tuning and storing copy of the large language model for each downstream task is expensive. Parameter-efficiemt fine-tuning (PEFT) techniques alleviate this problem: instead of tuning all model parameters, PEFT only tunes small number of additional parameters (usually less than 1%) and can often achieve comparable or better performance [57, 54, 30]. The ZO optimizer compatible with PEFT methods, since ZO operate any subset of the model parameters. We interested in the following two common PEFT methods, designed for transformers [96].
LoRA [46] adds tunable low-rank delta to linear layer during fine-tuning. Suppose linear layer performed War during pre-training with W E R When fine-tuning, LoRA introduces two smaller matrices A [ mx and B pg such that < min(m, n). The linear layer is then
26

---

Experiment
Hyperparameters
Values
McZO
Batch size Learning rate
16 fle-6, le-7]
1e-3
MeZO (prefix)
Batch size Learning rate < # prefix tokens
16
{ {1e-2, 1e-3) le-1
5
MeZO (LoRA)
Batch size Learning rate
16
11e-1,5e-5)
1e-2
(r,a)
(8, 16)
FT with Adam
Batch size
8 { f1e-5, 5e-5, 5e-5))
Learning Rates
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.
computed as
( + AA x b
(6)
where and a are hyperparameters. and are trained on the downstream task while W frozen at its pre-trained value. In transformers, this modification to the linear layer is applied to the query and value operations of each attention layer. Empirically, can be very small, so the number of trainable parameters during fine-tuning is small. We choose r = 8 and a = 16.
Prefix-tuning [57] adds prefix of m tunable representations 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 McZO 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 might expect from classical analyses (see Section 4). Accordingly, our theoretical analysis in Section suggests that the convergence rate of ZO-SGD does not depend on the parameter dimension during fine-tuning.
Task Type
SST-2 sentiment
SST-5
SNLI
MNLI
RTE
TREC topic
natural language inference
FT (prefix, random init) FT (prefix, real 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)
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-initiin 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 Fl objective, we use a batch size of 16. For MeZO, we use a learning rate of (1e- 5e-66 6, 1e-5) and = 1e-3. For MeZO (prefix), we use learning rate of {le-1, -2- 1e-2) and 1e-1.
27

---

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 to test the minimum hardware requirement to run the model with specific optimization algorithms. For multi-GPU backpropagation,
fully sharded data parallel (FSDP) [33] provided by PyTorch [76]. For multi- Dulti-GPU McZO, we use transformers multi- inference of large models. We use Nvidia's 's nvidia- a-mi command to monitor the GPU memory usage. We call run "successful" there no out memory error from GPUs for at least 100 steps. We also profile fine-tuning with LoRA, but find its memory usage similar that of fine-tuning with prefix-tuning, Hence here we only show the analysis with prefix-tuning.
F
More experiment results
F1
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
SST-5
SNLI
RTE
TREC topic
sentiment
natural language inference
Zero-shot
79.0
35.5
50.2
48.8
51.4
32.0
Gradient- free methods:
16
LP
76.0 (2.8)
40.3 (1.9)
66.0 (2.7)
56.5 (2.5)
59.4
(5.5)
MeZO
90.5 (1.2)
45.5 (2.0)
68.5 (3.9)
58.7 (2.5)
63.3)
(2.7)
MeZO (LoRA)
91.4 (0.9)
(1.6
69.7 (6.0)
(2.5)
64.9 (3.6)
(3.5)
MeZO (prefix)
90.8 (1.7)
45.8 (2.0)
(2.5)
63.4 (1.8)
63.4
80.3 (3.6)
Meam-Adam
90.4 (1.4)
(1.5)
74.1 (2.7)
64.3 (0.8)f
59.2 (11.1)
78.3 (1.4)
Gradient- based methods:
= 16
FT
91.9 (1.8)
47.5 (1.9)
77.5 (2.6)
70.0 (2.3)
66.4 (7.2)
85.0
FT (LoRA)
91.4 (1.7)
46.7 (1.1)
74.9
67.7 (1.4)
66.1 (3.5)
82.7 (4.1)
FT (prefix)
91.9 (1.0)
47.7 (1.1)
(1.3)
66.5 (2.5)
66.6 (2.0)
85.7 (1.3)
Gradient
Cradient-free free methods:
512
LP
(0.5)
(0.5)
(0.0)
71.5 (1.1)
(1.5
89.4
McZO
(0.7)
(1.4
(1.0)
78.3 (0.5)
72.6
94.3 (1.3)
McZO (LoRA)
93.4
52.4
84.0
77.9 (0.6)
71.6
90.7)
MeZO (prefix)
93.3 (0.1)
53.6
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 (9.4)
(1.2
95.1 (0.3)
bradient-based methods:
512
FT FT (LoRA)
93.9 (0.7)
(0.9)
88.7 (0.8)
84.4 (0.8)
82.7 (1.4)
97.3 (0.2)
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
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: 70, 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 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 le-4 to 0.01) in the scipy solver. Experiments on a few settings show that M.-M can sometimes improve performance without increasing the memory consumption (see Table 19). However, sometimes, linear probing first can severely hurt performance.
28

---

Task
SST-2
SST-5
SNLI
TREC
Zero-shot
79.0
35.5
50.2
32.0
FT
91.9
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 16 for these experiments.
F.2
OPT experiments
Table 20 present the full results of OPT-30B and OPT-66B, with detailed MeZ0 numbers.
Task
SST-2
RTE
BoolQ
WSC
WIC
SQuAD
30B zero-shol
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 Me70
90.6
66.4
67.2
63.5
56.3
85.2
30B McZ0 (prefix)
87.5
72.6
73.5
55.8
59.1
83.9
66B zero-shol
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 MeZ0 (prefix)
93.6
66.4
73.7
57.7
58.6
85.0
Table 20: Experiments OPT-30B and OPT-66B (with 1,000 examples). *; McZO requires further
tuning to successfully optimize.
F.3 Convergence of MeZO with full-parameter and PEFT
We demonstrate the convergence rate of MeZO, McZO (LoRA) and MeZO (prefix) on SST-2 and SNLI for the first 5,000 steps Figures 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.
SNLI
SST-2
0.55
MeZO Train
2.0
MeZO Train
MeZO (LoRA) Train
MeZO (LoRA) Train
0.50
MeZO (prefix) Train
1.8
MeZO (prefix) Train
0.45
1.4
Loss
0.40
Loss
1.2
0.35
1.0
0.30
0.8
0.25
0.6
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Steps
Steps
Figure 5: Me7. does not optimize significantly faster when tuning fewer parameters, agreeing with our theory in Section 4.
29

---

F.4 ZO vs BBTv2
We compare ZO with BBTv2 [90] on mutually assessed tasks in Table 21. ZO significantly outperform BBTv2. Furthermore, BBTv2 is limited to optimize in low-dimensional space and requires prefix- tuning and down- projeccoctttonn reduce the number of optimized parameters. BBTv2 also employs an iterative scheme which only optimizes one layer at time. In contrast, 70 works with both full-parameter tuning and PEFT, as shown in our experiments (Section 3) and theory (Section 4).
Task
SST-2
SNLI natural language inference
RTE
Task type
sentiment
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)
Table 21: 7O vs BBTv2 with RoBERTa-large. BBTv2 performance is from Sun et al. [90].
E5
Memory profiling
We show the detailed numbers of memory profiling results Table 22, which also corresponds to Figure 3. For how profile the memory usage, please refer Appendix E.7.
Method
Zero-shot / MeZO
ICL
Prefix FT
Full- parammerre FT
1.3B
1xA100 (4GB)
1xA100 (6GB)
1xA100 (19GB)
1xA100 (27GB)
2.7B
1xA100 (7GB)
1xA 1x0000 (8GB)
1xA 100 (29GB)
1xA100 (55GB)
6.7B
IxA100 (14GB)
1xA100 (16GB)
IxA100 (46GB)
2xA100 (156GB)
13B
1xA100 (26GB)
IxA100 (29GB)
2xA100 (158GB)
4xA 100 (316GB)
30B
1xA100 (58GB)
1xA100 (62GB)
4xA100 (31SGB)
8xA100 (633GB)
66B
2xA100 (128GB)
2xA100 (134GB)
8xA 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-arr solutions for distributed training. As shown in Table 23, on the MultiRC datasets, training with MeZO brings 7.74x speedup per step compared to full-parameter FT 30B model. This is due to (I) McZO does not require costly backpropagation and (2) McZO requires fewer GPUs and reduces the multi-GPU communication overhcad. We can that MeZO has 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 32x more steps than standard FT, while FT takes 8x more GPUs and 7.74x more time per step. Overall, MeZO requires only half as many GPU-hours as FT for a 30B model.
30

---

1.3B
2.7B
13B
30B
66B
Mc70 (hsz=16)
0.815s (1)
1.400s (1)
2.702s (1)
5.896s (1)
12.438s (4)
Mc70 (hsz=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)
hspd=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 numbers of GPUs required. It is measured on 80GB A 1000s with NVLink and InfiniteBand connections. The wallelock time is averaged over 100 steps. It is measured the MultiRC task with the OPT family. We use batch size ("bsz") also add MeZO with batch size of 8. For FT (FSDP), we show the
for FT and 16 for MeZO (consistant with our main experiment
setting). For comparison following additional information. "bspd":: batch size per device. "ga": gradient accumulation steps. The effective batch size is bspdxgax #GPUs. Note that for FT 66B, the effective batch size 20.
31

---

G Proofs
Proof of Lemma' 2. We first note that in the 0 limit, we have
O) = Bn
Σ
Σ z
VC(0:((2.)))
(z. ie(n]
Taking expectation over the batch and the 2;, we have E[V = V2(0), so V2(0; B) an unbiased estimator of the gradient.
Computing the second moment, we get
E
Σ
E E [[ (I VC(( ((22, (()))))))))
B
Let u, v two arbitrary vectors. have that
Bz, [z1z] uw 2zz, uo
when * j. and
Ez, ( uo ]= 04 (u, v)
3d
d Sym (I (02) (u,
d d +
2d
I
uv
d 2
Therefore
E
n
2d
E [[[(0; ((x1, (())) ((22 "2
B²
n
n(d + 2)
(ar Mn
BB
E [C(0; (x (Ix1,1))) 91))) L(0; (100((())) (102, ya)h) I. nld+ +
Next, note that when (x1, y1) * (22, 122). we have
= E [C(0; ((x1 ((0; =2C000000
and when (æ1, y1) (æ2, y2) have
[C: ((x) 91) ((; [ +MM=())))))))))
EM)(O).
Therefore
Σ
E ((0 -VC00 =(0))
(1L BI JEJE
and plugging this yields
d-2 (71(0)2110) + 1 I
E [
1+
n
n(d (2)+
B
d
(IV (
+
I
n(d +2)
Finally, we have
(1+ d + d ++d-2)
[[
1 (<(((0))
E
n(d + 2)
d d n - E |VC(0:))
32

---

Proof of Theorem 1. By Taylor's theorem with remainder, we have that
[(0++)) = ((),) (((0.)"(00::) - 0.)
+ . 10 -0,+1-0,)T0 -10
-0,) da
Next, note that
00+1 - 0ell = 1) [)) < <nvd
Bn Σ z VC ((x, y))e - ndG(()).
Therefore
||0 + (1 - X)0, - - S nd)((,)) By the assumption we have the upper bound V2L(X0++1 + (1 x)0,) ~ H())). and thus
[(0(+1) ((0,) + VL(0.) ((((O,)((0))))))) [0.+1 0.) + (0.+1 0.) 0.) ((.))(00+ 0:) = =(0,) - VV (L(0,; "
+ 2" " V (0,
Taking the conditional expectation with respect to 0, and plugging in (9), the formula for the covariance of our ZO estimate CL(,,; yields
1[((0++1) 0,1 < ((0,) - 1 [[ + 7/2
H(O,), E
2
(IV 1 ir( in(H(0,))
_(0.) [V +
2
n(d + 2)
B
+ n 1+ d 2-2
n(d m(d+2) 2) )
(VC(O)) H(O.)VC((0)) + (E())),
H(.))) H(O)))
2
By assumption, the Hessian upper bound H (0,) satisfies [H()))lopp S and tr H (0)) s fr. Thus
B[C(0++1) 0.] s L(0.) - [V)) + T
dr + +1) (IVC(O.)) + (r(=(0)))
n(d + 2)
(dr
= [(,,) - 7 ( +
[ndd+2 +1
[VC[0::
2
as desired.
G.1 Proofs of Global Convergence
Lemma 4. Let L(0) be μ-PL and let there exist such that tr tr(E(0)) I a(C(() L*) for all O.
Then after
t=0 (-0(("-
ta
log [(Op) L*
1²B
iterations of SGD we have B[C(O,)) < L +e.
Proof of Lemma 4. The descent lemma for SGD yields
EE(((0++1) | 0.] ((,,) < - -V + " B))
Plugging in B[] = [[V)) yields
ir()())) and selecting learning rate S
([C(O.+1) 0.1 S ((..) |VC(0.)11 T'C
tr()(0.))
2B
Since is p-PL, we get
B(C(0(41) 0,1 S L(O.) - ()(C((,) L*) (*)+ + 2f2
tr (EO)))
2B
33

---

Since tr(2()))) - <a((())) - L) we have
E[c(0:+1) | - (())) - m)(((()) - -L") + 22 (L(6)) L*).
2B
Altogether,
R[[(0++1)] - ( 12
(E([((.)) L*)
2B
Choosing min(? 2B) we obtain E
B[C(O++1]] L (1(1- min( 12B
(B|C((.)) L*).
28 28a
Therefore we reach solution with E[[(0.)) L S after
28 26ax
((00)
- -(00)-(() =
( 1
(
la
[(00) 6-
t t= max
-
log
log
1
12B
2²
iterations.
Proof of Lemma 3. By Corollary 1, ZO-SGD with 120 = = yields
[[((0.+1) 0.] - ((O,) - - [ TSGD |VC(0)) mane B[I
As in the proof for SGD, choosing 1JSGD S yields
[[C(O++1) 0.] C(0,) S--1
1SSD
[V(()
n
tr tr()(()))
2
2B
Therefore under y-PL and the tr(E((,)) S a(((0)) - L") assumption we obtain
B( 5--1 -TSCDDU + [
(B[[((,)] - L*)
2B
= = R[[(((++)) - f s (1 (1-7-1 1SSDD
1/ša0 &a ))
(B([(0,)) L²).
2B
Choosing 7JSGD min(? t, yields
B|C(0 - -- - (1-7-1_mm miinn 1 1²B)) (E(C((,)) - L²).
26 2Pa
Therefore we reach a solution with E[C(O.)] L- S after
22 24a p
[((0) -(0)))(() 0 ((1) m
L'
+1)
(L + la log L(Ou) L-
t 7 max
log
12B
iterations.
G.1.1 Verification of assumptions
We show that the tr (r,)) S <(C(0,)) L*) assumption holds for certain losses.
First, consider optimizing the model fa:; 0) with square loss, so that
((
1 N
Σ [E[N]
(f(a;; 0) y2)
One then has that
2 Σ (f(x;) 0) y:) ([(; -
VC(O)V2())T.
N ME[M
34

---

Therefore
2 E (f(a;; - (1; tr (T(E))) - N
ie[N]
- 22( KEIM]
Σ [V
Assume that the data can be interpolated, i.e., L 0. If the function L-Lipschitz, i.e., fx; 0)l I L, then the condition holds with 2NL2. If we are in the kernel regime, i.e, (a:;; 0) = ((x;) 0 for some feature map 6, then
V2C(0)
Σ f(z; ([=:0)V((a;0)) EEIN]
N
Thus
((() S Nu((V (()) S Nor C()).
So the condition holds for a Ndr.
Next, consider the cross entropy loss function, i.e
Σ ( N iE|M]
exp(-Uif yif (x;; 0))
One then has that
1 Σ( N EE(N]
Σ
fx: 0)), (V -((0))))
Assume that the targets Yi are bounded in [-1, (which true for binary classification tasks), and that L = 0 (which can be achieved if ff 0)1 can be sent to oc) we have that
Lr(E(0)) S N
Σ
exp( exp(-2y:J(ar) 0)) V (x,; 0)1
iE[N]
In the kernel regime, f( 0) ((x:;)
0, and thus
V2C(0)
Σ
N
KE[N]
Therefore
Ir (0) S NN((V L()) < Nor L())-
Therefore the condition holds with a = Nêr as well.
G.2 Proofs for Gaussian perturbations
The first lemma computes the second moment of the covariance estimate (c; when z drawn I).
Lemma 5. Let z; ~ (0, I) i.i.d. Then
E [C(..B\\000. =
(1+ (VCIOVVCII) + 1 IMBO()) - (IVC10)1)= + (r(ZMB(0)))
(8)
Proof. As the proof of Lemma have that in the limit
Bc B)
Σ
Σ E [Z,2TT]LL( ((,,, (m1,, (1))) VVC(0 ((02, 7
B
35

---

For vectors u, v, we have that
B2,,2, 2 uv 2Ø2, uv
when * J. and
Ex (1, = 3Sym(I )u,, = -u" v. ----- I+ +uv Ex, [= we 2,,2,11 "2,22,
Therefore
E
[[(0:B]V<(0:8]1
n-1 2 n1
Σ
E [2(0: (())
1
(0) Mi (22)
+ [c; ((0)) bii)) C(:; (122, y2)1) I.
In the proof of Lemma we showed that
Σ
E [c((; (1)
+ - 1(0)
B
(FI-01) BB
Plugging this yields
("+1) ( (VC(0)VZ(0) 1 B
210) 1 I. (IVC(0)) (r(E(()))
E
n
(9)
1
+
n
We can prove an analog to Theorem 1 in the case where the z; are Gaussian. One challenge that |0.+1 Oel| is 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 ||0+11 Oel being large.
Consider the following modified version the local r-effective rank assumption, where the upper bound on the Hessian is measured over ball of radius twice as large the one in Assumption 1.
Assumption 2 (Local r-effective rank, Gaussian). Let G(O.) (ax)) V L(O, ((x, y)
There exists a matrix H (0:) such that:
1. For all 0 such that |10 0,ll < 20,d we have V²C(0) (O.)).
2. The effective rank of H (0,), i.e., tr (O,))) (0))) is most
Theorem 2 (Dimension-Free Rate, Gaussian 2). Assume the loss exhibits local r-effective rank (Assumption 2). If 8.41 = 0, T20 ((,;
( B3) a single step of ZO-SGD using the n-SPSA estimate with minibatch of size B, then there exists O(r/n) such that the expected loss decrease can be bounded as
[E(C(0++)) 0.] C(0,)
- VL + .y - exp( exp(-S((nd))
Proof of Theorem 2. Let A be the event that ||0.1 - 0,11 s 2n.G(()). On A, we have that
C(0,+1) E ((0,) TC(0,)) 2 "
Likewise, since is f-smooth, we have that
((0,+1) s L(0,) 1)C(())
36

---

Therefore
E[((0 00) [(0++1) - 7lVv
+ 2" (B [C00:B)( 1(4)] .H(I))
1
1
|| [[(;;; B3 -1(-A)
+ ~" ATE
1, "²²²² ( [E .HO.))
= [(0++1) - ( [[
2 1"² E [CC0BB)C((0))) 11--0) ,1-
1
The latter term can be bounded follows
[.0 1(-A)), , e H(0))) (HT-H(O.)) < 1T
| 1(-AA)
2 r²
$ P
Pr|-LAj1/2.
The gradient estimate B) satisfies
Σ
2;
VL
1
iE|m]
The expectation term is upper bounded as
E
[
||]
Σ E 2> VC -|=11
n
sEin]
1/2
E $[12118] 1/2
s E > ():
S v10 6))
where we have plugged in explicit formulas for moments of Gaussian and x² random variables. Next, note that on the event -A, we have
(O.) S 10.11 Ocll - CL Sn Σ ||= G(O.)..
vE(n)
Therefore
Pr[-A] Pr
1 >
2nd
sEim]
Lemma (Standard x2 tail bound). Let be a x² random variable with k degrees of freedom. Then
(-
u2 W
Pr[Z 2 k + u] & exp
(-min minn
16k 16)
Since Eiejol ||=,|| is x² random variable with nd degrees of freedom, thus have that
nd
Pr[-A]] < exp
16
Altogether,
"" (B ( 1(-A) .11 11- ((H(()))) <
exp(= nd.
32
37

---

Finally, plugging in (8), along with the fact that || [H I and tr ((()) - lr,
E [[0 H(0)))
r +++1 1
(VVC(0))) + tx(z((,))
7
+ 1
E [C(00:
Thus letting = r+=++ yields
S - (E((((+11) 0,] -[(,)
+ 2 7 -[I (G) eeppppp(000 (xp(-(((dd))
desired.
38


Reviews:
[{"review_id": "S7HI0Ep7Pf", "reviewer": "Reviewer_qocq", "summary": "This work introduced a memory-efficient zeroth-order optimizer that can fine-tune large language models with the same memory footprint as inference, using only forward passes and gradient estimates. Comprehensive experiments across model types, scales, and tasks, showing that MeZO outperforms zero-shot, in-context learning, and linear probing, and achieves comparable performance to fine-tuning with backpropagation, while reducing memory cost by up to 12 times. Non-differentiable objectives that MeZO can optimize, such as accuracy or F1 score, which are usually not amenable to backpropagation. Theoretical insights that explain why MeZO can optimize LMs with billions of parameters, despite classical zeroth-order analyses suggesting otherwise.\n", "questions": "How is the training stability of MeZO? Does it an advantage over the backpropagation-based methods, especially for the large models?", "limitations": "1. It would be better to have the experiment results on the generation tasks, e.g. translation and summarization.\n\n2. I suggest the authors to have a discussion about the scope of application for the proposed method.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. It proposes a novel and memory-efficient method to fine-tune large language models without backpropagation, which can save up to 12x memory compared to standard methods.\n\n2. It demonstrates that the proposed method can achieve comparable or superior performance to fine-tuning with backpropagation across various tasks, models, and tuning techniques.\n\n3. It shows that the proposed method can optimize non-differentiable objectives, such as accuracy or F1 score, which are useful for many applications.\n\n4. It provides theoretical insights on why the proposed method can overcome the classical limitations of zeroth-order optimization and leverage the benefits of pre-training and task prompts.\n", "weaknesses": "1. While the experiments demonstrate good performance on the language understanding tasks, it is unknown whether the method is also applicable to the generation tasks.\n\n2. It relies on the assumption of low effective rank of the Hessian matrix, which may not hold for all loss functions. It would be great to have a discussion about the scope of application for the proposed method.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Does MeZO work for generation tasks?**\n\nTable 1 and Figure 1 show the performance of MeZO on DROP and SQuAD, which are two question-answering tasks that are formatted as generation tasks in our experiments. For each task, given the question, we train the model to directly generate the answer text (please see our Table 12 for details). We leave the study of more generation tasks like summarization and translation to future work.\n\n**When does the assumption of the low effective rank of the Hessian hold?**\n\nPlease refer to our general response. We hypothesize that when using a good prompt, the Hessian of the downstream objective likely exhibits a low effective rank.\n\n**How stable is MeZO training? Does it have an advantage over backpropagation?**\n\nMeZO is not very sensitive to hyperparameter choices. As shown in Tables 13 (RoBERTa-large) and 14 (OPT), we restrict the grid searches to a very narrow range of hyperparameters and often test MeZO with fewer configurations than we use to test fine-tuning with backpropagation. Also, as a gradient estimate, MeZO avoids well-known issues with backpropagation such as vanishing and exploding gradients, though these rarely occur when training networks with residual connections. MeZO is unstable in other ways, like if $\\epsilon$ must be set very small. However, in practice, we find that MeZO succeeds with a relatively large $\\epsilon$ and reduces the loss consistently over the course of training.  \n"}, {"review_id": "KqMj1R9nkq", "reviewer": "Reviewer_x4kH", "summary": "The paper proposes an enhanced memory efficient zero-order optimization method named MeZO. MeZO only requires the same memory as inference time and thus can enable model tuning for large LMs with limited memory budget. The authors demonstrate the efficacy of MeZO on multiple NLP benchmarks compared with linear probing, in context learning and fine-tuning in few/low shot learning regime. The authors also provided detailed theoretical proof for MeZO. ", "questions": "- Can authors also report practical training time compared with standard fine-tuning (e.g., in terms of # steps/ second)?\n\n- Maybe report average numbers as well in experiment results (e.g., Table 1)\n\n- Interested to see how the performance of MeZO compares with fine-tuning in the scenario of full training data (rather than k = 100, 500, etc). ", "limitations": "The authors have discussed limitations in the conclusion section and aim to explore them in future work.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "0 - The authors targeted a not well understood domain (zero order optimization) and open up new opportunities for future works. In the era of LLMs, this method can enable many future work especially for those who don't have access to large-scale GPU clusters. \n1 - Comprehensive experiments and ablation studies on the proposed method.\n2 - Strong theoretical support on the proposed method.\n3 - Good writing and flow which makes the paper easy to follow and understand. \n4 - Well-articulated future work.", "weaknesses": "No major weaknesses. I left some comments in the questions section and hope  the authors can answer and address. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your suggestion, and we will report average numbers in the experiment results in the next revision.\n\n**What is the practical training time compared to standard fine-tuning?**\n\nPlease refer to our general response for a wall clock time analysis. In short, MeZO reduces the number of GPU-hours needed to train large models, leading to a 2x GPU-hour reduction on a 30B model compared to Adam fine-tuning.\n\n**How does MeZO perform with full training data?**\n\nWe choose the fixed number of training example setting because of compute limitations. Some datasets have millions of examples, so fine-tuning on the entire dataset can be very expensive for the model scale we are studying. We will include more ablations showing how MeZO performance changes with the dataset size in the next revision.\n"}, {"review_id": "1MnLUO8WoU", "reviewer": "Reviewer_unzy", "summary": "The paper present a new optimiser MeZo based on stochastic approximation using gradient perturbation. \nThis optimiser is very memory efficient as it only requires to perform 2 forward passes with different deltas/epsilons on the parameters and multiple gaussian samplings.   These algorithms allows \"finetunning\" large language models to specific tasks very efficiently ( up to 30B on a single A100) yielding between x4 to x12 memory reductions. Since the proposed algorithm is an optimizer it can be combined with other standard techniques such as LORA or prefix tunning. All this is applied in finetunning setups similar to ICL. ", "questions": "While I have many question none of them affect the paper directly. \nIt would be nice however if some of the weakness could be discussed by authors: \n* in which task will this technique work? do you have a characterisation ?  have you tried other tasks where it failed ? \n* why the Hessian hypotheses ? is this based on some experimental or preliminar analysis ?\n* when ft surpasses this techniques? when we have to ft on hundreds of thousands of examples ? \n* why it only works with prompting ? \n", "limitations": "While the author do not focus on the limitations,  they are more or less clear per previous analysis. The focus of the paper is more on the direction of stressing the surprise of the technique with all known drawbacks is actually workig. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The algorithm is a new application of well known stochastic gradient approximation but many times forgotten due to their slowness. \nIt is very surprising that this algorithm works, and the authors provide a theoretical justification of why this could be working in this case. \nThey also acknowledge that despite of what it might sound this approach only works in prompting fine-tunning scenario [Appendix B.2]. \n\nThe results are insight full , the baselines are fair and the theoretical analysis is correct to the best of my knowledge.\n\nThis can settle as an alternative to in context learning with prompting, by  fine-tuning with  a limited set of examples (512). \n\nThe proposed technique seems to work on the benchmarks used and seems to surpass other techniques such as Zero-shot, LP, ICL, and approaches very close performance to fine-tunning. This can be set a cheap alternative to ICL for some tasks.\n\n", "weaknesses": "While there was a huge and titanic effort in the paper there were many questions that steam from the technique.\nThe first area which has not been explored too much and given as true is the need of having a prompt to apply MeZO.\nThis key ingredient is not well studied but rather given on some preliminary experiments, e.g. Table 5. \nWhy is MeZO not working w/o prompts, even some very simple prompts ? \n\nThe need of the prompt also raises the question of how this is related to ICL , as there are some works that suggest ICL maybe doing some alike to fine-tuning or back-propagation though the attention weights. Is this combination of prompting and MeZo that is guiding the forward propagations ? is there a mixed cooperation between the prompt and the stochastic technique ? How much of the prompt is needed ?\n\nAnother question would be how the different techniques behave as a function of the number of examples k. It would have been nice to see a plot for some models at least between ICL , MeZo, and possible ft on the selected tasks. Why have authors stopped at 512 ? why only 16 or 512 ? there are some dataset that contain more training examples. Why didn't they compare fine-tuning and MeZo in other setups with larger examples ?\n\nThere is the relationship between the task itself and the optimiser. It is not clear to me, in which tasks this will work properly. I suspect given the prompting above that this might only work on low-perplexity tasks or task in which prompting or ICL can generate good results and not in other more complex tasks. \n\nClearly this is maybe too much to address in the paper, but all aspects above point toward the little understanding the reader is left with on under which conditions this technique can be applied. The future work seems to already assume the MeZO algorithm is working and proved, but there is just an hypotheses and a very low link between the experiment conditions and the theory. The link is stablished as \"We attribute these phenomena to the Hessian of the loss exhibiting  small local effective rank.\" . It would have been nice to strengthen this connection with some experiments or computation. Could this explain when or how this algorithm is applied ? if we remove the prompt does this increases the H effective rank ? would other tasks exhibit larger ranks ?  how can we reduce it for each of the tasks ?\n\n\nPlease, I would kindly ask the authors to read above questions and discussion as a signal of the interest the paper brought to the research field and not as criticism on their very interesting work.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Why does MeZO require using a prompt? What tasks can MeZO work on? Why do you need the Hessian hypothesis (Assumption 1)?**\n\nPlease refer to our general response. In short, we hypothesize that using a prompt makes the fine-tuning objective similar to the pre-training one, which likely has a Hessian with a low effective rank. Following this, we agree with you that a task that has a lower perplexity (i.e., a more “natural” prompt) will probably work better with MeZO. Once we added a simple prompt, we did not encounter any tasks that MeZO completely failed to train on.\n\n**How does the empirical success of MeZO interact with the theoretical hypothesis that transformers may simulate fine-tuning on a smaller, internal model during inference time?**\n\nMeZO is a useful tool for fine-tuning currently popular LLMs trained with standard pre-training practices. It is not guaranteed to work for models that are designed or trained in new ways, such as the scenario you mention. It may be the case that the internal model is not stable to fine-tuning the large model. Alternatively, it could be stable to fine-tuning in some way (e.g., analogous to noise-tolerant circuits) that allows it to not be destroyed during MeZO. Overall, we are not sure if currently existing pre-trained models are simulating and updating internal models, so we cannot be sure how such constructed models would behave during fine-tuning, whether it is done with backpropagation or MeZO. \n\n**How does MeZO behave with different numbers of examples?**\n\nThanks for your question. We will include experiments ablating against different dataset sizes in a subsequent revision. "}, {"review_id": "cBASPGbK7o", "reviewer": "Reviewer_Qhm1", "summary": "This paper proposes a new zeroth order optimizer, MeZO, for LM training.  This is proposed as an improvement to ZO-SGD.  The advantage of this approach is a 12x reduction in the amount of memory required for training compared to backpropagation.  This enables the training of much larger models.\n\nThe effectiveness of MeZO is shown across a range of benchmarks and model sizes.  The results compare favorably to linear probing and in context learning.", "questions": "Appendix A (and Section 3) demonstrate that promoting is crucial to MeZO performance.  Why is this? Much of the other behavior is supported by a theoretical treatment, but this observation stands out as relatively uninterrogated.", "limitations": "yes though fairly lightly.  The observation about prompts being critical for training may be a limitation for some (new) tasks or datasets.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The MeZO technique stands to unlock substantial capability for LM training.  This enables training of much larger networks.  The compatibility with LoRA, prefix tuning are important use cases for many LM users.  There is an ability for optimizing non-differentiable objectives which is compelling, and could be expanded in the future.\n\nThe empirical behavior are coupled with a section on theory which effectively describes the both the expected behavior, but elaborates on the expected slow convergence by expanding the theoretical analysis to address low effective rank networks.", "weaknesses": "While the analysis refers to the convergence rate of MeZO, there is a very brief treatment of convergence behavior in the paper (Appendix E.2) It might be helpful for this to be expanded and possibly compared to backpropagation, especially in the context of the presentation of Section 4 Theory.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Can the theoretical convergence analysis of MeZO be compared to backpropagation?**\n\nCorollary 1 directly compares the SGD convergence rate to the convergence rate of MeZO, since the term in brackets in equation 5 is the per-step loss decrease of SGD (see Lemma 1). Two factors make MeZO converge more slowly than standard backpropagation (lines 211-218): (1) MeZO has to be run with a smaller learning rate than SGD in order to reliably decrease the loss at each step, and (2) MeZO reduces the amount that the loss can decrease at each step. \n\n**Why is prompting crucial to MeZO?**\n\nPlease refer to our general response. In short, we hypothesize that using a prompt makes the fine-tuning objective similar to the pre-training one, which likely exhibits a Hessian with low effective rank. \n"}, {"review_id": "Xuec5j7u8v", "reviewer": "Reviewer_XvSx", "summary": "Fine-tuning with backpropogation becomes infeasible for very large language models because it uses too much memory. While zeroth-order optimization uses far less memory and could in principle fine-tune the model with just forward passes, past theory suggested that the learning rate must scale down with the number of parameters, making convergence prohibitively slow. However, this paper finds that zeroth-order optimization actually performs quite well and converges quickly even on very large language models. They provide theory to explain this fast convergence, where they show that under an assumption they call \"low effective rank,\" the learning rate scales down with the rank rather than the number of parameters. They also provide a memory-efficient implementation of zeroth-order optimization that they call MeZO, along with memory efficient zeroth-order versions of SGD with momentum and Adam. In experiments, the method performs similarly to backpropogation with 1/12 the memory usage, while outperforming in-context learning and linear probing.", "questions": "(1) While MeZO is much more memory efficient, is it slower than backprop in terms of wall-clock time? (The appendix does state that FT used 1K steps while MeZO used 100K steps in the experiments. How much faster is each step of MeZO compared to each step of backprop?)\n\n(2) It's a bit surprising to me that MeZO, MeZO-prefix, and MeZO-LoRA optimize at similar speeds (Figure 5 in the appendix). Do the backprop versions also optimize at similar speeds? And do the three parameterizations actually have similar effective ranks empirically?", "limitations": "n/a", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "(1) The paper is well-written.\n\n(2) The method is simple and easy to understand.\n\n(3) The theory provides useful insights into why zeroth-order optimization works for fine-tuning large pre-trained models.\n\n(4) The experimental results are strong, and the appendix contains thorough ablations.\n\n(5) The idea that zeroth-order optimization works well for fine-tuning LMs seems practically useful and addresses a pressing need in the community for memory-efficient methods.", "weaknesses": "The paper seems strong overall, and I support its acceptance regardless of whether the suggested experiments below are run or not during the rebuttal period.\n\n(1) From what I understand, the paper does not verify the low effective rank assumption empirically, nor is it verified in the papers cited (which either study the effective rank / Hessian spectra in non-LLMs, or study the LLMs but not the effective rank and instead study the intrinsic dimensionality of fine-tuning). Therefore, to justify the assumption, it seems useful to study the Hessian spectra of the downstream fine-tuning loss in LLMs, at whatever size is feasible.\n\n(2) Related to (1), to verify the theory and confirm that the effective rank is indeed the quantity that determines convergence rates, it seems useful to run simulated experiments where one constructs a synthetic model + data and varies the effective rank, and examines whether the convergence rate or gradient norm scales with the effective rank as predicted in the theory.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Can we study the effective rank assumption in some language models?**\n\nIt is difficult to translate results on very small models, on which we would be able to measure the effective rank, to the large ones that we would find MeZO useful for. We would also likely need to pre-train these very small models ourselves, which is expensive, and they may be so small that they do not achieve meaningful results on the benchmarks we study.\n\n**Can we use simulations to verify the theoretical convergence analysis?**\n\nThanks for the suggestion! We added the simulated experiment in our attached PDF and also reported the results in our general response. In short, we observed that the convergence rate of MeZO does depend on the effective rank in the simulated experiments.\n\n\n**What is the wall-clock efficiency of MeZO compared to fine-tuning with backpropagation?**\n\nPlease see our general response. In short, MeZO reduces the number of GPU-hours needed to train large models, leading to a 2x GPU-hour reduction on a 30B model compared to Adam fine-tuning.\n\n**Do the backpropagation versions of full fine-tuning, prefix tuning, and LoRA optimize at roughly the same rates, as was observed with MeZO? And do the three parameterizations actually have similar effective ranks empirically?**\n\nWe did not measure the empirical effective rank of different methods due to limited compute. We observe that with backpropagation, all three methods converge roughly at a similar speed, with LoRA and prefix-tuning slightly slower on some tasks. The interesting case with MeZO is that classical ZO analyses suggest that full-parameter MeZO would converge much more slowly, but it is not the case empirically. Our theory in Section 4 highlights why the convergence rate does not depend on the number of parameters.\n"}]

Author rebuttal:
We thank all reviewers for their valuable feedback. We address some shared questions here.

**When can MeZO succeed in fine-tuning? What losses satisfy Assumption 1 (i.e., the Hessian has a low effective rank)? Why is a prompt necessary for MeZO to be able to fine-tune the model? Can you verify the dependence of MeZO convergence rate on the effective rank?**

Our theory in Section 4 provides a sufficient (but not necessary) condition for MeZO to succeed: the Hessian should exhibit a small local effective rank during fine-tuning (Assumption 1). We hypothesize that the Hessian of the pre-training objective likely exhibits low effective rank, because the model has been trained for many steps during pre-training. Ample evidence suggests that training for many steps can make the Hessian have a low effective rank in the case of vision (see lines 222-228). Adding a prompt turns the downstream task into next-word prediction [1] (or masked language modeling [2]). So, the Hessian of the fine-tuning objective when using a prompt (similar to pre-training) likely exhibits a small effective rank like the pre-training one [1, 3]. There is additional empirical evidence that the Hessian of a language model during fine-tuning likely has low rank [4].

Additionally, per reviewer XvSx’s suggestion, we ran simulations in a simple setting to verify the dependence of MeZO convergence rate on the effective rank and reported the results in the attached PDF. We observed that the slowdown of the convergence scales with the effective rank. We will include these experiments in the next revision of the paper.

[1] Nikunj Saunshi, Sadhika Malladi, Sanjeev Arora. A Mathematical Exploration of Why Language Models Help Solve Downstream Tasks. ICLR 2021.

[2] Tianyu Gao, Adam Fisch, Danqi Chen. Making Pre-trained Language Models Better Few-shot Learners. ACL 2021.

[3] Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, Sanjeev Arora. A Kernel-Based View of Language Model Fine-Tuning. ICML 2023.

[4] Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoyer. Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning. ACL 2021.

**What is the wall-clock efficiency of MeZO compared to standard training with backpropagation?**

The attached PDF with this rebuttal shows the wall-clock efficiency of MeZO compared to fine-tuning. MeZO takes more steps to achieve similar performance than fine-tuning, but requires much less wall-clock time per step and requires fewer GPUs. The gains are more prominent for larger models, which are the ones that require more memory to fine-tune. For example, for a 30B model, we show that MeZO enjoys a 7.74x per-step speed up and a 2x total GPU-hour reduction compared to fine-tuning with Adam. We will include these results in the next revision of the paper.






Meta-review: This paper proposed a memory-efficient zeroth-order optimizer that can fine-tune large language models using only forward passes. All reviewers give positive ratings to this paper and mention some merits of this paper such as targeting a not well-understood domain (zero order optimization), with good experiment and theoretical support, and good paper writing. Although reviewers ask several questions including whether it can apply to the generation task, or when the assumption of the low effective rank of the Hessian matrix holds, the authors answer these questions well. Overall, it is a strong paper with a clear problem setting, nice experiment results, theoretical support, and practical application. I recommend accepting this paper as a spotlight.
Decision: Accept (oral)