Backpropagation: find which weights changed the loss

Backpropagation computes how a loss changes with the parameters that produced it. It follows the computation backward, applying the chain rule through each operation. The result is a gradient for every trainable parameter involved in that computation.

An optimizer then uses those gradients to update the weights. Backpropagation calculates the direction of change; the optimizer chooses how to take the step. They are separate operations.

A small next-token network, inside out

This NumPy example uses a token embedding followed by a vocabulary projection. It has no attention or context mixing: each input token predicts one next token. tokens and targets are equal-length ID arrays, embedding has shape vocabulary × width, and output is width × vocabulary.

The forward pass creates scores and probabilities. The backward pass differentiates mean cross-entropy, then sends gradients through the projection and into the selected embedding rows. Both weight arrays are updated in place.

import numpy as np

def train_step(tokens, targets, embedding, output, rate):
    hidden = embedding[tokens]
    logits = hidden @ output
    probabilities = np.exp(logits - logits.max(axis=1, keepdims=True))
    probabilities /= probabilities.sum(axis=1, keepdims=True)

    grad_scores = probabilities.copy()
    grad_scores[np.arange(len(tokens)), targets] -= 1
    grad_scores /= len(tokens)
    grad_output = hidden.T @ grad_scores
    grad_hidden = grad_scores @ output.T

    np.add.at(embedding, tokens, -rate * grad_hidden)
    output -= rate * grad_output

np.add.at matters when the same token occurs more than once: contributions to its shared embedding row must accumulate.

The same idea at a larger scale

Autodifferentiation libraries record the operations and perform this chain-rule bookkeeping automatically. A Transformer has many more operations, but training still propagates a loss gradient back through them.

Intermediate activations are needed to calculate those derivatives. Keeping or recomputing them is one reason training has different memory requirements from simply running a model.