Fine-Tuning Language Models with Just Forward Passes

By Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alexandru Damian, Jason D. Lee, Danqi Chen, Sanjeev Arora. Neural Information Processing Systems, 2023

422 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Electrical and Electronic Engineering

View Paper: https://www.semanticscholar.org/paper/ad4b365630f1c13d74d78f0f5d8cee87ef356d41

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Tianyi Zhou
2. Can Xu
3. Tao Shen
4. Dacheng Tao
5. Ming Li
6. Jinyang Li
7. Jeff Zhang
8. Xiaohan Xu

In Electrical and Electronic Engineering

1. Dacheng Tao

Citations in the corpus

• Parameter-Efficient Fine-Tuning for Large Models: A Comprehensive Survey (2024, methodology): “MeZO [131] introduces a memory-efficient zeroth-order (ZO) optimizer for LLMs. Unlike conventional PEFT techniques, which rely on backpropagation to compute gradients for updating model parameters, MeZO fine-tunes LLMs through only forward passes.”
• A Survey on Knowledge Distillation of Large Language Models (2024, methodology): “In the realm of Efficient Fine-Tuning, Parameter Efficient Fine-Tuning (Hu et al., 2023b; Liu et al., 2022b; Wang et al., 2022b; Hu et al., 2021; Li and Liang, 2021; Liu et al., 2022c), and Memory Efficient Fine-Tuning (Dettmers et al., 2023; Kim et al., 2023d; Malladi et al., 2024) are utilized.”