Optimizer state: the memory bank of deep learning

Training a neural network is an exercise in navigation. You are dropped into a massive, foggy landscape with millions of valleys and ridges, and your only goal is to find the lowest point, the absolute minimum loss. How you choose to take your steps matters just as much as where you are trying to go.

To appreciate why modern networks train successfully at all, we have to look closely at the mathematical memory that drives them: the optimizer state.

1Why vanilla SGD fails

In the early days of deep learning, models relied on vanilla stochastic gradient descent (SGD). It is a purely reactive algorithm. It looks at the slope right under its feet, takes a step down that slope, and immediately forgets everything else.

The formula is entirely memoryless:

θ_t+1 = θ_t − η · ∇_t

Where θ represents the model weights, η is the learning rate, and ∇_t is the gradient, the current slope.

The ravine problem

Imagine a landscape that looks like a long, narrow canyon. The floor slopes very gently toward the exit (the true minimum), but the walls of the canyon are incredibly steep. Drop a vanilla SGD optimizer into this canyon and here is what happens:

It calculates the gradient. The slope pointing down the steep canyon walls is massive compared to the gentle slope leading toward the exit.
SGD takes a massive step based on that steep slope. It violently hurtles across the canyon floor and smashes into the opposite wall.
On the next step, it sees a massive slope pointing back the other way. It bounds across the floor again.

Because SGD has no memory of its past trajectory, it spends all its energy oscillating wildly between the canyon walls. It makes almost zero progress down the length of the canyon toward the exit. In complex modern AI models, this exact behavior causes training to stall out completely.¹

fn 1The oscillation is not bad luck. A reactive rule always lets the steepest local direction dominate the step, even when that direction is the useless one.

Lesson
A memoryless optimizer pays for every oscillation twice: once in wasted motion across the canyon, once in progress never made along it.

2AdamW: the optimizer with history

To solve this lack of awareness, modern training relies on AdamW, Adam with decoupled weight decay. AdamW succeeds because it keeps records. For every single parameter in your model, the optimizer state tracks historical statistics and creates a custom, adaptive update using two historical metrics: the first momentum and the second momentum.²

fn 2β₁ = 0.9 and β₂ = 0.999 are the PyTorch defaults for AdamW. The mechanics hold for any choice in (0, 1); only the smoothing timescales change.

What the optimizer state holds per parameter (typical mixed-precision setup)
Record	Job	Precision	Bytes
Weights θ	what gets updated	bf16 / fp16	2
First moment m_t	smoothed direction	fp32	4
Second moment v_t	smoothed volatility	fp32	4
Master weight copy	full-precision target for tiny updates	fp32	4

3The mechanics of first and second momentum

The core of the optimizer state consists of two vectors that act as filters for the noisy gradients generated during training.

The first momentum (m_t): the direction manager

The first momentum tracks the directional trend of the gradients over time. It functions exactly like physical momentum:

m_t = β₁ · m_t−1 + (1 − β₁) · ∇_t

The work: it acts as an exponential moving average. If the gradient points in the same direction for several steps, the optimizer speeds up. If the gradient starts oscillating back and forth (like bouncing off canyon walls), the opposing signs cancel each other out in the running average.
The result: it flattens noisy, erratic bounces and keeps the model pushing steadily along the true path.

The second momentum (v_t): the scale manager

The second momentum tracks the volatility, or the sheer variance, of the gradients for each parameter:

v_t = β₂ · v_t−1 + (1 − β₂) · (∇_t)²

The work: by squaring the gradient, the optimizer strips away the directional sign and focuses purely on magnitude.
The result: it acts as an automatic speed regulator. The final update divides the step by √v_t. A wildly volatile parameter with massive gradients gets divided by a large √v_t, shrinking its step to keep it safe. A parameter with tiny, rare gradients gets divided by a tiny √v_t, boosting its step so it does not get left behind. This is what makes the learning rate adaptive.

4Bias correction and the full update

Because m_t and v_t are initialized at zero, the early estimates are artificially pulled toward zero. Bias correction compensates:

m̂_t = m_t1 − β₁^t , v̂_t = v_t1 − β₂^t

As t grows large, β₁^t → 0 and β₂^t → 0, so the corrections fade and the estimates converge to their true values. In the early steps, however, they provide an essential rescaling that prevents the optimizer from taking negligibly small steps at the start of training.³

fn 3The figure assumes a constant unit gradient so the effect is isolated: nothing about the model changes, only how fast the running average catches up to reality.

step t2500

raw v(t)0.92

corrected v̂(t)1.00

β₂0.999

bias correction restores full step size from step one; the raw estimate needs thousands of steps to catch up

Fig. 1. Constant unit gradient, AdamW default β₂ = 0.999. The dark curve is the raw running estimate v(t) = 1 − 0.999ᵣ, still below 0.92 after 2500 steps. The green line is the bias-corrected value, which equals 1 at every single step.

The full update rule

Combine the direction manager and the scale manager and you get the unified mathematical step that updates the weights:

θ_t+1 = θ_t − η · m̂_t√v̂_t + ε − ηλθ_t

η is your global learning rate.
m̂_t/(√v̂_t + ε) is the bias-corrected adaptive step: the smoothed direction divided by the smoothed volatility scale. The tiny constant ε prevents division by zero.
ηλθ_t is the decoupled weight decay term, the W in AdamW. Crucially, it is applied as a separate subtraction, not folded inside the adaptive fraction. In vanilla Adam, L2 regularization is mixed into the gradient before the adaptive scaling, which makes the effective decay rate vary per parameter. AdamW decouples it entirely, so every weight decays at a uniform, predictable rate regardless of its gradient history.

opt = torch.optim.AdamW(model.parameters(), lr=1e-3,
                        betas=(0.9, 0.999), weight_decay=0.01)   # defaults

Lesson
An optimizer is a tiny database with an update policy. What it remembers (direction, volatility) decides how it moves; make sure you know which memories your training run is paying for.

5FP32 master weights in mixed precision

When scaling up large language models, storing everything in 32-bit floating-point numbers (FP32) consumes too much GPU memory. Mixed-precision training runs the active forward and backward passes in 16-bit brain floats (BF16) or half-precision floats (FP16).

This introduces a massive problem called underflow. Because 16-bit floats have limited precision, they cannot represent incredibly small numbers. During an update, the step size is often a tiny fraction. Subtracting a tiny 16-bit fraction from a 16-bit weight rounds that fraction down to absolute zero. The weight never changes, and the model stops learning.⁴

fn 4BF16 trades mantissa bits for range, FP16 keeps more mantissa but risks overflow. Either way, increments far below the representable spacing near the weight's magnitude vanish when added.

The solution: the FP32 master weight copy

To bypass underflow, the optimizer state holds a hidden, full-precision FP32 copy of all the model weights:

The model does its fast calculations in 16-bit to save memory.
The gradients are computed and passed to the optimizer.
The optimizer calculates the first and second moments.
The resulting tiny update step is applied directly to the FP32 master weights. Full 32-bit precision preserves the tiny numbers instead of rounding them to zero.
The freshly updated FP32 weights are cast back down to 16-bit for the next round of training.

Lesson
Precision is a budget. Spend 16-bit where the math is bulk, keep 32-bit where tiny increments land. Without the FP32 master copy, low-precision training of modern architectures would be mathematically impossible.

6References

torch.optim.AdamW documentation: defaults for lr, betas, and weight decay.
Fixing Weight Decay Regularization in Adam (Loshchilov and Hutter, 2017): the decoupled decay argument behind AdamW.
Adam: A Method for Stochastic Optimization (Kingma and Ba, 2014): original bias correction derivation.
Mixed Precision Training (Micikevicius et al., 2017): FP16 compute with FP32 master weights.
Local source: bin/blogs/optimizer-state.md, the original markdown note this page was rewritten from.

0Training is navigation