Tensors and the training loop

01 · 1The one idea

Every model in this course — a digit classifier, a tiny GPT (generative pre-trained transformer), a speech recogniser — learns the same way. There are only three moving parts:

1. A pile of adjustable numbers, called parameters or weights.
2. A single number that says how wrong the model currently is, called the loss.
3. A rule that changes every parameter a little, so that the loss goes down.

That is the whole thing. Everything else — attention, convolutions, quantization — is detail about the first item. The learning itself never changes.

It is worth knowing where this goes. The course ends with an offline voice assistant that runs entirely on a phone, with a median voice-to-voice reply under 1.5 seconds. It listens, understands, answers and speaks, with no server. Every model inside it was trained by the rule in this lesson.

Start with the smallest possible version: one parameter. Call it w. For each value of w the model makes some error, so we can draw a curve of loss against w. We cannot see this curve — the model has millions of parameters and we can only afford to look at one point at a time. But we can always ask one question: which way is downhill from here?

Two things in that animation are worth saying out loud, because they explain most training problems you will ever have.

The steps get smaller on their own. Near the bottom the curve is almost flat, so the slope is small, so the step is small. Nobody programmed that. It is a consequence of the rule itself.

The step size is a choice, and it is the wrong choice by default. We multiply the slope by a small number called the learning rate. Too small and training takes a week. Too large and the model jumps over the valley and the loss grows instead of shrinking. Eventually the loss becomes nan.

nan is short for not a number. It is a value that floating-point arithmetic produces when a number grows past the largest one it can store, or when zero is divided by zero. It is contagious: anything you add to a nan or multiply by a nan becomes nan too. So one overflowing gradient poisons every parameter it touches, usually within a single step. When you see it, check the learning rate first.

01 · 2Tensors: shapes are the real API

A tensor is an array of numbers with a shape. That is all. A single number is shape (). A list of 10 numbers is (10,). A greyscale image is (28, 28). A batch of 32 such images is (32, 1, 28, 28): batch, channels, height, width.

Three of those four words are obvious. Channels is not. A channel is one full-size grid of numbers per image. A grey picture needs only one, holding brightness, which is the 1 in the shape above. A colour picture needs three — red, green and blue — so the same picture in colour would be (3, 28, 28), and a batch of 32 of them (32, 3, 28, 28).

When you write model code, you spend maybe 10% of your attention on the maths and 90% on shapes. Almost every error you will hit is a shape error. So build the habit now: after every line, know the shape.

The first dimension is nearly always the batch: how many independent examples travel through the network together. The network treats them in parallel and they never mix. Batching exists mainly for speed: hardware is far faster on 32 rows at once than on one row 32 times. It does have a small effect on the maths. A bigger batch averages more examples, so its gradient is smoother and less noisy. Lesson 04 returns to how a training batch size is really chosen — there the number that matters is tokens per step, not the batch on its own.

import torch

x = torch.randn(32, 1, 28, 28)      # 32 grey images, 28 x 28
print(x.shape)                       # torch.Size([32, 1, 28, 28])

flat = x.view(32, -1)                # -1 means "work it out": 1*28*28 = 784
print(flat.shape)                    # torch.Size([32, 784])

W = torch.randn(784, 128)
h = flat @ W                         # (32, 784) @ (784, 128) -> (32, 128)
print(h.shape)                       # torch.Size([32, 128])

The rule for @ (matrix multiply) is short: the last dimension of the left side must equal the first dimension of the right side, and those two disappear. (32, 784) @ (784, 128) becomes (32, 128).

A network stacks these layers, but stacking alone buys nothing. Multiply by A, then multiply by B, and the result is just a multiply by A @ B — one single matrix. Two layers would have exactly the power of one. Check it on the smallest case: with A = [[1, 2], [0, 1]] and B = [[1, 0], [3, 1]], the pair collapses to A @ B = [[7, 2], [3, 1]]. One matrix, same answers.

So we put a non-linear function between the layers. The simplest one is ReLU, the rectified linear unit: negative values become 0, positive values pass through unchanged. That single bend is enough to stop the collapse, because you cannot write "multiply, bend, multiply" as one matrix. Everything a network can express beyond a straight line comes from these bends. Lesson 02 builds one from scratch and shows the difference on real data.

01 · 3Autograd: the tape that remembers

We need the gradient of the loss with respect to every parameter. Doing that by hand for a million parameters is not possible. PyTorch does it for you, with a mechanism called autograd — short for automatic differentiation.

While you compute the forward pass, PyTorch writes down every operation you performed, in order, in a graph. When you then call loss.backward(), it walks that list backwards. Each operation knows its own local derivative, and the chain rule multiplies them together on the way back.

Notice the direction of the colours. The data colour moves forward, left to right. The gradient colour moves backward, right to left. Six colours carry meaning across all thirteen lessons, and this legend is the authority on them. In Parts IV and V nothing is being trained, so there the loss colour reads more broadly as "cost or mistake".

Activations share the data colour, because they are data: the intermediate tensors each layer produces on the way forward. They are not free. PyTorch has to keep them in memory until backward() has used them, which is why the memory cost of training depends on batch size and depth, not only on parameter count.

import torch

w = torch.tensor(3.0, requires_grad=True)   # a parameter we want to learn
x = torch.tensor(2.0)                       # data: no gradient needed

y = w * x          # 6.0     ... and PyTorch remembers "this was a multiply"
loss = (y - 10) ** 2   # 16.0 ... and "this was a subtract, then a square"

loss.backward()    # walk the tape backwards

print(w.grad)      # tensor(-16.)   dL/dw = 2*(y-10)*x = 2*(-4)*2 = -16

Where does -16 come from? The chain rule is three small numbers multiplied together, and this example is small enough to do in your head. Give the error its own name first: u = y - 10, which is 6 - 10 = -4. The loss L is then u², which is 16. Now walk backwards, one operation at a time.

1. The last operation was squaring: L = u². Nudge u up a little and L changes about 2u times as much. With u = -4, that local derivative is 2 × (-4) = -8.
2. Before that we subtracted 10: u = y - 10. Subtracting a constant does not change how fast things move, so its local derivative is 1, and ∂L/∂y = -8 × 1 = -8.
3. Before that we multiplied: y = w × x, with x = 2. Nudge w by a little and y moves twice as much, so multiply once more: -8 × 2 = -16.

That is the whole chain rule. Each operation knows only its own local derivative, and the answer is the product of them along the path. The gradient is -16: negative, so increasing w would decrease the loss. Gradient descent therefore moves w in the opposite direction of the gradient — that is where the minus sign in the update rule comes from.

Now use it. Take lr = 0.1 and the gradient we just computed, and turn one step of the handle. You can check every line of this on a phone calculator.

One step, and the loss fell from 16.00 to 0.64. That is training. Everything that follows in this course is this same step, repeated a few million times, on a few million parameters at once.

One more thing is visible in those numbers. The best w here is 5, and one step took us from 3 to 4.60 — most of the way, but not all of it. That undershoot is what a safe learning rate looks like. Each step with lr = 0.1 removes 80% of whatever error is left, so w creeps up on 5 and never jumps past it. Now push the rate up. At lr = 0.125 the step lands exactly on 5. At lr = 0.5 the same step sends w to 11, which is further from 5 than we started, and the loss rises from 16 to 144. That is the animation in section 1, in arithmetic.

01 · 4Loss: turning "wrong" into one number

The loss must be a single number, because we can only walk downhill on a single surface. Which number you choose defines what the model actually learns.

Cross-entropy is the one that matters for the rest of this course, because language models and speech models are both category predictors. A language model predicts “which token comes next” out of ~50,000 options, thousands of times per sentence.

It works in two moves. First softmax turns raw scores (called logits) into probabilities that add up to 1. Then the loss is simply −log(probability of the correct class).

One detail decides whether you can check the numbers: log here is the natural logarithm, base e, written ln on most calculators. Every deep learning framework means this one when it writes log. So −ln(0.10) = 2.30, not the 1.00 you would get from the base-10 button.

Read the loss values in that animation and the intuition comes with no extra work:

That last row is the important one. Cross-entropy punishes confident mistakes far more than uncertain ones. It also tells you where a language model's loss curve starts. A fresh model spreads its probability evenly over all V tokens in its vocabulary — V is the vocabulary size — so the loss begins at ln(V). That is about 10.8 for a production-size vocabulary of 50,000 tokens. The tiny GPT you train in Lesson 04 uses a smaller vocabulary of 4,096, so its curve starts near ln(4096) = 8.32 and falls from there.

01 · 5Optimizers: better steps, same idea

Plain gradient descent has a well-known weakness. Real loss surfaces are not round bowls; they are long narrow valleys. The gradient points mostly across the valley, not along it, so the parameters cross the valley again and again and creep forward slowly.

The plain rule from section 1 has a name: SGD, stochastic gradient descent. "Stochastic" means each step is measured on one random batch rather than on all the data, so every gradient is a noisy estimate of the true one. It is the baseline in the figure below, and everything else on this page is an improvement on it.

The figure needs one word of warning, because the picture changes. Up to now the loss has been a curve, because there was one parameter. With two parameters the loss is a landscape. The usual way to draw a landscape flat is to look down on it from above, joining every point of equal height with a ring. That is exactly a hiking map of a hill. So the rings below are lines of equal loss, and the centre of the rings is the bottom of the valley.

Momentum keeps a running average of past gradients. The side-to-side parts cancel out and the forward part adds up, like a heavy ball that does not react to every bump.

Adam adds a second idea. It tracks how big each parameter's recent gradients have been — their typical size, ignoring the sign — and divides the step by that. Parameters with tiny gradients get bigger steps; parameters with wild gradients get smaller ones. This is why Adam works with almost no tuning, and why it costs memory: it stores two extra numbers per parameter.

Those two extra numbers are not a small detail, and this is a good moment to see how quickly model memory adds up. To count bytes you need one table, which the rest of the course will lean on constantly.

Use AdamW as your default. It is Adam with weight decay applied correctly, and it is what nearly every model in this course was trained with. Weight decay is defined in section 7, with the other names for pulling a model back from memorising.

01 · 6The training loop

Here it is. Five lines, and you will write them hundreds of times. One word first: one epoch is one complete pass over the training data. Three epochs means the model has seen every training image three times.

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

device = ("cuda" if torch.cuda.is_available()
          else "mps" if torch.backends.mps.is_available()
          else "cpu")                    # section 8 explains this line

train_set = datasets.MNIST("data", train=True, download=True,
                           transform=transforms.ToTensor())
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)

model = nn.Sequential(
    nn.Flatten(),                # (64, 1, 28, 28) -> (64, 784)
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),          # ten scores, one per digit
).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
epochs = 3

for epoch in range(epochs):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()            # 1. forget last batch's gradients
        logits = model(images)           # 2. forward: build the tape
        loss = criterion(logits, labels) # 3. one number that says "how wrong"
        loss.backward()                  # 4. walk the tape back, fill .grad
        optimizer.step()                 # 5. nudge every parameter downhill

The line that ties it all together is model.parameters(). That call hands the optimizer a list of exactly the tensors that backward() will fill with .grad. From then on the optimizer holds those same tensors. That is why step() needs no arguments: it already knows every parameter it may move, and where to read each gradient. Get this handover wrong — build the optimizer from a different model — and training will run happily and change nothing.

The order of the five lines is not a style choice. backward() needs the graph that model(images) built. step() needs the .grad values that backward() filled in.

And zero_grad() is first for a reason that surprises everyone: PyTorch adds gradients instead of replacing them. If you forget the line, batch 2 is trained on the sum of batches 1 and 2, batch 3 on the sum of 1, 2 and 3, and your effective learning rate grows without limit. Training does not crash. It just quietly gets worse, which is much harder to debug.

Evaluation is a different mode

val_set = datasets.MNIST("data", train=False, download=True,
                         transform=transforms.ToTensor())
val_loader = DataLoader(val_set, batch_size=64)

model.eval()                       # layers that act differently in training switch over
correct = 0
with torch.no_grad():              # do not build the tape: less memory, faster
    for images, labels in val_loader:
        preds = model(images.to(device)).argmax(dim=1)
        correct += (preds == labels.to(device)).sum().item()
print(correct / len(val_loader.dataset))

Two separate switches, and people mix them up. model.eval() changes the behaviour of the few layer types that deliberately act differently while training. The one you will meet in this course is dropout, defined in section 7: it switches off random values while training, and it has to be off when you measure. The model above has no dropout, so the call changes nothing yet. Write it anyway, so that the line is already there on the day you add some. torch.no_grad() stops recording the tape. You almost always want both when measuring, and neither when training.

One line in there is the moment the model finally answers a question, and it is easy to read past. argmax(dim=1) says: along dimension 1 — the ten class scores, since dimension 0 is the batch — give me the position of the largest value. On a (64, 10) tensor that returns 64 numbers, each between 0 and 9, and that is the predicted digit. No softmax is needed here: softmax never changes which score is largest, so the biggest logit is always the biggest probability.

01 · 7Reading a loss curve

Training loss going down means the model is memorising the training data. It does not mean the model is any good. That is why you always hold out data the model never trains on.

Held-out data comes in two roles, and the names matter. Validation data is held out so you can decide things: when to stop, which learning rate, how many layers. You look at it often. Test data is held out and looked at once, at the very end, so that the number you report has not been tuned against. Every time you use a split to make a decision, it stops being a test set and becomes a validation set.

Four shapes, four diagnoses. Learn these and you will save yourself days:

That last column is full of words this course has not used yet. Here is what each one means, in one line, so you can act and not only diagnose:

• Stop early — keep the weights from before the validation curve turned upward, and throw away the later ones.
• Augmentation — make new training examples out of old ones by shifting, rotating or cropping them, so memorising gets harder.
• Dropout — during training, switch off a random fraction of the values inside the network on every batch. No single unit can be relied on, so the model has to spread the work out.
• Weight decay — pull every weight slightly toward zero on every step. Large weights have to earn their size.
• Clip gradients — cap the size of a gradient before the step, so one bad batch cannot throw the weights across the map.
• Capacity — how much the model is able to represent, which in practice means how many parameters it has. Raising capacity means a wider or deeper network.

You already used weight decay in this lesson without noticing: it is the W in AdamW. Dropout returns in Lesson 05, where the fine-tuning config sets lora_dropout=0.05. For now, recognising the names on the shelf is enough.

01 · 8Your machine

You do not need a data-centre GPU for this course. You need a device string that is chosen at runtime, so the same code runs on your laptop and on a rented GPU.

import torch

device = (
    "cuda" if torch.cuda.is_available()          # NVIDIA GPU
    else "mps" if torch.backends.mps.is_available()  # Apple Silicon
    else "cpu"
)
print(device)

torch.manual_seed(0)          # same random numbers every run
model = MyModel().to(device)  # move parameters to the device

mps is Apple's GPU backend. On an M-series Mac it is typically several times faster than the CPU for the models in this course, and everything in Lessons 01 to 07 will run on it. A few operations still fall back to the CPU; set PYTORCH_ENABLE_MPS_FALLBACK=1 if you hit one.

Both the data and the model must be on the same device. Half of all beginner PyTorch errors are a variant of "expected all tensors to be on the same device".

Check yourself

Go deeper