Title: Fine-Tuning Language Models with Just Forward Passes
Venue: NeurIPS 2023
Accepted: yes
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

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 inplace, 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. 1

1 Introduction

Fine-tuning pre-trained language models (LMs) has been the dominant methodology for solving many language tasks [28], adapting to specialized domains [42], or incorporating human instructions and preferences [73]. However, as LMs are scaled up [13, 72], computing gradients for backpropagation requires a prohibitive amount of memory - in our test, up to 12 × the memory required for inference because it needs to cache activations during the forward pass, gradients during the backward pass, and, in the case of Adam [52], also store gradient history (see Section 3.4 for a detailed analysis).

As a result, while it is possible to run inference with a 30-billion (30B) parameter LM on a single Nvidia A100 GPU (with 80GB memory), backpropagation with Adam is feasible only for a 2.7B LM. Parameter-efficient fine-tuning methods (PEFT [46, 57, 54]) update just a fraction of the network

SST-2

RTE

CB

BoolQ

WSC

WIC

MultiRC

Copa

ReCoRD

SQuAD

DROP

10

30

50

70

90

Accuracy/F1 (%)

Zero-shot

ICL

MeZO

FT (12x memory)

parameters, but still need to cache many activations, because the tuned parameters are scattered throughout the model. In our tests, fine-tuning an OPT-13B model with full parameter tuning or PEFT requires 12 × and 6 × more memory than inference respectively.

In-context learning (ICL [13]) has allowed solving many tasks with a single inference pass, during which the model processes labeled examples ( demonstrations ) in its context and then outputs a prediction on a test example. While this allows for quick adaptation of the model to specific use cases, current models allow a limited context size (and thus, limited demonstrations) and the performance is sensitive to the formatting and choice of demonstrations [60, 66]. ICL can slow with the number of demonstrations, and it often performs worse than fine-tuning of medium-sized models [13].

Backpropagation also cannot optimize non-differentiable criteria, which have gained popularity in fine-tuning LMs according to human preference scores or set safety standards [89, 73]. Typically, these adaptations involve expensive reinforcement learning from human feedback (RLHF [20]).

A classical zeroth-order optimization method, ZO-SGD [88], uses only differences of loss values to estimate the gradients. Thus, in principle, the method can update neural networks with just forward passes, though naive implementation still doubles the memory overhead and classical lower bounds [69, 32] suggest that convergence slows linearly with model size. As such, ZO methods have been applied in deep learning settings to find adversarial examples or tune input embeddings [91, 90] but not to directly optimize large-scale models (see Liu et al. [61] for a survey).

In this work, we propose a memory-efficient zeroth-order optimizer (MeZO), which adapts the classical ZO-SGD algorithm and reduces its memory consumption to the same as inference . We apply MeZO to fine-tune large LMs and show that, both empirically and theoretically, MeZO can successfully optimize LMs with billions of parameters. Specifically, our contributions are:

In MeZO, we adapt the ZO-SGD algorithm [88] and a number of variants to operate in-place on arbitrarily large models with almost no memory overhead (see Algorithm 1 and Section 2).

We conduct comprehensive experiments across model types (masked LM and autoregressive LM), model scales (from 350M to 66B), and downstream tasks (classification, multiple-choice, and generation). MeZO consistently outperforms zero-shot, ICL, and linear probing. Moreover, with RoBERTa-large, MeZO achieves performance close to standard fine-tuning within 5% gap; with OPT-13B, MeZO outperforms or performs comparably to fine-tuning on 7 out of 11 tasks, despite requiring roughly 12 × less memory (Figure 1 and Section 3). In our implementation, MeZO requires only half as many GPU-hours as Adam fine-tuning for a 30B model (see Appendix F.6).

We demonstrate MeZO's compatibility with full-parameter tuning and PEFT (e.g., LoRA [46] and prefix-tuning [57]) in Section 3.

Further exploration showcases that MeZO can optimize non-differentiable objectives such as accuracy or F1 score, while still requiring only the same memory as inference (Section 3.3).

Our theory suggests that adequate pre-training ensures the per-step optimization rate (Theorem 1) and global convergence rate (Lemma 3) of MeZO depend on a certain condition number of the landscape (i.e., the local effective rank, see Assumption 1) instead of numbers of parameters. This

Algorithm 1: MeZO

result is in sharp contrast to existing ZO lower bounds [69, 32] suggesting that the convergence rate can slow proportionally to the number of parameters (Section 4).

2 Zeroth-order optimization

Zeroth-order (ZO) optimizers have long been studied in the context of convex and strongly convex objectives. In the following, we first introduce a classical ZO gradient estimator, SPSA (Definition 1 [88]) and the corresponding SGD algorithm, ZO-SGD (Definition 2). Then we describe MeZO, our in-place implementation that requires the same memory as inference in Section 2.1 and Algorithm 1. We highlight that SPSA can also be used in more complex optimizers, such as Adam, and we provide memory-efficient implementations for those algorithms too (Section 2.2).

Consider a labelled dataset D = { ( x i , y i ) } i ∈ [ |D| ] and a minibatch B ⊂ D of size B , we let L ( θ ; B ) denote the loss on the minibatch. We introduce a classical ZO gradient estimate in this setting. 2

Definition 1 (Simultaneous Perturbation Stochastic Approximation or SPSA [88]) . Given a model with parameters θ ∈ R d and a loss function L , SPSA estimates the gradient on a minibatch B as

where z ∈ R d with z ∼ N (0 , I d ) and ϵ is the perturbation scale . The n -SPSA gradient estimate averages ̂ ∇L ( θ ; B ) over n randomly sampled z .

SPSA requires only two forward passes through the model to compute the gradient estimate (for n -SPSA, each estimate requires 2 n forward passes). As ϵ → 0 , the SPSA estimate can be understood as a rank-1 reconstruction of the gradient. During training, n can be treated as a hyperparameter and follow a schedule [11, 15], though in cursory experiments (Appendix A), n = 1 is the most efficient.

We use n = 1 as the default. It is widely known that the SPSA estimate can be used to replace the backpropagation gradient in any optimizer such as SGD.

Definition 2 (ZO-SGD) . ZO-SGD is an optimizer with learning rate η that updates parameters as θ t +1 = θ t -η ̂ ∇L ( θ ; B t ) where B t is the minibatch at time t and ̂ ∇L is the SPSA gradient estimate.

2.1 Memory-efficient ZO-SGD (MeZO)

The vanilla ZO-SGD algorithm costs twice the memory of inference, as it needs to store z ∈ R d . We propose a memory-efficient implementation of ZO-SGD called MeZO , as illustrated in Algorithm 1. At each step, we first sample a random seed s , and then for each of z 's four uses in Algorithm 1, we reset the random number generator by s and resample the relevant entry of z . Using this in-place implementation, MeZO has a memory footprint equivalent to the inference memory cost.

We note that Algorithm 1 describes perturbing each parameter separately, which may be timeconsuming for large models. In practice, we can save time by perturbing an entire weight matrix instead of each scalar independently. This incurs an additional memory cost as large as the largest weight matrix; usually, this is the word embedding matrix (e.g., 0.86GB for OPT-66B).

Storage Efficiency of MeZO. Parameter-efficient fine-tuning (PEFT) techniques fine-tune just a fraction of the network parameters and have thus been proposed as a way to reduce the storage costs of fine-tuned model checkpoints. Fine-tuning with MeZO reduces the storage cost of the resulting checkpoint far more than popular PEFT techniques (e.g., LoRA [46] and prefix tuning [57]). We reconstruct the MeZO trajectory using a single seed, which spawns step-wise seeds to sample z , and the projected_grad at each step. 3 As such, for fine-tuning a 66B model, MeZO requires saving the seed plus 20,000 (steps) × 2 bytes, which is less than 0.1MB. LoRA fine-tunes 19M parameters and requires 38MB storage, and prefix tuning fine-tunes 6M parameters and requires 12MB storage.

2.2 MeZO extensions

We note that SPSA is a popular ZO gradient estimator but not the only one. Many one-point gradient estimators have been proposed in past works [34, 87, 95], and using such estimators in place of SPSA would halve the training time. However, cursory experiments with one such promising estimator [113] reveal that these are not as efficient as SPSA when fixing the number of forward passes (Appendix B.5). As such, we implement MeZO with the SPSA estimator.

MeZO can also be combined with other gradient-based optimizers, including SGD with momentum or Adam. Though naive implementation would require additional memory to store the gradient moment estimates, MeZO-momentum and MeZO-Adam alleviate such overhead by recomputing the moving average of the gradients using saved past losses and z (see Appendix B for a full discussion).

We also note that all of the coordinates of the SPSA gradient estimate have the same scale, but deep Transformers can have gradients of different scales for each layer [59, 61]. As such, we draw inspiration from layerwise adaptive optimizers [109, 110] to design several MeZO variants. Cursory experiments showed that these algorithms are not more efficient (in terms of forward passes), but we nevertheless present them as potential optimizers for more complex objectives. See Appendix B.

Forward Auto-Differentiation Note that z ⊤ ∇L ( θ ; B ) is a Jacobian-vector product (JVP), which can be computed in parallel with an inference pass with excess memory consumption equivalent to that of the largest activation in the network [40]. In this case, z must be stored on the GPU in order to construct the gradient estimate, so this procedure requires slightly more than two times the memory needed for inference. We analyze this algorithm in detail in Appendix D. Note that using a non-zero ϵ in SPSA, which is not possible through the JVP method, may boost generalization by promoting a sharpness-minimizing term. Past works (e.g., Baydin et al. [9]) have also studied JVP-based training but achieved limited empirical success.

SST-2

SST-5

SNLI

MNLI

RTE

TREC

30

40

50

60

70

80

90

100

Accuracy (%)

k=16 RoBERTa-large

SST-2

SST-5

SNLI

MNLI

RTE

TREC

30

40

50

60

70

80

90

100

k=512 RoBERTa-large

Zero-shot

MeZO (prefix)

LP

FT

MeZO

FT (LoRA)

MeZO (LoRA)

FT (prefix)

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 a suitable prompt ensures the fine-tuning objective is closely related to the pre-training one. In Section 4, we extend these ideas to show how using a simple prompt simplifies the fine-tuning optimization procedure, thereby enabling zeroth order methods to work efficiently. All experiments below use prompts detailed in Appendix E.2. All fine-tuning with backpropagation (FT) experiments follow convention and use Adam, though we also report results when performing FT with SGD in Appendix F.

We conduct comprehensive experiments on both medium-sized masked LMs (RoBERTa-large, 350M [65]) and large autoregressive LMs (OPT-13B, 30B, 66B [112]) in few-shot and many-shot settings with prompts. We also explore both full-parameter tuning and PEFT including LoRA [46] and prefix-tuning [57] (see Appendix E.5 for details). We compare MeZO with zero-shot, in-context learning (ICL), linear-probing (LP), and fine-tuning with Adam (FT). MeZO uses substantially less memory than FT but requires significantly more training steps.

Wefirst show that MeZO improves substantially over zero-shot, ICL, and LP across model types, sizes, and task types. Moreover, MeZO performs comparably to FT over a number of tasks, while drastically reducing the memory cost by, for example, 12 × on OPT-13B. Further experiments demonstrate that MeZO can optimize non-differentiable objectives, such as accuracy and F1 score (Section 3.3). We compare the memory consumption of ICL, FT, LP, and MeZO in Figures 3 and 4.

3.1 Medium-sized masked language models

We conduct experiments with RoBERTa-large on sentiment classification, natural language inference, and topic classification tasks. We follow past works [35, 67] in studying the few-shot and many-shot settings, sampling k examples per class for k = 16 and k = 512 (details in Appendix E). We run MeZO for 100 K steps and fine-tuning for 1000 steps, noting that one MeZO step is substantially faster than one fine-tuning step (see Appendix F.6 for a comparison). We summarize the results from Figure 2 and Table 18 below.

MeZO works significantly better than zero-shot, linear probing, and other memory-equivalent methods. On all six diverse tasks, MeZO can optimize the pre-trained model and consistently perform better than zero-shot and linear probing. We also show for several tasks that MeZO can outperform another ZO algorithm, BBTv2 [90], by up to 11% absolute (Appendix F.4). 4

With enough data, MeZO achieves comparable performance (up to 5% gap) to FT. MeZO achieves close-to-fine-tuning performance on k = 16 , with some tasks only having 2% gaps. When using k = 512 data, the gap between MeZO and FT further reduced to within 5% across all tasks.

MeZO works well on both full-parameter tuning and PEFT. Full-parameter tuning (MeZO) and PEFT (MeZO with LoRA and prefix-tuning) achieve comparable performance, while MeZO (prefix)

sometimes outperforms MeZO. We also show in Appendix F.3 that the three variants converge at similar rates, agreeing with our theory in Section 4, which shows that MeZO converges at a rate independent of the number of parameters being optimized.

We show additional results with more FT and MeZO variants in Appendix F.1. We see that (1) ZO-Adam sometimes outperforms ZO-SGD but is not consistent across tasks; (2) LP and then MeZO, as suggested for fine-tuning [53], can sometimes improve the performance.

3.2 Large autoregressive language models

With the promising results from RoBERTa-large, we extend MeZO to the OPT family [112], on a scale of 13B (Table 1), 30B, and 66B (Table 2). We select both SuperGLUE [98] tasks 5 (including classification and multiple-choice) and generation tasks. We randomly sample 1000 , 500 , and 1000 examples for training, validation, and test, respectively, for each datset. We run MeZO for 20 K steps and fine-tuning for 5 epochs, or 625 steps, noting that each step of MeZO is substantially faster than fine-tuning (see Appendix F.6 for a comparison). Please refer to Appendix E for details. Table 1 yields the following observations.

MeZO outperforms memory-equivalent methods and closely approaches fine-tuning results. We see that on a 13B-parameter scale, MeZO and its PEFT variants outperform zero-shot, ICL, and LP across almost all tasks. When comparing to FT, which costs 12 × more memory (Section 3.4), MeZO achieves comparable (within 1%) or better performance on 7 out of the 11 tasks.

MeZO exhibits strong performance across classification, multiple-choice, and generation tasks. We investigate MeZO on generation tasks, which are regarded as more intricate than classification or multiple-choice tasks. We evaluate on two question answering datasets, SQuAD [80] and DROP [31]. We use teacher forcing for training and greedy decoding for inference (details in Appendix E).

Table 1 shows that, on all generation tasks, MeZO outperforms zero-shot, ICL, and LP, and achieves comparable performance to FT. Considering that many applications of fine-tuning LMs - including instruction tuning or domain adaptation - target generation tasks, our results underscore the potential of MeZO as a memory-efficient technique to optimize large LMs for realistic and exciting applications.

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=14,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>

GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=14,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>

GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>

GLYPH<c=18,font=/AAAAAB+DejaVuSans>GLYPH<c=23,font=/AAAAAB+DejaVuSans>GLYPH<c=25,font=/AAAAAB+DejaVuSans>GLYPH<c=3,font=/AAAAAB+DejaVuSans>GLYPH<c=21,font=/AAAAAB+DejaVuSans>GLYPH<c=28,font=/AAAAAB+DejaVuSans> !#'GLYPH<c=3,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>GLYPH<c=18,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=28,font=/AAAAAB+DejaVuSans># ! GLYPH<c=7,font=/AAAAAB+DejaVuSans>$GLYPH<c=30,font=/AAAAAB+DejaVuSans>!%

GLYPH<c=19,font=/AAAAAB+DejaVuSans>GLYPH<c=16,font=/AAAAAB+DejaVuSans>GLYPH<c=20,font=/AAAAAB+DejaVuSans>

GLYPH<c=17,font=/AAAAAB+DejaVuSans> GLYPH<c=24,font=/AAAAAB+DejaVuSans>

GLYPH<c=17,font=/AAAAAB+DejaVuSans> GLYPH<c=24,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAB+DejaVuSans>"# GLYPH<c=28,font=/AAAAAB+DejaVuSans>GLYPH<c=29,font=/AAAAAB+DejaVuSans>GLYPH<c=31,font=/AAAAAB+DejaVuSans>&

GLYPH<c=26,font=/AAAAAB+DejaVuSans>GLYPH<c=22,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAC+DejaVuSans>GLYPH<c=8,font=/AAAAAC+DejaVuSans>GLYPH<c=12,font=/AAAAAC+DejaVuSans>GLYPH<c=15,font=/AAAAAC+DejaVuSans>

GLYPH<c=11,font=/AAAAAC+DejaVuSans>GLYPH<c=8,font=/AAAAAC+DejaVuSans>GLYPH<c=14,font=/AAAAAC+DejaVuSans>GLYPH<c=15,font=/AAAAAC+DejaVuSans>

GLYPH<c=13,font=/AAAAAC+DejaVuSans>GLYPH<c=8,font=/AAAAAC+DejaVuSans>GLYPH<c=14,font=/AAAAAC+DejaVuSans>GLYPH<c=15,font=/AAAAAC+DejaVuSans>

GLYPH<c=10,font=/AAAAAC+DejaVuSans>GLYPH<c=12,font=/AAAAAC+DejaVuSans>GLYPH<c=15,font=/AAAAAC+DejaVuSans>

GLYPH<c=12,font=/AAAAAC+DejaVuSans>GLYPH<c=9,font=/AAAAAC+DejaVuSans>GLYPH<c=15,font=/AAAAAC+DejaVuSans>

GLYPH<c=10,font=/AAAAAC+DejaVuSans>

GLYPH<c=10,font=/AAAAAC+DejaVuSans>GLYPH<c=9,font=/AAAAAC+DejaVuSans>

GLYPH<c=10,font=/AAAAAC+DejaVuSans>GLYPH<c=9,font=/AAAAAC+DejaVuSans>GLYPH<c=9,font=/AAAAAC+DejaVuSans>

GLYPH<c=10,font=/AAAAAC+DejaVuSans>GLYPH<c=9,font=/AAAAAC+DejaVuSans>GLYPH<c=9,font=/AAAAAC+DejaVuSans>GLYPH<c=9,font=/AAAAAC+DejaVuSans>

GLYPH<c=18,font=/AAAAAC+DejaVuSans>GLYPH<c=23,font=/AAAAAC+DejaVuSans>GLYPH<c=25,font=/AAAAAC+DejaVuSans>GLYPH<c=3,font=/AAAAAC+DejaVuSans>GLYPH<c=21,font=/AAAAAC+DejaVuSans>GLYPH<c=28,font=/AAAAAC+DejaVuSans> !#'GLYPH<c=3,font=/AAAAAC+DejaVuSans>GLYPH<c=5,font=/AAAAAC+DejaVuSans>GLYPH<c=18,font=/AAAAAC+DejaVuSans>GLYPH<c=15,font=/AAAAAC+DejaVuSans>GLYPH<c=6,font=/AAAAAC+DejaVuSans>

GLYPH<c=26,font=/AAAAAC+DejaVuSans>GLYPH<c=28,font=/AAAAAC+DejaVuSans># ! GLYPH<c=7,font=/AAAAAC+DejaVuSans>$GLYPH<c=30,font=/AAAAAC+DejaVuSans>!%

GLYPH<c=19,font=/AAAAAC+DejaVuSans>GLYPH<c=16,font=/AAAAAC+DejaVuSans>GLYPH<c=20,font=/AAAAAC+DejaVuSans>

GLYPH<c=17,font=/AAAAAC+DejaVuSans> GLYPH<c=24,font=/AAAAAC+DejaVuSans>

GLYPH<c=17,font=/AAAAAC+DejaVuSans> GLYPH<c=24,font=/AAAAAC+DejaVuSans>GLYPH<c=3,font=/AAAAAC+DejaVuSans>GLYPH<c=5,font=/AAAAAC+DejaVuSans>"# GLYPH<c=28,font=/AAAAAC+DejaVuSans>GLYPH<c=29,font=/AAAAAC+DejaVuSans>GLYPH<c=31,font=/AAAAAC+DejaVuSans>&GLYPH<c=6,font=/AAAAAC+DejaVuSans>

GLYPH<c=21,font=/AAAAAC+DejaVuSans>GLYPH<c=28,font=/AAAAAC+DejaVuSans>GLYPH<c=26,font=/AAAAAC+DejaVuSans>GLYPH<c=22,font=/AAAAAC+DejaVuSans>

7x

8x

11x

12x

11x

GLYPH<c=4,font=/AAAAAB+DejaVuSans>GLYPH<c=23,font=/AAAAAB+DejaVuSans>

GLYPH<c=27,font=/AAAAAB+DejaVuSans>#GLYPH<c=27,font=/AAAAAB+DejaVuSans> GLYPH<c=28,font=/AAAAAB+DejaVuSans>%GLYPH<c=28,font=/AAAAAB+DejaVuSans>#$

GLYPH<c=4,font=/AAAAAC+DejaVuSans>GLYPH<c=23,font=/AAAAAC+DejaVuSans>

GLYPH<c=27,font=/AAAAAC+DejaVuSans>#GLYPH<c=27,font=/AAAAAC+DejaVuSans> GLYPH<c=28,font=/AAAAAC+DejaVuSans>%GLYPH<c=28,font=/AAAAAC+DejaVuSans>#$

MeZO scales up to 66 billion parameter models. We demonstrate the efficacy of MeZO on even larger models, up to 66B, in Table 2. While directly fine-tuning models at such scales is extremely costly (Section 3.4), MeZO can effectively optimize these models and outperform zero-shot and ICL.

3.3 Training with non-differentiable objectives

We demonstrate the efficacy of MeZO for optimizing non-differentiable objectives through initial experiments. Accuracy and F1 are used as the respective objectives (details in Appendix E.6). Table 3 reveals that MeZO with accuracy/F1 successfully optimizes LMs with superior performance to zero-shot. Although minimizing cross entropy results in stronger performance, these preliminary findings highlight the promising potential of applying MeZO to optimize non-differentiable objectives without clear differentiable surrogates, such as human preferences [73].

3.4 Memory usage and wall-clock time analysis

In this section we profile the memory usage of zero-shot, ICL, FT, FT (prefix), and MeZO. We test OPT models of various sizes with Nvidia 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 3 (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, as illustrated in Figure 4. Specifically, using a single A100 GPU, MeZO allows for tuning a model that is 11 times larger than what is feasible with FT.

In Appendix F.6, we compare the wall-clock time efficiencies of our implementations of MeZO and Adam fine-tuning. MeZO achieves 7.74 × per-step speedup and requires 8 × fewer GPUs with a 30B model, but takes more steps to converge. Overall, MeZO only requires half as many GPU-hours to fine-tune a 30B model compared to full-parameter fine-tuning. The wall-clock benefit of MeZO is not inherent to the algorithm and is highly dependent on 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 of MeZO and backpropagation and find that MeZO is always more memory-efficient 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.

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 a simple prompt is crucial for MeZO to succeed (Appendix A). Past works [83, 67] have suggested that including such a prompt ensures that the fine-tuning objective is 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 MeZO is not catastrophically slower than SGD when fine-tuning. 6 For ease of illustration, we assume that z is sampled from a sphere with radius √ d , and in Appendix G.2, we derive the rate for a general Gaussian z , which was used in the experiments.

We follow classical analyses of SGD and replace the minibatch gradient estimate with SPSA (Definition 1). Consider the minibatch SGD update θ t +1 ← θ t -η ∇L ( θ ; B t ) where B t is a minibatch drawn uniformly from D B . Crucially, the SGD minibatch gradient estimate is unbiased.

Definition 3 (Unbiased Gradient Estimate) . Any minibatch gradient estimate g ( θ , B ) is said to be unbiased if E [ g ( θ , B )] = ∇L ( θ ) .

4.1 Per-step analysis

The classical descent lemma uses a 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 1 (Descent Lemma) . Let L ( θ ) be ℓ -smooth. 7 For any unbiased gradient estimate g ( θ , B ) ,

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

where n is the number of z sampled in n -SPSA (Definition 1) and d is the number of parameters.

Thus, in the usual case where n ≪ d , MeZO has a much larger gradient norm than SGD. 8 The descent lemma also shows that to guarantee loss decrease, one needs to choose the learning rate as

where η ZO and η SGD are the maximum permissible learning rates for MeZO and SGD respectively. Thus we see that without any further assumptions, MeZO can slow optimization by decreasing the largest permissible learning rate by a factor of d . Moreover, MeZO reduces the loss decrease that can be obtained at each step and, as a consequence, slows convergence by a 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 substantially 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 a large LM 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

effective rank [74, 75, 36, 107, 105, 82]. In particular, the bulk of the spectrum concentrates around 0 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 a very low dimensional subspace ( < 200 parameters), which further supports the below assumption. We formalize the assumption on the effective rank below. In particular, we require an upper bound on the Hessian in a neighborhood around the current iterate to have effective rank at most r .

Assumption 1 (Local r -effective rank) . Let G ( θ t ) = max ( x , y ) ∈D ∥∇L ( θ t ; { ( x , y ) } ) ∥ . There exists a matrix H ( θ t ) ⪯ ℓ · I d such that:

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). If θ t +1 = θ t -η ZO ̂ ∇L ( θ t ; B ) is a single step of ZO-SGD using the n -SPSA estimate with a minibatch of size B , then there exists a γ = Θ( r/n ) such that the expected loss decrease can be bounded as

By applying Equation (3), we can directly compare to the SGD descent lemma.

-1

Corollary 1. Choosing the learning rate η ZO = γ · η SGD, ZO-SGD obtains a loss decrease of

Here we see that comparing to 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 to an optimal value.

4.2 Global convergence analysis

We show that the global convergence rate also slows by a 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 of the iterate.

Definition 4 (PL Inequality) . Let L ∗ = min θ L ( θ ) . The loss L is µ -PL if, for all θ , 1 2 ∥∇L ( θ ) ∥ 2 ≥ µ ( L ( θ ) -L ∗ ) .

The PL inequality is not as strong as 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 is bounded, so noise does not disrupt the trajectory too drastically. Definition 5 (Gradient Covariance) . The SGD gradient estimate on a minibatch of size B has covariance Σ ( θ ) = B ( E [ ∇L ( θ ; B ) ∇L ( θ ; B ) ⊤ ] -∇L ( θ ) ∇L ( θ ) ⊤ ) .

As we show in Appendix G.1, 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 a slowdown proportional to the effective rank r , not the parameter dimension. Lemma 3 (Global Convergence of ZO-SGD) . Let L ( θ ) be µ -PL and let there exist α such that tr( Σ ( θ )) ≤ α ( L ( θ ) -L ∗ ) for all θ . Then after

iterations of ZO-SGD we have E [ L ( θ t )] ≤ L ∗ + ϵ.

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 Ω( sd 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 black-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-efficient backpropagation Several algorithms have been proposed to efficiently approximate backpropagation by sparsifying gradients [92, 102], approximating Jacobians [1, 19], and subsampling the computational graph [71, 2]. However, these methods may accrue large approximation 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 to only optimize a low-dimension projection of prefixes and they focus on RoBERTalarge 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.

6 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 a limitation, MeZO takes many steps in order to achieve strong performance, though we show that the per-step speedup in MeZO can often make fine-tuning with MeZO run faster than a standard implementation of fine-tuning with backpropagation. We did not explore combining MeZO with other memory-efficient 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 a number of areas, including but not limited to: pruning, distillation, saliency, interpretability, and dataset selection for fine-tuning. Non-differentiable objectives are a 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.

References

Hany S Abdel-Khalik, Paul D Hovland, Andrew Lyons, Tracy E Stover, and Jean Utke. A low rank approach to automatic differentiation. In Advances in Automatic Differentiation , pages 55-65, 2008.

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.

Alekh Agarwal, Peter L. Bartlett, Pradeep Ravikumar, and Martin J. Wainwright. Informationtheoretic lower bounds on the oracle complexity of stochastic convex optimization. IEEE Transactions on Information Theory , 58(5):3235-3249, May 2012.

Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoyer. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. 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 7319-7328, 2021.

Stephen H Bach, Victor Sanh, Zheng-Xin 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.

Krishnakumar Balasubramanian and Saeed Ghadimi. Zeroth-order (non)-convex stochastic optimization via conditional gradient and gradient updates. In Advances in Neural Information Processing Systems , volume 31, 2018.

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.

Roy Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and Idan Szpektor. The second PASCAL recognising textual entailment challenge. In Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment , 2006.

Atılım Güne¸ s Baydin, Barak A. Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without backpropagation, 2022.

Luisa Bentivogli, Peter Clark, Ido Dagan, and Danilo Giampiccolo. The fifth PASCAL recognizing textual entailment challenge. In TAC , 2009.

Raghu Bollapragada, Richard Byrd, and Jorge Nocedal. Adaptive sampling strategies for stochastic optimization. SIAM Journal on Optimization , 28(4):3312-3343, 2018.

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.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in neural information processing systems , volume 33, pages 1877-1901, 2020.

HanQin Cai, Yuchen Lou, Daniel McKenzie, and Wotao Yin. A zeroth-order block coordinate descent algorithm for huge-scale black-box optimization. In International Conference on Machine Learning , pages 1193-1203, 2021.

HanQin Cai, Daniel McKenzie, Wotao Yin, and Zhenliang Zhang. Zeroth-order regularized optimization (zoro): Approximately sparse gradients and adaptive sampling. SIAM Journal on Optimization , 32(2):687-714, 2022.

Yekun Chai, Shuohuan Wang, Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. Cliptuning: Towards derivative-free prompt learning with a mixture of rewards. arXiv preprint arXiv:2210.12050 , 2022.

Pin-Yu Chen, Huan Zhang, Yash Sharma, Jinfeng Yi, and Cho-Jui Hsieh. Zoo: Zeroth order optimization based black-box attacks to deep neural networks without training substitute models. In Proceedings of the 10th ACM workshop on artificial intelligence and security , pages 15-26, 2017.

Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 , 2016.

Krzysztof M Choromanski and Vikas Sindhwani. On blackbox backpropagation and jacobian sensing. In Advances in Neural Information Processing Systems , volume 30, 2017.

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.

Christopher Clark, Kenton Lee, Ming-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 1 (Long and Short Papers) , pages 2924-2936, 2019.

Ido Dagan, Oren Glickman, and Bernardo Magnini. The PASCAL recognising textual entailment challenge. In the First International Conference on Machine Learning Challenges: Evaluating Predictive Uncertainty Visual Object Classification, and Recognizing Textual Entailment , 2005.

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.

Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser. The commitmentbank: Investigating projection in naturally occurring discourse. In Sinn und Bedeutung , volume 23, pages 107-124, 2019.

Mingkai Deng, Jianyu Wang, Cheng-Ping Hsieh, 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.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. GPT3.int8(): 8-bit matrix multiplication for transformers at scale. In Advances in Neural Information Processing Systems , 2022.

Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via blockwise quantization. In International Conference on Learning Representations , 2022.

Jacob Devlin, Ming-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 1 (Long and Short Papers) , pages 4171-4186, 2019.

Shizhe Diao, Xuechun Li, Yong Lin, Zhichao Huang, and Tong Zhang. Black-box prompt learning for pre-trained language models. arXiv preprint arXiv:2201.08531 , 2022.

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.

Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. 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 2368-2378, 2019.

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 on Information Theory , 61(5):2788-2806, 2015.

FairScale authors. Fairscale: A general purpose modular pytorch library for high performance and large scale training, 2021.

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 0898715857.

Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better fewshot 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 Papers) , pages 3816-3830, 2021.

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-2241, 2019.

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.

Daniel Golovin, John Karro, Greg Kochanski, Chansoo Lee, Xingyou Song, and Qiuyi Zhang. Gradientless descent: High-dimensional zeroth-order optimization. In International Conference on Learning Representations , 2020.

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 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation . SIAM, 2008.

José Grimm, Lo¯ ıc Pottier, and Nicole Rostaing-Schmidt. Optimal time and minimum spacetime product for reversing a certain class of programs . PhD thesis, INRIA, 1996.

Suchin Gururangan, Ana Marasovi´ c, Swabha Swayamdipta, Kyle Lo, Iz 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 8342-8360, 2020.

Davood Hajinezhad and Michael M Zavlanos. Gradient-free multi-agent nonconvex nonsmooth optimization. In 2018 IEEE Conference on Decision and Control (CDC) , pages 4939-4944, 2018.

Geoffrey Hinton. The forward-forward algorithm: Some preliminary investigations. arXiv preprint arXiv:2212.13345 , 2022.

Bairu Hou, Joe O'Connor, Jacob Andreas, Shiyu Chang, and Yang Zhang. Promptboosting: Black-box text classification with ten forward passes. arXiv preprint arXiv:2212.09257 , 2022.

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.

Kevin G Jamieson, Robert Nowak, and Ben Recht. Query complexity of derivative-free optimization. In Advances in Neural Information Processing Systems , volume 25, 2012.

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.

Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance 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.

Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximalgradient methods under the polyak-łojasiewicz condition, 2020.

Daniel Khashabi, Snigdha Chaturvedi, Michael Roth, Shyam Upadhyay, and Dan Roth. Looking beyond the surface: A 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 1 (Long Papers) , pages 252-262, 2018.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Finetuning can distort pretrained features and underperform out-of-distribution. In International Conference on Learning Representations , 2022.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 3045-3059, 2021.

Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning , 2012.

Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations , 2018.

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.

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.

Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar. Robust training of neural networks using scale invariant architectures. In International Conference on Machine Learning , pages 12656-12684, 2022.

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 100-114, 2022.

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.

Sijia Liu, Bhavya Kailkhura, Pin-Yu 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.

Sijia Liu, Pin-Yu Chen, Xiangyi Chen, and Mingyi Hong. signSGD via zeroth-order oracle. In International Conference on Learning Representations , 2019.

Sijia Liu, Pin-Yu Chen, Bhavya Kailkhura, Gaoyuan Zhang, Alfred O Hero III, and Pramod K Varshney. A primer on zeroth-order optimization in signal processing and machine learning: Principals, recent advances, and applications. IEEE Signal Processing Magazine , 37(5):43-54, 2020.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

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 Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 8086-8098, 2022.

Sadhika Malladi, Alexander Wettig, Dingli Yu, Danqi Chen, and Sanjeev Arora. A kernel-based view of language model fine-tuning. arXiv preprint arXiv:2210.05643 , 2022.

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.

Arkadij Semenoviˇ c Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.

Yurii Nesterov and Vladimir Spokoiny. Random gradient-free minimization of convex functions. Foundations of Computational Mathematics , 17:527-566, 2017.

Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P Adams. Randomized automatic differentiation. arXiv preprint arXiv:2007.10412 , 2020.

OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

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.

Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint arXiv:1811.07062 , 2018.

Vardan Papyan. Traces of class/cross-class structure pervade deep learning spectra. Journal of Machine Learning Research , 21(252):1-64, 2020.

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, highperformance deep learning library. In Advances in Neural Information Processing Systems 32 , pages 8024-8035. 2019.

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.

Archiki Prasad, Peter Hase, Xiang Zhou, and Mohit Bansal. Grips: Gradient-free, edit-based instruction search for prompting large language models. arXiv preprint arXiv:2203.07281 , 2022.

Maxim Raginsky and Alexander Rakhlin. Information-based complexity, feedback and dynamics in convex programming. IEEE Transactions on Information Theory , 57(10):70367056, 2011.

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.

Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. 2011.

Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv:1706.04454 , 2017.

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.

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.

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-1713, 2017.

Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing , 2013.

James C Spall. A one-measurement form of simultaneous perturbation stochastic approximation. Automatica , 33(1):109-112, 1997.

J.C. Spall. Multivariate stochastic approximation using a simultaneous perturbation gradient approximation. IEEE Transactions on Automatic Control , 37(3):332-341, 1992.

Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. In Advances in Neural Information Processing Systems , volume 33, pages 3008-3021, 2020.

Tianxiang Sun, Zhengfu He, Hong Qian, Yunhua Zhou, Xuanjing Huang, and Xipeng Qiu. BBTv2: Towards a gradient-free future with large language models. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 3916-3930, 2022.

Tianxiang Sun, Yunfan Shao, Hong Qian, Xuanjing Huang, and Xipeng Qiu. Black-box tuning for language-model-as-a-service. In International Conference on Machine Learning , pages 20841-20855, 2022.

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.

Yujie Tang and Na Li. Distributed zero-order algorithms for nonconvex multi-agent optimization. In 2019 57th Annual Allerton Conference on Communication, Control, and Computing (Allerton) , pages 781-786, 2019.

Zhiwei Tang, Dmitry Rybin, and Tsung-Hui Chang. Zeroth-order optimization meets human feedback: Provable learning via ranking oracles, 2023.

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:1827-1835, 2009.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems , volume 30, 2017.

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 in information retrieval , 2000.

Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. In Advances in neural information processing systems , volume 32, 2019.

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.

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.

Zhongruo Wang, Krishnakumar Balasubramanian, Shiqian Ma, and Meisam Razaviyayn. Zeroth-order algorithms for nonconvex minimax problems with improved complexities. arXiv preprint arXiv:2001.07819 , 2020.

Bingzhen Wei, Xu Sun, Xuancheng Ren, and Jingjing Xu. Minimal effort back propagation for convolutional neural networks. arXiv preprint arXiv:1709.05804 , 2017.

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 Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers) , 2018.

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: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, 2020.

Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Understanding common structure of hessian in neural networks. arXiv preprint arXiv:2010.04261 , 2020.

Jiaxi Yang, Binyuan Hui, Min Yang, Binhua Li, Fei Huang, and Yongbin Li. Iterative forward tuning boosts in-context learning in language models, 2023.

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.

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:1812.11377 , 2018.

Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv:1708.03888 , 2017.

Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv:1904.00962 , 2019.

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:1810.12885 , 2018.

Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022.

Yan Zhang, Yi Zhou, Kaiyi Ji, and Michael M Zavlanos. A new one-point residual-feedback oracle for black-box learning and control. Automatica , 136:110006, 2022.

A Algorithmic Ablations

We perform a 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 of 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 is a small budget and is only used as a setting to compare these ZO algorithms to each other. We find that using a 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 a constant learning rate schedule with no warmup for all of the below experiments. We perform few-shot experiments with k = 16 and average the results across 5 seeds.

A.1 Prompting

Westudy if adding a prompt is crucial to the ability of MeZO to optimize the network. We use prompts from Gao et al. [35]. Malladi et al. [67] claimed the prompt makes the optimization trajectory wellbehaved, though we note that the current paper considers RoBERTa-large and large autoregressive models while the previous work only studied RoBERTa-base. We note the similarity between kernel behavior and our theoretical setting in Section 4. MeZO succeeds on tasks that are reported to not exhibit kernel behavior in Malladi et al. [67], so we investigate whether or not the prompt is necessary.

Both experiments followed the grid in Table 4, but we also expanded the grid to include a learning rate of 1e -4 for the no prompt case. As a 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 n t noise vectors at the t th step and use n t -SPSA 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 is that the optimizer may need a higher fidelity gradient as it approaches the minimum. Increasing the number of z 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 a trade-off to study. We note that increasing the number of z s should be accompanied by

a proportional scaling of the learning rate, analogous to the linear scaling rule proposed in [39] (theoretical justification can follow the SDE technique [58]). Table 6 shows no consistent benefit in one schedule over the other, and it demonstrates that increasing the n in n -SPSA while fixing the number of forward passes allowed results in only marginal gains at best.

B MeZO Variants

There is a rich history of transferring ideas from first order optimization to enhance ZO algorithms. Below, we highlight several variants of MeZO that did not achieve as high performance as the algorithm presented in Algorithm 1.

B.1 Memory-efficient n-SPSA

We highlight how MeZO can perform n -SPSA (Definition 1) efficiently for n > 1 in Algorithm 2. In particular, if sampling n z vectors and averaging the projected gradients, we require storing 2 n 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 MeZO with Gradient History

The n -SPSA algorithm merely provides a gradient estimate that can subsequently be used in place of the gradient in any gradient-based optimizer. Many popular optimizers, such as Adam and SGD with momentum, require storing some historical information about gradients (e.g., a moving average). This requirement causes such algorithms to require 2 × or 3 × 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. In reference to Algorithm 1, note that the gradient only needs projected_grad and the random seed s used to compute the perturbation z , so we need to only store 2 scalars per step to reproduce the gradient history (i.e., up to 2 T scalars during training). This is a 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 to future work to investigate when MeZO-Adam may be more useful than MeZO.

Algorithm 2: MeZO with n > 1

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 a variance reduced ZO algorithm, analogous to 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 in a memory-efficient 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 a general SPSA estimate that has the same expectation (i.e., the true gradient) but has a scaled variance.

Definition 6 (Variance-Modified SPSA) . Given a matrix D = diag( d ) , the variance modified SPSA computes

where d ∈ R d has nonzero entries and d -1 denotes the coordinatewise reciprocal.

The above SPSA variant is an unbiased estimator of the gradient, because E [ ˜ ∇L ( θ ; B )] = E [ D -1 zz ⊤ D ∇L ( θ ; B )] = E [ ∇L ( θ ; B )] . We will draw inspiration from classical methods (i.e., 'control variates') and choose d to be a block vector with gradient norms or parameter norms [99].

To select the parameter groups, we split the model by layer, keeping the embedding and the head separate (i.e., RoBERTa-large has 24 + 2 = 26 parameter groups). It is straightforward to measure the parameter norms without consuming additional memory. We can measure the gradient norms without performing backpropagation, as shown below.

Proposition 1 (ZO Estimate of Gradient Norm of ℓ th Layer) . Define z ℓ to have z ∼ N (0 , 1) in each coordinate corresponding to parameters in the ℓ th layer and 0 everywhere else. Then, we can estimate the norm of the gradient of the loss w.r.t. the ℓ th layer ∇ θ ℓ as

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 , E z [( ⟨ v , z ⟩ ) 2 ] = ∥ v ∥ 2 2 for Gaussian z . It is clear that this estimate can be computed in a memory efficient way, although it requires 2 L forward passes to compute gradient norms for L parameter groups.

We show the experimental results for modifying the variance below. We follow the ablation setting and use a fixed budget of 10 , 000 steps (Appendix A). Generally, using the gradient norm to reduce the variance substantially hurts performance (Table 8). If we 'cheat' and allow one backpropagation through the network to 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 to layerwise adaptive rate methods, does not substantially impact the performance of MeZO (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 a dependence on 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 1 and Lemma 3 is likely a better characterization of fine-tuning dynamics. We leave it to future work to explore if these methods can be useful for other more complex objectives.

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 MeZO as well. This can be used to efficiently estimate coordinate-wise normalized gradient-based optimizer updates (e.g., Adam).

Definition 7 (Expectation-Modified SPSA) . Given a matrix D = diag( d ) , the variance modified SPSA computes

where d ∈ R d .

Now, we see that ˜ ∇L ( θ ; B ) = E [ D -1 zz ⊤ ∇L ( θ ; B )] so the SPSA estimate is no longer an unbiased estimator for ∇L ( θ ) . 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 our setting). We find that estimating the normalized gradient does not perform as well as directly estimating the gradient (Table 10). Regardless, we present this algorithm as a way to highlight that any coordinate-wise operation to the gradient can be applied in a memory-efficient manner.

B.5 One-point estimate

Here, we investigate the efficacy of one-point gradient estimators in place of the two-point SPSA method. Using a 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 on one estimator [113] that has a form reminiscent of SPSA but requires only one forward pass to estimate the gradient at each step.

Definition 8 (One-Point Gradient Estimate, Zhang et al. [113]) . For a loss function L evaluated on a batch B t with parameters θ t at step t , we can draw random noise z t ∼ N (0 , I d ) and compute the gradient estimate using hyperparameter ϵ as written below.

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 one forward pass at each iterate to estimate the gradient. For well-behaved loss functions and slow-moving optimization, these two formulas are intuitively similar. However, Table 11 finds this estimator to be much less efficient than SPSA when fixing the number of forward passes.

C Memory Analysis

The compute-memory tradeoff of backpropagation is complex to analyze. Griewank and Walther [40] provides a rigorous theoretical treatment of the problem. We empirically measure the memory

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 MeZO 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.

Proposition 2 (Time-Memory Tradeoff for Backpropagation, Griewank and Walther [40]) . Consider a network containing N bits. For any time-memory tradeoff hyperparameter c = O (1) , there exists a backpropagation algorithm that runs in time O ( cN ) and consumes memory proportional to O ( N 1 /c ) .

Grimm et al. [41] also gave sharp bounds for the memory-time product. Note that the popular gradient checkpointing [18] method allows one to tune c with limited precision (i.e., one cannot always further split a differentiable block and observe savings). Experiments in Chen et al. [18] choose c = 2 to achieve O ( √ N ) memory while consuming O (2 N ) computation. In the extreme case, gradient checkpointing allows one to use O ( N log N ) computation and O (log N ) memory.

MeZO always consumes 2 N compute and O (1) memory, so it is more compute-efficient 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, we do not consider that variant here because it is somewhere in the middle of the compute-memory 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 is easily differentiable, so it is hard to compare the two algorithms along the entire pareto curve.

We also compare backpropagation for the c = 1 case (i.e., storing everything during the forward pass). When storing everything, backpropagation consumes O ( N ) time and O ( N ) 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 of those algorithms compare to MeZO.

D Forward Auto-Differentiation

We discuss the merits of using forward auto-differentiation instead of two forward passes to construct a gradient estimate for fine-tuning. As ϵ → 0 , the SPSA gradient estimate (Definition 1) can be written as zz ⊤ ∇L ( θ ; B ) . The term z ⊤ ∇L ( θ ; B ) is a Jacobian-vector product (JVP), and it is well-known that this can be computed in parallel with a single forward pass while consuming additional memory equivalent to that of the largest activation in 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 RoBERTa-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-differentiation are reported in Table 12.

We see that forward auto-differentiation is substantially more memory efficient than backpropagation but less memory efficient than inference. Furthermore, forward auto-differentiation selects ϵ = 0 , which removes potentially beneficial third-and-higher order Taylor expansion terms from the estimate.

E Experiment setup

E.1 Datasets

For RoBERTa-large, we 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 set to 1 , 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, we consider the SuperGLUE dataset collection [98], including: BoolQ [21], CB [24], COPA [81], MultiRC [51], ReCoRD [111], 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 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 on 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 1K steps and MeZO experiments use

Table 14: The prompts of the datasets we used in our OPT experiments. There are three types of tasks: classification (cls.), multiple-choice (mch.), and question answering (QA). Prompts are adopted from GPT-3 [13] and PromptSource [5] with minor changes. <text> represents input from the dataset and Yes represents label words. For inference on multiple choice tasks, we put in different candidates in 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.

100 K 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 5 epochs and all MeZO experiments use 20 K 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-tuning paradigm for masked language models. Please refer to the original paper for more details.

In OPT experiments, for classification tasks, we train the model similarly to [35], i.e., we 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 to 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-likelihood. 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 but this does not improve performance and sometimes leads to unstable results. Thus we keep the 32-example results.

For linear probing of classification tasks, we take the output feature and use scipy package to 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, we do head-tuning, where the whole model is fixed except for the LM projection head. We use a batch size of 8 and a learning rate of { 1e -4 5e -4 } , and train the head for 5 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 a copy of the large language model for each downstream task is expensive. Parameter-efficient fine-tuning (PEFT) techniques alleviate this problem: instead of tuning all model parameters, PEFT only tunes a small number of additional parameters (usually less than 1%) and can often achieve comparable or better performance [57, 54, 30]. The ZO optimizer is compatible with PEFT methods, since ZO can operate on any subset of the model parameters. We are interested in the following two common PEFT methods, designed for transformers [96].

LoRA [46] adds a tunable low-rank delta to a linear layer during fine-tuning. Suppose a linear layer performed Wx + b during pre-training with W ∈ R m × n . When fine-tuning, LoRA introduces two smaller matrices A ∈ R m × r and B ∈ R r × n such that r ≪ min( m,n ) . The linear layer is then

computed as

where r and α are hyperparameters. A and B are trained on the downstream task while W is 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, r can be very small, so the number of trainable parameters during fine-tuning is small. We choose r = 8 and α = 16 .

Prefix-tuning [57] adds a prefix of m tunable representations at each layer and freezes the rest of the model. The representations are added as new keys and values and treated as additional context during the attention operation. We initialize these tunable representations by randomly sampling tokens from the vocabulary and passing them through the LLM to get their keys and values at different attention layers. We found this crucial to make prefix tuning stable with MeZO, and this trick additionally boosts the performance of prefix tuning with backpropagation, as shown in Table 17. We also tried the reparameterization trick in [57], which does not help MeZO training. In our experiments, we find m = 5 to be sufficient to achieve good performance on most tasks.

We also show that MeZO is compatible with parameter-efficient fine-tuning methods, such as prefix tuning and LoRA. Surprisingly, the performance of MeZO does not improve substantially when tuning much fewer parameters, as one might expect from classical analyses (see Section 4). Accordingly, our theoretical analysis in Section 4 suggests that the convergence rate of ZO-SGD does not depend on the parameter dimension during fine-tuning.

E.6 Training with non-differentiable objectives

The experiments maximizing the accuracy of a RoBERTa-large model were all conducted using the same grid as MeZO in Table 15.

For OPT experiments on SQuAD with F1 as objective, we use a batch size of 16 . For MeZO, we use a learning rate of { 1e -6 , 5e -6 , 1e -5 } and ϵ = 1e -3 . For MeZO (prefix), we use a learning rate of { 1e -1 , 5e -2 , 1e -2 } and ϵ = 1e -1 .

E.7 Memory profiling

In memory profiling, we use standard implementation with Huggingface's transformers [104] package. We did not turn on any advance memory-saving options, e.g., gradient checkpointing. We set the per-device batch size as 1 to test the minimum hardware requirement to run the model with specific optimization algorithms. For multi-GPU backpropagation, we use fully sharded data parallel (FSDP) [33] provided by PyTorch [76]. For multi-GPU MeZO, we use transformers multi-GPU inference of large models. We use Nvidia's nvidia-smi command to monitor the GPU memory usage. We call a run 'successful' if there is no out of memory error from GPUs for at least 100 steps. We also profile fine-tuning with LoRA, but find its memory usage similar to that of fine-tuning with prefix-tuning. Hence here we only show the analysis with prefix-tuning.

F More experiment results

F.1 RoBERTa-large experiments

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 1e -4 to 0 . 01 ) in the scipy solver. Experiments on a few settings show that LP-MeZO can sometimes improve performance without increasing the memory consumption (see Table 19). However, sometimes, linear probing first can severely hurt performance.

F.2 OPT experiments

F.3 Convergence of MeZO with full-parameter and PEFT

We demonstrate the convergence rate of MeZO, MeZO (LoRA) and MeZO (prefix) on SST-2 and SNLI for the first 5,000 steps in Figures 5. We see that despite the different number of parameters they optimize, MeZO demonstrates similar training speed on full parameter and PEFT. This agrees with our theory in Section 4, which shows that MeZO's optimization speed is independent of the number of parameters.

SNLI

0

500

1000

1500

2000

2500

3000

Steps

0.25

0.30

0.35

0.40

0.45

0.50

0.55

Loss

SST-2

MeZO Train

MeZO (LoRA) Train

MeZO (prefix) Train

0

500

1000

1500

2000

2500

3000

Steps

0.6

0.8

1.0

1.2

1.4

1.6

1.8

2.0

Loss

MeZO Train

MeZO (LoRA) Train

MeZO (prefix) Train

F.4 ZO vs BBTv2

Wecompare 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 prefixtuning and a down-projection to reduce the number of optimized parameters. BBTv2 also employs an iterative scheme which only optimizes one layer at a time. In contrast, ZO works with both full-parameter tuning and PEFT, as shown in our experiments (Section 3) and theory (Section 4).

F.5 Memory profiling

We show the detailed numbers of memory profiling results Table 22, which also corresponds to Figure 3. For how we profile the memory usage, please refer to Appendix E.7.

F.6 Wallclock time efficiency

In this section, we measure the wallclock time efficiency of MeZO compared to full-parameter FT, with respect to different model sizes. We conduct our experiments with 80GB A100s connected by NVLink and InfiniteBand, which are state-of-the-art solutions for distributed training. As shown in Table 23, on the MultiRC datasets, training with MeZO brings 7.74 × speedup per step compared to full-parameter FT on a 30B model. This is due to (1) MeZO does not require costly backpropagation and (2) MeZO requires fewer GPUs and reduces the multi-GPU communication overhead. We can see that MeZO has a bigger advantage when training larger models-the multi-GPU overhead for fine-tuning is larger.

Note that even though MeZO has better per-step wallclock efficiency, it requires significantly more steps than standard FT. Taking our OPT-30B experiments as an example: MeZO takes 32 × more steps than standard FT, while FT takes 8 × more GPUs and 7.74 × more time per step. Overall, MeZO requires only half as many GPU-hours as FT for a 30B model.

G Proofs

Proof of Lemma 2. We first note that in the ϵ → 0 limit, we have

Taking expectation over the batch B and the z i , we have E [ ̂ ∇L ( θ ; B )] = ∇L ( θ ) , so ̂ ∇L ( θ ; B ) is an unbiased estimator of the gradient.

Computing the second moment, we get

Let u , v be two arbitrary vectors. We have that

̸

when i = j , and

Therefore

̸

Next, note that when ( x 1 , y 1 ) = ( x 2 , y 2 ) , we have

and when ( x 1 , y 1 ) = ( x 2 , y 2 ) we have

Therefore

and plugging this yields

Finally, we have

Proof of Theorem 1. By Taylor's theorem with remainder, we have that

Next, note that

Therefore ∥ λ θ t +1 +(1 -λ ) θ t -θ t ∥ ≤ ηdG ( θ t ) . By the assumption we have the upper bound ∇ 2 L ( λ θ t +1 +(1 -λ ) θ t ) ⪯ H ( θ t ) , and thus

Taking the conditional expectation with respect to θ t and plugging in (9), the formula for the covariance of our ZO estimate ̂ ∇L ( θ t ; B ) , yields

By assumption, the Hessian upper bound H ( θ t ) satisfies ∥ H ( θ t ) ∥ ≤ ℓ and tr( H ( θ t )) ≤ ℓr . Thus

as desired.

G.1 Proofs of Global Convergence

Lemma 4. Let L ( θ ) be µ -PL and let there exist α such that tr( Σ ( θ )) ≤ α ( L ( θ ) -L ∗ ) for all θ . Then after

iterations of SGD we have E [ L ( θ t )] ≤ L ∗ + ϵ.

Proof of Lemma 4. The descent lemma for SGD yields

Plugging in E [ ∥∇L ( θ t ; B ) ∥ 2 ] = ∥∇L ( θ t ) ∥ 2 + 1 B tr( Σ ( θ t )) and selecting a learning rate η ≤ 1 ℓ yields

Since L is µ -PL, we get

Since tr( Σ ( θ t )) ≤ α ( L ( θ t ) -L ∗ ) , we have

Altogether,

Choosing η = min( 1 ℓ , µB ℓα ) , we obtain

Therefore we reach a solution with E [ L ( θ t )] -L ∗ ≤ ϵ after

iterations.

Proof of Lemma 3. By Corollary 1, ZO-SGD with η ZO = γ -1 η SGD yields

As in the proof for SGD, choosing η SGD ≤ 1 ℓ yields

Therefore under µ -PL and the tr( Σ ( θ t )) ≤ α ( L ( θ t ) -L ∗ ) assumption we obtain

Choosing η SGD = min( 1 ℓ , µB ℓα ) yields

Therefore we reach a solution with E [ L ( θ t )] -L ∗ ≤ ϵ after

iterations.

G.1.1 Verification of assumptions

We show that the tr( Σ ( θ t )) ≤ α ( L ( θ t ) -L ) assumption holds for certain losses.

∗

First, consider optimizing the model f ( x ; θ ) with square loss, so that

One then has that

Therefore

Thus

One then has that

Assume that the targets y i are bounded in [ -1 , 1] (which is true for binary classification tasks), and that L ∗ = 0 (which can be achieved if | f ( x ; θ ) | can be sent to ∞ ) we have that

In the kernel regime, f ( x i ; θ ) = ϕ ( x i ) ⊤ θ , and thus

Therefore

Therefore the condition holds with α = Nℓr as well.

G.2 Proofs for Gaussian perturbations

The first lemma computes the second moment of the covariance estimate ̂ ∇L ( θ ; B ) when z is drawn N (0 , I ) .

Lemma 5. Let z i ∼ N (0 , I ) i.i.d. Then

Proof. As in the proof of Lemma 2, we have that in the ϵ → 0 limit

Assume that the data can be interpolated, i.e., L ∗ = 0 . If the function is L -Lipschitz, i.e., ∥∇ f ( x ; θ ) ∥ ≤ L , then the condition holds with α = 2 NL 2 . If we are in the kernel regime, i.e., f ( x i ; θ ) = ϕ ( x i ) ⊤ θ for some feature map ϕ , then

So the condition holds for α = Nℓr .

Next, consider the cross entropy loss function, i.e

For vectors u , v , we have that

̸

when i = j , and

Therefore

In the proof of Lemma 2 we showed that

Plugging this yields

We can prove an analog to Theorem 1 in the case where the z i are Gaussian. One challenge is that ∥ θ t +1 -θ t ∥ is no longer bounded; instead we the r -local effective rank assumption only holds with high probability, and thus to bound the expected loss decrease we must control the probability of the ∥ θ t +1 -θ t ∥ being large.

Consider the following modified version of the local r -effective rank assumption, where the upper bound on the Hessian is measured over a ball of radius twice as large as the one in Assumption 1.

Assumption 2 (Local r -effective rank, Gaussian) . Let G ( θ t ) = max ( x , y ) ∈D ∥∇L ( θ t ; { ( x , y ) } ) ∥ . There exists a matrix H ( θ t ) such that:

For all θ such that ∥ θ -θ t ∥ ≤ 2 ηdG ( θ t ) , we have ∇ 2 L ( θ ) ⪯ H ( θ t ) .

The effective rank of H ( θ t ) , i.e., tr( H ( θ t )) / ∥ H ( θ t ) ∥ op , is at most r .

Theorem 2 (Dimension-Free Rate, Gaussian z ) . Assume the loss exhibits local r -effective rank (Assumption 2). If θ t +1 = θ t -η ZO ̂ ∇L ( θ t ; B ) is a single step of ZO-SGD using the n -SPSA estimate with a minibatch of size B , then there exists a γ = Θ( r/n ) such that the expected loss decrease can be bounded as

Proof of Theorem 2. Let A be the event that ∥ θ t +1 -θ t ∥ ≤ 2 ηdG ( θ t ) . On A , we have that

Likewise, since L is ℓ -smooth, we have that

Therefore

The latter term can be bounded as follows

The gradient estimate ̂ ∇L ( θ ; B ) satisfies

The expectation term is upper bounded as

where we have plugged in explicit formulas for moments of Gaussian and χ 2 random variables. Next, note that on the event ¬A , we have

Therefore

Lemma 6 (Standard χ 2 -tail bound) . Let Z be a χ 2 random variable with k degrees of freedom. Then

Since ∑ i ∈ [ n ] ∥ z i ∥ 2 is a χ 2 random variable with nd degrees of freedom, we thus have that

Altogether,

Finally, plugging in (8), along with the fact that ∥ H ( θ t ) ∥ op ≤ ℓ and tr( H ( θ t )) ≤ ℓr ,

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFine-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.\n\nSTRENGTHS:\n(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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n(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?\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

comments: 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.

The effectiveness of MeZO is shown across a range of benchmarks and model sizes.  The results compare favorably to linear probing and in context learning.

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.

The 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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nWhile 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\nQUESTIONS:\nWhile 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\nLIMITATIONS:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: 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.

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. 
1 - Comprehensive experiments and ablation studies on the proposed method.
2 - Strong theoretical support on the proposed method.
3 - Good writing and flow which makes the paper easy to follow and understand. 
4 - Well-articulated future work.

WEAKNESSES:
No major weaknesses. I left some comments in the questions section and hope  the authors can answer and address.

QUESTIONS:
- Can authors also report practical training time compared with standard fine-tuning (e.g., in terms of # steps/ second)?

- Maybe report average numbers as well in experiment results (e.g., Table 1)

- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\nHow is the training stability of MeZO? Does it an advantage over the backpropagation-based methods, especially for the large models?\n\nLIMITATIONS:\n1. 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}