Optimization behaviour

What you will be able to do

Read a training loss curve and say what is wrong with a run. The learning rate can be too small or too large, or the run can stop improving. These are the first questions to ask when a model does not learn.

NotebookOptional: check the numbers in a notebookColab · Kaggle · GitHub

The notebook is optional: this lesson works with a calculator. It runs gradient descent on the ticket model with any learning rate and draws each step as a text trace. You see a run converge, zig-zag and diverge to inf and nan, and a ball that stops in a local minimum. It also runs stochastic gradient descent with the same batches as the batch explorer on this page. The guided practice and the module practice have checks. It uses only Python’s standard library, so it needs no install.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 50 minutes

What you need to know

• The update rule, w − learning rate × gradient, and one step by hand, from Loss and gradients: The update rule and Worked example: two steps by hand.
• The numbers of the ticket model: the gradient is 15w − 94, the best w is about 6.27, and the lowest loss is about 1.22. They are in Loss and gradients: The loss curve and The gradient.
• Each ticket’s part of the gradient at w = 6 (−2, −4, 6 and −16), from Loss and gradients: Where the gradient comes from. The section on stochastic updates uses this table.

What you need to have ready (the notebook needs none of it)

• A calculator. A phone calculator is enough. The module practice has seven steps of the update rule.
• Your c03-notes.md file, with the heading ## How models learn. You added the heading in the “Before you start” box of Functions and change. If you cannot find the file, make a new empty file with that name and add the heading: this lesson does not use earlier answers.
• About 20 minutes at the end for the module practice and the module check.

The notebook was tested in Jupyter with Python 3.14.6 (standard library only).

Where you meet this in AI

• Training settings. A training script has lines such as learning_rate = 3e-4 and epochs = 3. The learning rate is often the most important setting of a run. People change it before anything else.
• Loss curves on a dashboard. Tools such as TensorBoard draw the loss after each step. Engineers look at the shape of this curve every day: steep, flat, noisy, or going up.
• “Loss is NaN”. A run stops, or the log prints loss: nan. The usual cause is a learning rate that is too large. This lesson shows why.
• Fine-tuning a language model. Fine-tuning changes the parameters of a trained model on new examples. A learning rate that is too large can destroy what the model already knew in a few steps.

Learning rate: too small, good, too large

In the last lesson, Priya used a learning rate of 0.05 and reached w = 6.2 in three steps. Tomás asks: “Why not 0.01, to be careful? Or 0.15, to be fast?”

Here is the key fact for this model. The gradient is 15 × (w − 6.27), so it is proportional to the distance from the bottom. Each step therefore multiplies that distance by the same number, 1 − 15 × learning rate:

Learning rate	1 − 15 × rate	What each step does	Behaviour
0.01	0.85	Keeps 85% of the distance	Slow: many small steps
0.05	0.25	Keeps 25% of the distance	Fast and smooth
0.0667 (= 1/15)	0	Lands on the bottom	One step, only for this exact curve
0.1	−0.5	Jumps over the bottom, half as far on the other side	Zig-zag, but it still reaches the bottom
0.13	−0.95	Jumps over, almost as far on the other side	Slow zig-zag
0.1333 (= 2/15)	−1	Jumps exactly as far on the other side	Bounces forever
0.14	−1.1	Jumps over, further each time	Diverges: the loss grows without limit

So, for this model:

• Below about 0.067, steps go smoothly downhill. Smaller rates are slower.
• Between 0.067 and 0.133, each step jumps over the bottom, but the jumps shrink. The path zig-zags to the bottom.
• Above 0.133, the jumps grow. The model diverges: it moves further from the bottom at every step.

The numbers 0.067 and 0.133 belong to this curve only. A steeper loss curve has lower limits. In real training, you cannot calculate them, so you try a few rates and watch the loss.

Explore

Find the learning rates that fail

Step

8

w

5.104

Loss

11.354

Slope (gradient)

−17.439

Next move = −rate × slope

0.174

Moving downhill. Each step is smaller, because the slope gets smaller near the bottom.

The ticket model from w = 2. Change the learning rate, then press Step or Play. Read the loss after each step in the readout and in the lower chart.

Text description of the explorer

The upper chart is the loss curve of the ticket model. It is a bowl with its bottom at w ≈ 6.27 and a loss of about 1.217. The ball starts at w = 2. The slider offers 13 learning rates from 0.005 to 0.16, starting at 0.01. Step makes one update, Play repeats it, and Reset goes back to w = 2. The readout shows the step, w, the loss, the gradient and the next move. At 0.01, w moves 2, 2.64, 3.18, 3.65: slow, small steps. At 0.06, w goes 2, 5.84, 6.22, 6.26: smooth and fast. At 0.12, w jumps to 9.68, then 3.54, then 8.45: a zig-zag that slowly closes in on 6.27. At 0.14, w jumps to 10.96, then 1.10, then 11.95, and the loss grows from 137.75 to 166.4, 201.1 and 243.1. After about 20 steps, the explorer stops the run. The status line names each behaviour. The lower chart has a fixed top. It draws a higher loss on the top edge, so a diverging run looks flat there. Read the loss in the readout.

What a training loss curve looks like

A training log does not show the bowl. It shows only the loss after each step. Here are three runs of the ticket model from w = 2, with three learning rates:

• Learning rate 0.01
• Learning rate 0.05
• Learning rate 0.13

Loss after each step for three learning rates. The good rate drops at once. The too-small and the almost-too-large rates both fall slowly.

Text description of the chart

All three runs start at a loss of 137.75. With learning rate 0.05, the loss is 9.75 after one step, 1.75 after two, and stays at 1.217 from step 5. With learning rate 0.01, the loss falls slowly: 99.9 after one step, 38.4 after four, 6.5 after ten and 1.42 after twenty. With learning rate 0.13, the loss also falls slowly: 124.4 after one step, 50.2 after ten and 18.8 after twenty.

Look at the runs with 0.01 and 0.13. Both curves fall slowly, but for opposite reasons. At 0.01, the steps are too short. At 0.13, the steps are too long: w jumps from one side of the bowl to the other, and the zig-zag closes in slowly. The loss curve alone does not tell you which one it is. If you increase the learning rate and the run gets faster, the steps were too short. If it gets worse or diverges, they were too long.

When the rate is too large, the curve goes up:

• Learning rate 0.12
• Learning rate 0.14

A small change in the learning rate, from 0.12 to 0.14, turns a run that converges into a run that diverges.

Text description of the chart

Both runs start at a loss of 137.75. With learning rate 0.12, the loss falls: 88.6, 57.1, 37.0, and 1.86 after twelve steps. With learning rate 0.14, the loss grows at every step: 166.4, 201.1, 243.1, and 1,346 after twelve steps. It grows by about 21% per step and never stops.

A diverging loss grows by the same factor at every step, so it quickly becomes very large. At learning rate 0.16, the loss of the ticket model grows past the largest number that a 32-bit number can store, about 3.4 × 10³⁸. That takes about 125 steps. The computer then stores it as inf (infinity). One more calculation with inf, such as inf − inf, gives nan: “not a number”. That is the loss: nan in a training log.

Convergence

A run converges when its steps become so small that the parameters and the loss stop changing. On the loss curve, the ball reaches the bottom. In the training log, the loss curve becomes flat. With learning rate 0.05, the ticket model converges after about 5 steps. With 0.01, it needs more than 30.

A flat curve does not always mean the best possible model. It means that this method, with these settings, has stopped improving. A team usually stops training when the loss on separate validation tickets stops falling. Many setups also lower the learning rate during training (a learning-rate schedule): long steps early, short steps near the bottom.

Local minima: the start matters

The bowl of the ticket model has one bottom. Loss curves of larger models are not bowls. A local minimum is a point where the loss is lower than everywhere near it, but not the lowest anywhere. At a local minimum, the gradient is 0, so gradient descent stops there. It cannot see the deeper valley.

The next explorer has a made-up loss curve with two valleys. You can change where the ball starts.

Explore

Two valleys: where does the ball stop?

Step

8

w

1.944

Loss

1.922

Slope (gradient)

−0.798

Next move = −rate × slope

0.08

Moving downhill. Each step is smaller, because the slope gets smaller near the bottom.

A made-up loss curve with a local minimum on the right and a deeper minimum on the left. Change the start, then press Play.

Text description of the explorer

The chart shows a made-up loss curve for w from −3.6 to 3.6. It has a deep valley near w = −2.5, a hump near w = 0.2, and a shallower valley near w = 2.3. The loss is about 0.1 in the deep valley and about 1.8 in the shallow one. The ball starts at w = 1. The learning rate starts at 0.1, and a second slider sets the start, from −3.4 to 3.4. From w = 1, the ball rolls right and stops in the shallow valley near w = 2.3: a local minimum, where the gradient is 0. From any start left of the hump, for example w = −0.5, it rolls into the deeper valley. The status line says when the ball stops at the local minimum.

What does this mean in practice? Models with millions of parameters have loss surfaces with very many dimensions, and a one-dimensional picture exaggerates the problem. Teams seldom fail because of a bad local minimum. But two runs that start from different random values can end at different models, with different results. So teams:

• fix the random seed, the number that sets the random start, so that they can repeat a run;
• train more than once with different seeds when a decision depends on a small difference between models.

Stochastic updates: why real loss curves are noisy

In the last lesson, the gradient was a mean over all four tickets. Larkfield has about 1,200 new tickets a week, and real models train on millions of examples. Computing the full gradient for every step is too slow.

So training uses a mini-batch: a small random group of examples, for example 32. Each step computes the gradient on one batch only, then makes the update. One pass through all training examples is an epoch. This method is stochastic gradient descent (SGD). “Stochastic” means “with randomness”.

A batch gradient is an estimate. Remember each ticket’s part of the gradient at w = 6:

Batch	Gradient on this batch	True gradient (all four)
T-3101 only	−2	−4
T-3102 only	−4	−4
T-3103 only	+6	−4
T-3104 only	−16	−4
T-3101 and T-3102	−3	−4
T-3103 and T-3104	−5	−4

A batch of one ticket can give a gradient four times too large (T-3104), or one that points the wrong way (T-3103). Together, the batches are right: the mean of the four one-ticket gradients is −4. Larger batches give estimates closer to the true gradient, but each step costs more.

Here is a run of SGD on the ticket model. Each step uses one ticket, the learning rate is 0.01, and each epoch puts the tickets in a new random order. A training log prints the loss on the batch, before the update:

• Loss on the batch (what the log prints)
• Loss on all four tickets

Steps 13 to 40 (epochs 4 to 10) of a stochastic run. The batch loss jumps up and down at every step, while the loss on all tickets stays close to its lowest value.

Text description of the chart

The horizontal axis shows steps 13 to 40. The batch loss, the loss on the one ticket of each step, jumps between about 0.37 and 8.7. For example, it is 2.58 at step 13, 8.72 at step 14 and 1.00 at step 15. The loss on all four tickets, after each step, is 3.13 at step 13 and then stays between about 1.22 and 1.95. Its lowest possible value is 1.217.

The batch loss jumps because each ticket has its own error, not because the model gets worse. T-3103 always has a large error with this simple model, so its steps show a high batch loss. That is why training dashboards show a noisy line, and often draw a smoothed version over it. Read the trend over many steps, not one step.

The noise has a second effect: w never settles exactly at the bottom. It keeps moving a little around 6.27. In the next explorer, you choose the batch size and watch both effects.

Explore

Change the batch size

Step

20

Epoch

5

Tickets in this batch

T-3104

Gradient on this batch

−11.959

True gradient (all four)

−2.106

w after the step

6.365

Step 20 used T-3104. The batch gradient is −11.96; the true gradient is −2.11. It is off by 9.85.

• Loss on the batch (what the log prints)
• Loss on all four tickets, and w
• Lowest loss (1.217) and best w (6.27)

Stochastic gradient descent on the four tickets, from w = 5.5 with learning rate 0.02. Each step uses one batch. Compare a batch of 1 ticket with a batch of all 4.

Text description of the explorer

The explorer runs stochastic gradient descent on the ticket model, from w = 5.5 with learning rate 0.02. Buttons choose 1, 2 or 4 tickets per batch; it starts at 1. Step makes one update, Play repeats it, and Reset goes back to w = 5.5. The readout shows the step, the epoch and the tickets in the batch. It also shows the gradient on the batch, the true gradient on all four tickets, and w after the step. The upper chart shows the loss after each step: the batch loss as dots, and the loss on all four tickets as a line. A dashed line marks the lowest loss, 1.217. The lower chart shows w after each step, with a dashed line at the best w, 6.27. With 1 ticket per batch, the first step uses T-3104 only: its gradient is −32, but the true gradient is −11.5. At step 3, T-3103 gives +9.14 while the true gradient is −1.38: the step goes the wrong way. From step 21 to 40, w keeps moving between about 6.01 and 6.49, and the batch loss jumps between 0.04 and 6.1. With 2 tickets per batch, w stays between about 6.16 and 6.37. With all 4 tickets, every step is exact, and w settles at 6.27.

Larger batches give smoother steps, but each step costs more computation. Real training uses batches of tens to thousands of examples, so it gets most of the smoothness at a small part of the cost. Many setups also lower the learning rate near the end of training, which makes the remaining noise smaller.

Worked example: Priya reads a training log

Tomás trains the real ticket classifier, a model with about 2 million parameters. After a few minutes, he shows Priya the log (fictional output):

Training log, run 1 (fictional)

epoch 1  step 100   lr 0.01   loss 0.693
epoch 1  step 200   lr 0.01   loss 0.912
epoch 1  step 300   lr 0.01   loss 4.85
epoch 1  step 400   lr 0.01   loss 961.2
epoch 1  step 500   lr 0.01   loss inf
epoch 1  step 600   lr 0.01   loss nan

Priya reads it from the top:

1. The loss grows at every report, and faster each time. This is the shape of divergence, as in the chart with learning rate 0.14.
2. inf, then nan. The loss became too large to store, and the next calculation gave “not a number”. From here, every parameter is probably nan too. Training cannot recover.
3. Diagnosis: the learning rate, 0.01, is too large for this model.

Tomás lowers the learning rate ten times, to 0.001, and starts again from the beginning:

Training log, run 2 (fictional)

epoch 1  step 100   lr 0.001  loss 0.693
epoch 1  step 200   lr 0.001  loss 0.481
epoch 1  step 300   lr 0.001  loss 0.402
epoch 1  step 400   lr 0.001  loss 0.377
epoch 2  step 500   lr 0.001  loss 0.351
epoch 2  step 600   lr 0.001  loss 0.358

The loss falls fast at first, then more slowly. It goes up a little at step 600. Priya says that is normal batch noise, not divergence: one higher value, not a growing trend. The run is healthy.

Guided practice: diagnose three runs

Each log shows the loss every 100 steps (fictional output). Write your diagnosis for each run in c03-notes.md.

Run A

loss 0.693  0.691  0.689  0.687  0.685  0.683

Run B

loss 0.693  0.512  0.388  0.341  0.330  0.327

Run C

loss 0.693  0.590  0.731  0.552  0.802  0.517

1. For each run, write whether the loss goes down, stays flat, goes up, or jumps.
2. For each run, write the most likely cause: learning rate too small, good, or too large.
3. For each run, write one action for Tomás.

Show an example answer

• Run A: the loss goes down by only 0.002 per report, a straight, slow line. The learning rate is probably too small. Action: try a learning rate 3 or 10 times larger and compare the first few hundred steps.
• Run B: the loss falls fast, then flattens near 0.33. A good run that is close to convergence. Action: keep it; check the loss on validation tickets to decide when to stop.
• Run C: the loss jumps up and down by a lot, and the high values get higher (0.731, then 0.802). The steps are too large: this is a zig-zag that may soon diverge. Action: lower the learning rate. Batch noise alone usually makes smaller jumps around a falling trend.

Check: for Run A, you did not write “the model has converged”. A slow, steady fall is not a flat curve.

Independent variation: a new learning rate

Tomás now trains with learning rate 0.003. The log shows a loss that falls fast, then a few higher values, then nan at step 2,000. Write two possible causes in c03-notes.md, and for each one, a check that tells them apart.

Show an example answer

1. The learning rate is still too large. The higher values before nan are a zig-zag that grows. Check: run again with 0.001 from the same start (same seed). If the nan goes away, the rate was the cause.
2. Bad data in one batch, for example a ticket with an empty or extremely long text that gives an inf value. Check: find the batch at step 2,000 and look at its examples. If the nan comes at the same step with every learning rate, the data is the likely cause.

Common mistake

The log shows loss: nan

Symptom: the training log shows loss: nan (or inf), and every prediction of the saved model is nan or the same value. Several hours of compute time are lost, and a model saved after that point is broken.

Cause: most often, a learning rate that is too large. Each step jumps further past the bottom, the loss grows until the computer cannot store it, and the next calculation gives nan. Sometimes, the data contains a value that is already inf or missing.

Fix: stop the run. Lower the learning rate, often 10 times, and restart from the start or from the last saved model without nan. If it happens again at the same step, check the batch at that step for broken data. Many teams also limit the size of each step, a method called gradient clipping.

Check your understanding

Gradient descent stops at a point where the gradient is 0. Explain why that point is not always the best model.

Show the answer

A gradient of 0 only means that the curve is flat here. It can be a local minimum: lower than everything near it, but higher than a deeper valley somewhere else. Gradient descent only sees the slope at its current point, so it cannot know about the other valley. The start decides which valley it reaches.

Check your understanding

In a stochastic run, the batch loss goes from 0.31 at one step to 0.52 at the next. Tomás wants to stop the run. What do you tell him?

Show the answer

One higher value is not a reason to stop. Each batch has different examples, so the batch loss jumps from step to step even when the model improves. Look at the trend over many steps, or at a smoothed curve. Stop only if the trend goes up, or if the loss on validation examples stops falling.

Module practice

You trace gradient descent on the ticket model by hand. For this model, the gradient is 15w − 94, and the loss is (1183 − 376w + 30w²) ÷ 4.

Worked example: learning rate 0.1

Priya starts at w = 2 with learning rate 0.1. For each step, she calculates the gradient at the current w, then the move (−0.1 × gradient), then the new w, then the loss at the new w:

Step	w before	Gradient 15w − 94	Move −0.1 × gradient	w after	Loss after
1	2	−64	+6.4	8.4	35.35
2	8.4	32	−3.2	5.2	9.75
3	5.2	−16	+1.6	6.8	3.35
4	6.8	8	−0.8	6.0	1.75

What the table shows:

• The sign of the gradient changes at every step, so w jumps from one side of 6.27 to the other: 2, 8.4, 5.2, 6.8, 6.0. This is a zig-zag.
• Each jump is half as long as the one before: 6.4, 3.2, 1.6, 0.8. The factor is 1 − 15 × 0.1 = −0.5.
• The loss still falls at every step: 137.75, 35.35, 9.75, 3.35, 1.75. A zig-zag in w can hide behind a smooth loss curve.

Your task: two other learning rates

Make a table like Priya’s in c03-notes.md for each run. Start at w = 2 each time.

1. Learning rate 0.02: make three steps.
2. Learning rate 0.15: make four steps.
3. For each run, write one sentence: is the learning rate too small, good, or too large, and what is the evidence in your table?
4. Check your tables in the first explorer of this lesson, at the same learning rates.

Show the model answer

Learning rate 0.02:

Step	w before	Gradient	Move	w after	Loss after
1	2	−64	+1.28	3.28	68.12
2	3.28	−44.8	+0.896	4.176	34.00
3	4.176	−31.36	+0.6272	4.8032	17.28

Too small, or at least slow. Every step goes in the right direction, but each one keeps 70% of the distance (1 − 15 × 0.02 = 0.7). After three steps, w is still far from 6.27, and the loss is still 17.28.

Learning rate 0.15:

Step	w before	Gradient	Move	w after	Loss after
1	2	−64	+9.6	11.6	214.55
2	11.6	80	−12	−0.4	334.55
3	−0.4	−100	+15	14.6	522.05
4	14.6	125	−18.75	−4.15	815.02

Too large. Each move is longer than the one before (9.6, 12, 15, 18.75). w jumps further from 6.27 on each side, and the loss grows at every step. The factor is 1 − 15 × 0.15 = −1.25, so each jump is 25% longer. The run diverges. A real run would reach inf and nan.

Run Python in your browser

Optional: check your traces

This box prints four steps of gradient descent on the ticket model. Run it with rate = 0.1 and compare with Priya’s table. Then change the rate to 0.02 and 0.15 (and the range for three or four steps) and check your answers.

Python code (you can change it)def gradient(w): """Slope of the loss of the ticket model at w.""" return 15 * w - 94 def loss(w): return (1183 - 376 * w + 30 * w ** 2) / 4 w = 2 # start rate = 0.1 # learning rate: try 0.02 and 0.15 print("step w gradient loss") for step in range(1, 5): g = gradient(w) w = w - rate * g print(f"{step:<6} {w:<9.4f} {g:<10.2f} {loss(w):.2f}")

Module check

Check your understanding

What does one step of gradient descent do? Write the update rule in words and name each part.

Show the answer

One step changes each parameter a little, to make the loss smaller. The rule: new w = w − learning rate × gradient. The gradient is the slope of the loss at the current w: it gives the direction and the steepness. The minus sign makes the step go downhill. The learning rate, chosen by a person, sets how long the step is.

Check your understanding

Larkfield trains a classifier on 1,200 tickets with mini-batches of 40 tickets. How many steps are one epoch? After 3 epochs, the batch loss still jumps between 0.30 and 0.45 at each step. Predict: is something wrong?

Show the answer

One epoch is 1,200 ÷ 40 = 30 steps, so 3 epochs are 90 steps. Jumps in the batch loss are normal in stochastic gradient descent, because each batch has different tickets. Nothing is wrong if the trend over many steps goes down or stays flat at a low value. Check a smoothed curve, or the loss on all validation tickets, before you decide.

Check your understanding

Completion check: explain why a step size can be too small or too large. Use the ticket model as your example.

Show the answer

The update rule moves w by the learning rate times the gradient. The learning rate sets the step size.

• Too small (for example 0.01): each step goes the right way but covers only a small part of the distance to the bottom (15% here). The run needs many steps, and costs time and compute. In a log, the loss falls slowly and steadily.
• Too large (above 0.133 for this model, for example 0.15): each step jumps over the bottom and lands further away on the other side. The loss grows at every step, the run diverges, and the log ends with inf and nan. Between 0.067 and 0.133, the steps also jump over the bottom, but they shrink, so the run zig-zags to the bottom.
• Good (for example 0.05): large steps far from the bottom, where the gradient is large, and small steps near it. The loss reaches its lowest value in a few steps.

The limits depend on how steep the loss curve is. For a real model, you cannot calculate them, so you try several rates and read the loss curves.

Recap

In this module, you watched a model learn:

• A model is a function with parameters. The slope of a function is the change in output per unit of input, and the derivative is the slope at one point.
• The loss measures how wrong the model is. The gradient is the slope of the loss, and gradient descent moves each parameter against its gradient, one step at a time.
• The learning rate sets the step size. Too small is slow; too large jumps over the bottom, zig-zags or diverges to nan.
• A training loss curve shows the result. Steep then flat is healthy. Slow and straight suggests a small rate. Rising suggests a rate that is too large.
• Gradient descent can stop in a local minimum. The random start decides which valley it finds.
• Real training uses mini-batches. Batch gradients are noisy estimates, so real loss curves are noisy. Read the trend.

Next, in Module 5, you check whether a result measured on a sample of tickets is real, or only chance.

Next: Samples and populations