Why training needs an optimizer — from gradients to AdamW

In earlier lessons, we followed the forward pass through a Transformer and looked closely at inference: KV caches, batching, memory bandwidth, and scheduling.

Today we'll move back to training and focus on the moment immediately after backpropagation.

Backpropagation gives us gradients:

\[ g_t = \nabla_\theta L_t. \]

But a gradient is not yet a parameter update.

It tells us:

If I nudge each parameter a tiny amount, which direction locally increases or decreases the loss?

The optimizer decides what to do with that information.

The central idea is:

\[ \boxed{\text{backprop computes the signal; the optimizer turns it into learning}} \]

Modern LLMs are usually trained with variants of Adam, commonly AdamW, rather than plain gradient descent.

Understanding why gives us a useful window into both the mathematics and the enormous memory cost of LLM training.

Start with the simplest possible optimizer¶

Suppose the model has parameters:

\[ \theta. \]

After a training batch, backpropagation gives us:

\[ g_t = \nabla_\theta L_t. \]

Plain stochastic gradient descent performs:

\[ \theta_{t+1} = \theta_t-\eta g_t \]

where:

\[ \eta \]

is the learning rate.

The minus sign matters because the gradient points in the direction of increasing loss.

So SGD says:

measure downhill direction
        ↓
take a small step downhill
        ↓
repeat

This is conceptually beautiful.

But LLM optimization is much noisier than the smooth bowl often drawn in textbooks.

A one-parameter example¶

Imagine our entire model has one parameter:

\[ \theta=5. \]

Suppose the current gradient is:

\[ g=2. \]

With learning rate:

\[ \eta=0.1, \]

SGD gives:

\[ \theta_{new} = 5-0.1(2) = 4.8. \]

If instead:

\[ g=-3, \]

then:

\[ \theta_{new} = 5-0.1(-3) = 5.3. \]

The gradient determines direction and magnitude.

The learning rate globally scales the step.

Now imagine doing this not for one parameter, but for hundreds of billions.

The problem: gradients are noisy¶

An LLM does not compute the gradient over its entire training dataset before every update.

It trains on minibatches.

One batch might contain lots of:

Python

while another contains:

French prose

and another:

mathematics

So the gradient changes from step to step.

For one parameter we might see:

step 1:  +0.8
step 2:  +1.2
step 3:  -0.4
step 4:  +0.9
step 5:  +0.3

The underlying useful signal might be positive, but individual batches are noisy.

It would be helpful to remember some recent history.

Momentum: a rolling direction of travel¶

A simple improvement is momentum.

Instead of using only the current gradient, maintain an exponentially weighted moving average:

\[ m_t = \beta_1m_{t-1} + (1-\beta_1)g_t. \]

If:

\[ \beta_1=0.9, \]

then roughly speaking, recent gradients matter most while older gradients fade away.

The update can use:

\[ \theta_{t+1} = \theta_t-\eta m_t. \]

Suppose gradients repeatedly point in approximately the same direction.

Momentum accumulates that evidence.

If one batch produces an unusual gradient in the opposite direction, the historical average prevents the optimizer from instantly reversing course.

A good mental model is a heavy ball rolling downhill.

SGD behaves more like:

look at slope
take step
forget everything

Momentum behaves more like:

look at slope
combine with existing velocity
keep rolling

But different parameters have very different gradient scales¶

Now consider two parameters.

Suppose their gradients are:

\[ g_A=0.001 \]

and:

\[ g_B=10. \]

Using one global learning rate means B's raw update is:

\[ 10{,}000\times \]

larger than A's.

Sometimes that is appropriate.

But neural networks contain parameters with dramatically different scales and gradient statistics.

We would like the optimizer to adapt somewhat to each parameter.

This is where Adam's second moving average enters.

Adam tracks squared gradients too¶

Adam maintains:

\[ m_t = \beta_1m_{t-1} + (1-\beta_1)g_t \]

and also:

\[ v_t = \beta_2v_{t-1} + (1-\beta_2)g_t^2. \]

The square is elementwise.

So for every parameter Adam tracks two pieces of history:

m → recent average gradient
v → recent average squared gradient

The first asks roughly:

Which direction have gradients consistently pointed?

The second asks:

How large have gradients for this parameter typically been?

Adam then forms an update approximately like:

\[ \Delta\theta = -\eta \frac{m_t}{\sqrt{v_t}+\epsilon}. \]

This is the key equation.

Why divide by the square root of \(v\)?¶

Suppose one parameter consistently gets gradients around:

\[ 100 \]

while another gets gradients around:

\[ 0.01. \]

For the first:

\[ v\approx100^2=10{,}000 \]

so:

\[ \sqrt v\approx100. \]

Its normalized gradient is roughly:

\[ \frac{100}{100}=1. \]

For the second:

\[ v\approx0.01^2=0.0001 \]

so:

\[ \sqrt v\approx0.01. \]

Its normalized gradient is also roughly:

\[ \frac{0.01}{0.01}=1. \]

Adam has partially normalized away the difference in gradient scale.

This does not mean every parameter gets exactly the same update.

The moving averages, signs, history, epsilon term, and changing gradients all matter.

But the basic effect is:

\[ \boxed{\text{parameters with persistently large gradients get automatically smaller effective steps}} \]

relative to their raw gradient magnitude.

A concrete two-parameter example¶

Suppose:

\[ m_A=0.2,\qquad v_A=0.04 \]

and:

\[ m_B=20,\qquad v_B=400. \]

Ignoring \(\epsilon\):

\[ \frac{m_A}{\sqrt{v_A}} = \frac{0.2}{0.2} = 1. \]

For B:

\[ \frac{m_B}{\sqrt{v_B}} = \frac{20}{20} = 1. \]

Despite a 100× difference in first-moment magnitude, their normalized Adam updates are equal.

This per-parameter adaptivity is one reason Adam became so useful for large neural networks.

The bias-correction wrinkle¶

At the beginning of training:

\[ m_0=0,\qquad v_0=0. \]

That means early moving averages are biased toward zero.

Adam corrects this:

\[ \hat m_t = \frac{m_t}{1-\beta_1^t} \]

and:

\[ \hat v_t = \frac{v_t}{1-\beta_2^t}. \]

Then the update is:

\[ \theta_{t+1} = \theta_t - \eta \frac{\hat m_t} {\sqrt{\hat v_t}+\epsilon}. \]

After many steps:

\[ \beta_1^t\rightarrow0 \]

and:

\[ \beta_2^t\rightarrow0, \]

so the correction becomes negligible.

It mainly matters early in training.

Where AdamW comes in¶

Large models also commonly use weight decay.

The idea is to gently shrink parameters toward zero.

A naive way to combine this with Adam is to add an L2 penalty to the loss.

But because Adam rescales gradients parameter by parameter, that interacts with regularization in a somewhat unintuitive way.

AdamW instead decouples weight decay from the gradient update.

A simplified AdamW step is:

\[ \theta_{t+1} = \theta_t - \eta \frac{\hat m_t} {\sqrt{\hat v_t}+\epsilon} - \eta\lambda\theta_t. \]

The first term is learning from the gradient.

The second is explicit weight decay.

Equivalently:

\[ \theta_{t+1} = (1-\eta\lambda)\theta_t - \eta \frac{\hat m_t} {\sqrt{\hat v_t}+\epsilon}. \]

So AdamW conceptually separates:

learn from data
+
gently shrink weights

That decoupling is the "W" in AdamW.

Now comes the systems surprise¶

Adam is not free.

For every model parameter, training may need to store:

parameter
gradient
first moment m
second moment v

And some mixed-precision training setups also keep a higher-precision master copy of parameters.

Suppose we make a deliberately simplified estimate using FP32 optimizer states.

For one parameter:

BF16 parameter     2 bytes
BF16 gradient      2 bytes
FP32 m             4 bytes
FP32 v             4 bytes
--------------------------
                  12 bytes

Already:

\[ 12\text{ bytes/parameter}. \]

A 70B-parameter model would therefore require roughly:

\[ 70\times10^9\times12 = 840\text{ GB} \]

just for those states.

And that excludes activations and other training memory.

If an FP32 master weight copy is also stored, add another:

\[ 4\text{ bytes/parameter}, \]

taking this simplified accounting to:

\[ 16\text{ bytes/parameter} \]

or:

\[ 1.12\text{ TB} \]

for 70B parameters.

Exact implementations differ, but the lesson is important:

\[ \boxed{\text{training memory can be many times larger than inference weight memory}} \]

Compare training and inference¶

At inference, we may only need something like:

model weights
KV cache
temporary activations

Training needs substantially more:

model weights
forward activations
gradients
optimizer first moments
optimizer second moments
possibly master weights
temporary buffers

So a model that comfortably fits for inference may be nowhere close to fitting for training on the same hardware.

This is one reason distributed training techniques are essential.

Why optimizer sharding matters¶

Suppose we train across:

\[ N=8 \]

GPUs.

If every GPU stores a full copy of Adam's optimizer states, enormous amounts of memory are duplicated.

Instead, systems such as ZeRO-style training can shard states across devices.

Conceptually:

GPU 0 → optimizer states for parameters 0–12.5%
GPU 1 → optimizer states for parameters 12.5–25%
GPU 2 → optimizer states for parameters 25–37.5%
...

Rather than every device holding all optimizer state.

At different sharding stages, systems may distribute:

optimizer states,
gradients,
parameters.

This trades memory duplication for communication.

Again we see a recurring theme in large-model systems:

\[ \boxed{\text{save memory by moving data between machines}} \]

The difficult engineering question becomes whether the interconnect is fast enough.

Why the learning rate still matters with Adam¶

Because Adam adapts each parameter's effective step, it is tempting to think the global learning rate no longer matters much.

It matters enormously.

The learning rate:

\[ \eta \]

still scales the entire update.

Too large:

loss spikes
training destabilizes
parameters overshoot useful regions

Too small:

training progresses slowly
compute is wasted
model may undertrain

Large LLM runs therefore use learning-rate schedules.

A common shape is:

learning rate
     ^
     |        /\
     |       /  \
     |      /    \
     |     /      \
     |____/        \________
     +------------------------> training steps
         warmup      decay

Why warmup exists¶

At the very beginning, the optimizer has little history.

The model is also far from a well-adapted region of parameter space.

Immediately applying the full learning rate can destabilize training.

So the learning rate often starts near zero and increases over some number of steps:

\[ \eta_t\uparrow \]

during warmup.

After warmup, it usually decays gradually.

One intuition is:

early training:
learn cautiously while optimizer statistics stabilize

middle training:
take larger productive steps

late training:
take increasingly fine steps

The exact schedule varies across training runs.

Gradient clipping is another safety mechanism¶

Occasionally a batch can produce unusually large gradients.

Define the global gradient norm:

\[ \|g\|_2. \]

If it exceeds some threshold:

\[ c, \]

we can rescale:

\[ g \leftarrow g \frac{c}{\|g\|_2}. \]

So if:

\[ \|g\|_2\le c, \]

nothing changes.

If:

\[ \|g\|_2>c, \]

the direction stays the same but the magnitude is reduced.

Conceptually:

normal gradient → leave alone
huge gradient   → cap its size

This helps prevent rare batches from producing catastrophic parameter jumps.

Put the whole training step together¶

A simplified LLM training iteration now looks like:

tokens
  ↓
forward pass
  ↓
loss
  ↓
backpropagation
  ↓
gradients
  ↓
gradient clipping
  ↓
AdamW
  ├─ update first moments
  ├─ update second moments
  ├─ normalize gradient update
  └─ apply weight decay
  ↓
new parameters

Mathematically:

\[ g_t=\nabla_\theta L_t \]

then perhaps clip \(g_t\), then:

\[ m_t = \beta_1m_{t-1} + (1-\beta_1)g_t, \]

\[ v_t = \beta_2v_{t-1} + (1-\beta_2)g_t^2, \]

bias-correct:

\[ \hat m_t=\frac{m_t}{1-\beta_1^t}, \qquad \hat v_t=\frac{v_t}{1-\beta_2^t}, \]

and finally:

\[ \theta_{t+1} = (1-\eta_t\lambda)\theta_t - \eta_t \frac{\hat m_t} {\sqrt{\hat v_t}+\epsilon}. \]

Then the next batch begins.

This happens millions of times during a large training run.

The deeper mental model: hiking through fog¶

Imagine trying to descend a mountain in dense fog.

You cannot see the valley.

You can only measure the local slope under your feet.

That local slope is the gradient:

\[ g_t. \]

Plain SGD says:

Step downhill according to the slope you see right now.

Momentum says:

Also remember which direction we have consistently been traveling.

Adam adds:

And remember how volatile the slope has historically been in each direction, so don't overreact along dimensions that routinely produce huge gradients.

AdamW adds:

While doing all of that, gently keep the parameter vector from drifting unnecessarily large.

The learning-rate schedule controls how ambitious your steps are over the course of the journey.

Gradient clipping says:

If one measurement suddenly claims the mountain has become a cliff, don't take an enormous leap.

So training is not simply:

\[ \text{compute gradient}\rightarrow\text{change weights}. \]

It is:

\[ \boxed{ \text{measure noisy local evidence} \rightarrow \text{combine it with historical statistics} \rightarrow \text{choose a carefully scaled update} } \]

And this reveals an important difference between training and inference.

Inference mostly asks:

How efficiently can we execute the parameters we already have?

Training asks something harder:

How do we repeatedly estimate tiny useful changes to billions of parameters, maintain enough historical state to make those changes stable, and distribute all of that state across enormous machines?

The optimizer sits directly at that boundary between learning theory and systems engineering.

AdamW is only a handful of equations.

At LLM scale, those equations can require hundreds of gigabytes or even terabytes of state.