Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AGZO: Activation-Guided Zeroth-Order Optimization for LLM Fine-Tuning

Wei Lin, Yining Jiang, Qingyu Song, Qiao Xiang, Hong Xu

TL;DR

This work tackles memory-constrained fine-tuning of large language models by improving zeroth-order optimization. It introduces Activation-Guided Zeroth-Order Optimization (AGZO), which constructs activation-informed low-rank subspaces from the forward pass and restricts perturbations to these subspaces, while nonlinear layers revert to Gaussian perturbations. Theoretical results show AGZO estimates a subspace-smoothed gradient with higher expected cosine similarity to the true gradient than isotropic baselines, and empirical results on Qwen3 and openPangu models demonstrate stronger gradient alignment, better downstream performance, and comparable memory footprints to traditional ZO methods. Practically, AGZO narrows the performance gap between ZO and first-order fine-tuning under memory constraints, enabling more effective large-model adaptation on commodity hardware.

Abstract

Zeroth-Order (ZO) optimization has emerged as a promising solution for fine-tuning LLMs under strict memory constraints, as it avoids the prohibitive memory cost of storing activations for backpropagation. However, existing ZO methods typically employ isotropic perturbations, neglecting the rich structural information available during the forward pass. In this paper, we identify a crucial link between gradient formation and activation structure: the gradient of a linear layer is confined to the subspace spanned by its input activations. Leveraging this insight, we propose Activation-Guided Zeroth-Order optimization (AGZO). Unlike prior methods, AGZO extracts a compact, activation-informed subspace on the fly during the forward pass and restricts perturbations to this low-rank subspace. We provide a theoretical framework showing that AGZO optimizes a subspace-smoothed objective and provably yields update directions with higher cosine similarity to the true gradient than isotropic baselines. Empirically, we evaluate AGZO on Qwen3 and Pangu models across various benchmarks. AGZO consistently outperforms state-of-the-art ZO baselines and significantly narrows the performance gap with first-order fine-tuning, while maintaining almost the same peak memory footprint as other ZO methods.

AGZO: Activation-Guided Zeroth-Order Optimization for LLM Fine-Tuning

TL;DR

This work tackles memory-constrained fine-tuning of large language models by improving zeroth-order optimization. It introduces Activation-Guided Zeroth-Order Optimization (AGZO), which constructs activation-informed low-rank subspaces from the forward pass and restricts perturbations to these subspaces, while nonlinear layers revert to Gaussian perturbations. Theoretical results show AGZO estimates a subspace-smoothed gradient with higher expected cosine similarity to the true gradient than isotropic baselines, and empirical results on Qwen3 and openPangu models demonstrate stronger gradient alignment, better downstream performance, and comparable memory footprints to traditional ZO methods. Practically, AGZO narrows the performance gap between ZO and first-order fine-tuning under memory constraints, enabling more effective large-model adaptation on commodity hardware.

Abstract

Zeroth-Order (ZO) optimization has emerged as a promising solution for fine-tuning LLMs under strict memory constraints, as it avoids the prohibitive memory cost of storing activations for backpropagation. However, existing ZO methods typically employ isotropic perturbations, neglecting the rich structural information available during the forward pass. In this paper, we identify a crucial link between gradient formation and activation structure: the gradient of a linear layer is confined to the subspace spanned by its input activations. Leveraging this insight, we propose Activation-Guided Zeroth-Order optimization (AGZO). Unlike prior methods, AGZO extracts a compact, activation-informed subspace on the fly during the forward pass and restricts perturbations to this low-rank subspace. We provide a theoretical framework showing that AGZO optimizes a subspace-smoothed objective and provably yields update directions with higher cosine similarity to the true gradient than isotropic baselines. Empirically, we evaluate AGZO on Qwen3 and Pangu models across various benchmarks. AGZO consistently outperforms state-of-the-art ZO baselines and significantly narrows the performance gap with first-order fine-tuning, while maintaining almost the same peak memory footprint as other ZO methods.
Paper Structure (38 sections, 13 theorems, 110 equations, 4 figures, 4 tables, 2 algorithms)

This paper contains 38 sections, 13 theorems, 110 equations, 4 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background: Zeroth-Order Fine-Tuning
  3. Memory-Efficient Full-Space Perturbations
  4. Low-Rank Zeroth-Order Perturbations
  5. Limitations of Existing ZO Baselines
  6. Structural Analysis of Gradients and Activations
  7. Gradient Confinement in Activation Subspaces
  8. Low-Rank Structure of Gradients and Activations
  9. Activation-Guided Zeroth-Order Optimization
  10. AGZO Algorithm
  11. Memory Usage Analysis
  12. Provable Superiority of AGZO
  13. AGZO is A Projected Gradient Estimator
  14. AGZO is A High-Precision Gradient Estimator
  15. Experiments
  16. ...and 23 more sections

Key Result

Proposition 5.1

Assume $F$ has $L$-Lipschitz gradient. The AGZO estimator satisfies, for each linear layer $\ell$, In particular, AGZO estimates the gradient of the subspace-smoothed objective eq:subspace-smooth, projected onto $\mathrm{span}(A_\ell)$.

Figures (4)

  • Figure 1: Structural analysis of gradients and activations. (a) Cosine similarity between the true gradient and its projection onto the activation subspace. (b) & (c) Singular value spectra of gradients and activations.
  • Figure 2: Gradient alignment during fine-tuning.
  • Figure 3: Peak GPU memory usage when fine-tuning Qwen3-0.6B on DROP.
  • Figure 4: Peak GPU memory usage when fine-tuning Pangu-1B on DROP.

Theorems & Definitions (25)

  • Proposition 5.1
  • proof
  • Proposition 5.2
  • proof
  • Corollary 5.3
  • Theorem 5.4
  • proof
  • Corollary 5.5
  • Theorem 5.6
  • proof
  • ...and 15 more