AGZO: Activation-Guided Zeroth-Order Optimization for LLM Fine-Tuning
Wei Lin, Yining Jiang, Qingyu Song, Qiao Xiang, Hong Xu
ICML 2026
cs.LG
2026-01-24
AGZO confines rank-1 ZO noise to the forward activation subspace. Qwen3-0.6B averages 0.768 on eight tasks, above MeZO at 0.736, with the same peak memory.
Full-parameter fine-tuning hits GPU memory because backpropagation keeps activations, and those tensors grow with sequence length and batch size. Zeroth-order (ZO) updates skip the backward pass and estimate a step from loss values. MeZO perturbs every weight with isotropic Gaussian noise, takes two forwards, and rebuilds the noise from a seed, so peak memory stays close to inference. LOZO uses low-rank noise, but the factors are still drawn without looking at the batch.
On a linear layer the weight gradient is the upstream gradient times the transpose of the input activations. Its row space sits inside the column space of those activations. Isotropic noise spends a large share of queries where gradient energy is near zero.
During the forward pass AGZO runs three steps of power iteration on the current activations, keeps an orthonormal basis, and drops the activations. Experiments fix the rank at r = 1, so only the dominant direction is retained. Linear-layer perturbations are confined to that direction. The Gaussian factor can be regenerated from a seed; the basis depends on the batch and has to be stored. Nonlinear layers fall back to ordinary Gaussian noise.
The estimator is a one-sided finite difference, the same shape MeZO uses. A clean forward and a perturbed forward yield a scalar step. The smoothing constant µ is 1e-7 on Qwen3 and 1e-4 on Pangu-1B.
Rank 1 minimizes storage and forces the random search onto one axis. A spectrum check on GPT-2 fine-tuned on SST-2 points elsewhere: the cosine between the true gradient and its projection onto the activation subspace is already near 1 for rank at least 10, and the full subspace there is rank 750. The main runs do not choose rank from that curve. When activation energy is concentrated and the leading directions hold at least average gradient energy, the expected cosine with the true gradient is strictly higher than MeZO's.
The GPU testbed is two RTX 3090s. ZO methods train for 20,000 steps and full first-order fine-tuning (FO) for 1,000. Averages below are arithmetic means of the per-task scores.
| Setting | FO | AGZO | MeZO | LOZO |
| Qwen3-0.6B, 8 tasks | 0.816 | 0.768 | 0.736 | 0.736 |
| Qwen3-4B, 7 tasks | OOM | 0.835 | 0.825 | 0.818 |
| Pangu-1B GPU, 6 tasks | 0.751 | 0.720 | 0.707 | 0.688 |
The 0.6B tasks are SST-2, COPA, CB, BoolQ, MultiRC, RTE, WiC, and SQuAD. AGZO is about 3 points above MeZO and LOZO, and the gap to FO shrinks from about 8 points to about 5. SST-2 is 0.904 / 0.877 / 0.858 for FO / AGZO / MeZO. CB is 0.892 versus 0.803. WiC is 0.595 against FO at 0.675, and SQuAD is 0.790 against 0.871. BoolQ is the exception: MeZO 0.730, AGZO 0.724.
On Qwen3-4B, FO does not fit, and AGZO's edge over MeZO shrinks to about 1 point. Pangu-1B is the GPU build of openPangu-Embedded-1B, scored on SST-2, COPA, CB, BoolQ, RTE, and WiC. The same weights on Ascend 910B2 NPUs average 0.709 for AGZO, 0.703 for MeZO, and 0.667 for LOZO. CB falls from 0.732 on GPU to 0.696 on NPU.
On SST-2 the cosine between the ZO estimate and the true same-batch gradient is on the order of 1e-4, and AGZO stays above MeZO. In the DROP memory sweep the three ZO methods stay within 0.02 GB of each other, while first-order SGD runs out of memory as sequence length or batch size grows.
If the run is already zeroth-order and updates every weight, AGZO barely changes peak memory. It only retargets the noise onto the dominant activation direction. A 4B dense backward pass does not fit on a 3090; the ZO run still does.
The accuracy gain is incremental. It is about 3 points over MeZO at 0.6B, about 1 point at 4B, and 0.006 on the NPU rescore. Each step is still two forwards, and ZO takes 20 times as many steps as FO.
Rank 1 conflicts with the GPT-2 spectrum, where rank 10 is what captures the gradient. There is no rank ablation, and the energy left outside the rank-1 subspace is not measured. The observed cosine is about 1e-4. Many task gaps are under a point, the main tables do not report run variance, and the 9-point CB gap is a single run.
ZO trains for 20,000 steps and FO for 1,000, so the budgets differ by 20 times. LoRA, QLoRA, and gradient checkpointing are absent. FO running out of memory at 4B does not show that AGZO beats an adapter. With µ at 1e-7, GPU training followed by NPU eval already moves scores, and CB drops from 0.732 to 0.696. The rank study is on GPT-2. The theorem is an expectation for noiseless, vanishing-µ differences, while the training curves also carry minibatch noise.