Forward pass, loss, backward, step in PyTorch - Deep Dive

Overview - Forward pass, loss, backward, step

What is it?

In machine learning with PyTorch, training a model involves four main steps: the forward pass, loss calculation, backward pass, and optimizer step. The forward pass means sending input data through the model to get predictions. The loss measures how far these predictions are from the true answers. The backward pass calculates how to change the model to improve it, and the step updates the model using this information.

Why it matters

These steps let a model learn from data by adjusting itself to make better predictions. Without this process, models would not improve and remain useless. This training loop is the core of teaching machines to recognize patterns, make decisions, or generate content, impacting fields like healthcare, self-driving cars, and language translation.

Where it fits

Before learning this, you should understand basic Python programming and what a neural network is. After mastering these steps, you can explore advanced topics like different loss functions, optimization algorithms, and model evaluation techniques.

Mental Model

Core Idea

Training a model is like repeatedly guessing answers, checking mistakes, learning from them, and improving guesses step by step.

Think of it like...

Imagine learning to shoot basketball hoops: you throw the ball (forward pass), see if it went in or missed (loss), think about how to adjust your throw (backward pass), and then try again with a better aim (step).

Input Data ──▶ [Model] ──▶ Predictions
                      │
                      ▼
                 Calculate Loss
                      │
                      ▼
                Backpropagation
                      │
                      ▼
                Optimizer Step
                      │
                      ▼
                Updated Model

Build-Up - 7 Steps

1

Foundation

Understanding the Forward Pass

Concept: The forward pass is how input data moves through the model to produce predictions.

In PyTorch, the forward pass means calling the model with input data. The model applies its layers and functions to transform inputs into outputs. For example, if the model is a simple neural network, it multiplies inputs by weights, adds biases, and applies activation functions to get predictions.

Result

You get predictions from the model based on current parameters.

Understanding the forward pass is key because it shows how the model uses its current knowledge to make guesses.

2

Foundation

Calculating the Loss

Concept: Loss measures how wrong the model's predictions are compared to true answers.

After the forward pass, you compare predictions to the actual labels using a loss function like Mean Squared Error or Cross Entropy. This gives a number representing the error size. Lower loss means better predictions.

Result

A single number quantifying prediction error.

Knowing loss lets you measure progress and guides how to improve the model.

3

Intermediate

Performing the Backward Pass

🤔Before reading on: do you think the backward pass changes model weights directly or just calculates information to help change them? Commit to your answer.

Concept: The backward pass computes gradients that show how to change each model parameter to reduce loss.

PyTorch uses automatic differentiation to calculate gradients of the loss with respect to each parameter. Calling loss.backward() triggers this process. These gradients tell us the direction and amount to adjust weights to improve predictions.

Result

Gradients stored in each parameter's .grad attribute.

Understanding backward pass reveals how models learn by knowing exactly how each parameter affects errors.

4

Intermediate

Updating Parameters with Optimizer Step

🤔Before reading on: does the optimizer step recalculate gradients or use existing ones to update parameters? Commit to your answer.

Concept: The optimizer uses gradients to adjust model parameters, making the model better.

After gradients are computed, calling optimizer.step() changes each parameter by moving it opposite to the gradient direction, scaled by a learning rate. This step is what actually updates the model's knowledge.

Result

Model parameters are changed to reduce future loss.

Knowing the optimizer step completes the learning loop by applying calculated improvements.

5

Intermediate

Clearing Gradients Before Next Step

Concept: Gradients accumulate by default, so they must be cleared before the next backward pass.

In PyTorch, gradients are added up each time backward() is called. To avoid mixing old and new gradients, call optimizer.zero_grad() before the forward pass of the next batch. This resets gradients to zero.

Result

Fresh gradients for each training iteration.

Understanding gradient clearing prevents bugs where updates become incorrect due to accumulated gradients.

6

Advanced

Putting It All Together in Training Loop

🤔Before reading on: do you think the order of zero_grad, forward, loss, backward, and step matters? Commit to your answer.

Concept: The training loop repeats zeroing gradients, forward pass, loss calculation, backward pass, and optimizer step for many batches.

A typical PyTorch training loop looks like this: for data, labels in dataloader: optimizer.zero_grad() # Clear old gradients predictions = model(data) # Forward pass loss = loss_fn(predictions, labels) # Calculate loss loss.backward() # Backward pass optimizer.step() # Update parameters This cycle repeats many times to improve the model.

Result

Model gradually learns to make better predictions over epochs.

Knowing the full loop order is crucial because changing it breaks training or causes wrong updates.

7

Expert

Surprising Effects of Gradient Accumulation

🤔Before reading on: do you think calling backward multiple times without zero_grad accumulates gradients or overwrites them? Commit to your answer.

Concept: Gradients accumulate by default, which can be used intentionally or cause subtle bugs.

If you call loss.backward() multiple times before optimizer.step(), gradients add up. This can simulate larger batch sizes or cause errors if unintended. Experts use this to train with limited memory but must carefully manage zero_grad calls.

Result

Controlled gradient accumulation can improve training efficiency or cause silent bugs.

Understanding gradient accumulation unlocks advanced training tricks and prevents hard-to-find bugs.

Under the Hood

PyTorch builds a computation graph dynamically during the forward pass, recording operations on tensors with requires_grad=True. When loss.backward() is called, it traverses this graph backward, applying the chain rule to compute gradients for each parameter. These gradients are stored in the .grad attribute of parameters. The optimizer then uses these gradients to update parameters according to its algorithm, like SGD or Adam.

Why designed this way?

Dynamic computation graphs allow flexibility to change model structure on the fly, which is useful for research and debugging. Automatic differentiation saves developers from manually calculating gradients, reducing errors and speeding up development. This design balances ease of use with powerful customization.

Input Data
   │
   ▼
[Dynamic Computation Graph]
   │
   ▼
Forward Pass (record ops)
   │
   ▼
Loss Computation
   │
   ▼
Backward Pass (auto diff)
   │
   ▼
Gradients in Parameters
   │
   ▼
Optimizer Step (update params)

Myth Busters - 4 Common Misconceptions

Quick: Does calling optimizer.step() automatically clear gradients? Commit yes or no.

Common Belief: Calling optimizer.step() clears gradients automatically.

Reality: optimizer.step() updates parameters but does NOT clear gradients; you must call optimizer.zero_grad() explicitly.

Why it matters: If you forget zero_grad(), gradients accumulate and cause incorrect parameter updates, leading to training failure.

Quick: Does loss.backward() change model weights directly? Commit yes or no.

Common Belief: loss.backward() updates model weights immediately.

Reality: loss.backward() only computes gradients; weights are updated later by optimizer.step().

Why it matters: Confusing these steps can cause misunderstanding of training flow and debugging difficulties.

Quick: Can you call backward() multiple times without zero_grad() safely? Commit yes or no.

Common Belief: Calling backward() multiple times without zero_grad() is safe and resets gradients each time.

Reality: Gradients accumulate by default; multiple backward calls add gradients unless zero_grad() is called.

Why it matters: Unintended gradient accumulation can silently break training or cause unexpected behavior.

Quick: Is the forward pass only about prediction, not learning? Commit yes or no.

Common Belief: The forward pass only predicts and does not affect learning.

Reality: The forward pass builds the computation graph needed for gradient calculation, so it is essential for learning.

Why it matters: Ignoring the forward pass's role in learning can lead to confusion about how gradients are computed.

Expert Zone

1

Gradient accumulation can be used intentionally to simulate large batch sizes when memory is limited, but requires careful zeroing of gradients.

2

Some optimizers maintain internal states (like momentum in SGD or running averages in Adam) that affect how step updates parameters beyond simple gradient scaling.

3

The order of operations in the training loop is critical; swapping zero_grad and backward calls can cause silent bugs that are hard to detect.

When NOT to use

This standard training loop is not suitable for models requiring custom gradient computations or non-differentiable operations. Alternatives include reinforcement learning algorithms or gradient-free optimization methods.

Production Patterns

In production, training loops often include mixed precision for speed, gradient clipping for stability, and checkpointing to save progress. Distributed training splits batches across multiple devices but still follows the forward-loss-backward-step pattern.

Connections

Chain Rule in Calculus

The backward pass uses the chain rule to compute gradients through layers.

Understanding the chain rule explains how gradients flow backward through complex models.

Control Systems Feedback Loop

The training loop acts like a feedback system adjusting parameters based on error signals.

Seeing training as feedback control helps grasp stability and convergence concepts.

Human Learning by Trial and Error

Model training mimics how humans learn by trying, seeing mistakes, and adjusting behavior.

This connection highlights why iterative improvement is fundamental to intelligence.

Common Pitfalls

#1 Forgetting to clear gradients before backward pass.

Wrong approach:

for data, labels in dataloader:
predictions = model(data)
loss = loss_fn(predictions, labels)
loss.backward()
optimizer.step()

Correct approach:

for data, labels in dataloader:
optimizer.zero_grad()
predictions = model(data)
loss = loss_fn(predictions, labels)
loss.backward()
optimizer.step()

Root cause:

Misunderstanding that gradients accumulate by default in PyTorch.

#2 Calling optimizer.step() before loss.backward().

Wrong approach:

for data, labels in dataloader:
optimizer.zero_grad()
predictions = model(data)
loss = loss_fn(predictions, labels)
optimizer.step()
loss.backward()

Correct approach:

for data, labels in dataloader:
optimizer.zero_grad()
predictions = model(data)
loss = loss_fn(predictions, labels)
loss.backward()
optimizer.step()

Root cause:

Confusing the order of gradient calculation and parameter update.

#3 Not setting requires_grad=True for model parameters.

Wrong approach:

model = nn.Linear(10, 1)
for param in model.parameters():
param.requires_grad = False

Correct approach:

model = nn.Linear(10, 1)
# Default requires_grad=True for parameters

Root cause:

Disabling gradient tracking prevents learning.

Key Takeaways

Training a PyTorch model involves a cycle of forward pass, loss calculation, backward pass, and optimizer step.

The forward pass produces predictions and builds the computation graph for gradients.

Loss quantifies prediction errors and guides learning direction.

Backward pass computes gradients using automatic differentiation, which optimizer.step() uses to update parameters.

Clearing gradients before each iteration is essential to prevent incorrect updates.