📊 statistics

Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers

This paper identifies and corrects two finite-sample biases in preconditioned language model optimizers—gradient-preconditioner coupling and nonlinear inversion bias—by proposing a single-batch framework that combines cross-fitted preconditioning with variance-corrected inversion, thereby reducing pretraining loss and improving performance across models like AdamW, Sophia, and Shampoo.

Original authors: Nikhil Nayak, Julia White, Urchade Zaratiana, Kelton Zhang, Henrijs Princis, Dhruv Atreja, Henry Fawcett, Matthew Thomas, George Hurn-Maloney, Ash Lewis

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nikhil Nayak, Julia White, Urchade Zaratiana, Kelton Zhang, Henrijs Princis, Dhruv Atreja, Henry Fawcett, Matthew Thomas, George Hurn-Maloney, Ash Lewis

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to write poetry. To do this, you show it thousands of examples of poems and ask it to adjust its "brain" (its parameters) to get better. The robot uses a special tool called an optimizer to figure out which way to nudge its brain.

Most modern optimizers are "smart." They don't just look at the mistake the robot made; they also look at the shape of the mistake to decide how big a step to take. This is called preconditioning.

However, the authors of this paper discovered that these smart optimizers have a hidden flaw: they are slightly biased because they are trying to make decisions based on small, noisy samples of data. It's like trying to judge the weather by looking out the window for only five seconds.

Here is the problem and their solution, explained through simple analogies.

The Two Hidden Flaws

The paper identifies two specific ways the optimizer gets the math wrong when working with small batches of data:

1. The "Self-Serving" Bias (Coupling Bias)

The Analogy: Imagine a student taking a test. To grade their own work, they use the same test paper they just took to calculate their score. Naturally, they might be tempted to grade themselves too harshly or too leniently because the questions and the answers are perfectly matched in their mind.
The Reality: In the optimizer, the robot calculates the "direction to move" (the gradient) and the "size of the step" (the preconditioner) using the exact same small group of data. Because they come from the same source, they are statistically "coupled." This creates a subtle distortion in the update rule, making the robot's step slightly off-target.

2. The "Non-Linear" Bias (Inverse Bias)

The Analogy: Imagine you are trying to guess the average speed of a car. You know the average time it takes to drive a mile is 1 minute. You might think, "Okay, so the average speed is 60 mph." But math doesn't work that simply! If you average the times first and then calculate the speed, you get a different answer than if you calculate the speed for every single trip and then average them. This is because "speed" is the inverse of "time," and doing math with inverses is tricky.
The Reality: The optimizer needs to divide by a number (the preconditioner) to figure out the step size. Even if the robot estimates the preconditioner perfectly on average, the act of inverting it (turning it upside down to divide) introduces a systematic error. It's like trying to guess the average of a group of people's heights by averaging their shoe sizes first; the math gets warped.

The Solution: A "Double-Check" System

The authors propose a clever fix that doesn't require the robot to re-read the data or slow down significantly. They call it a Single-Batch Bias Correction.

They use two main tricks, which they apply at the same time:

1. Cross-Fitting (The "Independent Judges")

How it works: Instead of using the same group of data for both the "direction" and the "step size," the robot splits its current batch of data into two separate groups.
- Group A calculates the direction.
- Group B calculates the step size.
The Result: By keeping the judges independent, you remove the "self-serving" bias. The step size is no longer influenced by the specific direction of that exact same group of data.

2. Variance Correction (The "Noise Meter")

How it works: The robot looks at how much the "step size" estimate varies across tiny sub-groups of the data. If the estimate is shaky (high variance), the robot knows the math is likely to be biased due to that non-linear inversion problem.
The Result: The robot uses a statistical formula (based on the "Delta Method") to subtract that expected error. It's like a mechanic who knows that a specific tool tends to measure 2% too high when the air is humid, so they automatically subtract 2% from the reading to get the true value.

What Happened When They Tried It?

The team tested this fix on three different types of "smart" optimizers used to train large language models (like the ones powering chatbots):

AdamW: The standard, workhorse optimizer.
Sophia: An optimizer that tries to guess the "curvature" of the problem (like knowing if a hill is steep or flat).
Shampoo: A very complex optimizer that uses matrix math to handle huge amounts of data.

The Results:

Pretraining (Learning from scratch): When training a model from zero, the bias correction worked very well. It lowered the error rate (loss) significantly for all three optimizers. It's as if the robot learned the poetry faster and more accurately because it wasn't being misled by its own noisy calculations.
Instruction Tuning (Fine-tuning an already smart model): When they tried to teach an already-trained model new tasks, the results were "neutral-to-positive." The correction didn't hurt, and in some cases, it gave a tiny boost, but the gains were smaller. This makes sense because the model was already stable, so the "noise" was less of a problem.

The Bottom Line

The paper argues that we have been treating these optimizers as if they are perfect mathematical machines, but they are actually making small, consistent mistakes because they are working with limited data.

By simply splitting the data to keep the direction and step-size independent, and measuring the noise to subtract the inversion error, the authors created a "bias-correction layer." This layer makes the training process more efficient, allowing language models to learn slightly better and faster, especially when they are starting from scratch. It's a small statistical tweak with a surprisingly big impact on how well these AI models learn.

Technical Summary: Correcting Stochastic Update Bias in Preconditioned Language Model Optimizers

1. Problem Statement

Preconditioned optimizers (e.g., AdamW, Sophia, Shampoo) are central to training large language models (LLMs). These methods estimate a descent direction and a geometry term (preconditioner) from minibatches to approximate the ideal population update:
$u^\star(\theta) = P(\theta)^{-1}\nabla f(\theta)$
In practice, optimizers apply a stochastic update using minibatch estimates $\hat{g}$ and $\hat{P}$ :
$\hat{u} = \hat{P}^{-1}\hat{g}$

The paper identifies that even when $\hat{g}$ and $\hat{P}$ are unbiased estimators of the population gradient and preconditioner, the product $\hat{P}^{-1}\hat{g}$ is generally a biased estimator of the ideal update. This finite-sample bias arises from two distinct sources:

Gradient–Preconditioner Coupling Bias: When $\hat{g}$ and $\hat{P}$ are estimated from the same minibatch, they are statistically dependent. This coupling introduces a cross-moment term that deviates from the target population vector field.
Inverse-Preconditioner Bias: Matrix inversion (or inverse-root operations) is a nonlinear transformation. By Jensen's inequality, even if $\hat{P}$ is unbiased for $P$ , the expectation of the inverse $\mathbb{E}[\hat{P}^{-1}]$ is not equal to the inverse of the expectation $(\mathbb{E}[\hat{P}])^{-1}$ . This bias is particularly pronounced when preconditioner estimates are noisy, aggressive, or matrix-valued.

These biases are not implementation errors but inherent consequences of applying nonlinear adaptive rules to finite stochastic estimates.

2. Methodology

The authors propose a single-batch bias-correction framework that addresses both sources of bias without requiring revisiting the same parameter value (which is impractical in neural network training). The method operates at a fixed parameter step $\theta_t$ by splitting the current training batch into microbatch groups.

2.1 Cross-Fitted Preconditioning (Addressing Coupling)

To eliminate gradient–preconditioner coupling, the method estimates the numerator (gradient) and denominator (preconditioner) from independent microbatch groups.

Let $A$ and $B$ be two independent groups sampled from the same training distribution.
The gradient is estimated from $A$ : $\hat{g}_A$ .
The preconditioner is estimated from $B$ : $\hat{P}_B$ .
The update becomes $\hat{u}_{cf} = \hat{P}_B^{-1}\hat{g}_A$ .
Due to independence, $\mathbb{E}[\hat{P}_B^{-1}\hat{g}_A] = \mathbb{E}[\hat{P}_B^{-1}]\mathbb{E}[\hat{g}_A]$ , effectively removing the coupling bias term.

2.2 Variance-Corrected Inversion (Addressing Nonlinearity)

To reduce the bias introduced by the nonlinearity of inversion, the method utilizes the variability within the preconditioner estimate group $B$ .

Group $B$ is split into $m$ microbatches $B_1, \dots, B_m$ .
Diagonal Case (AdamW, Sophia): The authors apply a delta-method correction. For a scalar preconditioner entry $p$ , the inverse functional $T(p) = (p+\lambda)^{-1}$ is expanded to second order. The leading bias term, proportional to the variance of the microbatch mean, is estimated and subtracted:
$\tilde{T}(\bar{p}) = \frac{1}{\bar{p} + \lambda} - \frac{\widehat{\text{Var}}(\bar{p})}{(\bar{p} + \lambda)^3}$
This correction is applied elementwise to diagonal preconditioners.
Matrix Case (Shampoo): For matrix-valued preconditioners, the method applies the correction in the eigenbasis of the averaged preconditioner $\bar{P}$ . Microbatch preconditioners are projected into this basis, and the scalar delta-method correction is applied to the eigenvalues (or inverse-root eigenvalues) before reconstructing the matrix.
Jackknife Alternative: A more general leave-one-out (LOO) jackknife correction is also described, which cancels the leading $O(1/m)$ bias term, though it may be computationally expensive for matrix preconditioners.

3. Key Contributions

Statistical Decomposition: The paper explicitly decomposes the stochastic preconditioned update error into two finite-sample biases: gradient–preconditioner coupling and the nonlinearity of inversion.
Correction Framework: It proposes a practical, single-step framework combining cross-fitted preconditioning (to remove coupling) and variance-corrected inversion (to remove inverse bias) using microbatch variability.
Empirical Instantiation: The framework is instantiated for three major optimizer families:
- AdamW: Diagonal moment-based (using LOO cross-fitting and Jensen/delta correction).
- Sophia: Diagonal curvature-based.
- Shampoo: Matrix-preconditioned (using eigenbasis correction).

4. Experimental Results

The framework was evaluated on Qwen2.5-0.5B under two regimes:

Pretraining: From random initialization on FineWeb-Edu sequences (clean and mixed-quality/noisy).
Instruction Tuning: Supervised fine-tuning (SFT) on Alpaca-style data.

Key Findings:

Pretraining Performance: Bias correction consistently reduced held-out pretraining loss (cross-entropy) across all three optimizer families:
- AdamW: Reduced loss by 0.149 nats (from 4.836 to 4.687).
- Sophia: Reduced loss by 0.070 nats (from 6.665 to 6.595).
- Shampoo: Reduced loss by 0.110 nats (from 5.792 to 5.681).
- In mixed-quality (noisy) pretraining, the corrected AdamW variant showed a smaller but positive gain (0.019 nats) on clean evaluation data, while fitting corrupted spans less aggressively.
Instruction Tuning: Results were "neutral-to-positive."
- AdamW: Showed small gains (e.g., -0.001 nats) after retuning learning rates.
- Sophia: Nearly tied with standard optimization on small sets, slightly better on larger 5K sets.
- Shampoo: Effectively neutral in the attention-only SFT setting.
Ablations: The paper notes that naive two-fold splitting (using half the batch for the denominator) can be too conservative due to increased noise. The proposed Leave-One-Out (LOO) construction for AdamW preserves the denominator scale while removing coupling, proving critical for performance.

5. Significance and Claims

The paper positions this work as a practical mechanism for improving preconditioned optimizers by addressing statistical biases inherent in finite-sample updates, rather than proposing a new optimizer architecture or curvature approximation.

Theoretical Insight: It establishes that standard adaptive updates are biased estimators of the population preconditioned direction due to coupling and nonlinearity.
Practical Utility: The correction is lightweight for diagonal optimizers and extensible to matrix optimizers. It improves sample efficiency and stability, particularly in pretraining regimes where estimates are noisy and training runs are long.
Modest Scope: The authors are modest about the generalizability of the effect size. They acknowledge that results are based on a single model scale (0.5B) and single-seed runs. They note that in short instruction-tuning runs starting from strong pre-trained checkpoints, the variance cost of sample splitting can offset the bias-reduction benefits, leading to neutral results.
Conclusion: The primary contribution is demonstrating that correcting finite-sample bias in stochastic updates can make preconditioned language model optimization more effective, particularly in pretraining, without requiring fundamental changes to the optimizer's memory footprint or update structure.