Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Prior-Informed Zeroth-Order Optimization with Adaptive Direction Alignment for Memory-Efficient LLM Fine-Tuning

Feihu Jin, Shipeng Cen, Ying Tan

TL;DR

This work tackles the memory-intensive bottleneck of backpropagation in fine-tuning large language models by introducing prior-informed zeroth-order optimization (GV-ZO). It presents two plug-and-play components—Guiding Vector augmentation (MeZO-GV) and Greedy Perturbation (MeZO-Greedy)—that direct perturbations toward informative directions, yielding stronger alignment with the true gradient and faster convergence. Theoretical results show improved directional alignment under prior-informed perturbations, and extensive experiments across OPT-13B and Llama2 scales demonstrate superior performance and memory efficiency compared with standard ZO methods and many gradient-based baselines. The approach is architecture-agnostic and readily integrates with existing optimization pipelines, offering a practical, scalable solution for memory-constrained fine-tuning of modern LLMs.

Abstract

Fine-tuning large language models (LLMs) has achieved remarkable success across various NLP tasks, but the substantial memory overhead during backpropagation remains a critical bottleneck, especially as model scales grow. Zeroth-order (ZO) optimization alleviates this issue by estimating gradients through forward passes and Gaussian sampling, avoiding the need for backpropagation. However, conventional ZO methods suffer from high variance in gradient estimation due to their reliance on random perturbations, leading to slow convergence and suboptimal performance. We propose a simple plug-and-play method that incorporates prior-informed perturbations to refine gradient estimation. Our method dynamically computes a guiding vector from Gaussian samples, which directs perturbations toward more informative directions, significantly accelerating convergence compared to standard ZO approaches. We further investigate a greedy perturbation strategy to explore the impact of prior knowledge on gradient estimation. Theoretically, we prove that our gradient estimator achieves stronger alignment with the true gradient direction, enhancing optimization efficiency. Extensive experiments across LLMs of varying scales and architectures demonstrate that our proposed method could seamlessly integrate into existing optimization methods, delivering faster convergence and superior performance. Notably, on the OPT-13B model, our method outperforms traditional ZO optimization across all 11 benchmark tasks and surpasses gradient-based baselines on 9 out of 11 tasks, establishing a robust balance between efficiency and accuracy.

Prior-Informed Zeroth-Order Optimization with Adaptive Direction Alignment for Memory-Efficient LLM Fine-Tuning

TL;DR

This work tackles the memory-intensive bottleneck of backpropagation in fine-tuning large language models by introducing prior-informed zeroth-order optimization (GV-ZO). It presents two plug-and-play components—Guiding Vector augmentation (MeZO-GV) and Greedy Perturbation (MeZO-Greedy)—that direct perturbations toward informative directions, yielding stronger alignment with the true gradient and faster convergence. Theoretical results show improved directional alignment under prior-informed perturbations, and extensive experiments across OPT-13B and Llama2 scales demonstrate superior performance and memory efficiency compared with standard ZO methods and many gradient-based baselines. The approach is architecture-agnostic and readily integrates with existing optimization pipelines, offering a practical, scalable solution for memory-constrained fine-tuning of modern LLMs.

Abstract

Fine-tuning large language models (LLMs) has achieved remarkable success across various NLP tasks, but the substantial memory overhead during backpropagation remains a critical bottleneck, especially as model scales grow. Zeroth-order (ZO) optimization alleviates this issue by estimating gradients through forward passes and Gaussian sampling, avoiding the need for backpropagation. However, conventional ZO methods suffer from high variance in gradient estimation due to their reliance on random perturbations, leading to slow convergence and suboptimal performance. We propose a simple plug-and-play method that incorporates prior-informed perturbations to refine gradient estimation. Our method dynamically computes a guiding vector from Gaussian samples, which directs perturbations toward more informative directions, significantly accelerating convergence compared to standard ZO approaches. We further investigate a greedy perturbation strategy to explore the impact of prior knowledge on gradient estimation. Theoretically, we prove that our gradient estimator achieves stronger alignment with the true gradient direction, enhancing optimization efficiency. Extensive experiments across LLMs of varying scales and architectures demonstrate that our proposed method could seamlessly integrate into existing optimization methods, delivering faster convergence and superior performance. Notably, on the OPT-13B model, our method outperforms traditional ZO optimization across all 11 benchmark tasks and surpasses gradient-based baselines on 9 out of 11 tasks, establishing a robust balance between efficiency and accuracy.
Paper Structure (20 sections, 4 theorems, 29 equations, 6 figures, 11 tables, 4 algorithms)

This paper contains 20 sections, 4 theorems, 29 equations, 6 figures, 11 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Memory-Efficient ZO-SGD (MeZO)
  4. Parameter-Efficient Fine-Tuning (PEFT)
  5. Gradient-free Optimization of LLMs
  6. Our Proposed Method
  7. Memory-efficient ZO with Guiding Vector
  8. Memory-efficient ZO with Greedy Perturbation
  9. Theory Analysis
  10. Per-Step Decrease Analysis with Prior-Informed ZO
  11. Experiments and Analysis
  12. Medium-sized Language Models
  13. Large Language Models
  14. MeZO with Greedy Strategy
  15. Single Step Analysis for Different Models
  16. ...and 5 more sections

Key Result

Lemma 1

Under the ZO setting, assume the optimization problem has dimension $d$, and the sampling number is $k$, where $z_1, z_2,\ldots, z_k \sim \mathcal{N}(0, I_d)$. Let $\epsilon > 0$ and $\delta > 0$, and define When $k = O\!\left(\frac{1}{\epsilon^2} \log\!\left(\frac{d}{\delta}\right)\right)$, with probability at least $1 - \delta$, we have $\|S_k - I_d\| \leq \epsilon$. Note that in practice, the

Figures (6)

  • Figure 1: The training loss curves for the WSC, SST-2, and BoolQ tasks are evaluated using the OPT-1.3B model. Our proposed methods (MeZO-Greedy and MeZO-GV) are compatible with MeZO. For full fine-tuning, a learning rate of 2e-7 is employed. All experiments are conducted with a consistent batch size of 16 to ensure uniformity across evaluations.
  • Figure 2: Validation Accuracy on SST2 and BoolQ Tasks for Llama2-7B and Llama2-13B. All experiments are conducted with a batch size of 16. For LoRA-based methods, the learning rate is set to 1e-4, while for full-parameter methods, the learning rate is set to 5e-7.
  • Figure 3: Training loss on SST2, BoolQ, and CB Tasks for OPT-1.3B/13B Models. We employ a learning rate of 2e-7. All experiments are conducted with a consistent batch size of 16.
  • Figure 4: Training loss on BoolQ and RTE Tasks with Llama2-7B Model. We employ a learning rate of 5e-7. All experiments are conducted with a consistent batch size of 16.
  • Figure 5: Performance of OPT-13B Model Across three Datasets as a Function of Prior-Estimated Times
  • ...and 1 more figures

Theorems & Definitions (8)

  • Lemma 1
  • proof
  • Lemma 2
  • proof
  • Lemma 3
  • proof
  • Lemma 4
  • proof