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

TL;DR

AGZO introduces an activation-guided approach to zeroth-order optimization for fine-tuning large language models, leveraging activation structure to improve efficiency and performance under memory constraints.

Contribution

It proposes a novel activation-informed subspace method for ZO optimization, with theoretical guarantees and empirical improvements over existing approaches.

Findings

AGZO outperforms state-of-the-art ZO baselines.

It narrows the performance gap with first-order fine-tuning.

Maintains similar memory footprint as other ZO methods.

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.