PaperlayerView exact revision on arXiv

Source-linked AI summary

Sample-Adaptive Zeroth-Order Post-Training

Haiyuan Wan, Yanyi Li, Weicheng Lin, Cong Fang

arXiv:2610.12566v1cs.LG

TL;DR

ZO fine-tuning offers memory savings, but uniform example sampling overlooks evolving sample importance and leaves avoidable estimation variance. SAZO estimates importance from existing directional differences and tracks it across sparse minibatch feedback. The paper reports improved average downstream performance and up to 6× faster query-count convergence to the same training loss, while its convergence analysis matches the oracle's leading term.

  • Problem

    Uniform ZO sampling overlooks differences and changes in example importance, leaving avoidable estimation variance.

    Source: §1 · p. 2

  • Method

    SAZO estimates sample importance from existing directional differences and tracks sparse minibatch feedback with per-example exponential moving averages.

    Source: §1 · p. 2

  • Results

    Experiments report improved average downstream performance and up to 6× faster convergence to the same training loss in query count than MeZO.

    Source: §5 · p. 9

  • Takeaways & Limitations

    SAZO asymptotically matches the minimum-variance dynamic oracle's leading convergence term and has a smaller convergence coefficient than uniform sampling when sample importance is heterogeneous.

    Source: §5 · p. 9

    The convergence comparison fixes a common parameter trajectory and does not rank the separate trajectories generated by different samplers.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

Zeroth-order (ZO) optimization offers a memory-efficient alternative to backpropagation for fine-tuning large language models. However, existing ZO methods typically sample training examples uniformly, overlooking their evolving importance and leaving avoidable estimation variance. To address this limitation, we propose Sample-Adaptive Zeroth-Order Optimization (SAZO), which learns sample importance online and adaptively allocates queries to accelerate convergence. SAZO extracts gradient-norm information from existing directional differences to estimate sample importance without access to true gradients. Furthermore, to track evolving importance from sparse minibatch feedback, SAZO maintains per-example exponential moving averages that aggregate observations across updates. The method requires no additional model evaluations and only $O(N)$ additional scalar state compared with MeZO. We prove that SAZO asymptotically matches the leading convergence term of the minimum-variance sampling oracle, and achieves a strictly faster convergence rate than standard methods. Extensive experiments across four model families demonstrate improved average downstream performance and up to $6\times$ faster convergence to the same training loss in terms of query count.

1 Introduction

Zeroth-order fine-tuning reduces activation-memory demands but retains noisy gradient estimates, and uniform example sampling overlooks importance that varies and evolves during training. SAZO estimates sample importance from existing directional differences, tracks it across sparse minibatch feedback, and uses it to allocate queries more effectively.

1 Introduction

  • ZO queries reveal scalar directional projections, making unrestricted d-dimensional adaptation require Ω(d) observations and potentially offsetting its acceleration.The paper therefore asks whether sample-level information in ZO queries can instead guide adaptive allocation.Source: §1 · pp. 1–2
  • Per-example gradient norms differ and evolve during training, so uniform sampling overlooks sample heterogeneity and leaves avoidable estimation variance.Figure 1 shows norm changes across checkpoints and lower sampling regret for SAZO than uniform or fixed initial-oracle sampling.Source: §1 · p. 2, Figure 1 · p. 2
  • Inverse-probability weighting preserves the expected full-data finite-difference update while learned scores guide sampling.SAZO incurs lower relative cumulative sampling regret than uniform sampling and an oracle allocation fixed at initialization.Source: §1 · p. 2
  • SAZO estimates sample importance from existing directional differences and tracks sparse minibatch observations with per-example exponential moving averages.It requires no additional model evaluations and only O(N) scalar state.Source: §1 · p. 2
  • SAZO asymptotically matches the minimum-variance sampling oracle's leading convergence term and achieves a strictly faster convergence rate than standard methods.Source: §1 · p. 3
  • Experiments report improved average downstream performance and up to 6× faster convergence to the same training loss in query count than MeZO.The method retains MeZO-like memory and per-update cost.Source: §1 · p. 3

2 Related work

Related work covers adaptive sampling in first-order optimization and several ways to accelerate zeroth-order optimization. SAZO instead adapts query allocation across examples in a setting where gradients are unavailable and directional feedback is noisy.

Sample-adaptive first-order optimization

  • First-order sample-adaptive methods learn importance distributions from sampled feedback, with work spanning regret guarantees, Adam connections, federated convergence, joint convex optimization, and gradient-continuity guarantees.Source: §2 · p. 3

Accelerating ZO post-training

  • Prior ZO methods accelerate optimization through curvature-based preconditioning, low-rank perturbations, variance reduction, directional alignment, update normalization, optimizer adaptivity, and adaptive layer selection.Source: §2 · p. 3
  • Where prior ZO approaches adapt parameter updates or perturbations, SAZO adapts query allocation across examples despite unavailable gradients and noisy directional feedback.The passage notes that reliable parameter-wise adaptation is intrinsically difficult in high dimensions.Source: §2 · p. 3

3 Sample-adaptive zeroth-order optimization

SAZO replaces equal example sampling with learned importance allocation in zeroth-order updates, estimating evolving sample scores from existing directional feedback. Its analysis shows asymptotic minimum-variance sampling with a lower-order adaptation correction under stated conditions.

3 Sample-adaptive zeroth-order optimization

  • The base ZO update samples an example uniformly and uses an independent Gaussian direction for a finite-difference estimator whose expectation is the gradient plus O(µ^2).The finite-difference error is subject to the stated Hessian regularity conditions.Source: §3 · p. 3

3.1 Sample-adaptive updates

  • Importance weighting assigns examples different query probabilities while preserving the expected full-data ZO update up to finite-difference error.The probabilities are learned from the changing importance estimates.Source: §3.1 · pp. 3–4

3.2 Estimating sample importance from directional differences

  • SAZO estimates per-example gradient norms from existing directional differences, tracks sparse observations with per-example EMAs, and refreshes stale scores through exploration.It reuses existing queries and requires O(N) additional scalar state; larger β responds faster, while smaller β suppresses directional noise.Source: §3.2 · p. 4

3.3 Convergence analysis

  • Under the stated assumptions and schedules, SAZO’s leading optimization term matches the stepwise minimum-variance oracle, with adaptation correction o(T^-1/2).Uniform sampling’s horizon-averaged second moment exceeds the oracle when example gradient norms differ; this comparison uses a common parameter trajectory.Source: §3.3 · pp. 5–6

4 Experiments

Experiments cover masked and generative language models, convergence across scales, training-population scaling, computational efficiency, and sampling-hyperparameter robustness. SAZO improves downstream accuracy and convergence while retaining zeroth-order methods’ memory and per-update-cost advantages.

4 Experiments

  • The evaluation spans downstream performance on masked and generative models, convergence, population scaling, computational efficiency, and sampling-hyperparameter robustness.Detailed settings and additional diagnostics are provided in Appendix E.Source: §4 · p. 6

4.1.1 Masked language models

  • SAZO raised average accuracy over the stronger of MeZO and HiZOO by 1.6 points (full-parameter) and 1.1 points (PEFT).These comparisons cover RoBERTa-large tasks; Table 2 reports mean test accuracy and identifies the best zeroth-order result.Source: §4.1.1 · p. 6, Table 2 · p. 7

4.1.2 Generative language models

  • SAZO improved average performance over the stronger of MeZO and HiZOO by 1.2, 1.7, and 3.9 percentage points on OPT-1.3B, Phi-2, and Llama3-8B, respectively.Table 3 compares full-parameter generative-LM results.Source: §4.1.2 · p. 7, Table 3 · p. 7

4.1.3 Convergence across model scales

  • 4K, 6K, and 6K forward passes let SAZO reach MeZO’s endpoint losses on RoBERTa-large, Phi-2, and Llama 3, versus HiZOO’s 10.5K, 9K, and 6K.Figure 2 compares training-loss trajectories at a common 24K forward-pass budget.Source: §4.1.3 · p. 7, Figure 2 · p. 8

4.2 Scaling with training-population size

  • SAZO outperformed MeZO on SST-2 test accuracy at every population size, with gains of 0.6, 1.1, 1.4, and 1.5 percentage points from N through 4N.Figure 3 plots test accuracy as the RoBERTa-large training population grows from N = 1,024 to 4N.Source: §4.2 · p. 7, Figure 3 · p. 8

4.3 Memory and wall-clock efficiency

  • SAZO matched MeZO’s measured memory across OPT configurations and kept per-update time comparable across evaluated models.Compared with first-order methods, memory use was lower by 6.8–12.2× for full FT and 3.3–6.1× for prefix FT.Source: §4.3 · p. 8, Figure 4 · p. 8, Table 4 · p. 9

4.4 Ablation studies

  • Across β values 0.005, 0.05, and 0.5, loss trajectories stayed stable with closely clustered endpoints; β = 1 diverged.At β = 1, the update discards accumulated history and uses only the latest directional observation; SAZO was also robust to ε over a broad range.Source: §4.4 · p. 9, Figure 5 · p. 9

5 Conclusion

SAZO estimates sample importance from existing directional differences, tracks its evolution with per-example moving averages, and requires only O(N) scalar state. The analysis and experiments report improved convergence and downstream performance.

5 Conclusion

  • SAZO asymptotically matches the dynamic oracle’s leading convergence term and achieves a smaller convergence coefficient than uniform sampling when sample importance is heterogeneous.Experiments show improved average downstream performance and up to 6× faster convergence to the same training loss by query count.Source: §5 · p. 9

AI use statement

The authors used generative AI for research organization, literature retrieval, mathematical derivation, drafting, LaTeX preparation, numerical identity checks, and a separate adversarial review. It did not generate the reported training measurements.

AI use statement

  • Generative AI assisted with research tasks and manuscript preparation, but did not generate the reported training measurements.The listed assistance included mathematical derivation, numerical identity checks, and a separate adversarial review.Source: AI use statement · p. 10

A Proofs of the main results

The appendices justify SAZO’s sampling and tracking guarantees, then qualify how those guarantees depend on trajectory quantities, finite-difference radius, and comparator choice. They also establish bounds for sparse EMA feedback and distinguish dynamic from fixed sampling costs.

The corresponding update satisfies

A.2 Gaussian identities and the sampling oracle

  • For a fixed parameter vector and selected index, Gaussian directional observations yield the leading sampling moment, with equality in the oracle allocation at pi ∝ ai on positive-norm support.Zero-norm terms contribute zero, and boundary infima are understood by limits.Source: §A.2 · p. 13

A.3 Stationary variance of the EMA

  • At a fixed target, the sparse EMA’s visit-time variance converges to βvµ_i/(2 − β), reducing observation variance while retaining finite-difference bias.This fixed-target result does not remove drift on an evolving optimization trajectory.Source: §A.3 · p. 13

A.4 A global normalization bound

  • The global normalization bound covers S = 0 and does not require closeness between M and S.Source: §A.4 · p. 13

A.5 Sparse EMA with correlated target drift

  • The sparse-EMA tracking analysis permits correlated target drift and does not assume independence between observation noise and target drift.The proof retains a drift constant of 12 for consistency with the finite-difference extension, although its argument gives 6.Source: §A.5 · p. 14

A.6 Closing the drift and proving stationarity

  • The descent analysis yields Theorem 1 with an additive finite-difference remainder, while its leading term avoids the coarse moment bound’s factor of two.The proof controls the additional terms by Ωµ,T and establishes Lemma 1 without multiplying its leading term by that factor.Source: §A.6 · p. 14
  • Theorem 2’s bound follows after substituting the horizon-dependent schedules, subject to κ ≤ 1/2 and a conditional parameter balance.The balance depends on quantities from the resulting trajectory, so it is not an online tuning rule or an existence guarantee for self-consistent parameters.Source: §A.6 · p. 15, B Parameter balance and comparator… · p. 15

B.1 Three fixed and dynamic costs

  • For positive moments, the sampling-cost hierarchy is eDT ≤ eHT ≤ eGT ≤ eUT; differing example norms can make eDT strictly smaller than eUT.The two-step example has eDT = 1/4 while eHT, eGT, and eUT are each 1/2; the hierarchy does not imply that hindsight sampling is implementable.Source: §B.1 · p. 15

C Finite-difference error

The finite-difference analysis bounds the error in SAZO’s convergence guarantee under stated smoothness and trajectory conditions. The minibatch extension reuses evaluations for repeated draws while preserving their multiplicities.

C Finite-difference error

  • The explicit finite-radius convergence bound retains a radius-dependent remainder along the actual finite-difference trajectory.The theorem uses with-replacement updates and once-per-included-example EMA updates; its bound includes regret and finite-difference terms.Source: C Finite-difference error · p. 16
  • With fixed problem and schedule coefficients, bounded trajectory moments, and horizon-dependent radii µt = µ0T^-α, α > 3/10 makes the stated radius errors o(T^-3/5).This conclusion concerns a horizon-dependent radius, not a schedule µt = µ0t^-α with fixed initial radius; evaluation randomness and numerical rounding are not represented by the Taylor bounds.Source: C Finite-difference error · p. 17

D.1 Minibatch update and reused evaluations

  • With-replacement minibatches reuse one two-sided difference per distinct included example, preserving repeated-draw multiplicities without additional evaluations for duplicates.The distribution stays fixed within a batch, and each included example updates its score once.Source: §D.1 · p. 17

D.2 Effect of the shared direction on variance

Under a shared Gaussian direction, norm-proportional allocation reduces the leading sampling-variance term, and minibatching improves the tracking contribution to convergence. The remaining sections specify experiment settings for RoBERTa-large, GLMs, and training-population scaling.

D.2 Effect of the shared direction on variance

  • Sharing a direction and reusing duplicate evaluations preserve the expected full-data finite-difference update, even when example differences are dependent.The shared-direction minibatch moments are derived under Assumption 1.Source: §D.2 · p. 18
  • Unequal gradient norms yield a strictly positive leading variance reduction; for learned allocation, the leading gain is the oracle gain minus a tracking-error term.The shared direction term cancels in the comparison with MeZO, so sharing a direction adds no variance penalty relative to MeZO.Source: §D.2 · p. 18
  • For batch size B, both methods’ sampling variance falls by 1/B, while norm-proportional allocation minimizes this term at every batch size.SAZO uses the same two perturbed model states as MeZO, and duplicate evaluation reuse reduces the work for its B draws.Source: §D.2 · p. 18

D.3 Consequences for tracking and convergence

  • When Cdη(1 − 1/B) < 2, the descent coefficient is positive and shared with uniform sampling; SAZO instead contributes a dynamic-oracle regret term.The bound contrasts uniform sampling’s eUT/B term with SAZO’s (eDT + RegretT/T)/B term.Source: §D.3 · p. 19
  • For small Bε/N, the regret-dependent update-moment coefficient satisfies κB/κ = B^-3[1 + O(Bε/N)], combining a 1/B sampling reduction with an approximately B^2 tracking improvement.Here κ is the B = 1 coefficient at the same η, ε, and β.Source: §D.3 · p. 19

E.1 Detailed experiments on RoBERTa-large

  • RoBERTa-large SAZO experiments use a constant learning rate, β = 0.05, ε = 0.1, and four burn-in passes.For Figure 2(a), each method’s learning rate is selected from {1, 5} × {10^-7, 10^-6, 10^-5}.Source: §E.1 · p. 19

E.2 Detailed results on various LLMs

  • GLM experiments use the hyperparameter grids in Table 6 and select constant learning rates for each model in Figures 2(b,c).Other settings match the corresponding Table 3 experiments; Table 6 specifies zero weight decay and SAZO’s 20K steps.Source: §E.2 · p. 19, Table 6 · p. 20

E.3 Training-population scaling settings

  • The training-population scaling experiment compares MeZO and SAZO on SST-2 at N = 1,024 through 4N = 4,096, using batch size 16 and 100K steps.Test accuracy is averaged over seeds 13, 21, and 42; the test set is not used for selection.Source: §E.3 · p. 20

E.4 Ablation experiment details

The ablations test SAZO’s sensitivity to β and ε on SST-2 under a fixed RoBERTa-large setup. Optimization remains stable across the tested ε values below 1, while ε = 1 becomes uniform sampling and diverges.

E.4 Ablation experiment details

  • The ablations use one seed of full-parameter RoBERTa-large on SST-2 with k = 512, a fixed 1,024-example training population, batch size 64, constant learning rate 3×10−6, perturbation radius 10−3, and zero weight decay.Source: §E.4 · p. 20
  • SAZO remains stable for ε ∈ {0.05, 0.1, 0.5}, but ε = 1 reduces to uniform sampling and becomes unstable.The prose reports that the ε = 1 trajectory eventually diverges under the same learning rate.Source: §E.4 · p. 20, Figure 6 · p. 21

E.5 Sample-importance diagnostics

Exact-gradient audits show that sample importance is heterogeneous and changes across checkpoints, while additional analyses compare group-restricted and adaptive sampling allocations.

E.5 Sample-importance diagnostics

  • Figure 1(a) tracks exact per-example gradient norms at initialization, 10K, and 30K updates, with each example retaining its initial horizontal position.Norms are divided by the median initial norm; every example is shown without smoothing or aggregation.Source: §E.5 · p. 21
  • Across six tasks, the top 10% of examples account for an increasing share of total gradient norm, above the approximately 10% uniform reference.The top set is recomputed at each checkpoint; audits use frozen RoBERTa-large MeZO checkpoints.Source: Figure 7 · p. 21
  • Length and label controls restrict allocation probabilities to be equal within label groups or joint prompt-length-and-label groups.The group-restricted allocation is evaluated by the fraction of the uniform-to-oracle variance gap remaining at 20K updates.Source: §E.5 · pp. 21–22
  • The SST-2 comparison evaluates uniform, fixed initial oracle-proportional, and SAZO sampling over 30K updates using scanned learning rates.SAZO uses β in {0.025, 0.05, 0.075} and ε = 0.1; the fixed sampler uses initial exact gradient-norm proportions without an exploration mixture.Source: §E.5 · p. 22
  • Figure 1(b) reports relative cumulative regret, computed as a ratio of separately trapezoid-integrated numerator and denominator across 61 audits.The audits span 0–30K updates at 500-update intervals; each method’s current parameters define its exact-gradient reference oracle, and loss uses the complete training population.Source: §E.5 · p. 22

E.6 Interpretable sample-level structure

Sample-importance rankings shift substantially during training, and selected SST-2 examples show these shifts alongside changing predictions. Together, the population and example-level patterns indicate that initial importance can become stale, motivating continued tracking during optimization.

E.6 Interpretable sample-level structure

  • Initial-to-20K Spearman correlations are 0.52 on SST-2 and 0.12 on TREC, showing that initial rankings do not reliably preserve later allocation priorities.Figure 8 compares each example’s normalized importance at initialization and 20K updates; points above or below the diagonal gain or lose relative importance.Source: §E.6 · p. 22, Figure 8 · p. 22
  • Among four SST-2 cases examined at initialization, 5K, and 20K updates, B (“over-dramatic at times”) becomes correct as normalized importance falls from 3.113 to 0.230.Source: §E.6 · p. 22
  • The “smug grin” example remains incorrect while its importance rises from 2.863 to 11.170.The Figure 9 cases connect sample-level predictions with normalized gradient norms.Source: §E.6 · p. 22, Figure 9 · p. 23
  • In the representative example ending “a minor treat,” correctness is lost as importance rises from 0.805 to 10.930; for D, the absolute gradient norm also increases.Source: §E.6 · p. 23
  • With words, length, and labels fixed, predictions and gradient contributions still change, so allocation based only on initial importance can become stale.The cases include an example that stays correct with low importance throughout training.Source: §E.6 · pp. 22–23

Related papers

Loading 2610.12566v1…