PaperlayerView exact revision on arXiv

Source-linked AI summary

SubZero+: Efficient Zeroth-Order LLM Fine-Tuning via Large Learning Rates

Ziming Yu, Shuyao Xiao, Xingyu Zhao, Sike Wang, Pan Zhou, Peiyu Zang, Xiangda Yan, Yongjie Yang, Jia Li

arXiv:2608.15665v1cs.LG

TL;DR

Zeroth-order LLM fine-tuning is hindered by high-variance gradient estimates and learning-rate sensitivity. SubZero+ combines multi-query low-rank subspace estimation, subspace Adam, and QR sign correction, consistently outperforming prior ZO baselines while improving learning-rate robustness and convergence speed.

  • Problem

    Existing zeroth-order methods produce high-variance gradient estimates for large models, making convergence unstable and learning-rate selection sensitive.

  • Method

    SubZero+ combines multi-query estimation in low-rank subspaces, adaptive Adam updates within those subspaces, and sign-corrected QR-based subspace construction.

  • Results

    Across 1.3B–32B models, tasks, and full-parameter or LoRA tuning, SubZero+ consistently outperforms ZO baselines with improved learning-rate robustness and convergence speed.

  • Takeaways & Limitations

    SubZero+ narrows the gap to first-order methods while retaining inference-level memory for zeroth-order fine-tuning.

  • Takeaways & Limitations

    SubZero+ uses data-independent random subspaces that may discard important gradient information despite preserving memory efficiency.

Abstract

from arXiv · show

Zeroth-order (ZO) optimization enables backpropagation-free fine-tuning of large language models, but existing ZO methods suffer from high-variance gradient estimators, making convergence unstable and highly sensitive to learning rates. We propose SubZero+, an improved SubZero framework that improves stability in three complementary ways: (i) multi-query gradient estimation within layer-specific low-rank subspaces to reduce variance without exhibiting the multi-query paradox; (ii) a subspace Adam optimizer that performs adaptive updates using in-subspace multi-query gradient statistics; and (iii) a sign correction for QR-based subspace construction to ensure Haar-distributed projection matrices, eliminating implementation-dependent orientation ambiguity. Experiments on models from 1.3B to 32B across SuperGLUE, under both full-parameter tuning and LoRA, show that SubZero+ consistently outperforms prior ZO baselines, enlarges the stable learning-rate range, and narrows the gap to first-order methods with minimal extra memory overhead.

1 INTRODUCTION

SubZero+ targets the high-variance, learning-rate-sensitive convergence of zeroth-order LLM fine-tuning by combining low-rank subspaces with multi-query estimation and adaptive in-subspace optimization. It is designed to improve stability and convergence while retaining inference-comparable memory usage.

  • Motivation: Zeroth-order optimization avoids backward passes, but its gradient-estimation variance grows with parameter dimension, creating a convergence and performance gap versus first-order methods.This challenge is especially severe for large language models with extremely high-dimensional parameters.
  • Prior approach: SubZero reduces estimation difficulty and variance by projecting gradients into layer-specific low-dimensional subspaces that preserve informative low-rank optimization directions.The smaller search space can outweigh reduced gradient fidelity, producing practical improvements.
  • Empirical motivation: ∼0.6GB of GPU memory overhead versus MeZO and SubZero accompanies smoother training curves, larger supported learning rates, faster convergence, and improved stability and performance.These observations are reported for full-parameter tuning of OPT-1.3B on SST-2.
  • Problem: ZO fine-tuning remains highly learning-rate-sensitive because larger rates amplify estimator noise, causing oscillations or divergence; MeZO requires an unusually small stable learning rate.This sensitivity limits convergence speed and stability compared with the desired behavior.
  • SubZero+ contributions: SubZero+ averages multiple independent perturbations within one low-rank subspace, reducing gradient variance while mitigating the multi-query paradox.All statistics are computed in-subspace, keeping memory usage comparable to inference-time.
  • SubZero+ contributions: SubZero+ applies Adam entirely within the low-rank subspace, using multi-query gradient statistics for adaptive parameter updates.The method is intended to enable more effective training with larger learning rates.

2 RELATED WORK

Prior zeroth-order LLM fine-tuning methods reduce memory through finite-difference estimation, sparse perturbations, or low-rank subspaces, but remain limited by estimator variance, single-query updates, and unresolved optimizer or multi-query trade-offs.

  • ZO Gradient Estimation: MeZO pioneered SPSA-based zeroth-order LLM fine-tuning using only forward passes, achieving inference-level memory but suffering variance that scales with parameter dimension d.ZO replaces backpropagation with finite-difference gradient estimates.
  • ZO Gradient Estimation: Sparse perturbations and low-rank projections reduce zeroth-order estimation costs, while LOZO uses low-rank perturbations and SubZero constructs column-orthonormal projections with QR decomposition.Both LOZO and SubZero use a single query per iteration, limiting variance reduction.
  • ZO Parameter Update: MeZO and SubZero use vanilla SGD, LOZO applies SGDM in the low-rank subspace with negligible memory overhead, and TeZO exploits temporal low-rankness across gradients.TeZO can be extended to Adam, but that extension does not consistently improve over its SGD counterpart.
  • Multi-Query ZO: Averaging K > 1 independent perturbation evaluations reduces per-step gradient variance by a factor of K, but the multi-query paradox cancels this benefit under a fixed forward-pass budget.Using K queries also reduces the number of training steps by a factor of K, yielding no net gain over single-query training.

3 PRELIMINARIES

This section formulates zeroth-order fine-tuning for layered LLM parameters and explains how MeZO and SubZero estimate gradients. SubZero reduces variance through layer-wise low-rank subspaces, but remains limited by single-query estimation, vanilla SGD, and implementation-dependent QR construction.

  • Problem Setup: LLM fine-tuning represents each layer’s trainable parameters as a matrix, flattens all matrices into w ∈ R^d, and minimizes loss L(W) over dataset D.The total parameter count is d = Σ_i m_i n_i.
  • ZO Gradient Estimation (MeZO): MeZO uses a two-forward-pass central-difference SPSA estimator with Gaussian perturbations, but its variance scales as O(d), producing extremely noisy estimates for large d.The estimator is unbiased for the smoothed objective E_z[L(w + εz)].
  • Subspace ZO Optimization (SubZero): SubZero constructs per-layer column-orthonormal projections through QR decomposition and restricts perturbations to rank-r subspaces with r ≪ min{m_i, n_i}.Each layer uses U_i ∈ R^{m_i×r} and V_i ∈ R^{n_i×r}, with low-dimensional random matrices Z_i ∈ R^{r×r}.
  • Limitations of SubZero: SubZero’s limitations are single-query estimation with K = 1, vanilla SGD without per-layer adaptivity, and implementation-dependent QR-based subspace construction.These limitations motivate the SubZero+ improvements.

4 THE SUBZERO+ FRAMEWORK

SubZero+ extends SubZero with multi-query estimation in low-rank subspaces, subspace Adam, and QR sign correction. Together, these components reduce gradient noise, enable curvature-aware layer adaptation, stabilize subspaces, and preserve low memory overhead.

  • Framework overview: SubZero+ combines multi-query subspace gradient estimation, low-rank subspace Adam, and QR-based projection construction with sign correction.The components form a dependency chain: sign correction improves subspace stability, multi-query estimation supplies low-noise gradient norms, and subspace Adam leverages them.
  • Multi-query gradient estimation: Multi-query estimation resolves the full-space multi-query paradox by reducing variance within SubZero’s subspace structure without proportionally increasing query cost.Single-query variance is O(r2) per layer, while naive full-space averaging incurs a K× query-cost increase.
  • Subspace Adam: Subspace Adam stores moment buffers in Rr×r, requiring only 2r2 floats per layer and assigning smaller effective steps to layers with larger gradient magnitudes.The optimizer operates entirely in the low-rank space, adapting updates to steep and flat regions.
  • Memory efficiency: 27.9 GB is SubZero+’s memory use for multi-query Adam on OPT-13B, versus 97.2 GB for MeZO, while MeZO’s increase is 3.8× from single-query SGD.SubZero+ confines auxiliary states to r × r subspace buffers; all four SubZero+ variants consume identical memory.
  • Subspace Adam: With K = 100, multi-query gradient norms faithfully reflect local curvature, enabling subspace Adam to assign small steps on steep walls and large steps in flat regions.With K = 1, the gradient norm field is dominated by noise; adaptive steps enable near-monotone descent.
  • QR subspace construction: QR sign correction enforces positive diagonal entries in R, restoring sign symmetry so projection bases are Haar-distributed and invariant to implementation-dependent QR conventions.The correction costs O(r) column sign flips per layer, negligible relative to the O(minir) QR cost.

5 EXPERIMENTS · 5.1 EXPERIMENTAL SETUP

SubZero+ is evaluated against established zeroth-order baselines across diverse decoder-only LLMs, SuperGLUE tasks, and both full-parameter and LoRA fine-tuning. The experimental setup fixes query, optimizer, batching, checkpoint-selection, and matched-forward-pass-budget procedures.

  • 5 EXPERIMENTS: Experiments evaluate SubZero+ against state-of-the-art zeroth-order baselines across diverse models, tasks, and fine-tuning schemes.Additional hyperparameter configurations and results are provided in Appendix A.2.
  • 5.1 EXPERIMENTAL SETUP: The model suite includes OPT-1.3B, OPT-13B, OPT-30B, LLaMA3.1-8B, and Qwen2.5-32B.All listed models are autoregressive decoder-only LLMs.
  • 5.1 EXPERIMENTAL SETUP: Evaluation spans SuperGLUE classification, multiple-choice, and generation tasks under FT and LoRA schemes.The LoRA configuration uses r = 8 and α = 16.
  • 5.1 EXPERIMENTAL SETUP: The comparisons include MeZO and SubZero, with Ada included as a first-order reference.The supplied passage truncates the first-order reference description after “Ada”.
  • 5.1 EXPERIMENTAL SETUP: K = 99, β1 = 0.9, β2 = 0.95, and ϵadam = 10−8 are SubZero+’s default query and Adam settings.Subspace rank r and update frequency F are specified in Appendix A.2; learning rates and perturbation scales are tuned by benchmark grid search.
  • 5.1 EXPERIMENTAL SETUP: All zeroth-order methods use batch size 16 and select the best checkpoint by validation loss.These procedures follow Malladi et al. (2023), while MeZO and SubZero use two-point gradient estimation.
  • 5.1 EXPERIMENTAL SETUP: SubZero+ constructs perturbations from layer-specific projection sets and random r × r matrices, with subspaces changed according to frequency F.The implementation includes Haar projection construction with QR-based sign correction and seeded random perturbations.
  • 5.1 EXPERIMENTAL SETUP: 400 steps, consuming 99 + 1 = 100 forwards per step, give SubZero+ the same 40K forward-pass budget as MeZO and SubZero.The matched baseline budget is 2 × 20K = 40K total forwards.

5.2 MAIN RESULTS · 5.3 QUERY COUNT ABLATION · 5.4 MEMORY ANALYSIS

SubZero+ consistently improves zeroth-order fine-tuning results across model scales and tuning schemes, while multi-query estimation strengthens learning-rate robustness. Query-count effects vary by task, and memory overhead remains minimal.

  • 5.2.1 AUTO-REGRESSIVE LLMS UNDER FT SCHEME: SubZero+ achieves 66.3% average accuracy on OPT-1.3B, surpassing SubZero at 65.5% and MeZO at 64.6% across 9 tasks.All ZO methods use FT, and Table 1 gives them a shared 40K forward-pass budget.
  • 5.2.2 AUTO-REGRESSIVE LLMS UNDER LORA SCHEME: Under LoRA, SubZero+ outperforms both MeZO and SubZero on OPT-13B, LLaMA3.1-8B, and OPT-30B.The comparison uses r = 8, α = 16 and each method’s best result from hyperparameter searches.
  • 5.2.3 SCALING TO QWEN2.5-32B: On Qwen2.5-32B, SubZero+ reaches 87.4% average accuracy, outperforming SubZero by +2.3%, MeZO by +3.8%, MeZO-LoRA by +10.7%, and zero-shot by +11.4%.SubZero+ records the highest score on every task, with particularly large gains on SST-2 and WIC.
  • 5.2.4 LEARNING RATE ROBUSTNESS: SubZero+ reaches 86.5% training accuracy at LR = 2 × 10−2 and stays above 79.0% across 10−2 ∼ 3 × 10−2, whereas MeZO peaks at 83.3% at LR = 1 × 10−4.The experiment fine-tunes Qwen3-0.6B with LoRA on SST-2 across four orders of magnitude.
  • 5.2.5 HAAR-CORRECTED QR-BASED SUBSPACE CONSTRUCTION: Haar-corrected QR subspaces improve average accuracy by +0.7% under Adam and +0.3% under SGD, from 61.5% to 62.2% and 65.1% to 65.3%, respectively.MultiRC gains are +2.8% under Adam and +2.6% under SGD; RTE benefits depend on the optimizer.
  • 5.4 MEMORY ANALYSIS: Under FT on OPT-13B, SubZero+ incurs approximately 1% peak GPU-memory overhead over MeZO, while under LoRA all three methods consume identical memory.The memory comparison uses batch size 16 and fp16 across four representative tasks.
  • 5.3 QUERY COUNT ABLATION: SST-2 accuracy rises from 50.9% at K = 1 to 92.2% at K = 100, a gain of +41.3 percentage points, while RTE peaks at 63.9% at K = 50.With a fixed 30,000 forward-pass budget, larger K reduces training steps; MultiRC varies only from 52.3% to 61.1%.

6 CONCLUSION

SubZero+ is a memory-efficient zeroth-order fine-tuning framework that combines multi-query subspace estimation, low-rank subspace Adam, and sign-corrected QR construction. Across 1.3B–32B models and FT/LoRA settings, it outperforms existing ZO baselines with improved learning-rate robustness and convergence speed.

  • 6 CONCLUSION: SubZero+ integrates multi-query subspace gradient estimation, low-rank subspace Adam, and QR-based subspace construction with sign correction.The framework is designed for memory-efficient zeroth-order fine-tuning.
  • 6 CONCLUSION: Across 1.3B–32B models, tasks, and FT/LoRA schemes, SubZero+ consistently outperforms existing ZO baselines while maintaining inference-level memory.Experiments also show substantially improved learning-rate robustness and convergence speed.
  • Limitations and Future Work: Random subspaces are data-independent, preserving memory efficiency but potentially discarding important gradient directions.This is identified as a fundamental trade-off of SubZero+.
  • Limitations and Future Work: The subspace is regenerated from scratch every F steps rather than adapted from prior optimization history.Data-dependent subspace construction and adaptive subspace updating are proposed as future directions.
  • Broader Impact: SubZero+ lowers the barrier to LLM customization by reducing memory cost and hyperparameter sensitivity.The paper reports no specific negative societal impacts beyond those generally associated with LLM fine-tuning.

A APPENDIX · A.1 IMPLEMENTATION DETAILS

SubZero+ constructs sign-corrected layer-specific subspaces using Gaussian matrices and thin QR decomposition, while retaining memory-efficient regeneration and compact optimizer-state storage. Multi-query perturbations are generated sequentially to keep peak memory at inference levels.

  • A.1 IMPLEMENTATION DETAILS: SubZero+ draws Gaussian matrices ΩU and ΩV for each layer’s subspace construction.The matrices have dimensions mi×r and ni×r, respectively.
  • A.1 IMPLEMENTATION DETAILS: Thin QR decomposition is computed with PyTorch’s torch.linalg.qr during subspace construction.The implementation uses standard thin QR decomposition before applying sign correction.
  • A.1 IMPLEMENTATION DETAILS: Sign correction multiplies each column of Ui and Vi by the sign of the corresponding diagonal entry of RU and RV.This follows the procedure described in §4.3 for the QR-based construction.
  • A.1 IMPLEMENTATION DETAILS: SubZero+ inherits SubZero’s random seed trick and per-layer parameter update strategy for memory-efficient perturbation regeneration.These implementation choices are attributed to Yu et al. (2025).
  • A.1 IMPLEMENTATION DETAILS: Projection matrices Ui and Vi are regenerated every F steps and cached between regenerations.This strategy reduces repeated subspace construction while preserving the stated per-layer update approach.
  • A.1 IMPLEMENTATION DETAILS: Adam moment buffers are stored in Rr×r, contributing only 2r2 floats per layer.The storage requirement is expressed in terms of the subspace rank r.
  • A.1 IMPLEMENTATION DETAILS: Multi-query perturbations are generated sequentially within each iteration.Each perturbation is regenerated from its random seed before the forward pass and immediately subtracted after loss computation.
  • A.1 IMPLEMENTATION DETAILS: Sequential perturbation handling keeps peak memory at inference levels.The immediate subtraction after loss computation prevents perturbations from accumulating in memory.

A.2 ADDITIONAL EXPERIMENTAL DETAILS

The experiments use specified single-GPU software configurations, tune SubZero+ beyond prior ZO baselines, and equalize forward-pass budgets across methods. Full-parameter tuning uses model-dependent subspace ranks, while LoRA fixes rank and scaling.

  • Hyperparameter Configurations: Experiments run on a single NVIDIA H200 GPU, except OPT-1.3B on a single NVIDIA RTX 4090, using PyTorch 2.1.0 and CUDA 12.1.The H200 has 141GB memory and the RTX 4090 has 24GB.
  • Hyperparameter Configurations: SubZero+ tunes subspace rank r, update frequency F, query count K, and Adam parameters in addition to prior-work search grids used for MeZO and SubZero.The prior grids follow Malladi et al. (2023) and Yu et al. (2025).
  • Hyperparameter Configurations: For full-parameter tuning, SubZero+ uses r = 16 for OPT-1.3B and r = 128 for OPT-13B and LLaMA3.1-8B, matching SubZero’s model-dependent scheme.The rank selection is model-dependent and follows the same scheme as SubZero.
  • Forward-Pass Budget: MeZO and SubZero use 2 forward passes per step for a 40K total budget, while SubZero+ uses K + 1 passes and adjusts its step count to match 40K.With default K = 99, SubZero+ runs 400 steps at 100 forward passes per step.
  • Hyperparameter Configurations: LoRA uses rank 8 and α = 16 for all methods.These LoRA settings are fixed across the compared methods.

A.3 EMPIRICAL VALIDATION OF HAAR-CORRECTED QR SUBSPACES

Empirical validation shows that uncorrected torch.linalg.qr produces non-Haar-distributed subspaces, whereas sign correction yields Q distributions closely matching the Haar reference. This result is consistent across random seeds, hardware backends, and QR variants.

  • Distributional validation: Raw torch.linalg.qr Q deviates systematically from the Haar distribution, while sign-corrected Q closely matches normalized-Gaussian reference vectors.The experiment uses m = 50, r = 10 Gaussian matrices and averages results over 2000 trials.
  • Statistical validation: The raw Q is strongly rejected by the KS test (p ≪0.001), while the sign-corrected Q passes (p > 0.05).Table 9 defines D as the KS statistic, with lower values preferred, and p > 0.05 indicating failure to reject identical distributions.
  • Robustness: The same distributional result holds across different random seeds, CPU and GPU backends, and both standard and column-pivoted QR.The sign-corrected Q remains virtually indistinguishable from the Haar reference, while raw Q is concentrated closer to zero with lighter tails.
Loading 2608.15665v1…