Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Bayesian Subspace Gradient Estimation for Zeroth-Order Optimization of Large Language Models

Jian Feng, Zhihong Huang

TL;DR

This work addresses the memory bottleneck in fine-tuning extremely large language models by adopting zeroth-order optimization. It introduces Bayesian Subspace Zeroth-Order Optimization (BSZO), which reframes gradient estimation as Bayesian inference over a projected gradient in a k-dimensional subspace and uses Kalman filtering to fuse multiple finite-difference measurements. The method includes a residual-based adaptive mechanism for perturbation scales and proves a convergence-rate improvement by a factor of $k/\gamma$, with only $O(k^2)$ additional memory. Empirically, BSZO outperforms MeZO, MeZO-Adam, and HiZOO across masked and decoder-only LLM benchmarks, achieving up to $6.67\%$ absolute average gains (e.g., OPT-13B) while keeping memory close to inference-only baselines, and showing stability across reduced-precision settings.

Abstract

Fine-tuning large language models (LLMs) with zeroth-order (ZO) optimization reduces memory by approximating gradients through function evaluations, but existing methods rely on one-step gradient estimates from random perturbations. We introduce Bayesian Subspace Zeroth-Order optimization (BSZO), a ZO optimizer that applies Kalman filtering to combine finite-difference information across multiple perturbation directions. By treating each finite-difference measurement as a noisy observation, BSZO builds a posterior distribution over the projected gradient and updates it through Bayesian inference, with a residual-based adaptive mechanism to adjust perturbation scales. Theoretical analysis shows that BSZO improves the convergence rate by a factor of $k/γ$ compared to standard ZO methods. Experiments on RoBERTa, Mistral, and OPT models show that BSZO outperforms MeZO, MeZO-Adam, and HiZOO across various tasks, achieving up to 6.67\% absolute average improvement on OPT-13B while keeping memory usage close to inference-only baselines (1.00$\times$--1.08$\times$ of MeZO).

Bayesian Subspace Gradient Estimation for Zeroth-Order Optimization of Large Language Models

TL;DR

This work addresses the memory bottleneck in fine-tuning extremely large language models by adopting zeroth-order optimization. It introduces Bayesian Subspace Zeroth-Order Optimization (BSZO), which reframes gradient estimation as Bayesian inference over a projected gradient in a k-dimensional subspace and uses Kalman filtering to fuse multiple finite-difference measurements. The method includes a residual-based adaptive mechanism for perturbation scales and proves a convergence-rate improvement by a factor of , with only additional memory. Empirically, BSZO outperforms MeZO, MeZO-Adam, and HiZOO across masked and decoder-only LLM benchmarks, achieving up to absolute average gains (e.g., OPT-13B) while keeping memory close to inference-only baselines, and showing stability across reduced-precision settings.

Abstract

Fine-tuning large language models (LLMs) with zeroth-order (ZO) optimization reduces memory by approximating gradients through function evaluations, but existing methods rely on one-step gradient estimates from random perturbations. We introduce Bayesian Subspace Zeroth-Order optimization (BSZO), a ZO optimizer that applies Kalman filtering to combine finite-difference information across multiple perturbation directions. By treating each finite-difference measurement as a noisy observation, BSZO builds a posterior distribution over the projected gradient and updates it through Bayesian inference, with a residual-based adaptive mechanism to adjust perturbation scales. Theoretical analysis shows that BSZO improves the convergence rate by a factor of compared to standard ZO methods. Experiments on RoBERTa, Mistral, and OPT models show that BSZO outperforms MeZO, MeZO-Adam, and HiZOO across various tasks, achieving up to 6.67\% absolute average improvement on OPT-13B while keeping memory usage close to inference-only baselines (1.00--1.08 of MeZO).
Paper Structure (27 sections, 13 theorems, 90 equations, 1 figure, 9 tables, 2 algorithms)

This paper contains 27 sections, 13 theorems, 90 equations, 1 figure, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Zeroth-Order Optimization for LLMs.
  4. Population-Based Gradient Estimation.
  5. Bayesian Inference for Optimization.
  6. Method
  7. Preliminaries
  8. Bayesian Gradient Estimation
  9. Posterior Update In Subspace
  10. Algorithm
  11. Theoretical Analysis
  12. Experiments
  13. Experimental Setup
  14. Performance in Masked Language Models
  15. Performance in decoder-only models
  16. ...and 12 more sections

Key Result

Lemma 3.5

For any direction $d\in\mathbb{R}^k$ of subspace $B$, the expectation of one-side difference $\hat{y}(d)$ satisfies:

Figures (1)

  • Figure 1: GPU memory usage comparison across different models. BSZO maintains memory consumption comparable to MeZO, while MeZO-Adam and HiZOO require significantly more memory due to storing optimizer states or Hessian estimates.

Theorems & Definitions (26)

  • Definition 3.3
  • Definition 3.4
  • Lemma 3.5
  • Corollary 3.6
  • Theorem 3.7
  • Definition 3.8
  • Theorem 3.9
  • Corollary 3.10
  • Lemma 2.1: Expectation of Direction Derivative
  • proof
  • ...and 16 more