Training: loss, backpropagation and optimisers

The training loop

1. Forward pass: push a mini-batch of examples through the network to get predictions.
2. Loss: compare predictions with the true labels. Mean squared error for numbers; cross-entropy for classes.
3. Backward pass (backpropagation): compute the gradient of the loss with respect to every weight.
4. Optimiser step: update every weight a little in the direction that lowers the loss (gradient descent from module 5).
5. Repeat for every mini-batch; one pass over all the data is an epoch.

Cross-entropy loss

For classification, the network outputs probabilities with softmax. Cross-entropy loss for one example is −ln(probability given to the correct class). If the right class gets 0.9, the loss is about 0.105. If it gets 0.05, the loss is about 3.0. So the network is pushed to put high probability on the right answer. Binary cross-entropy (log loss, module 5) is the two-class version.

Backpropagation in plain words

A deep network has millions of weights, and changing any one of them changes the loss a tiny bit. Backpropagation computes all these effects efficiently by going backwards from the loss, layer by layer, using the chain rule of calculus: the effect of a weight on the loss equals its effect on the next layer times that layer's effect on the loss, and so on. You never do this by hand in practice; frameworks like PyTorch record the forward pass and compute gradients automatically (autograd).

One tiny worked step

Take one neuron with no activation: prediction = w·x + b, with w = 0.5, b = 0, input x = 2 and target y = 3. Prediction is 1, error is 1 − 3 = −2, and the squared-error loss is 4. The gradients are: d(loss)/dw = 2 × error × x = −8, and d(loss)/db = 2 × error = −4. With learning rate 0.05: w becomes 0.5 − 0.05 × (−8) = 0.9 and b becomes 0 − 0.05 × (−4) = 0.2. The new prediction is 0.9 × 2 + 0.2 = 2.0, closer to 3.

Optimisers

Plain stochastic gradient descent (SGD) uses the same learning rate for every weight. SGD with momentum keeps a running direction, like a ball rolling downhill, which speeds up and smooths training. Adam adapts the step size for each weight based on recent gradients and is a popular default. Learning-rate schedules lower the rate as training goes on.

import torch
from torch import nn

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 3))
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    for X_batch, y_batch in train_loader:
        optimizer.zero_grad()            # clear old gradients
        loss = loss_fn(model(X_batch), y_batch)
        loss.backward()                  # backpropagation
        optimizer.step()                 # update the weights

Keeping networks from overfitting

• Dropout: during training, randomly switch off a share of neurons (say 20%) in each step, so the network cannot rely on any single path.
• Weight decay (L2 regularisation): penalise large weights.
• Early stopping: watch the validation loss and stop when it stops falling.
• Data augmentation: for images, train on flipped, cropped or slightly rotated copies, so the network sees more variety.
• More data, or a smaller network.

Common mistakes

• Forgetting optimizer.zero_grad(), so gradients pile up from previous batches.
• Applying softmax before nn.CrossEntropyLoss, which already includes it.
• Only watching training loss. Plot validation loss too; when it rises while training loss falls, you are overfitting.
• Vanishing gradients in very deep networks with sigmoid activations; ReLU and careful initialisation help.