Probe-Space Preconditioning for Fast and Stable Zero-Order Training
Not on the map. Why?
Abstract (arXiv)
Backpropagation (BP) dominates deep learning but imposes a massive memory tax. For example, training OPT-30B with Adam requires ≈ 600GB of GPU memory (assuming batch size 8 and sequence length 2048). Alternatively, zero-order optimization (ZOO) trains in inference-mode (requiring only ≈ 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).
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 10 of its 20 references that are in the map. First posted on arXiv 29 Sep 2026.
Builds on
10 of 20 references.
References · 10
- 2017Evolution Strategies as a Scalable Alternative to Reinforcement Learningcited by 574
- 2023Fine-Tuning Language Models with Just Forward Passescited by 345
- 2019An Investigation into Neural Net Optimization via Hessian Eigenvalue Densitycited by 292
- 2017Empirical Analysis of the Hessian of Over-Parametrized Neural Networkscited by 290
- 2014Identifying and attacking the saddle point problem in high-dimensional non-convex optimizationcited by 289
- 2024StableToolBench: Towards Stable Large-Scale Benchmarking on Tool Learning of Large Language Modelscited by 132
- 2019ZO-AdaMM: Zeroth-Order Adaptive Momentum Method for Black-Box Optimizationcited by 76
- 2023DeepZero: Scaling up Zeroth-Order Optimization for Deep Model Trainingcited by 76
- 2024Variance-reduced Zeroth-Order Methods for Fine-Tuning Language Modelscited by 54
- 2024Simultaneous Computation and Memory Efficient Zeroth-Order Optimizer for Fine-Tuning Large Language Modelscited by 9
Built on by
Not cited yet.