PaperlayerView exact revision on arXiv

Source-linked AI summary

StructZO: Structured Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning

Xuanlin Yang, Jia Li, Zhouchen Lin

arXiv:2610.16451v1cs.LGcs.AI

TL;DR

Memory-efficient zeroth-order fine-tuning still trails first-order optimizers, motivating a forward-only method that recovers useful optimization structure. StructZO estimates updates in activation-induced subspaces and spectrally shapes them; across the reported benchmarks it outperforms ZO baselines and remains close to FO-Adam.

  • Problem

    Zeroth-order fine-tuning reduces memory use through forward-only estimates, but remains slower and less effective than first-order methods in high-dimensional, ill-conditioned parameter spaces.

    Source: §1 · p. 1

  • Method

    StructZO aggregates multiple ZO directions in an activation-induced informative subspace and applies spectral shaping to the reduced-space update.

    Source: §1 · p. 2, §3.3.1 · p. 3

  • Results

    Across OPT-1.3B and Llama3.1-8B tasks, StructZO achieves up to 8.7% average improvement over the strongest ZO baseline and about a 1.5% average gap to FO-Adam.

    Source: §4.1 · p. 6

  • Takeaways & Limitations

    The analysis reduces estimation dimension from mn to mr, with remaining projection bias controlled by the activation Gram matrix’s spectral tail.

    Source: §1 · p. 2

    StructZO’s benefit may decrease when the forward-induced subspace captures optimization-relevant directions weakly; experiments focus on limited-data GLUE/SuperGLUE-style fine-tuning.

    Source: §5 · p. 8

Abstract

from arXiv · show

Fine-tuning large language models at scale is increasingly limited by the memory cost of backpropagation. Zeroth-Order (ZO) optimizers reduce this burden by relying only on forward function evaluations, but they still exhibit a substantial performance gap to first-order (FO) optimizers such as Adam. We argue that this gap does not arise solely from the absence of explicit gradients: FO fine-tuning also benefits from informative low-rank optimization geometry, large-scale signal aggregation, and well-conditioned update directions. Motivated by this, we introduce StructZO, a structured Zeroth-Order optimization framework for memory-efficient LLM fine-tuning. StructZO performs large-query ZO estimation within an activation-induced informative subspace and applies spectral control to the resulting reduced-space update before mapping it back to the original parameters. This lets ZO aggregate useful directional signals in a compact, geometry-aware space while avoiding the activation and optimizer-state memory costs of backpropagation. We show theoretically that, when the induced subspace captures concentrated optimization-relevant directions, StructZO reduces the effective ZO estimation dimension from $mn$ to $mr$, with the remaining projection error governed by the spectral tail of the activation Gram matrix. Empirically, across OPT-series and Llama3.1-8B on diverse downstreamtasks, StructZO improves over the strongest existing ZO baselines by up to $8.7%$ relative on average and up to $26.9%$ relative on individual tasks, while remaining within $1.5%$ of FO-Adam on average. These results suggest that structured and spectrally controlled ZO updates can substantially narrow the gap between ZO and FO optimization.

1 Introduction

The introduction frames a persistent performance gap between memory-efficient zeroth-order fine-tuning and first-order methods, then asks whether forward-only optimization can recover structural benefits of first-order updates. It presents StructZO as a coupled framework for structured signal estimation and update conditioning, with theoretical and empirical support.

Motivation

  • The paper identifies noisy estimates in high-dimensional, ill-conditioned parameter spaces as a source of ZO’s slow convergence and performance gap to first-order methods.Source: §1 · pp. 1–2

Proposed framework

  • StructZO combines large-query directional signal aggregation in activation-induced informative subspaces with Muon-style spectral shaping in reduced coordinates.The framework maps the shaped update back to the original parameters without storing activations for backpropagation.Source: §1 · p. 2
  • The framework couples where ZO queries are spent, how their signals are aggregated, and where spectral shaping is applied before lifting the update.The paper characterizes the resulting estimator through a variance–bias trade-off.Source: §1 · p. 2

Theoretical contribution

  • The analysis reduces the ZO estimation dimension from mn to mr, with remaining projection bias controlled by the activation Gram matrix’s spectral tail.Source: §1 · p. 2

Experiments

  • On Llama3.1-8B, StructZO raises average score from 76.3 to 82.9 over the strongest ZO baseline and improves CB from 73.2 to 92.9.The introduction also reports that StructZO remains close to FO-Adam while retaining forward-only optimization’s memory advantage.Source: §1 · p. 2

2 Related Work

Related work addresses memory-efficient fine-tuning, zeroth-order optimization, and subspace-based updates. The paper positions StructZO as a study of how activation-induced spaces, multi-query aggregation, and spectral conditioning interact.

Fine-tuning LLMs

  • Parameter-efficient methods such as adapters, prefix tuning, and LoRA reduce trainable parameters, but still rely on backpropagation and retain activation-memory costs.Source: §2 · p. 2

Zeroth-Order Optimization

  • ZO methods such as MeZO enable LLM fine-tuning with memory close to inference, while later work explores structured perturbations, variance reduction, momentum, and Hessian-informed updates.Source: §2 · p. 2

Subspace Optimization

  • Prior subspace methods include random low-rank projections and activation-guided AGZO; the paper motivates studying their interaction with query aggregation and update conditioning.Source: §2 · p. 2

3 Method

The method is motivated by first-order updates’ use of informative geometry, aggregated derivative signals, and conditioning, which standard ZO methods typically lack. Its preliminaries explain why full-parameter multi-query estimation is costly and motivate an informative lower-dimensional search space.

3.1 Motivation

  • StructZO aims to recover first-order optimization’s informative geometry, signal aggregation, and well-conditioned transformations within a forward-only pipeline.Source: §3.1 · p. 3

3.2 Preliminaries

  • Random Gradient Estimation uses objective values from perturbations instead of explicit gradients, and its estimate is unbiased as the smoothing parameter approaches zero.Source: §3.2 · p. 3
  • Because LLM parameter dimension is extremely large, directly applying multi-query ZO estimation in full parameter space is computationally inefficient, motivating an informative lower-dimensional search space.Source: §3.2 · p. 3

3.3 StructZO

StructZO constructs a reduced search space from forward activations, estimates updates within that space, and applies spectral shaping before lifting updates back to the parameters. Its diagnostics compare the informative basis and resulting updates with random-space and raw-gradient alternatives.

3.3.1 Activation-Induced Informative Subspaces

  • When output-side covariance is relatively isotropic or slowly varying, the activation Gram matrix’s dominant eigenspace provides informative coordinates for zeroth-order perturbations.Source: §3.3.1 · p. 3
  • StructZO obtains a rank-r basis from forward activations using a Gaussian sketch and QR decomposition, without storing activations for backward propagation.The resulting projector captures dominant activation directions and defines the reduced optimization space.Source: §3.3.1 · p. 3

3.3.4 Empirical Geometry Diagnostics

  • The activation-induced basis aligns more strongly with the gradient’s right singular subspace than a random basis, while spectral shaping moves the final update toward the conditioned direction.Before shaping, the structured estimator is closer to the raw gradient than its random-space counterpart.Source: Figure 1 · p. 4, §3.3.4 · p. 5

3.3.2 Multi-Query Structured Zeroth-Order Estimation

  • Restricting perturbations to activation-derived coordinates reduces their dimension from mn to mr and improves query efficiency relative to random low-dimensional perturbations.The reduced coordinates are constructed from current forward activations and are intended to align better with optimization-relevant directions.Source: §3.3.2 · pp. 3–4

3.3.3 Spectral Shaping via Reduced-Space Muon

  • Muon-style spectral shaping equalizes active singular directions in the informative subspace, so the final update approximates a conditioned direction rather than merely the raw gradient.Source: §3.3.3 · pp. 4–5

3.4 MSE Analysis of StructZO

StructZO constructs an activation-based reduced space for each mini-batch and analyzes a multi-query estimator in those coordinates. Its error reflects both query variance, reduced from full-space dimension mn to mr, and projection bias governed by the activation spectrum.

3.4 MSE Analysis of StructZO

  • StructZO builds P from each mini-batch’s activation Gram matrix and uses the resulting projector as its informative reduced space.The subspace is statistically coupled with the same mini-batch randomness used in the estimator.Source: §3.4 · pp. 5–6
  • The MSE bound separates query-induced variance from subspace-approximation bias, under stated assumptions and standard smoothness conditions.The analysis concerns the pre-Muon multi-query estimator.Source: §3.4 · p. 6
  • The reduced-space query variance has dimension factor mr instead of mn, yielding a variance reduction controlled by r/n when r ≪ n.The remaining projection bias is controlled by the spectral tail of the activation Gram matrix.Source: §3.4 · p. 6
  • When the activation spectrum has low effective rank, the informative subspace can retain most useful gradient directions while avoiding full-dimensional ZO perturbations.The estimator also includes an O(µ^2) smoothing error.Source: §3.4 · p. 6

4 Experiments

Across OPT-1.3B and Llama3.1-8B tasks, StructZO outperforms the ZO baselines and narrows the average gap to FO-Adam. Ablations and efficiency results examine its components, memory cost, and training-time trade-offs.

4 Experiments

  • The evaluation covers GLUE and SuperGLUE tasks on OPT-1.3B and Llama3.1-8B, with scalability tests on OPT-13B and OPT-30B.StructZO uses q = 100 queries for 400 steps, compared with 20K steps for single-query ZO baselines; methods share preprocessing and evaluation protocols.Source: §4 · p. 6

4.1 Main Results

  • StructZO is best among ZO methods across classification, multiple-choice, and generation tasks, with up to 8.7% average improvement and about 1.5% average gap to FO-Adam.The CB result on Llama3.1-8B improves over the best ZO baseline by 26.9%; Table 2 extends comparisons to OPT-13B and OPT-30B.Source: §4.1 · p. 6, Table 1 · p. 7, Table 2 · p. 7

4.2 Ablation Study

  • The full StructZO outperforms both the random-subspace StructZO-R and SGD-style StructZO-S variants across the evaluated models and tasks.Increasing q also leads to faster and lower loss reduction when queries are collected in the informative reduced space.Source: §4.2 · p. 7, Table 3 · p. 8

4.3 Memory and Wall-Clock Efficiency

  • StructZO’s basis adds O(rn) storage per weight matrix, while resampling regenerates query directions; it achieves up to 11.2× memory reduction over FO-Adam.Despite higher per-step cost at q = 100, total training time is lower or comparable to MeZO and LOZO and substantially faster than HiZOO/HiZOO-L. Figure 2 compares training loss against forward passes on BoolQ and CB.Source: §4.3 · pp. 7–8, Figure 2 · p. 8

5 Conclusion

StructZO combines large-query estimation in an activation-induced subspace with spectral shaping, while avoiding backpropagation. The conclusion notes dependence on the subspace capturing useful directions and limits evaluation to GLUE/SuperGLUE-style limited-data fine-tuning.

5 Conclusion

  • StructZO aggregates directional signals in an informative subspace and applies spectral shaping in reduced coordinates to improve forward-only optimization.Its benefit may decrease if the forward-induced subspace does not capture optimization-relevant directions.Source: §5 · p. 8
  • Figure 3 plots training loss against forward passes for different query numbers on BoolQ with OPT1.3B and CB with Llama3.1-8B.The caption identifies the comparisons but does not report their outcomes.Source: Figure 3 · p. 9
  • Figure 4 compares memory overhead across models, ZO algorithms, and FO-Adam.The caption directs readers to Appendix B for details.Source: Figure 4 · p. 9

A Proof

The proof establishes an identity for Muon-style matrix-sign updates when the reduced basis matches the gradient’s right singular space. It frames activation-induced bases as approximations rather than exact recovery.

A Proof

  • When the reduced basis matches the gradient’s right singular space, the reduced-coordinate matrix-sign update lifted back to the full space equals the full-space Muon update.The proof uses a polar-factor representation and the gradient’s singular value decomposition.Source: A Proof · p. 12
  • In practice, the activation-induced basis only approximates the gradient’s right singular space, so the paper does not claim exact recovery of the full-space Muon update.The cited interpretation points to empirical alignment with gradient-relevant structure over a random basis.Source: A Proof · p. 12
  • Theorem 2’s assumptions include bounded stochastic gradient noise and a sketchability condition involving the activation Gram matrix’s eigenvalues.The supplied assumption fragments identify the eigenvalues in descending order.Source: A Proof · p. 12

A.1 A Coupled MSE Decomposition

This proof decomposes the coupled multi-query estimator’s error by analyzing its first-order part, conditional variance, and smoothing remainder. Under standard smoothness conditions, the squared-error contribution of the finite-difference remainder is O(µ^2).

A.1 A Coupled MSE Decomposition

  • Proposition 7 frames the result as a coupled MSE decomposition for independent Gaussian query directions.The directions are drawn in the reduced mr-dimensional space and are independent of the batch and subspace randomness.Source: §A.1 · p. 13
  • The finite-difference estimator is split into a first-order part and a smoothing remainder, whose squared-error contribution is O(µ^2) under standard smoothness conditions.The difference between bG and bG(1) is identified as the finite-difference smoothing error.Source: §A.1 · p. 13
  • The proof conditions on the batch and subspace randomness to analyze the first-order estimator and its conditional variance.It states that the cross term vanishes under this conditioning and that the two error components are orthogonal in Frobenius inner product.Source: §A.1 · pp. 13–14
  • The proof takes expectation over the batch and sketch randomness and adds the smoothing remainder to complete the decomposition.The final step uses the O(µ^2) smoothing bound.Source: §A.1 · pp. 13–14

A.2 Proof of Theorem 2

The proof relates activation-Gram sketchability to projection bias and contrasts activation-induced ZO estimation with full-space and random-subspace alternatives. Its dimension and bias advantages depend on the activation geometry and stated assumptions.

A.2 Proof of Theorem 2

  • For a linear layer, the gradient’s right-side energy is controlled by the activation Gram matrix, motivating the activation-gradient alignment assumption.The input activation and output-side error signal form the gradient expression used in this argument.Source: §A.2 · p. 15
  • The activation-Gram construction is a randomized range finder, and sketchability bounds how well its subspace captures the dominant eigenspace.When the activation Gram matrix has fast spectral decay, the residual projection error is small.Source: §A.2 · p. 15
  • StructZO reduces the ZO search dimension from mn to mr, with the remaining activation-subspace projection bias potentially smaller when the activation Gram matrix has low effective rank.The bias is controlled by the activation Gram matrix’s spectral tail.Source: §A.2 · p. 15
  • For an isotropic random rank-r projector, the proof compares projection bias with the activation-induced projector, whose bias can be much smaller under low effective rank.The argument frames the comparison in terms of retained gradient energy and activation-Gram structure.Source: §A.2 · p. 15

B Experimental details

Experiments use GLUE and SuperGLUE tasks, with memory, wall-clock, and ablation analyses. StructZO’s reduced query count offsets its higher per-step cost, while query number and subspace rank affect performance.

Datasets and setup

  • Evaluation covers classification, multiple-choice, and generation tasks from GLUE and SuperGLUE, using 1,024 training examples and 512 validation examples per task where available.CB and COPA reserve 100 examples for validation and use the remainder for training; experiments run on a single NVIDIA H20 96GiB GPU.Source: B Experimental details · pp. 16–17

Memory comparison

  • StructZO reduces peak GPU memory from 27.14 GiB to 4.29 GiB on OPT-1.3B and from 85.29 GiB to 20.80 GiB on Llama3.1-8B versus FO-Adam.It also reduces memory from 316 GiB to 28.33 GiB on OPT-13B and from 633 GiB to 63.67 GiB on OPT-30B; its overhead over MeZO and LOZO is small.Source: B Experimental details · p. 17, Table 6 · p. 17

Wall-clock time

  • On Llama3.1-8B, StructZO reduces estimated total time versus MeZO from 3.656h to 2.954h on SST-2, 9.79h to 9.17h on RTE, and 14.99h to 14.76h on BoolQ.StructZO uses 50× fewer update steps, trading higher per-step cost for comparable or lower total time than MeZO and LOZO in most settings.Source: B Experimental details · p. 17, Table 7 · p. 18

Time breakdown

  • ZO forward evaluations account for 83.91%–93.32% of StructZO’s per-step time, while basis construction contributes 4.95%–6.37%.The final update takes less than 2.3% on Llama3.1-8B across the three tasks.Source: B Experimental details · pp. 17–18, Table 8 · p. 19

Effect of query number

  • Increasing queries from 4 to 100 improves accuracy from 63.5 to 68.5 on OPT-1.3B BoolQ and from 64.3 to 92.9 on Llama3.1-8B CB.Larger query budgets beyond a point do not necessarily improve performance under a fixed computational budget; q = 100 is the chosen balance.Source: B Experimental details · p. 18, Table 10 · p. 20

Effect of subspace rank

  • Increasing rank from 16 to 64 improves OPT-1.3B RTE accuracy from 63.2 to 69.0, while further increases do not necessarily improve results.Performance on Llama3.1-8B SST-2 remains relatively stable for r = 16, 32, 64, and 128, with a mild drop at r = 256; r = 64 is the default trade-off.Source: B Experimental details · p. 19, Table 9 · p. 19

Related papers

Loading 2610.16451v1…