PaperlayerView exact revision on arXiv

Source-linked AI summary

Fine-Tuning Language Models with Just Forward Passes

Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, Sanjeev Arora

arXiv:2305.17333v3cs.LGcs.CL

TL;DR

Scaling language models makes backpropagation memory-prohibitive, while zeroth-order optimization is expected to become slow for large models. MeZO adapts ZO-SGD to operate in place at inference-level memory. Across models and tasks, it outperforms zero-shot and in-context learning, often approaches backpropagation performance, supports parameter-efficient tuning and non-differentiable objectives, and has author-stated limitations concerning optimization steps and untested memory-saving combinations.

  • Problem

    As language models grow, backpropagation requires prohibitive memory, while classical zeroth-order analyses suggest optimization slows with parameter count.

  • Method

    MeZO adapts ZO-SGD to operate in place, using forward-pass loss differences for memory-efficient language-model fine-tuning.

  • Results

    Across model types, scales, and tasks, MeZO outperforms zero-shot, in-context learning, and linear probing, while achieving performance comparable to fine-tuning on multiple tasks and optimizing accuracy or F1.

  • Takeaways & Limitations

    MeZO extends memory-efficient tuning to full parameters and PEFT methods, including LoRA and prefix-tuning, and scales to models up to 66B.

  • Takeaways & Limitations

    MeZO takes many optimization steps, and the study does not explore combining it with methods such as gradient checkpointing, FlashAttention, or quantization.

Abstract

from arXiv · show

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 12x memory reduction and up to 2x 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 Introduction

As language models scale, backpropagation and parameter-efficient tuning become memory-intensive, while zeroth-order methods face concerns about optimization speed. MeZO adapts ZO-SGD to operate in place, enabling large-language-model tuning with inference-level memory and competitive performance.

  • Motivation: Backpropagation requires prohibitive memory at scale because it caches activations, gradients, and Adam’s gradient history.In tests, backpropagation required up to 12× inference memory.
  • Motivation: On one 80GB A100, inference supports a 30B LM, whereas Adam fine-tuning supports only a 2.7B LM.
  • Motivation: Classical ZO-SGD uses loss differences from forward passes, but naive implementations still add memory overhead and may slow convergence with model size.
  • MeZO: MeZO adapts ZO-SGD to operate in place on arbitrarily large models with almost no memory overhead.
  • Results: Across model types, scales, and task types, MeZO outperforms zero-shot, ICL, and linear probing, while often matching fine-tuning with much less memory.With OPT-13B, it is comparable to fine-tuning on 7 of 11 tasks while using roughly 12× less memory.
  • Results: MeZO supports full-parameter tuning, LoRA, and prefix-tuning, and can optimize non-differentiable objectives such as accuracy and F1.

2 Zeroth-order optimization

MeZO adapts zeroth-order optimization to fine-tune large language models in-place with inference-level memory, while retaining strong performance across model scales and tasks. Experiments show competitive results against backpropagation, compatibility with parameter-efficient tuning, and support for non-differentiable objectives.

  • Classical zeroth-order optimization: SPSA estimates gradients from loss differences using two forward passes, while n-SPSA averages estimates from n random perturbations.ZO-SGD updates parameters using the SPSA estimate; the paper uses n = 1 by default.
  • MeZO: MeZO adapts ZO-SGD to operate in-place with almost no memory overhead, matching inference memory requirements.The method is implemented with the SPSA estimator and can also be combined with optimizers such as Adam or momentum.
  • Empirical performance: MeZO consistently outperforms zero-shot, in-context learning, and linear probing across model types, model sizes, and downstream task types.The experiments cover masked and autoregressive LMs, classification, multiple-choice, and generation tasks.
  • Empirical performance: Within 5% of fine-tuning performance, MeZO matches standard fine-tuning on RoBERTa-large when k = 512, while using much less memory.On OPT-13B, MeZO is comparable within 1% or better than fine-tuning on 7 of 11 tasks.
  • Extensions and theory: MeZO supports both full-parameter tuning and PEFT variants, whose performance is generally comparable and whose convergence rates are similar.The accompanying theory states that MeZO’s convergence rate is independent of the number of optimized parameters.
  • Scaling and efficiency: MeZO optimizes models up to 66B parameters and non-differentiable objectives such as accuracy and F1 while retaining inference-level memory use.On a single A100 GPU, MeZO enables tuning a model 11 times larger than full fine-tuning, and its implementation uses half as many GPU-hours for a 30B model.

4 Theory

The theory explains why MeZO can avoid the parameter-dimension slowdown predicted by classical zeroth-order analyses: under favorable pre-training and loss-landscape assumptions, its rates depend on local effective rank instead. Without those assumptions, MeZO can suffer a slowdown proportional to the number of parameters.

  • 4 Theory: The paper attributes MeZO’s practical optimization of billion-parameter models to adequate pre-training, task prompts, and favorable low-effective-rank landscapes.The theory assumes the model has already been trained for many steps on the fine-tuning objective; experiments use n = 1.
  • 4.1 Per-step analysis: Classical analysis predicts that MeZO’s permissible learning rate and per-step loss decrease can worsen by a factor of d.Here, d denotes the number of model parameters.
  • 4.1 Per-step analysis: Under a local r-effective-rank assumption, ZO-SGD’s convergence rate does not depend on the number of parameters.The resulting slowdown is governed by the local effective rank r rather than parameter dimension d.
  • 4.1 Per-step analysis: With a learning rate scaled by γ^-1, where γ = Θ(r/n), ZO-SGD’s loss decrease can be compared directly with SGD’s descent rate.The comparison uses the dimension-free per-step bound from the effective-rank analysis.
  • 4.2 Global convergence: The global convergence result requires a µ-PL loss landscape and a gradient-covariance trace bounded by suboptimality.Under these stronger assumptions, the slowdown is proportional to effective rank rather than parameter dimension.

5 Related work

Related work places MeZO among zeroth-order optimization, memory-efficient backpropagation, and gradient-free language-model adaptation. These approaches reduce gradients or memory in different ways, but prior methods generally face approximation, dimensionality, or model-scale limitations described by the paper.

  • Zeroth-order optimization: Classical zeroth-order lower bounds often depend on parameter count, while newer results exploit low-dimensional gradient structure.The newer bounds scale with intrinsic dimension and logarithmically with parameter count, but gradient estimation can remain costly.
  • Memory-efficient backpropagation: Memory-efficient backpropagation methods sparsify gradients, approximate Jacobians, subsample computation, checkpoint activations, or recompute attention quantities.The paper notes that some approximation-based methods may incur large errors, while checkpointing trades memory for recomputation.
  • Gradient-free adaptation of large language models: Gradient-free LM adaptation methods such as BBT optimize low-dimensional prefix projections, while black-box prompt tuning updates discrete prompts without updating model parameters.The cited BBT work focuses on RoBERTa-large-scale models and few-shot settings.

6 Conclusion

MeZO optimizes large language models across tasks and scales while matching inference memory, though it requires many optimization steps and has not yet been combined with other memory-saving methods.

  • MeZO effectively optimizes large language models across many tasks and model scales.
  • MeZO can optimize non-differentiable objectives that backpropagation usually cannot handle.
  • MeZO requires many steps for strong performance, although per-step speedups can sometimes make it faster than standard backpropagation fine-tuning.
  • The work does not explore combining MeZO with FlashAttention or quantization.

A Algorithmic Ablations

The ablations compare zeroth-order algorithms under controlled forward-pass budgets and examine prompts, schedules, and optimization settings for MeZO.

  • A Algorithmic Ablations: The ablations treat forward passes as the main computational cost, affected by gradient steps, accumulation, and noise samples.
  • A Algorithmic Ablations: MeZO performance improves monotonically with more steps, with no apparent overfitting in the reported experiments.
  • A Algorithmic Ablations: With a fixed 10,000-forward-pass budget, the ablations compare algorithms on SST-2, SNLI, and TREC.
  • A Algorithmic Ablations: Learning-rate schedules and warmup produce identical or non-improving results in the reported three-task ablations, so later experiments use a constant rate without warmup.
  • A Algorithmic Ablations: The study fixes prompt-based fine-tuning for subsequent experiments after comparing prompted and unprompted settings.

A.2 Sample schedules

The sample-schedule ablations study the trade-off between gradient-estimate fidelity and forward-pass cost, alongside memory-efficient optimizer-history mechanisms.

  • A.2 Sample schedules: Increasing the number of noise vectors reduces gradient variance but increases forward passes per optimization step, creating a computation–variance trade-off.
  • A.2 Sample schedules: Under a fixed forward-pass budget, increasing n in n-SPSA yields at best marginal gains, with no consistently superior sampling schedule.
  • B.2 Augmenting MeZO with Gradient History: MeZO can recompute gradient history using only the projected gradient and random seed, storing 2 scalars per step instead of full optimizer history.
  • B.2 Augmenting MeZO with Gradient History: MeZO-Adam can sometimes improve performance, but each gradient step requires additional computation without additional forward passes.

B.3 Modifying the Variance of MeZO

The variance-modification experiments test unbiased and expectation-modified SPSA variants, finding no consistent convergence benefit and weaker performance for normalized-gradient estimation.

  • B.3 Modifying the Variance of MeZO: Variance-Modified SPSA preserves the gradient expectation while scaling its variance through a diagonal parameter transformation.
  • B.3 Modifying the Variance of MeZO: Estimating layer gradient norms without backpropagation requires 2L forward passes for L parameter groups.
  • B.3 Modifying the Variance of MeZO: Gradient-norm-based variance reduction substantially hurts performance, while parameter-norm modification does not yield a consistently reported acceleration.
  • B.3 Modifying the Variance of MeZO: Expectation-modified SPSA is no longer unbiased for the original gradient and can estimate a normalized gradient when d is chosen as the gradient norm.
  • B.3 Modifying the Variance of MeZO: The normalized-gradient variant performs worse than directly estimating the gradient in the reported experiments.
  • B.5 One-point estimate: A one-point estimator halves forward passes per step but is much less efficient than SPSA when methods are compared at equal forward-pass budgets.

C Memory Analysis

This section compares backpropagation and zeroth-order methods through their compute–memory tradeoff. MeZO uses inference-level memory while avoiding the extreme computation required by aggressive checkpointing, but its gradient is approximate and prompt-dependent.

  • Backpropagation tradeoff: Backpropagation trades memory for recomputation by caching or regenerating outputs of differentiable blocks.The tradeoff depends on how the network is decomposed and on the hyperparameter c.
  • Checkpointing: O(N log N) computation and O(log N) memory are achievable with extreme gradient checkpointing.Checkpointing reduces memory by increasing recomputation.
  • MeZO tradeoff: MeZO uses 2N compute and O(1) memory, making it more compute-efficient than gradient checkpointing at the same memory cost.The comparison concerns the stated compute–memory Pareto curve.
  • Backpropagation tradeoff: O(N) memory and O(N) time describe backpropagation when all forward-pass activations are stored.At this endpoint, SPSA uses slightly more time and substantially less memory.
  • Scope boundary: MeZO computes an approximate gradient that is useful for prompt-based fine-tuning, making it less broadly useful than gradient checkpointing.The passage also notes that comparisons with other low-memory gradient approximations remain unclear.

D Forward Auto-Differentiation

Forward auto-differentiation computes a zeroth-order estimate using a Jacobian-vector product during inference. It is more memory-efficient than backpropagation but remains less memory-efficient than inference and excludes nonzero-ϵ effects.

  • JVP construction: A Jacobian-vector product computes z⊤∇L(θ; B) during a forward pass and requires memory equivalent to the largest activation.The method still stores z to construct the full gradient estimate.
  • Implementation: PyTorch support was memory-inefficient, while JAX enabled profiling of forward auto-differentiation on RoBERTa-large.The implementation discussion motivates using JAX for the reported comparison.
  • Memory comparison: Forward auto-differentiation is substantially more memory-efficient than backpropagation but less memory-efficient than inference.The comparison uses RoBERTa-large with batch size 16 on MultiRC; reported memory excludes model storage.
  • Approximation: ϵ = 0 removes potentially beneficial third-and-higher-order Taylor terms from the gradient estimate.This limitation follows from the JVP formulation.

E.1 Datasets

The experiments cover masked and autoregressive language models across classification, multiple-choice, generation, and question answering. They use prompt-based objectives, specified datasets, and distinct optimization schedules for MeZO and backpropagation.

  • Datasets: RoBERTa-large experiments use SST-2, SST-5, TREC, MNLI, SNLI, and RTE classification datasets with k = 16 or k = 512 examples per class.The test set is limited to 1,000 examples for fast iteration.
  • Datasets: OPT experiments use SuperGLUE tasks, SST-2, SQuAD, and DROP across classification, multiple-choice, and question-answering settings.The sampling scheme uses 1,000 training, 500 validation, and 1,000 test examples.
  • Prompting: RoBERTa prompts use templates and label words that fill the [MASK] token, while OPT prompts are adapted from GPT-3 and PromptSource.The prompt constructions are described in Tables 13 and 14.
  • Training schedules: MeZO uses constant learning rates and longer runs than backpropagation: 100K versus 1K steps for RoBERTa, and 20K steps versus 5 epochs for OPT.Validation checkpoints are selected at different intervals for the two methods.
  • Objectives: OPT classification uses label-word logits with cross-entropy, while multiple-choice and QA training uses teacher forcing on candidate tokens only.Prompt tokens are excluded from the loss.
  • Parameter-efficient tuning: MeZO is compatible with PEFT methods including LoRA and prefix-tuning, which update only a small subset of parameters.LoRA uses low-rank trainable matrices, while prefix-tuning adds trainable key and value representations; real-word activation prefixes outperform random initialization.

F.1 RoBERTa-large experiments

The experiments evaluate MeZO against linear probing, fine-tuning, other zeroth-order methods, and parameter-efficient variants. Across the reported settings, MeZO approaches or matches backpropagation performance while reducing memory and, for large models, GPU-hour cost.

  • RoBERTa-large results: RoBERTa-large results show MeZO outperforming zero-shot and linear probing while approaching Adam fine-tuning with much lower memory cost.Table 18 reports averaged accuracy across the evaluated tasks.
  • LP-MeZO: Linear probing before MeZO can improve performance without increasing memory overhead, but can also severely hurt performance.The effect is setting-dependent.
  • Parameter count: MeZO shows similar training speed with full-parameter tuning, LoRA, and prefix-tuning despite optimizing different numbers of parameters.This agrees with the stated theory that optimization speed is independent of parameter count.
  • ZO comparisons: ZO significantly outperforms BBTv2 on their mutually assessed tasks while supporting both full-parameter tuning and PEFT.BBTv2 is limited to low-dimensional, layer-wise optimization.
  • Wallclock efficiency: 7.74× speedup per step is achieved by MeZO over full-parameter fine-tuning on MultiRC with an OPT-30B model.The reported comparison uses 80GB A100 GPUs connected by NVLink and InfiniteBand.
  • Wallclock efficiency: Half as many GPU-hours are required by MeZO as fine-tuning for the OPT-30B example, despite MeZO taking 32× more steps.Fine-tuning uses 8× more GPUs and 7.74× more time per step in this comparison.
  • Theory: MeZO’s theoretical analysis includes unbiased gradient estimation and convergence results under stated smoothness, Hessian, and PL-condition assumptions.The proof passages derive the estimator’s expectation and convergence lemmas rather than report new empirical measurements.

G.2 Proofs for Gaussian perturbations

The Gaussian-perturbation analysis establishes a dimension-free ZO-SGD rate under a local effective-rank assumption. Because Gaussian step lengths are unbounded, the proof controls large-step events using χ2 tail bounds before bounding the expected loss decrease.

  • Moment calculations: The Gaussian fourth-moment identity used in the proof is Ez[z⊗4](u, v) = u⊤v · I + 2uv⊤.This follows from the Gaussian fourth-moment expression 3Sym(I⊗2)(u, v).
  • Proof strategy: Gaussian perturbations make the update norm unbounded, so the proof separately controls the event of an unusually large step.The analysis introduces an event A where the step norm is at most 2ηdG(θt), then treats its complement.
  • Theorem 2: Theorem 2 gives a dimension-free rate for one n-SPSA ZO-SGD step under the local r-effective-rank assumption.The assumption bounds the effective rank of H(θt) by r.
  • Expected decrease: The resulting expected loss decrease is bounded using γ = Θ(r/n) with minibatch size B, together with operator-norm and trace bounds on the Hessian.The proof plugs in ∥H(θt)∥op ≤ ℓ and tr(H(θt)) ≤ ℓr.
  • Proof strategy: A χ2 tail bound controls the probability of large Gaussian perturbation norms because the summed squared norms follow a χ2 distribution with nd degrees of freedom.The proof uses standard χ2 concentration and Gaussian moment formulas.
Loading 2305.17333v3…