Kahibaro
Discord Login Register

6.1 Forward Pass and Backward Pass

What the forward pass is in practice

In a PyTorch training step, the forward pass is the part where you feed a batch of inputs into the model to produce predictions, then compute a scalar loss that measures how wrong those predictions are. Concretely, you call the model like a function, such as y_pred = model(x), and then compute something like loss = loss_fn(y_pred, y). The forward pass is also where PyTorch records the operations needed for automatic differentiation, as long as gradient tracking is enabled.

A useful mental model is that the forward pass builds a computation history that connects the loss back to the model parameters. In supervised learning you typically want a single scalar loss, because that is what you will differentiate with respect to the parameters.

What the backward pass is in practice

The backward pass is where you ask PyTorch to compute gradients of the loss with respect to all trainable parameters. You do this by calling loss.backward(). Under the hood, PyTorch applies the chain rule to propagate derivatives from the loss back through the operations performed during the forward pass, accumulating a .grad tensor for each parameter.

If $L$ is the loss and $\theta$ represents a parameter tensor, the backward pass computes $\frac{\partial L}{\partial \theta}$ and stores it in theta.grad. Those gradients are then consumed by the optimizer step, which is covered elsewhere.

Rule: loss.backward() computes gradients for the current computation graph, and stores them in each parameter’s .grad field. Gradients accumulate by default if you call backward() multiple times without clearing .grad.

A minimal training step with forward and backward

A typical single iteration looks like this:

python
model.train()
x, y = batch
x = x.to(device)
y = y.to(device)
y_pred = model(x)          # forward pass
loss = loss_fn(y_pred, y)  # still part of forward
loss.backward()            # backward pass
optimizer.step()           # parameter update, covered in the optimizer chapter

Everything up to and including loss = ... is the forward computation. The gradients are not computed until loss.backward() is called. After backward, each parameter in model.parameters() that has requires_grad=True will typically have a non None .grad tensor.

How autograd knows what to differentiate

PyTorch autograd tracks operations on tensors that have requires_grad=True, and it tracks any tensor that is produced from them through differentiable operations. Model parameters created by nn.Module layers usually have requires_grad=True automatically. Your input data x usually does not need gradients, so it typically has requires_grad=False, which is fine because gradients are needed for parameters, not for inputs, in standard training.

During the forward pass, the intermediate tensors store references to how they were computed. When you call loss.backward(), autograd traverses those references backward to compute gradients.

Rule: If you accidentally disable gradient tracking during training, for example by wrapping the forward pass in torch.no_grad(), then loss.backward() will fail or produce no gradients because the graph was not recorded.

Scalar losses, gradient outputs, and common shape expectations

Most beginner training loops use a scalar loss such as mean squared error or cross entropy reduction, so loss has shape [], meaning a zero dimensional scalar tensor. In that common case, loss.backward() needs no arguments.

If your loss is not a scalar, you must provide the gradient of some scalar objective with respect to that tensor, using loss.backward(gradient=...). In practice, beginners should keep the loss reduced to a scalar by using the default reduction or explicitly averaging or summing.

Rule: Prefer a scalar loss. If loss is not scalar, loss.backward() requires an explicit gradient argument, and mistakes here often lead to confusing errors.

Retaining graphs and multiple backward passes

By default, after loss.backward(), PyTorch frees the saved intermediate buffers for that computation graph to save memory. If you need to call backward more than once on the same graph, you must set retain_graph=True on the first backward call. This is uncommon in basic training, but can appear in advanced setups like certain multi loss objectives or higher order derivatives.

A typical beginner error is calling loss.backward() twice on the same loss without retaining the graph, which triggers an error about trying to backward through the graph a second time.

Rule: You normally call backward() once per forward pass. If you truly need multiple backward calls on the same forward, you must use retain_graph=True, and you should expect higher memory use.

Gradient accumulation is a feature, not a bug

Each call to loss.backward() adds gradients into existing .grad buffers. This is why gradient zeroing is necessary and is handled in the separate chapter on zeroing and accumulation. Here, the key point is conceptual: backward does not overwrite gradients by default, it accumulates them.

This accumulation can be used intentionally, for example to simulate a larger batch by running several forward and backward passes before applying an optimizer step, but if you do not plan for it, it will silently change training behavior.

Training mode, inference mode, and where they matter

Forward and backward passes occur during training, but some layers behave differently depending on whether the model is in training mode or evaluation mode. You set this with model.train() and model.eval(). The backward pass itself only cares that a graph exists, but the forward pass can change depending on mode, which affects the loss and thus the gradients.

For evaluation or inference, you typically use model.eval() and disable gradient tracking to save memory and compute, but you should not do that during training forward passes that need gradients.

Rule: During training, do the forward pass with gradient tracking enabled, and use model.train() so train time layer behavior is correct. During evaluation, use model.eval() and usually torch.no_grad().

Views: 70

Comments

Please login to add a comment.

Don't have an account? Register now!