Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers
Not on the map. Why?
Abstract (arXiv)
Preconditioned optimizers are central to language model training, but their stochastic update rules are usually treated as direct approximations to population preconditioned descent. We show that this view misses two finite-sample biases. First, the gradient and preconditioner are typically estimated from the same minibatch, introducing gradient--preconditioner coupling bias. Second, even when the preconditioner estimate is unbiased, its inverse or inverse-root is generally biased because inversion is nonlinear. We propose a single-batch bias-correction framework that addresses both effects: cross-fitted preconditioning estimates the numerator and preconditioner from independent microbatch groups, while variance-corrected inversion uses microbatch variability to subtract the leading delta-method bias term. The framework applies to diagonal moment, diagonal curvature, and matrix preconditioning methods, instantiated in AdamW, Sophia, and Shampoo. Bias correction reduces held-out pretraining loss on Qwen2.5-0.5B by 0.15, 0.07, and 0.11 nats, respectively; the effects on mixed-quality pretraining and downstream instruction tuning are consistently neutral-to-positive. Together, these results establish bias correction as a practical mechanism for reducing finite-sample update bias and improving the performance of preconditioned optimizers.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 11 of its 27 references that are in the map. First posted on arXiv 20 May 2026.
Builds on
11 of 27 references.
References · 11
- 2020Sharpness-Aware Minimization for Efficiently Improving Generalizationcited by 1,417
- 2018Adafactor: Adaptive Learning Rates with Sublinear Memory Costcited by 905
- 2019On the Convergence of Adam and Beyondcited by 841
- 2015Optimizing Neural Networks with Kronecker-factored Approximate Curvaturecited by 796
- 2013Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programmingcited by 724
- 2019Large Batch Optimization for Deep Learning: Training BERT in 76 minutescited by 674
- 2018Shampoo: Preconditioned Stochastic Tensor Optimizationcited by 426
- 2014SAGA: A Fast Incremental Gradient Method With Support for Non-Strongly Convex Composite Objectivescited by 402
- 2019Lookahead Optimizer: k steps forward, 1 step backcited by 336
- 2023Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-trainingcited by 265
- 2020On the Convergence of SGD with Biased Gradientscited by 83
Built on by
Not cited yet.