Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Low-Rank Curvature for Zeroth-Order Optimization in LLM Fine-Tuning

Hyunseok Seung, Jaewoo Lee, Hyunsuk Ko

TL;DR

This work tackles the memory bottleneck of fine-tuning large language models by developing LOREN, a curvature-aware zeroth-order optimizer. It unifies adaptive perturbation sampling via natural evolution strategies, a low-rank Kronecker-factored covariance to capture heterogeneous loss-landscape curvature, and variance reduction through REINFORCE leave-one-out. The method yields higher accuracy and faster convergence than state-of-the-art ZO baselines across masked and autoregressive LLMs, while keeping memory usage modest. These results demonstrate that curvature-aware ZO optimization can enable scalable, memory-efficient fine-tuning of very large models. The approach has practical implications for deploying memory-efficient training on expensive LLMs and sets a new benchmark for ZO methods in NLP applications.

Abstract

We introduce LOREN, a curvature-aware zeroth-order (ZO) optimization method for fine-tuning large language models (LLMs). Existing ZO methods, which estimate gradients via finite differences using random perturbations, often suffer from high variance and suboptimal search directions. Our approach addresses these challenges by: (i) reformulating the problem of gradient preconditioning as that of adaptively estimating an anisotropic perturbation distribution for gradient estimation, (ii) capturing curvature through a low-rank block diagonal preconditioner using the framework of natural evolution strategies, and (iii) applying a REINFORCE leave-one-out (RLOO) gradient estimator to reduce variance. Experiments on standard LLM benchmarks show that our method outperforms state-of-the-art ZO methods by achieving higher accuracy and faster convergence, while cutting peak memory usage by up to 27.3% compared with MeZO-Adam.

Low-Rank Curvature for Zeroth-Order Optimization in LLM Fine-Tuning

TL;DR

This work tackles the memory bottleneck of fine-tuning large language models by developing LOREN, a curvature-aware zeroth-order optimizer. It unifies adaptive perturbation sampling via natural evolution strategies, a low-rank Kronecker-factored covariance to capture heterogeneous loss-landscape curvature, and variance reduction through REINFORCE leave-one-out. The method yields higher accuracy and faster convergence than state-of-the-art ZO baselines across masked and autoregressive LLMs, while keeping memory usage modest. These results demonstrate that curvature-aware ZO optimization can enable scalable, memory-efficient fine-tuning of very large models. The approach has practical implications for deploying memory-efficient training on expensive LLMs and sets a new benchmark for ZO methods in NLP applications.

Abstract

We introduce LOREN, a curvature-aware zeroth-order (ZO) optimization method for fine-tuning large language models (LLMs). Existing ZO methods, which estimate gradients via finite differences using random perturbations, often suffer from high variance and suboptimal search directions. Our approach addresses these challenges by: (i) reformulating the problem of gradient preconditioning as that of adaptively estimating an anisotropic perturbation distribution for gradient estimation, (ii) capturing curvature through a low-rank block diagonal preconditioner using the framework of natural evolution strategies, and (iii) applying a REINFORCE leave-one-out (RLOO) gradient estimator to reduce variance. Experiments on standard LLM benchmarks show that our method outperforms state-of-the-art ZO methods by achieving higher accuracy and faster convergence, while cutting peak memory usage by up to 27.3% compared with MeZO-Adam.

Paper Structure

This paper contains 34 sections, 4 theorems, 56 equations, 4 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Notations
  5. Zeroth-Order Gradient Estimates
  6. REINFORCE with Leave-One-Out Baseline
  7. Methods
  8. Preconditioning via Evolution Strategies
  9. Low-Rank Structured Covariance Matrices
  10. Search Distribution Parameter Updates
  11. Algorithm
  12. Convergence Analysis
  13. Experiments
  14. Masked Language Models
  15. Autoregressive Language Models
  16. ...and 19 more sections

Key Result

Proposition 3.3

For $f:\mathbb{R}^{d}\to \mathbb{R}$, the gradient of Gaussian smoothed $f$ is given by

Figures (4)

  • Figure 1: (a) Mean squared errors of ZO gradient estimates, with and without RLOO, relative to the true gradient on the 1,000-dimensional Sphere, Rastrigin, and Rosenbrock functions. (b) Optimization trajectories of FO-SGD and ZO optimizers on the monkey saddle function, all initialized at (2.9, -0.01). Accuracy curves for (c) GPT-2-XL fine-tuned on QNLI and (d) OPT-13B fine-tuned on CB, using early stopping.
  • Figure 2: Training loss curves for different ZO optimizers when fine-tuning OPT-13B on SuperGLUE tasks.
  • Figure 3: Training loss curves for different ZO optimizers when fine-tuning GPT-2-XL on GLUE tasks.
  • Figure 4: Fine-tuning results of QNLI task on RoBERTa-large with varying (Left) covariance learning rate, (Center) damping, and (Right) number of forward passes.

Theorems & Definitions (7)

  • Definition 3.1: SPSA Spall1992MultivariateSA
  • Definition 3.2: Generalized Gaussian smoothing
  • Proposition 3.3
  • Proposition 4.1
  • Theorem 4.2
  • Lemma D.3
  • proof