Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Towards Fast LLM Fine-tuning through Zeroth-Order Optimization with Projected Gradient-Aligned Perturbations

Zhendong Mi, Qitao Tan, Grace Li Zhang, Zhaozhuo Xu, Geng Yuan, Shaoyi Huang

TL;DR

This work tackles the memory bottleneck of fine-tuning very large language models by adopting zeroth-order optimization. It introduces P-GAP, which first estimates a low-dimensional gradient space and then uses projected gradient-aligned perturbations within that space to substantially reduce gradient-estimation variance and accelerate convergence. The authors provide variance and convergence analyses and validate the approach across RoBERTa-large, OPT, and LLaMA models, showing up to 12% gains in generation tasks and substantial reductions in training iterations and GPU hours. Overall, P-GAP offers a practical, scalable path for fast, resource-efficient ZO fine-tuning of very large models, with strong compatibility with LoRA-style PEFT.

Abstract

Fine-tuning large language models (LLMs) using zeroth-order (ZO) optimization has emerged as a promising alternative to traditional gradient-based methods due to its reduced memory footprint requirement. However, existing ZO methods suffer from high variance in gradient estimation, leading to slow convergence and suboptimal performance on large-scale models. In this work, we propose P-GAP, a fast LLM fine-tuning approach through zeroth-order optimization with Projected Gradient-Aligned Perturbations. Specifically, we first estimate a low-dimensional gradient space and then align perturbations in projected gradients' direction within the space. This approach enables reduced the number of perturbed parameters and decreased variance, therefore accelerated convergence for LLM fine-tuning. Experiments on LLMs show that P-GAP consistently surpasses the baselines, achieving up to 6% increase in accuracy on classification tasks and up to 12% higher accuracy on generation tasks, with up to about 81% less training iterations and 70% less GPU hours. These results demonstrate that P-GAP enables fast, scalable, and resource-efficient ZO LLM fine-tuning.

Towards Fast LLM Fine-tuning through Zeroth-Order Optimization with Projected Gradient-Aligned Perturbations

TL;DR

This work tackles the memory bottleneck of fine-tuning very large language models by adopting zeroth-order optimization. It introduces P-GAP, which first estimates a low-dimensional gradient space and then uses projected gradient-aligned perturbations within that space to substantially reduce gradient-estimation variance and accelerate convergence. The authors provide variance and convergence analyses and validate the approach across RoBERTa-large, OPT, and LLaMA models, showing up to 12% gains in generation tasks and substantial reductions in training iterations and GPU hours. Overall, P-GAP offers a practical, scalable path for fast, resource-efficient ZO fine-tuning of very large models, with strong compatibility with LoRA-style PEFT.

Abstract

Fine-tuning large language models (LLMs) using zeroth-order (ZO) optimization has emerged as a promising alternative to traditional gradient-based methods due to its reduced memory footprint requirement. However, existing ZO methods suffer from high variance in gradient estimation, leading to slow convergence and suboptimal performance on large-scale models. In this work, we propose P-GAP, a fast LLM fine-tuning approach through zeroth-order optimization with Projected Gradient-Aligned Perturbations. Specifically, we first estimate a low-dimensional gradient space and then align perturbations in projected gradients' direction within the space. This approach enables reduced the number of perturbed parameters and decreased variance, therefore accelerated convergence for LLM fine-tuning. Experiments on LLMs show that P-GAP consistently surpasses the baselines, achieving up to 6% increase in accuracy on classification tasks and up to 12% higher accuracy on generation tasks, with up to about 81% less training iterations and 70% less GPU hours. These results demonstrate that P-GAP enables fast, scalable, and resource-efficient ZO LLM fine-tuning.
Paper Structure (24 sections, 13 theorems, 93 equations, 5 figures, 10 tables, 3 algorithms)

This paper contains 24 sections, 13 theorems, 93 equations, 5 figures, 10 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Methodology
  4. Projected Gradient-Aligned Perturbation
  5. Low-dimensional Gradient Space Design
  6. Adapting Projected Gradient-Aligned Perturbation in Low-dimensional Gradient Space
  7. Experiments
  8. Results on medium-sized model
  9. Results on large autoregressive Models
  10. Performance on LLMs with Various Scales
  11. Convergence and Wall-clock Time Analysis
  12. Memory analysis
  13. Related Work
  14. Conclusion
  15. Appendix
  16. ...and 9 more sections

Key Result

Lemma 1

Let $P\in\mathbb R^{d\times q}$ satisfy $P^{\top}P=I_q$, and sample $z\sim\mathcal{N}(0,I_q)$. Define the two-point estimator Let $\nabla f=\nabla f(x)$ and $u:=P^\top \nabla f\in\mathbb R^q$. Then: (A) Quadratic objective (exact formula). If $f(x)=x^\top Hx$ is quadratic, then so the variance grows linearly in the perturbation dimension $q$. (B) General $L$-smooth objective (upper bound). If $f

Figures (5)

  • Figure 1: Estimation of directional derivative magnitudes on the $V$ matrix from the first Transformer layer of the OPT-2.7B model, using perturbations sampled from a standard Gaussian distribution and from P-GAP
  • Figure 2: Accuracy comparison of MeZO and P-GAP (Ours) on LLaMA3-3B and LLaMA3-8B
  • Figure 3: Training loss comparison with iterations of MeZO and P-GAP on RoBERTa-large
  • Figure 4: Comparison of GPU hours for full FT across different datasets on OPT-2.7B between MeZO and P-GAP. Results are presented as normalized time (numbers in red indicate speedup)
  • Figure 5: Wall clock time for OPT-2.7B on CB datasets

Theorems & Definitions (28)

  • Lemma 1
  • proof
  • Corollary 1: Full-space perturbation
  • Lemma 2: Kronecker projection
  • proof
  • Lemma 3: Block diagonal model projection
  • proof
  • Lemma 4: Rotation invariance
  • proof
  • Lemma 5: Moments of $\mathcal{N}(0,I_n)$
  • ...and 18 more