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 沿激活主方向做秩 1 扰动，Qwen3-0.6B 八任务平均 0.768，超过 MeZO 的 0.736。
全参微调的显存墙多半来自反传要留的激活。零阶方法避开反传，MeZO 用各向同性高斯扰动配两次前向，噪声只留种子，峰值显存接近推理。LOZO 把扰动做成低秩，因子仍与数据无关。
线性层的权重梯度等于上游梯度乘输入激活的转置，行空间锁在激活张成的列空间里。满空间随机方向里，大量查询落在梯度能量接近零的地方。
AGZO 在前向里用 3 步幂迭代，从当前 batch 的激活抽出正交基底，随即释放激活。实验取秩 r = 1，只留能量最大的方向。线性层扰动的行空间被限制在这个方向上，随机因子仍可由种子重放；非线性层退回高斯扰动。
更新沿用单侧差分：未扰动损失和加扰后的损失相减，除以平滑系数 µ，得到步长标量。Qwen3 上 µ 为 10 的负 7 次方，Pangu-1B 上为 10 的负 4 次方。
r = 1 把基底存到最小，探索也压进主方向。GPT-2 在 SST-2 上，梯度投到激活子空间的余弦在 r 不小于 10 时已接近 1，满秩对应 r = 750。主实验没有顺着这条谱选秩。激活谱集中、且头部方向的梯度能量不低于平均时，期望余弦严格高于 MeZO。
两张 RTX 3090。ZO 跑 20000 步，一阶全参（FO）跑 1000 步。下表平均由各任务分数算术平均得到。
| 设置 | FO | AGZO | MeZO | LOZO |
| Qwen3-0.6B，八任务 | 0.816 | 0.768 | 0.736 | 0.736 |
| Qwen3-4B，七任务 | 显存不足 | 0.835 | 0.825 | 0.818 |
| Pangu-1B GPU，六任务 | 0.751 | 0.720 | 0.707 | 0.688 |
0.6B 的任务是 SST-2、COPA、CB、BoolQ、MultiRC、RTE、WiC、SQuAD。AGZO 比 MeZO、LOZO 高约 3 个点，距 FO 从约 8 个点收到约 5 个点。SST-2 为 0.904 / 0.877 / 0.858（FO / AGZO / MeZO），CB 为 0.892 对 0.803。WiC 0.595、SQuAD 0.790，对应 FO 0.675 和 0.871。BoolQ 上 MeZO 0.730，略高于 AGZO 的 0.724。
4B 上 FO 显存不够，AGZO 比 MeZO 只高约 1 个点。Pangu-1B 是 openPangu-Embedded-1B 的 GPU 版，六任务为 SST-2、COPA、CB、BoolQ、RTE、WiC。同一权重在昇腾 910B2 上，表内平均分为 AGZO 0.709、MeZO 0.703、LOZO 0.667，CB 从 GPU 的 0.732 掉到 0.696。
SST-2 上，估计梯度与真梯度的余弦在 10 的负 4 次方，AGZO 整段高于 MeZO。DROP 显存扫描里，三种 ZO 的峰值差不超过 0.02 GB，一阶 SGD 在更长序列或更大 batch 上 OOM。
打算用零阶更新全部权重时，AGZO 几乎不加显存，只把噪声换成当前激活的主方向。4B 的全参反传在 3090 上放不下，零阶仍跑得动。
相对 MeZO 的优势随规模变小：0.6B 约 3 个点，4B 约 1 个点，NPU 复测只剩 0.006。步数仍是一阶的 20 倍。
r = 1 和「r 不小于 10 才罩住梯度」拧在一起。正文没有秩消融，漏掉的梯度能量也没量化。实测余弦只有 10 的负 4 次方。主表没有多次运行的方差，CB 上大约 9 个点的单次差距站不稳。
ZO 跑 20000 步，FO 跑 1000 步，更新预算差着 20 倍。LoRA、QLoRA 和梯度检查点没有进对照。4B 上 FO 显存不够，说明不了 AGZO 比适配器更值。µ 取到 10 的负 7 次方后，GPU 训练换 NPU 评测，平均分会掉，CB 从 0.732 掉到 0.696。秩扫描做在 GPT-2 上。定理给的是无噪声、µ 趋于 0 的期望余弦，训练曲线里还有 minibatch 噪声。