cs.LGSep 28, 2026

AIM-ZO: Activation-Informed Subspace Maintenance for Zeroth-Order LLM Fine-Tuning

Authors: Yue Xie, Zhi Zheng, Yunpeng Ba, Xuyang Wu, Xialiang Tong, Zhichao Lu, Tao Zhong, Zhenkun Wang

Organizations: Southern University of Science and Technology · National University of Singapore · Huawei Technologies Ltd. · City University of Hong Kong

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Zeroth-order (ZO) optimization offers a memory-efficient alternative for LLM fine-tuning by estimating updates only from forward evaluations of perturbed parameters, without backpropagation or activation storage. However, in billion-parameter LLMs, isotropic perturbations often waste many forward evaluations on weakly informative directions. To make these evaluations more informative, existing ZO methods restrict perturbations to low-dimensional subspaces. Yet the quality of these subspaces is critical: overly compressed or poorly maintained spaces can miss useful update directions. To obtain a high-quality subspace for ZO updates, this paper proposes AIM-ZO, a ZO fine-tuning method based on Activation-Informed Subspace Maintenance. AIM-ZO uses forward activations as local directional information and continuously integrates them into a broad, evolving subspace over training. To access broader gradient-relevant structure while keeping individual perturbations low-dimensional, AIM-ZO activates only a smaller set of shared and sampled directions, decoupling the maintained width from the active width. We evaluate AIM-ZO across 5 LLMs and 11 downstream tasks under matched forward-evaluation budgets; its six-task average exceeds the strongest fully evaluated ZO baseline by 1.26 percentage points on OPT-2.7B and MeZO by 2.85 percentage points on OPT-30B. Our code is available at https://github.com/EkkoXy/AIM-ZO

Figures & tables

Figure 1: Perturbation spaces. MeZO explores the full parameter space; AGZO constructs a subspace from current activations. AIM-ZO maintains a width- KK subspace across training and activates a width- kk subspace with hh shared and k−hk-h sampled directions per perturbation.
MethodRTEBoolQSST-2WiCWSCSQuADAvg.
OPT-2.7B
Zero-shot55.2352.8756.6554.8636.5426.9247.18
MeZO65.13±1.6765.13_{\pm 1.67}66.10±2.2666.10_{\pm 2.26}92.50±0.5492.50_{\pm 0.54}58.53±0.4158.53_{\pm 0.41}54.81±1.5254.81_{\pm 1.52}80.99±1.4280.99_{\pm 1.42}69.68
CurvZO64.53±4.8364.53_{\pm 4.83}68.01±0.6268.01_{\pm 0.62}92.96±0.31\mathbf{92.96}_{\pm 0.31}57.96±2.0757.96_{\pm 2.07}46.47±7.2846.47_{\pm 7.28}58.71±2.4758.71_{\pm 2.47}64.77
HiZOO60.29±2.9460.29_{\pm 2.94}67.31±1.3467.31_{\pm 1.34}91.77±0.6791.77_{\pm 0.67}58.88±0.8058.88_{\pm 0.80}53.53±2.4253.53_{\pm 2.42}73.81±0.9073.81_{\pm 0.90}67.60
AGZO64.40±3.1064.40_{\pm 3.10}66.26±1.4666.26_{\pm 1.46}92.59±0.5192.59_{\pm 0.51}57.02±1.7357.02_{\pm 1.73}47.76±2.2247.76_{\pm 2.22}29.82±5.5429.82_{\pm 5.54}59.64
Table 1: Multi-method comparison on OPT-2.7B and OPT-13B. Scores are accuracy (%) except SQuAD (F1). Avg. is the unweighted mean over six tasks. Best mean scores are bold; subscripts report available sample standard deviations. Evaluation details are in Appendix H .
OPT-30B
MethodRTEBoolQSST-2WiCWSCSQuADAvg.
Zero-shot53.7940.9058.8352.9837.5046.4848.41
MeZO63.25±1.7963.25_{\pm 1.79}70.74±1.1770.74_{\pm 1.17}91.08±0.3691.08_{\pm 0.36}55.96±1.8155.96_{\pm 1.81}58.08±2.8558.08_{\pm 2.85}80.22±0.9980.22_{\pm 0.99}69.89
AIM-ZO (Ours)68.16±2.66\mathbf{68.16}_{\pm 2.66}74.90±1.53\mathbf{74.90}_{\pm 1.53}93.47±0.56\mathbf{93.47}_{\pm 0.56}57.71±0.91\mathbf{57.71}_{\pm 0.91}59.23±2.77\mathbf{59.23}_{\pm 2.77}82.96±0.69\mathbf{82.96}_{\pm 0.69}72.74\mathbf{72.74}
Qwen3-8B-Base
MethodRTEBoolQSST-2WiCMultiRCSQuADAvg.
Table 2: Scaling and model-family comparisons on OPT-30B and Qwen3-8B (%). SQuAD uses F1; other tasks use accuracy. Avg. is the unweighted mean over six tasks. Best mean scores are bold; subscripts report standard deviations. Qwen3-8B entries use three seeds and development-selected checkpoints.
Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

ModelLearning rate η\etaInitial scale ϵ0\epsilon_{0}
OPT-2.7B1.5×10−31.5\times 10^{-3}8×10−58\times 10^{-5}
OPT-13B1×10−31\times 10^{-3}8×10−58\times 10^{-5}
OPT-30B5×10−45\times 10^{-4}6×10−56\times 10^{-5}
Qwen3-0.6B-Base1×10−31\times 10^{-3}8×10−58\times 10^{-5}
Qwen3-8B-Base1×10−31\times 10^{-3}8×10−58\times 10^{-5}
Appendix
Table 3: AIM-ZO learning rates and initial perturbation scales.
MethodCalls/updateLearning ratePerturbation scale
MeZO2(1–5)×10−7(1\text{--}5)\times 10^{-7}10−310^{-3}
CurvZO2(1–8)×10−7(1\text{--}8)\times 10^{-7}5×10−4–10−35\times 10^{-4}\text{--}10^{-3}
AGZO35×10−75\times 10^{-7}10−4–10−310^{-4}\text{--}10^{-3}
HiZOO35×10−75\times 10^{-7}5×10−4–10−35\times 10^{-4}\text{--}10^{-3}
LoZO210−7–10−610^{-7}\text{--}10^{-6}10−310^{-3}
ZO-Muon510−210^{-2}10−310^{-3}
Appendix
Table 4: Baseline hyperparameter search ranges and calls per update for OPT-13B and Qwen3-0.6B, with OPT-2.7B MeZO also included.
ModelLayerr=8r=8r=32r=32r=64r=64r=128r=128
Qwen3-4BEarly.622.820.879.923
 Middle.911.964.980.990
 Late.992.994.995.995
OPT-2.7BEarly.443.678.784.868
 Middle.954.980.987.992
 Late.923.966.981.991
Appendix
Table 5: Temporal held-out shared-space capture. The rank- rr space is fitted on eight steps and evaluated on the next eight.
ModelLayerK=8K=8K=16K=16K=32K=32K=64K=64K=128K=128
Qwen3-4BEarly.2710.1242.0523.0491.0462
 Middle.1120.0909.0636.0469.0370
 Late.2101.1572.0636.0344.0225
OPT-2.7BEarly.2439.1550.0526.0295.0209
 Middle.0394.0223.0142.0064.0046
 Late.2557.2252.1571.1398.1184
Appendix
Table 6: Temporal held-out activation alignment residual eact2e_{\rm act}^{2} at fixed shared rank r=8r=8 .
ModelLayerOja meanOja rangeCurrent SVD mean
Qwen3-4BEarly.522.476–.628.695
 Middle.862.791–.951.890
 Late.878.699–.957.926
OPT-2.7BEarly.635.587–.675.717
 Middle.951.932–.966.964
 Late.657.537–.719.770
Appendix
Table 7: Capture of the finite-training-set mean gradient over 16 steps at K=128K=128 ; the current-batch reference uses exact SVD.
  Wide subspace ( K=128K=128 )Active space ( k=64k=64 )
ModelLayerDifference (pp)Win rateDifference (pp)Win rate
Qwen3-4BEarly-11.063.9%-10.3111.7%
 Middle+4.7198.8%+4.0694.5%
 Late+1.0249.6%+2.3864.5%
OPT-2.7BEarly+8.9596.1%+6.2983.2%
 Middle+1.5480.9%+2.1983.6%
Appendix
Table 8: Oja minus PI-5 current-batch capture in percentage points (pp), and Oja win rates. Pairs share checkpoints and are not independent training runs.
ModelLayerηq=.03\eta_{q}=.03.1.1.3.3.6.6
Qwen3-4BEarly(.943,.140)(.926,.315)(.892,.414)(.837,.549)
 Middle(.527,.898)(.330,.901)(.338,.894)(.410,.886)
 Late(.533,.924)(.618,.923)(.657,.922)(.670,.921)
OPT-2.7BEarly(.161,.713)(.263,.677)(.349,.646)(.384,.635)
 Middle(.349,.943)(.417,.940)(.559,.935)(.617,.933)
 Late(.440,.704)(.572,.662)(.663,.648)(.697,.644)
Appendix
Table 9: Stationary replay: (normalized Φ\Phi , gradient capture). Qwen and OPT use steps 500 and 2000, respectively.
WidthQwen3-4BOPT-2.7B
8.7022 / .0020 / .002272.8137 / .0026 / .002855
16.7758 / .0068 / .001615.8458 / .0043 / .002052
32.8243 / .0121 / .001072.8736 / .0082 / .001374
64.8624 / .0277 / .000829.8975 / .0183 / .000982
128.8866 / .0510 / .000564.9166 / .0416 / .000752
Appendix
Table 10: Full active-space width scan using independent Gaussian coefficient matrices. Each cell gives learned capture / random capture / mean learned single-direction first-order cosine.
ModelSpaceOP MSETP MSEReductionCosine gap
Qwen3-4BLearned50,35794,00246.4%1.192×10−31.192\times 10^{-3}
 Random1,0421,68038.0%1.911×10−41.911\times 10^{-4}
 Dense1,825,7523,487,91547.7%1.647×10−41.647\times 10^{-4}
OPT-2.7BLearned30,75565,77753.2%1.129×10−31.129\times 10^{-3}
 Random7641,36944.2%2.530×10−42.530\times 10^{-4}
 Dense1,115,0831,998,32544.2%2.280×10−42.280\times 10^{-4}
Appendix
Table 11: Matched normalized Gaussian control against a fixed batch gradient. MSE is normalized by gradient energy; cosine gap is OP minus TP.
ModelSpaceOP MSETP MSEReductionCosine gap
Qwen3-4BLearned1,995,8353,743,20146.7%+7.47×10−5+7.47\times 10^{-5}
 Random45,46291,24250.2%+3.06×10−5+3.06\times 10^{-5}
 Dense96,707,409192,548,22549.8%−3.72×10−5-3.72\times 10^{-5}
OPT-2.7BLearned1,374,6882,154,53536.2%+8.44×10−5+8.44\times 10^{-5}
 Random31,73153,72340.9%+7.29×10−5+7.29\times 10^{-5}
 Dense63,415,059116,824,52845.7%−2.92×10−5-2.92\times 10^{-5}
Appendix
Table 12: Finite-training-set reference at ϵ=8×10−5\epsilon=8\times 10^{-5} , K=64K=64 , and 16 loss evaluations.
Modelϵ\epsilonLearnedRandomDense
Qwen3-4B6×10−56\times 10^{-5}46.7 / +7.6450.2 / +2.9949.8 / -3.69
 7.43×10−57.43\times 10^{-5}46.7 / +7.4850.2 / +3.0449.9 / -3.72
 10−310^{-3}39.1 / -1.1750.3 / +3.0750.6 / -3.70
OPT-2.7B2.57×10−52.57\times 10^{-5}36.5 / +7.6641.1 / +7.4945.7 / -2.68
 6×10−56\times 10^{-5}36.3 / +8.5141.2 / +7.2845.7 / -2.98
 10−310^{-3}33.9 / +2.5540.9 / +7.2145.7 / -2.82
Appendix
Table 13: Perturbation-scale control against the exact training gradient. Entries are OP MSE reduction (%) / OP–TP cosine gap in units of 10−510^{-5} .
ModelSpaceOPTPDifference95% interval
Qwen3-4BLearned38.7232.935.791[1.96,9.63]
 Dense2.7742.443.3307[-.514,1.18]
OPT-2.7BLearned40.4832.607.882[4.17,11.6]
 Dense5.4824.2461.236[.536,1.94]
Appendix
Table 14: BF16 implementation-level cosine. All numeric entries are in units of 10−510^{-5} ; intervals are paired 95% intervals conditional on the checkpoint and batch.
VariantRTEBoolQSST-2SQuAD
MeZO65.13±1.6765.13\pm 1.6766.10±2.2666.10\pm 2.2692.50±.5492.50\pm.5480.99±1.4280.99\pm 1.42
Common active basis66.14±.7866.14\pm.7866.84±.6366.84\pm.6392.48±.7892.48\pm.7881.06±1.0281.06\pm 1.02
AIM-ZO67.51±3.1867.51\pm 3.1867.22±1.6767.22\pm 1.6792.87±.4792.87\pm.4781.19±.4981.19\pm.49
Appendix
Table 15: Common active basis versus per-perturbation resampling on OPT-2.7B. Classification scores are accuracy; SQuAD uses F1 (%).
Maintenance ruleAccuracy
Fixed subspace68.40±2.0468.40\pm 2.04
Oja every ten steps68.92±3.3868.92\pm 3.38
Oja every step70.20±1.9570.20\pm 1.95
Appendix
Table 16: OPT-2.7B RTE maintenance-frequency ablation (five seeds; final development accuracy, %).
ModelK=64, k=h=64K=64,\ k=h=64K=128, k=h=64K=128,\ k=h=64AIM-ZO
OPT-2.7B64.26±3.8264.26\pm 3.8266.55±1.1066.55\pm 1.1067.51±3.18\mathbf{67.51}\pm 3.18
Qwen3-0.6B77.38±1.7177.38\pm 1.7177.38±.4277.38\pm.4277.62±.92\mathbf{77.62}\pm.92
Appendix
Table 17: RTE accuracy (%) with fixed-prefix activation and the AIM-ZO main-result reference.
Maintained width KKAccuracy
6465.70±1.2565.70\pm 1.25
12865.46±3.6265.46\pm 3.62
25667.15±0.7267.15\pm 0.72
Appendix
Table 18: OPT-2.7B RTE maintained-width scan at k=64k=64 and h=48h=48 (three seeds; official-validation accuracy, %).
Shared width hhSampled tail k−hk-hAccuracy
06460.77±1.9960.77\pm 1.99
481665.46±3.6265.46\pm 3.62
64066.55±1.1066.55\pm 1.10
Appendix
Table 19: OPT-2.7B RTE shared-width scan at K=128K=128 and k=64k=64 (three seeds; official-validation accuracy, %).
Population size NNTraining forwardsAccuracy
839,99661.01±3.7761.01\pm 3.77
1540,00065.46±3.6265.46\pm 3.62
Appendix
Table 20: OPT-2.7B RTE population-size ablation (three seeds; official-validation accuracy, %).
Coefficient formAccuracy (%)Step time (s)Memory (GiB)
Rank-one ab⊤Q⊤ab^{\top}Q^{\top}65.46±3.6265.46\pm 3.6210.3511.214±0.15811.214\pm 0.158
Rank-four ABQ⊤ABQ^{\top}64.98±2.0164.98\pm 2.019.7711.216±0.16011.216\pm 0.160
Full RQ⊤RQ^{\top}65.58±2.2165.58\pm 2.2112.2411.153±0.10511.153\pm 0.105
Appendix
Table 21: Coefficient-rank ablation on OPT-2.7B RTE over three seeds. Accuracy is official-validation accuracy (%); time is mean per-run median non-evaluation step time; memory is peak CUDA allocated memory.
EstimatorBest checkpointFinal update
OP15-RLOO71.40±3.0271.40\pm 3.0271.12±3.2971.12\pm 3.29
Raw TP870.04±4.3870.04\pm 4.3869.28±4.6269.28\pm 4.62
Appendix
Table 22: Online OPT-2.7B RTE estimator ablation over five matched seeds (development accuracy, %).
 Qwen3-4BOPT-2.7B
ssCaptureΔcos⁡\Delta_{\cos}CaptureΔcos⁡\Delta_{\cos}
0.0046.0795.0020.0668
.10.0857.365.0931.448
.25.2072.567.2304.705
.50.4094.798.4598.996
.75.6116.975.68941.220
Appendix
Table 23: Fixed-width space-quality control with an exact linear oracle. Alignment gains are reported in units of 10−310^{-3} .
Modelϵ\epsilonCorrelationMonotoneOP MSE wins
Qwen3-4B6×10−56\times 10^{-5}.95212/1284/84
 10−310^{-3}.95212/1284/84
OPT-2.7B2.57×10−52.57\times 10^{-5}.92712/1284/84
 10−310^{-3}.92612/1284/84
Appendix
Table 24: Real-loss space-quality sweep. Correlation relates measured capture to the OP–TP alignment gain; monotone counts refer to checkpoint/batch settings.
ModelTaskAIM-ZOMeZORatio
Qwen3-0.6BRTE4.3391.8712.319
OPT-2.7BRTE7.3633.5082.099
OPT-13BBoolQ26.60321.9191.214
Appendix
Table 25: Training-step hours at 40,000 forward evaluations. Ratio is AIM-ZO time divided by MeZO time.
ModelMethodEvaluations/steps/steps/eval.GiB
Qwen3-8BMeZO2.653.32719.929
 AIM-ZO167.396.46219.937
 AGZO31.196.39918.772
OPT-30BMeZO21.810.90558.552
 AIM-ZO1614.052.87860.034
 AGZO32.727.90958.087
Appendix
Table 26: Large-model short-run time and memory.
Caches/stepAllocated GiBReserved GiB
None9.16410.25119.111
Factor restoration7.27310.28719.143
Factor restoration + full noise7.24810.57620.516
Appendix
Table 27: Cache ablation for AIM-ZO . Time averages all ten steps; allocated and reserved memory are reported separately.
DatasetMetricMeZOAIM-ZOW/T/L
RTEAccuracy65.13±1.6765.13\pm 1.6767.51±3.18\mathbf{67.51\pm 3.18}4/1/0
BoolQAccuracy66.10±2.2666.10\pm 2.2667.22±1.67\mathbf{67.22\pm 1.67}5/0/0
SST-2Accuracy92.50±0.5492.50\pm 0.5492.87±0.47\mathbf{92.87\pm 0.47}5/0/0
WiCAccuracy58.53±0.4158.53\pm 0.4160.13±1.70\mathbf{60.13\pm 1.70}5/0/0
MultiRCAccuracy61.96±2.0461.96\pm 2.0463.94±1.59\mathbf{63.94\pm 1.59}4/0/1
MultiRCF1a34.71±10.8134.71\pm 10.8148.91±4.09\mathbf{48.91\pm 4.09}5/0/0
Appendix
Table 28: Five-seed official-validation results for OPT-2.7B under 40,000 matched training forwards. Higher is better. W/T/L counts paired-seed outcomes for AIM-ZO relative to MeZO. Scores are percentages.
MethodRTEBoolQSST-2WiCMultiRCSQuADAvg.
Zero-shot58.4862.2058.2651.2559.0856.2257.58
MeZO77.62±1.1177.62_{\pm 1.11}73.87±0.5773.87_{\pm 0.57}88.42±1.3588.42_{\pm 1.35}55.36±2.4455.36_{\pm 2.44}72.78±1.4972.78_{\pm 1.49}83.23±0.7683.23_{\pm 0.76}75.21
CurvZO77.14±1.0477.14_{\pm 1.04}68.28±0.8068.28_{\pm 0.80}82.87±1.1582.87_{\pm 1.15}49.74±1.3749.74_{\pm 1.37}71.72±2.8671.72_{\pm 2.86}71.29±3.4671.29_{\pm 3.46}70.17
HiZOO77.86±0.5577.86_{\pm 0.55}72.76±0.8472.76_{\pm 0.84}88.80±0.8888.80_{\pm 0.88}53.92±1.9653.92_{\pm 1.96}76.11±0.9376.11_{\pm 0.93}80.42±0.8980.42_{\pm 0.89}74.98
AGZO69.68±4.3969.68_{\pm 4.39}68.40±0.2668.40_{\pm 0.26}84.75±3.0484.75_{\pm 3.04}52.77±3.7052.77_{\pm 3.70}73.89±0.3873.89_{\pm 0.38}73.17±1.4373.17_{\pm 1.43}70.44
ZO-Muon75.69±0.2175.69_{\pm 0.21}72.73±0.7972.73_{\pm 0.79}88.76±0.7088.76_{\pm 0.70}57.79±2.0557.79_{\pm 2.05}75.10±0.6375.10_{\pm 0.63}81.10±1.3681.10_{\pm 1.36}75.20
Appendix
Table 29: Multi-method comparison on Qwen3-0.6B-Base (%). SQuAD uses F1; other tasks use accuracy. Avg. is the unweighted mean over six tasks; subscripts report available sample standard deviations.
MetricZero-shotMeZOAIM-ZO
MultiRC F1a54.7083.95±0.2583.95\pm 0.2584.48±0.3784.48\pm 0.37
MultiRC question EM28.2354.60±1.3754.60\pm 1.3756.31±0.4456.31\pm 0.44
SQuAD EM72.4078.20±2.2378.20\pm 2.2384.63±0.4284.63\pm 0.42
Appendix
Table 30: Additional Qwen3-8B-Base evaluation metrics (%; three seeds for trained methods).
MethodSubspace perturbationsInformation- guidedHistory- informedWidth decoupling
MeZO×\times×\times×\times×\times
HiZOO×\times✓\checkmark✓\checkmark×\times
CurvZO✓\checkmark✓\checkmark✓\checkmark×\times
LoZO✓\checkmark×\times×\times×\times
SubZero✓\checkmark×\times×\times×\times
ZO-Muon✓\checkmark×\times×\times×\times
Appendix
Table 31: Design properties of representative ZO methods. Sparse coordinate spaces count as subspace perturbations; guidance and history refer to information used in perturbation construction.

Explore similar work

CardsList

  1. ZO-Act: Efficient Zeroth-Order Fine-Tuning via One-Shot Activation-Informed Low-Rank Subspaces

    Jul 1, 2026Xun Dong, Yibo Xu, Naigang Wang +3Zeroth-Order OptimizationZeroth-Order

    University at Albany, SUNY · IBM T. J. Watson Research Center

  2. SubZero+: Memory-Efficient Adaptive Zeroth-Order LLM Fine-Tuning in Random Subspaces

    Aug 16, 2026Ziming Yu, Shuyao Xiao, Xingyu Zhao +6Zeroth-OrderLarge Language Model Fine-Tuning

    Beijing Normal University · Singapore Management University · Xiaomi Inc. +2

  3. Dominant-Layer ZO: A Single Layer Dominates Zeroth-Order Fine-Tuning of LLMs

    Jun 3, 2026Wanhao Yu, Ziyan Wang, Zheng Wang +7Continual Fine-TuningZeroth-Order

    University of North Carolina at Charlotte · University of Houston · State University of New York at Binghamton +2