Lesson 6: Training a network — Loss & Backpropagation

In Lesson 5 we assembled a complete multi-layer MLP by hand and ran a forward pass to produce raw predictions. Untrained, though, those predictions are meaningless. For the model to actually learn the pattern, we have to give it a measure of its error and a mechanism for correcting itself.

This lesson walks you through setting up the loss function, understanding the backpropagation algorithm properly — it rests on the chain rule — and the adaptive Adam optimiser, to complete a professional training loop.

6.1 The loss function — measuring how wrong we are

A short reminder from Lesson 2: a loss function takes the model's prediction ($\hat{y}$, read "y-hat") together with the true label ($y$), and returns one number — how wrong the model is. Lower is better.

What is new in this lesson is which function to choose. Lesson 2 used a simple squared function because it only needed to illustrate the principle. On a real problem, choosing the wrong loss makes the model learn very slowly or not at all, and this section shows why — with numbers, not advice.

Two functions cover almost every real case:

• MSE loss (mean squared error): for regression problems. It measures the mean squared distance: $$L = \frac{1}{N} \sum_{i=1}^N (y_i - \hat{y}_i)^2$$ Seen probabilistically, minimising MSE is equivalent to maximum likelihood estimation (MLE) under the assumption that the model's errors follow a Gaussian distribution with mean zero.
• Cross-entropy loss (BCE & categorical cross-entropy): for classification problems. This measures the distance between two probability distributions (the true distribution $y$ and the predicted $\hat{y}$): $$L = -\frac{1}{N} \sum_{i=1}^N \left( y_i \log(\hat{y}_i) + (1 - y_i) \log(1 - \hat{y}_i) \right)$$ The idea comes from Claude Shannon's information theory. By minimising cross-entropy we are in effect minimising the Kullback-Leibler divergence between the model's predicted distribution and the real data distribution.

Let's prove that with concrete numbers rather than trusting the algebra above:

import math

def sigmoid(z):
    return 1 / (1 + math.exp(-z))

# The situation: the true label is 1, but the logit is very negative, so the
# model is predicting almost exactly the WRONG answer.
y_true = 1
z = -8.0                  # raw logit, before Sigmoid
y_pred = sigmoid(z)       # ~0.000335 — confidently wrong

# Gradient of MSE loss w.r.t. z: dL/dz = (y_pred - y_true) * sigmoid'(z)
sigmoid_derivative = y_pred * (1 - y_pred)
grad_mse = (y_pred - y_true) * sigmoid_derivative

# Gradient of cross-entropy w.r.t. z: dL/dz = y_pred - y_true.
# Note what is NOT here: the sigmoid'(z) factor.
grad_ce = y_pred - y_true

print(f"almost entirely wrong: y_pred = {y_pred:.6f} (true label = {y_true})")
print(f"MSE gradient (dL/dz):           {grad_mse:.6f}")
print(f"cross-entropy gradient (dL/dz): {grad_ce:.6f}")
# MSE gradient           -0.000335  -> TINY, the model barely learns
# cross-entropy gradient -0.999665  -> LARGE, the model corrects fast

For the same near-total misprediction, MSE's gradient is about 0.03% of cross-entropy's — that is the concrete figure behind the "gradient saturation" described above, and the reason essentially no real classification problem still uses MSE as its main loss.

6.2 Backpropagation & gradient descent optimisation

Once we have the loss at the final layer, how do we carry that error information backwards to the deeper hidden layers so their weights can be updated?

The backpropagation algorithm solves this by applying the chain rule of differentiation.

Before the formulas, look at the shape of what is happening. The forward pass runs left to right to produce a prediction; the backward pass runs right to left carrying the error information home:

Read the diagram in both directions. Green: $x$ goes through a matrix multiply and then an activation, twice, producing $\hat{y}$, which is compared with the true label to give $L$. Red: starting from $L$ and going back, each layer receives the share of the error that belongs to it — that is the entire meaning of "backward propagation". The two yellow arrows show which weight matrix each share of the error is used to fix.

Now the same content written as formulas, for a 2-layer MLP (input $X$, hidden layer $Z^{[1]}$ with ReLU, output layer $Z^{[2]}$ with Sigmoid). You do not need to derive any of this yourself — PyTorch does it for you. Read it to convince yourself there is no magic underneath loss.backward():

1. Forward: $$Z^{[1]} = W^{[1]} X + b^{[1]}$$ $$A^{[1]} = \text{ReLU}(Z^{[1]})$$ $$Z^{[2]} = W^{[2]} A^{[1]} + b^{[2]}$$ $$A^{[2]} = \sigma(Z^{[2]}) = \hat{Y}$$
2. Backward: Error at the final layer: $$\delta^{[2]} = A^{[2]} - Y$$ Derivatives for layer 2's weights and bias: $$\frac{\partial L}{\partial W^{[2]}} = \frac{1}{m} \delta^{[2]} (A^{[1]})^T$$ $$\frac{\partial L}{\partial b^{[2]}} = \frac{1}{m} \sum_{\text{samples}} \delta^{[2]}$$ Carry the error back to the hidden layer: $$\delta^{[1]} = \left( (W^{[2]})^T \delta^{[2]} \right) \odot \sigma_{\text{ReLU}}'(Z^{[1]})$$ Derivatives for layer 1's weights and bias: $$\frac{\partial L}{\partial W^{[1]}} = \frac{1}{m} \delta^{[1]} X^T$$ $$\frac{\partial L}{\partial b^{[1]}} = \frac{1}{m} \sum_{\text{samples}} \delta^{[1]}$$

With the gradients computed ($\nabla_W L$), the optimiser updates the weights in the direction opposite the gradient, reducing the error step by step:

6.3 Better optimisers: Adam & the learning rate

Plain stochastic gradient descent (SGD) updates every weight with the same fixed learning rate $\alpha$. That makes it easy for the model to get stuck in local minima or at saddle points — places where the gradient is zero without being a minimum or maximum.

From SGD to Adam:

• Momentum: models a physical ball rolling downhill. It adds a fraction of the previous steps' gradients into the current step, letting the model cross flat regions or narrow valleys quickly.
• RMSprop: adjusts the learning rate automatically by dividing the gradient by the square root of a running average of squared gradients. A weight that swings violently gets slowed down; one moving sluggishly gets sped up.
• Adam (adaptive moment estimation): the combination of momentum and RMSprop, tracking both the first moment (momentum) and the second moment (adaptive variance) of the gradient:

To see exactly what those four steps do, here is Adam written in pure Python (purely to ILLUSTRATE the mechanism — torch.optim.Adam does this far more efficiently):

def adam_step(theta, grad, m, v, t, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8, correct_bias=True):
    # 1. First moment: a running average of the gradient (this is momentum).
    m = beta1 * m + (1 - beta1) * grad
    # 2. Second moment: a running average of the SQUARED gradient (this is RMSprop).
    v = beta2 * v + (1 - beta2) * (grad ** 2)
    # 3. Bias correction. m and v both start at 0, so early on they are far too
    #    small; dividing by (1 - beta**t) compensates. Switchable here so we can
    #    measure what happens without it.
    if correct_bias:
        m_hat, v_hat = m / (1 - beta1 ** t), v / (1 - beta2 ** t)
    else:
        m_hat, v_hat = m, v
    # 4. The update itself.
    theta_new = theta - lr * m_hat / (v_hat ** 0.5 + eps)
    return theta_new, m, v

# Same constant gradient of 0.5, same 4 steps, with and without bias correction.
for label, correct in [("with bias correction", True), ("without bias correction", False)]:
    print(f"--- {label} ---")
    theta, m, v = 1.0, 0.0, 0.0
    previous = theta
    for t in range(1, 5):
        theta, m, v = adam_step(theta, grad=0.5, m=m, v=v, t=t, correct_bias=correct)
        print(f"step {t}: theta = {theta:.6f}  (moved {previous - theta:.6f})")
        previous = theta

--- with bias correction ---
step 1: theta = 0.999000  (moved 0.001000)
step 2: theta = 0.998000  (moved 0.001000)
step 3: theta = 0.997000  (moved 0.001000)
step 4: theta = 0.996000  (moved 0.001000)
--- without bias correction ---
step 1: theta = 0.996838  (moved 0.003162)
step 2: theta = 0.992588  (moved 0.004250)
step 3: theta = 0.987638  (moved 0.004950)
step 4: theta = 0.982196  (moved 0.005442)

The real numbers above show something worth pausing on: dropping bias correction does not make the first update timid, as simple intuition suggests. With $\beta_1=0.9$ and $\beta_2=0.999$ (the actual defaults) it makes the step MUCH larger than designed right from step 1 — about three times — and it keeps swelling with every step instead of settling. That is why bias correction is always on by default in every real Adam implementation, torch.optim.Adam included, rather than being an optional detail you can skip.

6.4 The training loop in full

Below is the complete Python source that trains an MLP on the non-linear two-concentric-circles dataset. It runs 500 epochs straight through using nn.BCELoss and the Adam optimiser, and draws an ASCII loss chart in the terminal so you can watch the convergence.

# train_circles.py
# Lesson 6: Training a network — Loss & Backpropagation
# Practical AI Engineer series
#
# Run it with:  python train_circles.py
# Requires:     pip install torch numpy      (matplotlib optional)
#
# A complete training loop on the two-concentric-circles dataset from Lesson 5.
# The four lines that do the actual learning are marked STEP 1..4 below; every
# training loop you will ever write is those same four lines in that same order.

import math

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim


def generate_concentric_circles(n_samples=1200, noise=0.05, factor=0.5):
    """Two concentric rings: outer labelled 0, inner labelled 1.

    Not linearly separable — no straight line splits them — which is exactly why
    the network needs its non-linear activations from Lesson 5.
    """
    np.random.seed(42)
    n_out = n_samples // 2
    n_in = n_samples - n_out

    theta_out = np.linspace(0, 2 * np.pi, n_out)
    X_out = np.vstack(
        (
            np.cos(theta_out) + np.random.normal(0, noise, n_out),
            np.sin(theta_out) + np.random.normal(0, noise, n_out),
        )
    ).T

    theta_in = np.linspace(0, 2 * np.pi, n_in)
    X_in = np.vstack(
        (
            factor * np.cos(theta_in) + np.random.normal(0, noise, n_in),
            factor * np.sin(theta_in) + np.random.normal(0, noise, n_in),
        )
    ).T

    X = np.vstack((X_out, X_in))
    y = np.concatenate((np.zeros(n_out), np.ones(n_in)))

    # Shuffle, so the train/test split below does not put every inner-ring point
    # in the test set.
    idx = np.arange(n_samples)
    np.random.shuffle(idx)
    return X[idx], y[idx]


class SimpleMLP(nn.Module):
    def __init__(self, input_dim=2, hidden_dim=8, output_dim=1):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
            nn.Sigmoid(),  # squashes the output into 0..1 so BCELoss can read it
        )
        self._initialise_weights()

    def _initialise_weights(self):
        # Kaiming init, for the symmetry-breaking reason covered in Lesson 5.
        for m in self.network:
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0.0)

    def forward(self, x):
        return self.network(x)


def draw_ascii_loss_chart(losses, epochs):
    """Bar chart of the loss curve, on a LOGARITHMIC scale.

    A linear scale is useless here: the loss falls from 0.69 to 0.0003, so after
    the first two rows every bar rounds down to zero characters and the chart
    goes blank exactly where the interesting part is. Log scale keeps the whole
    decay visible.
    """
    print('\n=== Loss curve (log scale — each bar is an order of magnitude) ===')
    lo, hi = math.log10(min(losses)), math.log10(max(losses))
    span = hi - lo or 1.0
    for epoch, loss in zip(epochs, losses):
        frac = (math.log10(loss) - lo) / span
        bar = '█' * max(1, int(frac * 40))
        print(f'Epoch {epoch:4d} | loss {loss:.4f} | {bar}')


def train():
    # Seed BOTH generators. numpy seeds the data; torch seeds the initial weights.
    # Without the torch seed the loss numbers differ on every run, which makes the
    # output in the lesson impossible to compare against.
    torch.manual_seed(42)

    X_np, y_np = generate_concentric_circles(n_samples=1000)

    # Hold back 20% the model never trains on, so the final number means something.
    split = 800
    to_t = lambda a: torch.tensor(a, dtype=torch.float32)
    X_train, X_test = to_t(X_np[:split]), to_t(X_np[split:])
    # unsqueeze(1): (800,) -> (800, 1), the shape BCELoss expects
    y_train = to_t(y_np[:split]).unsqueeze(1)
    y_test = to_t(y_np[split:]).unsqueeze(1)

    model = SimpleMLP()
    criterion = nn.BCELoss()  # binary cross-entropy, for a two-class problem
    optimizer = optim.Adam(model.parameters(), lr=0.01)

    epochs = 500
    sampled_losses, sampled_epochs = [], []

    print('=== Training an MLP on the concentric-circles data ===')
    for epoch in range(1, epochs + 1):
        model.train()

        predictions = model(X_train)  # STEP 1: forward pass
        loss = criterion(predictions, y_train)  # STEP 2: how wrong are we
        optimizer.zero_grad()  # STEP 3a: clear last round's gradients
        loss.backward()  # STEP 3b: backpropagate
        optimizer.step()  # STEP 4: nudge every weight

        if epoch == 1 or epoch % 50 == 0:
            sampled_losses.append(loss.item())
            sampled_epochs.append(epoch)
            hits = ((predictions >= 0.5).float() == y_train).float().mean().item() * 100
            print(f'Epoch {epoch:4d}/{epochs} | loss {loss.item():.4f} | train accuracy {hits:.2f}%')

    # eval() + no_grad(): stop tracking gradients, and switch layers that behave
    # differently at inference time. Neither matters for this small model, but
    # leaving them out of a real model is a bug that is hard to spot.
    model.eval()
    with torch.no_grad():
        test_preds = model(X_test)
        test_loss = criterion(test_preds, y_test).item()
        test_acc = ((test_preds >= 0.5).float() == y_test).float().mean().item() * 100

    print('\n=== After training ===')
    print(f'test loss     {test_loss:.4f}')
    print(f'test accuracy {test_acc:.2f}%')

    draw_ascii_loss_chart(sampled_losses, sampled_epochs)

    try:
        import matplotlib.pyplot as plt

        plt.figure(figsize=(8, 5))
        plt.plot(sampled_epochs, sampled_losses, marker='o', color='gold', label='loss')
        plt.yscale('log')  # same reason as the ASCII chart above
        plt.title('Loss decay — MLP classifying concentric circles')
        plt.xlabel('epoch')
        plt.ylabel('loss (log scale)')
        plt.grid(True)
        plt.legend()
        plt.savefig('loss_chart.png')
        print("\n[note] also saved a high-resolution chart to 'loss_chart.png'")
    except ImportError:
        pass  # matplotlib is optional; the ASCII chart above is enough


if __name__ == '__main__':
    train()

Running it produces exactly this (middle trimmed):

=== Training an MLP on the concentric-circles data ===
Epoch    1/500 | loss 0.7111 | train accuracy 49.75%
Epoch   50/500 | loss 0.4149 | train accuracy 85.50%
Epoch  100/500 | loss 0.0194 | train accuracy 100.00%
...
Epoch  500/500 | loss 0.0004 | train accuracy 100.00%

=== After training ===
test loss     0.0008
test accuracy 100.00%

=== Loss curve (log scale — each bar is an order of magnitude) ===
Epoch    1 | loss 0.7111 | ████████████████████████████████████████
Epoch   50 | loss 0.4149 | █████████████████████████████████████
Epoch  100 | loss 0.0194 | ████████████████████
Epoch  150 | loss 0.0049 | ████████████
Epoch  200 | loss 0.0026 | █████████
Epoch  250 | loss 0.0017 | ███████
Epoch  300 | loss 0.0012 | █████
Epoch  350 | loss 0.0009 | ███
Epoch  400 | loss 0.0007 | ██
Epoch  450 | loss 0.0005 | █
Epoch  500 | loss 0.0004 | █

Four things worth reading out of this. One: at epoch 1 the accuracy is 49.75% — exactly the coin-flip rate between two classes. The network knows nothing at that point, and that is the correct starting position. Two: by epoch 100 it is at 100%. So what are the remaining 400 epochs doing? Not raising accuracy — raising confidence: the loss keeps falling from 0.0194 down to 0.0004. Telling those two apart is one of the easiest things to muddle — accuracy counts correct answers, loss measures certainty. Three: the test loss (0.0008) is close to the training loss, meaning the model has not memorised the 800 points but genuinely captured the pattern. Four: the chart uses a logarithmic scale, so equal spacing means equal tenfold reductions. On a linear scale every bar from epoch 100 onwards would be zero characters long, and the chart would be blank exactly where it gets interesting.

How to run this project on your own machine

1. Install the libraries: pip install torch numpy (add matplotlib if you want the image file).
2. Download train_circles.py at the end of the lesson, or retype the code above.
3. Run it: python3 train_circles.py. The first import torch is slow; don't worry.
4. You will get exactly the numbers above, because the script fixes both random generators — np.random.seed(42) for the data and torch.manual_seed(42) for the initial weights.
5. Then break it, which is where most of the learning is: delete the optimizer.zero_grad() line and run again. Gradients from previous rounds will pile into the next, and you will see the loss behave completely differently. Meeting this bug once on purpose means recognising it instantly the next time it appears in a real model.

Lesson summary & bridge to what's next

Bridge to the next lesson: an MLP handles flat data very well, but for spatial data such as images we need a better-suited architecture — the convolutional network, in Lesson 7.