Probe-Space Preconditioning for Fast and Stable Zero-Order Training

Abstract: Backpropagation (BP) dominates deep learning but imposes a massive memory tax. For example, training OPT-30B with Adam requires $\approx$ 600GB of GPU memory (assuming batch size 8 and sequence length 2048). Alternatively, zero-order optimization (ZOO) trains in inference-mode (requiring only $\approx$ 60GB for the same model): no stored activations, no gradients, and no optimizer states. However, ZOO convergence has lagged behind BP. In this work, we evaluate two methods to close this gap. First, we show that reallocating training compute budget from many steps to large effective batch sizes with many perturbations (or probes) but fewer steps, allows 1SPSA (Spall, 1992) to outperform zero order methods like MeZO (Malladi et al., 2023) with less training compute. Next, we introduce 1.5-SPSA, adding a single "clean" forward-pass per step to 1SPSA to calculate a cheap diagonal preconditioner in probe-space, which improves convergence rate and convergence by down-weighting high curvature directions. Benchmarking on 6 post-training datasets on both Qwen3 and OPT model families, we show that 1.5-SPSA achieves State-of-the-Art results over previous ZOO solvers with much less optimization steps. For example, we train OPT-13B (for direct comparison to MeZO) and find 1.5-SPSA achieves +3.1% accuracy on SST-2 over both MeZO and BP in only 70 steps vs. MeZO's 100,000 steps. Finally, we combine an 8-bit-packing random generator, triton fused unpack/apply kernels, and distributed parallelism to achieve fast and stable training of models as large as OPT-30B in-place on commodity GPUs (e.g. A100).

• The paper proposes 1.5-SPSA, a fast, stable zero-order training method that combines batch size enhancement, increased perturbations, and curvature reweighting, achieving over 94.5% accuracy on SST-2 in 70 OPT-13B optimization steps and significantly outperforming standard ZOO baseline on various post-training tasks.
• You utilize 1.5-SPSA by a smaller number of high-cost optimization steps with increased batch sizes and more perturbations per step, producing a more accurate update, 1.5-SPSA allows batch sizes of 128–256 with more perturbation probes to improve signal-to-noise ratio.
• Using curvature-determined weight rounding stabilizes training process for large language models, while avoiding numerical instabilities in ill-conditioned geometry.

Problem setting and contribution

“Probe-Space Preconditioning for Fast and Stable Zero-Order Training” (2609.38095) addresses the memory–compute trade-off in derivative-free optimization for large neural networks. Backpropagation with Adam requires stored activations, gradients, and optimizer states; for OPT-30B, the paper estimates approximately 600 GB of GPU memory under a batch size of 8 and sequence length 2048. Inference-mode zero-order optimization (ZOO) eliminates these requirements, reducing the corresponding memory footprint to approximately 60 GB, but traditionally requires many more forward passes and exhibits substantial estimator noise.

The paper proposes two complementary changes to Simultaneous Perturbation Stochastic Approximation (SPSA). First, it reallocates a fixed forward-pass budget from many optimization steps to larger effective batches and more perturbation directions per step. Second, it introduces 1.5-SPSA, which adds one unperturbed forward pass to estimate directional curvature and uses that estimate to reweight perturbation directions. The resulting method retains inference-mode memory use while improving stability and convergence in highly ill-conditioned post-training objectives.

The central empirical claim is strong but deliberately scoped: 1.5-SPSA outperforms prior ZOO baselines and, under the paper’s selected compute allocations, can outperform the reported BP+Adam baseline on several post-training tasks. The authors explicitly note that BP is not exhaustively retuned for the extreme large-batch, few-step regime used by the proposed methods, so the results do not establish universal superiority over backpropagation.

Zero-order optimization and compute allocation

Standard 1SPSA estimates a gradient from central differences along random Rademacher probes. For perturbations $z_i$ and radius $\epsilon$, the method evaluates the loss at $\theta+\epsilon z_i$ and $\theta-\epsilon z_i$, then averages the resulting directional estimates. The number of forward passes per optimization step is independent of the parameter dimension, making the method applicable to models with billions of parameters. However, the estimator combines minibatch noise, perturbation noise, finite-difference bias, and curvature-induced instability.

The paper’s first substantive result is that a fixed forward-pass budget should not necessarily be spent on a large number of low-cost optimization steps. Instead, increasing the effective batch size and the number of perturbations per step can produce a more accurate update and permit a much larger step size. The authors define the 1SPSA budget as proportional to the number of steps, accumulation steps, and perturbations:

$$F_{\mathrm{1SPSA}} = s \times a \times 2 \times n_{\mathrm{pert}}.$$

On OPT-13B fine-tuned on SST-2, 1SPSA with effective batch size 128 and 160 perturbations reaches 94.2% accuracy in 80 steps using approximately 205,000 forward passes. This exceeds the reported MeZO result of 91.4% using a comparable forward-pass budget. The implication is that ZOO performance depends critically on compute allocation, not merely on the nominal number of function evaluations.

The associated optimization regime uses a learning rate of $5\times 10^{-4}$, compared with $10^{-6}$ for the reported MeZO and BP runs. Thus, the proposed 1SPSA configuration uses a step size approximately 500 times larger. The paper attributes this to reduced estimator noise from larger batches and more perturbations. The result also exposes a practical constraint: large effective batches improve update reliability, but excessive averaging can reduce the stochastic variation that helps the optimizer escape poor local regions. The best reported configuration is not the largest tested batch size or perturbation count.

Figure 1: Convergence on stiff paraboloids as the condition number increases, showing progressively larger advantages for 1.5-SPSA over 1SPSA.

Directional curvature and 1.5-SPSA

The second contribution is a preconditioner defined in probe space rather than parameter space. For each perturbation $z_i$, the method estimates scalar directional curvature with a three-point finite-difference stencil:

$$\hat{c}_i = \frac{L(\theta+\epsilon z_i)-2L(\theta)+L(\theta-\epsilon z_i)} {\epsilon^2}.$$

This approximates $z_i^\top \nabla^2L(\theta)z_i$, but does not require constructing or storing a Hessian, Hessian-vector products, or an $\epsilon$0 optimizer state. The clean loss evaluation $\epsilon$1 is shared across all probes and is often already needed for monitoring training loss.

The curvature estimate determines a robust weight,

$\epsilon$2

with $\epsilon$3 and $\epsilon$4 in the principal experiments. The update therefore attenuates directions with large absolute curvature while avoiding the numerical instability of direct inverse-curvature scaling. The saturation exponent is essential: $\epsilon$5 would approximate a full inverse-curvature correction in probe space, whereas the chosen $\epsilon$6 applies a substantially milder transformation.

The resulting update is

$\epsilon$7

This is not a parameter-space diagonal preconditioner in the Adam sense. It is a per-probe scalar reweighting of the random subspace sampled at each step. The method therefore avoids storing per-parameter moments while still responding to local anisotropy.

The motivation is supported by direct measurements on Qwen3-8B fine-tuned on SST-2. One-dimensional loss profiles along random directions exhibit both strongly negative and strongly positive local curvature, with magnitudes reaching approximately $\epsilon$8.

Figure 2: Random one-dimensional loss profiles for Qwen3-8B, illustrating highly variable and indefinite local curvature.

A larger probe sweep produces three-point curvature estimates spanning approximately $\epsilon$9 to $\theta+\epsilon z_i$0. Applying a common step size to such directions necessarily under-steps flat directions and over-steps sharp directions. This observation directly motivates probe-specific attenuation rather than a global learning-rate reduction.

Figure 3: Distribution of directional curvature estimates for Qwen3-8B on SST-2, with variation across roughly eight orders of magnitude.

The paper’s interpretation relies partly on a Johnson–Lindenstrauss argument: if random projections preserve relevant geometric relations among a finite set of probes, then curvature-related inner products can also be approximately preserved. This provides conceptual support for operating in probe space, but it is not a complete convergence theory for 1.5-SPSA on nonconvex neural objectives. In particular, the relative curvature error becomes uncontrolled when $\theta+\epsilon z_i$1 and $\theta+\epsilon z_i$2 are nearly orthogonal, and the finite-difference estimate has bias governed by third derivatives. The algorithm’s regularization and saturation are therefore practical safeguards, not consequences of an exact Hessian approximation.

Noise, batch size, and perturbation count

The empirical analysis separates minibatch noise from perturbation noise. For a fixed perturbation, the variance of the finite-difference estimate decreases approximately as $\theta+\epsilon z_i$3 with batch size $\theta+\epsilon z_i$4, as expected from minibatch averaging.

Figure 4: Distribution of finite-difference gradient estimates across minibatch sizes for Qwen3-8B.

The signal-to-noise ratio becomes favorable around effective batch sizes of 128–256. Below this range, the estimator variance is comparable to or larger than the median gradient signal, producing unstable optimization.

Figure 5: Empirical batch-noise variance follows the expected inverse-batch-size scaling, with stable optimization emerging near signal-to-noise ratio one.

Perturbation averaging exhibits an analogous $\theta+\epsilon z_i$5 variance reduction. In the controlled Qwen3-8B experiment, the perturbation estimator enters the signal-dominated regime around 256 probes.

Figure 6: Finite-difference estimator distributions narrow as the number of perturbations increases.

Figure 7: Empirical perturbation variance decreases approximately inversely with the number of probes.

These measurements provide a practical tuning heuristic: choose the smallest batch size and perturbation count for which the estimator’s signal-to-noise ratio is approximately one or greater. This is more informative than simply maximizing either quantity, since the paper observes diminishing returns in final accuracy and possible loss of useful stochasticity at very large values.

Results on language-model post-training

On OPT-13B, 1.5-SPSA reaches 94.5% on SST-2 in 70 optimization steps using approximately 179,000 forward passes. This is 3.1 percentage points above the reported MeZO result and 2.5 points above the reported BP+Adam result. The result is especially notable in step count: MeZO uses approximately 100,000 optimization steps, whereas 1.5-SPSA reaches its result in fewer than 100.

Across the five reported OPT-13B tasks, 1.5-SPSA obtains the following accuracies:

Method	SST-2	RTE	BoolQ	WSC	WiC
MeZO	91.4	66.1	67.6	63.5	61.1
BP+Adam	92.0	70.8	77.1	63.5	70.1
1SPSA	94.2	63.3	76.5	65.4	61.8
1.5-SPSA	94.5	77.7	76.5	71.2	61.9

The improvement is not uniform. 1.5-SPSA substantially improves RTE and WSC, but it does not exceed BP+Adam on BoolQ or WiC. This heterogeneity is important: curvature reweighting improves the optimization trajectory, but it does not guarantee better task generalization across all datasets.

On OPT-30B, the method remains competitive while using fewer than 300 optimization steps. It reaches 94.5% on SST-2, 77.0% on RTE, 74.0% on BoolQ, 67.5% on WSC, and 59.3% on WiC. The experiments demonstrate that the method scales to a model size for which the authors report inference-mode training on commodity A100 hardware, although the paper does not provide a complete end-to-end wall-clock comparison against a carefully optimized distributed BP system.

The Qwen3 experiments test whether the effect is architecture-specific. On Qwen3-8B, 1.5-SPSA improves over 1SPSA on SST-2, BoolQ, WSC, and WiC, with the largest reported gain on WiC: 71.2% versus 64.6%. On Qwen3-1.7B and StableToolBench, 1.5-SPSA remains stable over a broader learning-rate range. At learning rate $\theta+\epsilon z_i$6, it reaches 79.0% in 24 steps, whereas 1SPSA diverges at $\theta+\epsilon z_i$7 and $\theta+\epsilon z_i$8. The implication is that preconditioning increases the usable step-size range, not merely the final accuracy.

The hyperparameter ablation supports a moderate saturation exponent. At $\theta+\epsilon z_i$9, the method reaches 94.5% and remains stable; larger values degrade accuracy to 89.2% at $\theta-\epsilon z_i$0. The result is consistent with the paper’s argument that direct or aggressive inverse-curvature weighting is too sensitive to noisy, heavy-tailed curvature estimates.

Controlled conditioning experiments

The paper isolates the proposed mechanism using a rotated stiff paraboloid whose Hessian condition number is controlled explicitly. As the condition number increases, 1.5-SPSA increasingly outperforms 1SPSA, reaching an average speedup of up to approximately seven times. In the illustrated $\theta-\epsilon z_i$1 case, 1.5-SPSA converges in six steps while 1SPSA requires approximately 2,000 steps.

Figure 8: On a stiff paraboloid with $\theta-\epsilon z_i$2, 1.5-SPSA converges in approximately six steps versus roughly 2,000 steps for 1SPSA.

This experiment establishes the condition under which the method should be expected to help: not generic optimization, but optimization dominated by anisotropic curvature. At low condition numbers, 1.5-SPSA approximately matches 1SPSA, indicating that its extra curvature calculation does not provide a consistent advantage in well-conditioned objectives.

The nonconvex DNC experiments extend the comparison to recurrent models with external memory. At equal perturbation counts, 1.5-SPSA reduces steps-to-near-zero-loss by as much as six times relative to 1SPSA, and at high perturbation counts it can outperform BPTT in optimization steps. However, the compute comparison is more qualified. When backward passes are converted to approximately two forward-pass equivalents, BPTT is generally more compute-efficient on the DNC task. At 1.1 billion parameters, BPTT cannot run in the reported hardware configuration because of memory limitations, whereas 1.5-SPSA remains executable. Thus, the DNC results support a memory advantage and a step-efficiency advantage over 1SPSA, but not a general forward-pass efficiency advantage over backpropagation.

Systems implementation

The proposed algorithm has a nontrivial systems requirement: each random probe must be regenerated several times without retaining all perturbations in memory. The implementation addresses this through three components.

Rademacher probes are bit-packed, reducing storage from one floating-point sign per parameter to one bit per parameter. For a 13-billion-parameter model, the paper estimates a reduction from approximately 1.6 GB for an unpacked bf16 perturbation to approximately 100 MB when packed.

Custom Triton kernels fuse bit unpacking, sign conversion, scaling, and in-place parameter updates. On OPT-13B, the reported implementation reduces probe-generation time from 48.0 seconds to 9.8 seconds for 96 perturbations and reduces end-to-end step time from 59.9 seconds to 21.7 seconds, corresponding to a 2.76-times total speedup relative to the baseline PyTorch implementation.

Distributed execution assigns perturbations across ranks, communicates only scalar losses, and regenerates the probes on the update rank. The method therefore keeps the optimization state small, but distributed model synchronization remains necessary after parameter updates. The claim of inference-like memory use should consequently be interpreted per rank and under the stated distributed execution model; aggregate cluster memory still scales with the number of ranks, microbatch size, and model size.

Limitations and open questions

The most consequential limitation is the comparison with BP. The paper explicitly states that BP+Adam is not exhaustively retuned for the large-batch, few-step setting used by 1SPSA and 1.5-SPSA. The reported accuracy improvements over BP therefore establish competitiveness under a particular allocation and tuning protocol, not an unconditional advantage in accuracy, compute, or wall-clock time.

The method also requires several sensitive hyperparameters: $\theta-\epsilon z_i$3, the tied learning rate $\theta-\epsilon z_i$4, batch size, perturbation count, and saturation exponent $\theta-\epsilon z_i$5. The paper reports that tying $\theta-\epsilon z_i$6 is particularly stable and that $\theta-\epsilon z_i$7 performs best in its principal setting, but the acceptable range is narrow. A systematic learning-rate and probe-radius selection procedure is not provided.

The curvature estimator is local, finite-difference-based, and potentially biased in the presence of large third derivatives. The curvature distribution is heavy-tailed and indefinite, while the proposed weight uses absolute curvature and therefore does not distinguish positive curvature from negative curvature. This is appropriate for suppressing sharp directions but leaves open how the method interacts with saddle escape, where negative-curvature directions can be algorithmically useful.

Finally, the JL argument concerns preservation over finite probe sets and does not by itself prove convergence of the stochastic, nonconvex algorithm. The empirical relationship between the advantage of 1.5-SPSA and the Hessian condition number is persuasive in the controlled experiments, but the paper leaves open whether a condition-number-based predictor can be estimated online reliably enough to automate the choice of batch size, perturbation count, or $\theta-\epsilon z_i$8.

Conclusion

The paper presents 1.5-SPSA as an inference-mode optimizer that combines aggressive forward-pass aggregation with a probe-space curvature correction. Its main empirical findings are that larger batches and more perturbations can make 1SPSA substantially more effective at fixed forward-pass budgets, and that a single clean loss evaluation is sufficient to stabilize the resulting large updates in ill-conditioned objectives.

The strongest results occur in the intended regime: memory-constrained training of large models, highly anisotropic post-training losses, and distributed hardware capable of parallelizing many forward evaluations. On OPT and Qwen3, 1.5-SPSA achieves competitive or superior reported task accuracy in tens of optimization steps, while retaining approximately inference-mode memory requirements. The results do not replace the need for stronger BP baselines or broader hyperparameter studies, but they establish probe-space preconditioning as a technically viable way to improve the stability and efficiency of large-scale zero-order training (2609.38095).