Neural Networks & Transformers

AI · Neural Networks & Transformers

How AI Models Train: Backpropagation to the Full Training Loop

The chain rule, run backward through a network, is the entire mechanism behind how a model learns — here's how that one idea scales into a real training loop.

10 min read
BackpropagationModel TrainingDeep Learning

Every neural network, from a small classifier to a frontier language model, learns the same way: it makes a prediction, measures how wrong that prediction was, and adjusts its internal weights to be slightly less wrong next time. Repeated billions of times, that loop is the entire mechanism behind "training." The part that makes it possible at scale — computing exactly how much each individual weight contributed to the error — is called backpropagation. Everything else in a modern training pipeline is built around that one calculation.

The forward pass: making a guess

Training starts with a forward pass. An input — a batch of images, tokens, or feature vectors — flows through the network's layers in order, each one transforming it a little further, until the final layer produces a prediction. That prediction is then compared against the true label or target using a loss function: a single number that quantifies how wrong the guess was. A model predicting the next word gets scored by how much probability it assigned to the actual next word; a model predicting a price gets scored by how far off its number was. Either way, the loss is the one quantity the rest of training exists to shrink.

Backpropagation: assigning blame efficiently

Knowing the loss tells you that the model was wrong, not which weights to blame or in which direction to adjust them. A network can have billions of weights spread across hundreds of layers, and each one nudges the final prediction in some small, entangled way. Backpropagation is the algorithm that untangles this: it applies the chain rule of calculus to compute the gradient of the loss with respect to every single weight in the network, all in one efficient backward sweep.

The "backward" part is what makes it efficient. Backprop starts at the output layer, where the gradient of the loss is easy to compute directly, and then walks backward through the network one layer at a time. Each layer only needs two things to compute its own weights' gradients: the error signal flowing in from the layer after it, and the values it itself saw during the forward pass. It never needs to recompute anything about layers further downstream — that work was already done and passed along as a single compact signal. This is exactly the chain rule in action: the derivative of a composition of functions decomposes into a product of each function's local derivative, computed one link at a time. Without this reuse, computing the gradient for a deep network would require redoing an enormous amount of duplicate work for every layer; with it, a full backward pass costs roughly the same as a forward pass, regardless of how deep the network is.

Gradients, not answers

Backpropagation doesn't tell the network what the right weights are — it only tells it, for the current weights, which direction increases the loss and by roughly how much, for every weight simultaneously. What to do with that information is a separate question, answered by the optimizer.

Gradient descent: taking the step

Once every weight has a gradient, the update rule at its simplest is gradient descent: move each weight a small amount in the direction opposite its gradient, since that's the direction that decreases the loss. How large that step is gets scaled by the learning rate, and getting it wrong in either direction is costly. Too large, and each step overshoots the bottom of the loss curve — the weights swing past the minimum, the loss gets worse instead of better, and training can diverge outright. Too small, and each step barely moves anything; the model technically keeps improving, but so slowly that training that should take hours stretches into weeks, and can stall in shallow regions of the loss landscape long before reaching a good solution.

Why optimizers do more than plain gradient descent

In practice, almost no modern network trains on plain gradient descent. Optimizers like Adam keep a running history of each parameter's past gradients and use it to adapt the effective step size per parameter, rather than applying the same learning rate everywhere. A weight whose gradient has pointed the same direction consistently builds up momentum and moves faster; a weight whose gradient keeps flipping sign gets damped down. This matters enormously at scale: different layers and parameters in a large network can have wildly different natural scales and sensitivities, and a single global learning rate that suits one part of the network is often badly wrong for another. Adaptive optimizers smooth over exactly that mismatch, which is why they, not plain SGD, are the default for training large modern networks.

Batches and epochs: how much data per step

The gradient described above is technically computed from a loss — but whose loss? Using the entire training set for every single update would be accurate but far too slow and memory-hungry to be practical at any real scale. Using one example at a time is cheap per step but produces a gradient that is extremely noisy — a single outlier can send the weights in a bad direction — and it wastes the parallel throughput of modern GPUs, which are built to crunch many examples at once. The compromise that training actually uses is the mini-batch: a modest slice of the training set, often somewhere from a few dozen to a few thousand examples, averaged into one gradient per step. This keeps memory use bounded, keeps GPUs busy, and — somewhat counterintuitively — the mild noise from batching often helps generalization, nudging the model out of narrow loss-surface crevices that would trap a perfectly smooth gradient. One full pass through the entire training dataset, batch by batch, is called an epoch; a real training run typically repeats this for many epochs, reshuffling the data each time.

Putting the four stages together

Strip away everything else and a single training step is only four operations, repeated over and over:

python
for batch in dataloader:          # one mini-batch at a time
    inputs, targets = batch

    # 1. Forward pass — run the batch through the network
    predictions = model(inputs)

    # 2. Loss — how wrong were the predictions?
    loss = loss_fn(predictions, targets)

    # 3. Backward pass — backprop the error through every layer
    optimizer.zero_grad()        # clear gradients from the last step
    loss.backward()              # populates .grad on every weight

    # 4. Optimizer step — nudge every weight to reduce the loss
    optimizer.step()

The forward pass (model(inputs)) runs the batch through the network to get predictions. The loss computation (loss_fn(...)) scores those predictions against the true targets. The backward pass (loss.backward()) is backpropagation itself — it walks the computation backward from the loss and fills in a gradient for every weight. And the optimizer step (optimizer.step()) uses those gradients to actually update the weights, following whatever rule the chosen optimizer implements.

The line easiest to overlook is optimizer.zero_grad(), called right before the backward pass. Frameworks accumulate gradients into each weight's .grad field by default, rather than overwriting them — a design choice that's genuinely useful for a few advanced techniques, but a footgun if you forget about it. Without explicitly zeroing them out first, each new batch's gradients would pile on top of the previous batch's, silently corrupting the update. Zeroing them at the start of every step is what keeps each iteration's update dependent only on the batch currently being processed.

The same loop, underneath everything

This four-step loop is the entire mechanism, and it doesn't get replaced by later stages like fine-tuning or RLHF-style alignment — it gets reused. Fine-tuning a pretrained model runs the exact same forward-loss-backward-optimizer cycle, just on a smaller, more targeted dataset and usually with a smaller learning rate. Even RLHF-style alignment, which changes what counts as the "loss" (a reward or preference signal rather than a next-token label), still updates weights via backpropagation and an optimizer step. Backprop isn't one phase among several in a model's lifecycle — it's the foundation every phase stands on.

Key takeaways

  • The forward pass produces a prediction; the loss function scores how wrong it was against the target.
  • Backpropagation uses the chain rule to compute the gradient for every weight in one backward sweep, reusing each layer's incoming error signal instead of recomputing everything from scratch.
  • Gradient descent nudges weights opposite their gradient, scaled by a learning rate — too large overshoots and can diverge, too small trains painfully slowly.
  • Adaptive optimizers like Adam track gradient history to adjust the step size per parameter, which is why plain SGD is rarely used as-is on large networks.
  • Mini-batches balance the noise of single examples against the cost of using the full dataset; one full pass over the data is an epoch.
  • Zeroing gradients before each backward pass matters because frameworks accumulate gradients by default rather than overwriting them.
2026 — Built by Arpit Mishra