Optimizers and the Learning Rate

In Module 2, the training loop contained:

optimizer.zero_grad()
loss = loss_fn(model(xb), yb)
loss.backward()
optimizer.step()

Let's separate the jobs:

1. zero_grad()

PyTorch stores gradients inside each parameter's .grad attribute.

Before calculating a new gradient, we normally clear the old one.

2. loss.backward()

Backpropagation calculates how the loss changes with respect to each parameter.

In simplified notation:

loss
  ↓
how much should each parameter change?
  ↓
gradient for each parameter

For a parameter p, PyTorch stores this result in:

p.grad

3. optimizer.step()

This is the part we are studying today.

The optimizer takes the gradients and decides how to change the parameters.

For plain stochastic gradient descent (SGD), the rule is:

where:

• p = the current parameter;
• L = the loss;
• ∇p L = the gradient of the loss with respect to p;
• η (eta) = the learning rate.

Why do we subtract the gradient?

Because the gradient points in the direction in which the loss increases fastest. Moving in the opposite direction therefore tends to reduce the loss.

For example, if:

gradient = +0.01
learning rate = 0.1

then:

parameter change = -0.1 × 0.01
                 = -0.001

So the parameter moves down by 0.001.

This is the simplest optimizer possible.

The experiment below replaces optimizer.step() with this arithmetic so we can verify that PyTorch's SGD is doing exactly what the equation says.

The two implementations agree up to tiny floating-point rounding differences.

That gives us a precise interpretation of the learning rate.

The learning rate is not just a vague "speed" setting. It is the factor that scales the gradient when SGD turns a gradient into a parameter update.

For example:

gradient ≈ 0.01
learning rate = 0.1

update ≈ 0.1 × 0.01
       = 0.001

The actual update therefore depends on both the learning rate and the size of the gradient.

This explains why a learning rate that works on one problem may be terrible on another.

The four update rules (Other optimers) in this lesson can be thought of as modifications of the same basic idea:

SGD
→ use the raw gradient

Momentum
→ smooth the gradient using recent history

Adam
→ smooth the gradient AND scale it according to each parameter's typical gradient size

AdamW
→ do Adam's adaptive update, while applying weight decay separately

So the training loop stays almost the same. What changes is the rule used by optimizer.step().

It is easier to understand an optimizer on a simple surface than inside a 5,377-parameter neural network.

Imagine a long, narrow valley:

        steep
          ↕
       \     /
        \   /
         \_/
          →
       shallow

The bottom of the valley contains the minimum.

The important detail is that the surface is:

• very steep in one direction;
• much flatter in another direction.

This can happen when different parameters have very different scales.

The gradient points in the direction of steepest increase.

Gradient descent moves in the opposite direction, so it always tries to move downhill.

The problem is that in a narrow valley, the downhill direction can point strongly across the valley rather than along the valley floor.

So the optimizer can do this:

left wall → right wall → left wall → right wall

while making only slow progress toward the actual minimum.

This is the zig-zag problem.

Momentum is designed to reduce that behaviour.

The left panel shows the problem with plain SGD.

The optimizer repeatedly crosses the valley in the steep y direction while making much slower progress along the shallow x direction.

Momentum adds memory.

Instead of using only the current gradient, it keeps a velocity that contains information from recent gradients:

and then updates the parameter using that velocity:

Here:

• v = accumulated direction of recent gradients;
• β = how much of the previous velocity we keep;
• η = learning rate.

In the steep direction, the gradient may keep changing sign:

step 1 → right
step 2 → left
step 3 → right
step 4 → left

Those opposing directions partially cancel inside the velocity.

Along the shallow direction, the gradient may keep pointing in roughly the same direction:

→ → → → →

Those contributions accumulate.

So momentum has a useful intuition:

It reduces movement that keeps changing direction and reinforces movement that remains consistently useful.

This is why momentum is sometimes described as a kind of smoothing or low-pass filtering of the gradient.

With β = 0.9, the optimizer keeps a large amount of the previous velocity. A rough intuition is that it incorporates information from around the previous 1/(1-β) = 10 steps, although this is not a literal ten-step window.

The demonstration deliberately uses a smaller momentum value than the common 0.9.

The learning rate in this toy example is already large enough to be close to the stability boundary.

The multiplier in the steep direction is:

1 - learning_rate × 2
= 1 - 0.85 × 2
= -0.70

The negative sign means the parameter changes sides on each step. The magnitude 0.70 means the oscillation is still large.

Momentum can make the effective update larger because it accumulates gradients.

So:

Learning rate and momentum are not independent knobs.

Increasing momentum can make an otherwise stable learning rate unstable.

This does not mean beta=0.9 is bad. It means the learning rate and momentum need to be considered together.

A practical intuition is:

more momentum
    ↓
more accumulated movement
    ↓
you may need a smaller learning rate

Momentum remembers recent gradients.

Adam goes one step further: it also keeps track of the typical size of the gradients for each parameter.

Why is that useful?

Different parameters in the same neural network can receive gradients with very different magnitudes.

For example, imagine two input features:

age
job=student

age is a dense numerical feature. Many examples contain a meaningful age value.

job=student is a one-hot feature. It is 1 only for examples belonging to that category and 0 otherwise.

The corresponding weights can therefore experience very different gradient patterns.

One global learning rate may not be ideal for both.

For every parameter, Adam maintains:

1. First moment (m)

A running average of the gradient.

This is closely related to momentum.

2. Second moment (v)

A running average of the squared gradient.

Because the gradient is squared, v tells us roughly how large the gradient has typically been, without positive and negative values cancelling each other.

Adam then uses approximately:

as the direction and scale of the adaptive update.

The important intuition is:

If a parameter normally receives very large gradients, Adam scales its update down. If it receives small gradients, Adam scales its update up.

The global learning rate then controls the overall size of these already-normalised updates.

That is why Adam's learning rate behaves differently from plain SGD's learning rate.

Adam stores two extra values for every trainable parameter:

parameter
   ↓
first-moment value (m)
second-moment value (v)

So if the model has N parameters, Adam needs roughly:

N parameter values
+ N first-moment values
+ N second-moment values

That is about 3× the parameter storage before considering gradients and other training memory.

For a 5,377-parameter model, this is tiny.

For a multi-billion-parameter model, it becomes significant.

This is one reason memory-efficient fine-tuning methods such as LoRA become important for large models.

The key lesson for this lesson is simply:

Adam is more adaptive, but that adaptiveness has a memory cost.

The experiment shows that different parameters can have very different typical gradient magnitudes.

This is exactly the situation Adam was designed to handle.

Under plain SGD:

update = learning_rate × raw_gradient

So if one parameter has a gradient 100 times larger than another, its update will also tend to be much larger.

Adam normalises the gradient using its running estimate of gradient magnitude:

adaptive update ≈ m / (sqrt(v) + eps)

So the learning rate now acts more directly as a control over the size of the parameter movement.

This is why the useful numerical learning rates are often very different:

SGD on this problem → much larger lr
Adam on this problem → much smaller lr

Do not interpret that as a universal conversion such as "Adam is always exactly 100× smaller."

The numbers depend on the problem and the optimizer state.

Adam divides by:

sqrt(v) + eps

eps is a very small number, commonly 1e-8.

Its job is to prevent division by zero or by an extremely tiny number.

If a parameter has received zero gradient for many steps, v can be zero. Without eps, the denominator could be zero and the update would be undefined.

So eps is mainly a numerical-stability safeguard.

Weight decay is a regularisation technique.

Its basic idea is:

Unless the data gives the model a reason to keep a weight large, gently encourage that weight to become smaller.

Why?

Very large weights can contribute to overly complex solutions. Regularisation can encourage simpler parameter values and can sometimes improve generalisation.

The important question in this section is not whether regularisation is useful. Lesson 3.2 will study that.

The question here is:

How should weight decay be combined with Adam's adaptive update?

A traditional approach is to add an L2 penalty to the loss:

The derivative of that extra term is:

So the total gradient becomes:

With ordinary SGD, this behaves like directly shrinking the weights.

With Adam, however, the extra λp term becomes part of the gradient that Adam subsequently rescales using sqrt(v).

That means the regularisation is no longer applied uniformly.

AdamW separates the two operations.

Conceptually:

Adam update
     +
direct weight decay

The update is:

The weight-decay term bypasses Adam's adaptive scaling.

That is what decoupled weight decay means.

The experiment below compares these two behaviours directly.

The experiment shows why putting the L2 penalty inside Adam's gradient is not equivalent to decoupled weight decay.

With L2-in-the-gradient, the decay term λp becomes part of the gradient and is then affected by Adam's adaptive denominator.

That means the effective regularisation can depend on the parameter's gradient history.

AdamW instead applies:

Adam's adaptive update
+
direct shrinkage of the parameter

so the decay term does not pass through the sqrt(v) scaling.

AdamW with weight_decay=0 behaves like Adam.

So simply changing:

torch.optim.Adam(...)

to:

torch.optim.AdamW(..., weight_decay=0)

does not meaningfully introduce weight decay.

When weight decay is non-zero, switching optimizers changes the meaning of the regularisation setting.

Therefore, if you switch from Adam to AdamW and the result changes, the first hyperparameters to reconsider are the learning rate and weight decay rather than assuming the optimizer itself is inherently worse.

For this curriculum, the practical rule is:

Use AdamW when you want Adam-style adaptive optimisation with decoupled weight decay. It is the default in essentially every modern training stack — Hugging Face Trainer, the reference implementations of every recent open-weight model, most vision pipelines.

Now we compare:

• plain SGD;
• SGD with momentum;
• Adam;
• AdamW.

Everything else is held constant:

same architecture
same seed
same data
same batch size
same number of epochs
same learning rate

This is useful because changing several things at once would make it impossible to know what caused a difference.

However, there is an important trap:

Holding the learning rate equal does not necessarily make the comparison fair.

Different optimizers interpret and transform the gradient differently.

The first comparison intentionally uses the same 1e-3 learning rate so that we can see this effect.

Read the first comparison carefully.

At lr=1e-3, Adam and AdamW perform much better than the SGD variants in this experiment.

But it would be incorrect to conclude:

"Adam is better than SGD."

Why?

Because the learning rate was shared.

1e-3 is a reasonable default for Adam, but it is not necessarily a good learning rate for plain SGD on this problem.

So the experiment really asks:

What happens when different optimizers are forced to use the same numerical learning rate?

That is a valid experiment, but it is not a complete optimizer comparison.

Also notice that Adam and AdamW are identical here because:

AdamW + weight_decay=0
        =
Adam

So now we give each optimizer its own reasonable learning-rate range and compare them again.

This is a more informative comparison because each optimizer gets an opportunity to operate in a useful regime.

Once each optimizer is given a learning rate that works reasonably well for it, the differences become much smaller.

That changes the interpretation of the first experiment.

The first plot mostly showed:

optimizer + learning rate combination

rather than optimizer quality by itself.

The practical lesson is:

An optimizer cannot be evaluated independently of its hyperparameters.

For this particular problem, AdamW's practical advantage is its robustness to the initial learning-rate choice. Its adaptive scaling often puts it into a useful range without as much manual tuning.

That does not mean SGD is intrinsically incapable of performing well. With a suitable learning rate, it can become competitive on this problem.

This is an important experimental-design lesson:

If you compare two methods, give each method a reasonable opportunity to work.

A method that was badly tuned should not be used as evidence that the underlying method is bad.

If you ever read a paper or a blog post claiming architecture A beats architecture B, the first question is whether both were tuned with equal effort.

Adam/AdamW is better by default. It reaches a good result without you having to find the right learning rate first, because its adaptive scaling absorbs most of the problem-specific magnitude. That robustness is worth real money in engineering time, and it is why AdamW is the 2026 default. But on a well-conditioned problem with a tuned learning rate, the gap narrows sharply.

Now we keep the optimizer fixed and change only the learning rate.

This is one of the most useful experiments you can do when learning optimisation because the learning rate leaves recognizable patterns in the training curves.

We test:

1e-5
1e-4
1e-3
1e-2
1e-1
0.5

These values differ by factors of 10, so we can see what happens when the learning rate moves from extremely small to extremely large.

The important idea is:

The learning rate controls how far the parameters move in response to the gradients.

If the step is too small, learning is painfully slow.

If it is reasonable, the model learns efficiently.

If it is too large, the optimizer can overshoot good regions.

If it is extremely large, training can become unstable or the network can collapse.

The goal is to learn the visual fingerprints of these cases.

Here are the five learning-rate patterns.

The loss decreases smoothly, but very slowly.

Nothing looks obviously broken.

That is precisely the problem: the run is wasting compute because the model has not had enough movement to reach a good solution.

Diagnosis: the curve is still improving substantially when training ends.

The behaviour is similar, but the model gets further.

It is learning, but it may still need more epochs than necessary.

The difference from the previous case is mostly quantitative.

The loss drops quickly and then flattens.

The model reaches a useful region with time left in the training run.

This is the pattern we generally want.

The model still learns, but the curve becomes noisier.

The optimizer can move past the best region and then move back again.

A common symptom is:

best validation epoch
      ↓
good performance

later epochs
      ↓
slightly worse performance

A learning-rate schedule can sometimes help here by reducing the step size later in training.

The model can enter a collapsed state.

In this experiment, the validation AUC falls to approximately 0.5, which is the level expected from random ranking.

The important warning is:

A flat loss curve does not automatically mean successful convergence.

A model can be stable because it has reached a useful minimum, or because it has become stuck in a useless solution.

Always examine a meaningful evaluation metric alongside the loss.

There is also a middle region between "works" and "completely collapses" where the model learns something but performs poorly. That region is often harder to diagnose because nothing looks catastrophically broken.

The lr=0.5 model has effectively learned the simplest possible prediction: the base rate.

The base rate is the proportion of positive examples in the training data.

If approximately 11.7% of training examples are positive, the model can minimise its binary cross-entropy loss surprisingly well by predicting approximately 0.117 for every example.

That gives:

example 1 → 0.117
example 2 → 0.117
example 3 → 0.117
...

The model is no longer using the input features to distinguish examples.

This explains the combination:

stable loss
+
AUC ≈ 0.5
+
almost identical predictions
=
model collapse to a constant prediction

This is a valuable diagnostic pattern that will appear again later in the curriculum.

The broader lesson is:

Never interpret a training curve using loss alone. Always check whether the model is actually making useful, varying predictions.

In the previous section, we tried several learning rates separately.

That is called a learning-rate sweep or grid search.

It works, but it requires multiple training runs.

For a small example, that is cheap.

For a large model, one complete training run could take hours or days.

The LR range test tries to get useful information from a single run.

Start with an extremely small learning rate.

Then increase it after every batch:

very small
   ↓
small
   ↓
useful
   ↓
too large
   ↓
unstable

While doing this, record the loss.

The resulting curve gives us a picture of how the model reacts as the learning rate changes.

Typical pattern:

learning rate very small
→ loss barely changes

learning rate enters useful range
→ loss falls quickly

learning rate becomes too large
→ loss stops improving and eventually rises

The test is therefore not another training strategy. It is a diagnostic experiment for choosing a starting learning-rate range.

The usual idea is to choose a learning rate somewhat below the point where the range test reaches its minimum loss.

A common heuristic is:

recommended lr ≈ minimum-loss lr / 10

The reason is that the minimum of the short range-test curve may already be close to the point where training becomes unstable. Choosing a somewhat smaller value provides a safety margin for a full training run.

The code searches for the steepest decrease before the loss minimum.

Why?

After the minimum, the loss starts increasing.

The steepest part of that increase can produce a large gradient too, but it is pointing in the wrong direction for our purpose. A naive search over the entire curve could therefore identify the instability region instead of the useful region.

The LR range test tells us about learning-rate behaviour during this particular short experiment.

It does not guarantee the perfect learning rate for the full training run.

The weights are changing throughout the test, and the learning rate itself is changing every batch.

So the loss at a particular learning rate is not equivalent to running a complete training job at that learning rate.

Instead, use the curve to answer:

"What range of learning rates appears promising?"

Then confirm that range with a normal training run.

A second limitation is smoothing. The raw batch loss is noisy, so the test smooths it. Too little smoothing leaves the curve noisy; too much smoothing can move the apparent transition point.

Treat the result as a starting point, not a final answer.

A learning-rate schedule changes the learning rate during training.

Why would we want that?

Early in training, the model is usually far from a useful solution. Larger steps can help it move through parameter space quickly.

Later, once it is near a good solution, smaller steps can help it settle rather than repeatedly jumping around.

So the basic idea is:

early training
→ larger learning rate
→ cover ground quickly

later training
→ smaller learning rate
→ make finer adjustments

This is called learning-rate decay or a learning-rate schedule.

Three common patterns are:

Multiply the learning rate by a fixed factor every few epochs.

Example:

start: 0.01
epoch 5: 0.001
epoch 10: 0.0001

It is simple and easy to inspect.

Reduce the learning rate smoothly according to a cosine-shaped curve.

Instead of sudden jumps, the learning rate gradually decreases.

Start with a very small learning rate, gradually increase it for the first part of training, and then reduce it.

Warmup is particularly common in Transformer training, where large early updates can destabilise a freshly initialised model.

We test schedules starting from lr=1e-2, which the previous experiment identified as somewhat too large for a constant learning rate.