Module 6: PyTorch Autograd and Backpropagation (Part 2)
Topics Covered
- Backpropagation in PyTorch
- The Chain Rule
Tensor.backward()- Accessing gradients with
.grad - Scalar and non-scalar gradient computation
- Gradient computation with multiple variables
- Gradient accumulation
- Clearing gradients with
zero_() optimizer.zero_grad()retain_graph=True- Manual backpropagation examples
- Practical gradient debugging
- PyTorch gradient workflow
Learning Objectives
After completing this chapter, you will be able to:
- Understand how backpropagation works.
- Explain the chain rule behind neural network gradients.
- Use
backward()to calculate derivatives. - Access gradients using
.grad. - Calculate gradients for multiple variables.
- Understand why PyTorch accumulates gradients.
- Reset gradients correctly during training.
- Understand when
retain_graph=Trueis required. - Distinguish scalar and non-scalar backward operations.
- Debug common gradient-related problems in PyTorch.
Introduction
In Part 1, you learned that PyTorch Autograd builds a computational graph while tensor operations are executed.
The next step is to use that graph to calculate derivatives.
This process is called backpropagation.
Backpropagation is one of the fundamental algorithms behind neural network training. It determines how much each trainable parameter contributed to the final loss and provides the gradients required by the optimizer.
A simplified training process is:
1Input Data 2 │ 3 ▼ 4Forward Pass 5 │ 6 ▼ 7Model Prediction 8 │ 9 ▼ 10Loss Calculation 11 │ 12 ▼ 13Backward Pass 14 │ 15 ▼ 16Gradients 17 │ 18 ▼ 19Optimizer 20 │ 21 ▼ 22Updated Parameters
In PyTorch, the backward pass is commonly started with:
1loss.backward()
What Is Backpropagation?
Backpropagation is an algorithm for efficiently computing derivatives of an output, usually a loss function, with respect to the parameters of a neural network.
The word "backpropagation" describes the direction in which gradient information moves through the computational graph.
During the forward pass:
1Input → Operations → Prediction → Loss
During the backward pass:
1Loss → Gradients → Earlier Operations → Parameters
For a neural network, this allows PyTorch to calculate values such as:
1∂Loss/∂Weight 2∂Loss/∂Bias 3∂Loss/∂Input
These gradients are then used by an optimizer such as SGD or Adam to update model parameters.
Forward Pass vs Backward Pass
The forward pass calculates predictions.
1prediction = model(x) 2loss = loss_function(prediction, target)
The backward pass calculates gradients.
1loss.backward()
The optimizer then updates the parameters.
1optimizer.step()
A typical training iteration is therefore:
1optimizer.zero_grad() 2 3prediction = model(x) 4 5loss = loss_function(prediction, target) 6 7loss.backward() 8 9optimizer.step()
This sequence is one of the most important patterns in PyTorch.
What Is the Chain Rule?
The Chain Rule is a calculus rule used to differentiate composite functions.
Suppose:
1u = 3x + 4 2 3y = u²
Therefore:
1y = (3x + 4)²
The derivative is:
1dy/dx = dy/du × du/dx
We know:
1dy/du = 2u 2 3du/dx = 3
Therefore:
1dy/dx = 2u × 3
Since:
1u = 3x + 4
we get:
1dy/dx = 6(3x + 4)
For:
1x = 2
the gradient is:
16(10) = 60
PyTorch performs this chain-rule computation automatically through Autograd.
First Gradient Example
Consider:
1y = x²
The derivative is:
1dy/dx = 2x
Create the tensor:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x ** 2
Now calculate the gradient:
1y.backward()
Read the gradient:
1print(x.grad)
Output:
1tensor(4.)
Because:
1dy/dx = 2x 2 3 = 2 × 2 4 5 = 4
Understanding backward()
The backward() method tells PyTorch to perform reverse-mode automatic differentiation from a tensor.
For a scalar output:
1loss.backward()
is the standard form.
Example:
1import torch 2 3x = torch.tensor(5.0, requires_grad=True) 4 5y = 4 * x 6 7y.backward() 8 9print(x.grad)
Output:
1tensor(4.)
Mathematically:
1y = 4x 2 3dy/dx = 4
Cubic Function Example
Consider:
1y = x³
The derivative is:
1dy/dx = 3x²
PyTorch:
1import torch 2 3x = torch.tensor(3.0, requires_grad=True) 4 5y = x ** 3 6 7y.backward() 8 9print(x.grad)
Output:
1tensor(27.)
Because:
13 × 3² = 27
Understanding .grad
The .grad attribute stores the accumulated gradient for a leaf tensor that participates in gradient tracking.
Example:
1import torch 2 3x = torch.tensor(4.0, requires_grad=True) 4 5y = x ** 2 6 7y.backward() 8 9print(x.grad)
Output:
1tensor(8.)
The mathematical derivative is:
1dy/dx = 2x 2 3 = 2 × 4 4 5 = 8
Before calling backward():
1print(x.grad)
normally produces:
1None
After:
1y.backward()
the gradient becomes available:
1tensor(8.)
Computational Graph Example
Consider:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x * 3 6 7z = y + 4 8 9loss = z ** 2
The graph can be represented as:
1x 2│ 3├── × 3 4│ 5▼ 6y 7│ 8├── + 4 9│ 10▼ 11z 12│ 13├── Square 14│ 15▼ 16loss
Mathematically:
1y = 3x 2 3z = y + 4 4 5loss = z²
Therefore:
1loss = (3x + 4)²
The derivative is:
1dloss/dx = 2(3x + 4) × 3
At:
1x = 2
we get:
12(10) × 3 = 60
PyTorch:
1loss.backward() 2 3print(x.grad)
Output:
1tensor(60.)
Multiple Variables
Autograd can calculate gradients with respect to multiple tensors.
Example:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = torch.tensor(3.0, requires_grad=True) 6 7z = x * y 8 9z.backward() 10 11print("dz/dx =", x.grad) 12print("dz/dy =", y.grad)
Output:
1dz/dx = tensor(3.) 2dz/dy = tensor(2.)
Because:
1z = xy
Therefore:
1∂z/∂x = y 2 3∂z/∂y = x
At:
1x = 2 2y = 3
we get:
1∂z/∂x = 3 2 3∂z/∂y = 2
Multiple Operations with Multiple Variables
Consider:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5w = torch.tensor(3.0, requires_grad=True) 6 7b = torch.tensor(1.0, requires_grad=True) 8 9y = x * w + b 10 11loss = y ** 2 12 13loss.backward() 14 15print("x gradient:", x.grad) 16print("w gradient:", w.grad) 17print("b gradient:", b.grad)
The computation is:
1x ──┐ 2 × ──► y ──► Square ──► loss 3w ──┘ 4 5b ───────────────► +
Every trainable leaf tensor receives its corresponding gradient.
This is exactly the concept used for neural network weights and biases.
Gradient Accumulation
A critical PyTorch concept is that gradients accumulate by default.
PyTorch does not automatically replace the existing value in .grad.
Instead:
1new gradient 2 + 3existing gradient 4 = 5updated gradient
Example:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x ** 2 6 7y.backward() 8 9print(x.grad)
Output:
1tensor(4.)
Now create another graph and perform another backward pass:
1y = x ** 2 2 3y.backward() 4 5print(x.grad)
Output:
1tensor(8.)
The first backward pass produced:
14
The second produced another:
14
Therefore:
14 + 4 = 8
This behavior is called gradient accumulation.
Why Does PyTorch Accumulate Gradients?
Gradient accumulation is useful in several situations.
For example, suppose a large batch cannot fit into GPU memory.
Instead of processing:
1Batch = 128
you can process:
1Batch = 32 2Batch = 32 3Batch = 32 4Batch = 32
and accumulate gradients before performing an optimizer update.
Conceptually:
1Mini-batch 1 2 ↓ 3Gradient 4 ↓ 5Accumulate 6 7Mini-batch 2 8 ↓ 9Gradient 10 ↓ 11Accumulate 12 13Mini-batch 3 14 ↓ 15Gradient 16 ↓ 17Accumulate 18 19Mini-batch 4 20 ↓ 21Gradient 22 ↓ 23Optimizer Step
This technique is commonly called gradient accumulation or gradient accumulation steps.
Resetting Gradients with zero_()
Because gradients accumulate, they must often be cleared before another independent gradient calculation.
Example:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x ** 2 6 7y.backward() 8 9print("Before reset:", x.grad) 10 11x.grad.zero_() 12 13print("After reset:", x.grad)
Output:
1Before reset: tensor(4.) 2After reset: tensor(0.)
The underscore in:
1zero_()
indicates an in-place operation.
optimizer.zero_grad()
When training neural networks, you normally do not manually clear every parameter.
Instead, use:
1optimizer.zero_grad()
Typical training code:
1for inputs, targets in dataloader: 2 3 optimizer.zero_grad() 4 5 outputs = model(inputs) 6 7 loss = loss_function(outputs, targets) 8 9 loss.backward() 10 11 optimizer.step()
The order is important:
1zero_grad() 2 ↓ 3forward pass 4 ↓ 5loss 6 ↓ 7backward() 8 ↓ 9optimizer.step()
Complete Training Example
Here is a minimal neural network training example:
1import torch 2import torch.nn as nn 3 4model = nn.Linear(1, 1) 5 6optimizer = torch.optim.SGD( 7 model.parameters(), 8 lr=0.01 9) 10 11x = torch.tensor([[1.0], [2.0], [3.0]]) 12 13target = torch.tensor([[2.0], [4.0], [6.0]]) 14 15for epoch in range(5): 16 17 optimizer.zero_grad() 18 19 prediction = model(x) 20 21 loss = ((prediction - target) ** 2).mean() 22 23 loss.backward() 24 25 optimizer.step() 26 27 print( 28 f"Epoch {epoch + 1}, " 29 f"Loss: {loss.item():.4f}" 30 )
This demonstrates the fundamental PyTorch training loop:
1Clear Gradients 2 ↓ 3Forward Pass 4 ↓ 5Calculate Loss 6 ↓ 7Backward Pass 8 ↓ 9Update Parameters
Understanding retain_graph=True
By default, PyTorch frees parts of the computational graph after backward() has completed.
This helps reduce memory usage.
If you need to perform another backward pass through the same graph, you may need:
1retain_graph=True
Example:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x ** 2 6 7y.backward(retain_graph=True) 8 9print(x.grad) 10 11y.backward() 12 13print(x.grad)
Output:
1tensor(4.) 2tensor(8.)
The first backward pass calculates:
14
The second calculates another:
14
Since gradients accumulate:
14 + 4 = 8
When Should You Use retain_graph=True?
Use it only when you genuinely need to reuse the same computational graph for another backward computation.
For example:
1loss1.backward(retain_graph=True) 2 3loss2.backward()
However, unnecessary use of retain_graph=True can increase memory consumption.
In normal neural network training, you usually do not need it.
Scalar vs Non-Scalar Tensors
A common beginner mistake is assuming that every tensor can directly call:
1tensor.backward()
without additional information.
For a scalar tensor, PyTorch can implicitly start the backward pass with a gradient of 1.
Example:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x ** 2 6 7y.backward() 8 9print(x.grad)
This works because y is scalar.
But consider:
1x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) 2 3y = x ** 2 4 5y.backward()
This raises an error because y is non-scalar.
You must provide the gradient argument:
1y.backward(torch.ones_like(y)) 2 3print(x.grad)
Output:
1tensor([2., 4., 6.])
This is an important concept when working with vector and matrix outputs.
Why Does backward() Need a Gradient for Non-Scalar Outputs?
For a scalar:
1y = scalar
the derivative of the scalar output can be calculated directly.
For a vector:
1y = [y₁, y₂, y₃]
PyTorch needs to know which scalar combination of these outputs should be differentiated.
The supplied tensor acts as the vector in a vector-Jacobian product.
For beginners, the practical rule is:
1scalar_output.backward()
works directly.
For non-scalar output:
1output.backward(gradient)
requires a gradient tensor with a compatible shape.
Gradient of a Vector Example
1import torch 2 3x = torch.tensor( 4 [1.0, 2.0, 3.0], 5 requires_grad=True 6) 7 8y = x ** 2 9 10gradient = torch.ones_like(y) 11 12y.backward(gradient) 13 14print(x.grad)
Output:
1tensor([2., 4., 6.])
Because:
1y = [x₁², x₂², x₃²]
and:
1dy/dx = [2x₁, 2x₂, 2x₃]
Therefore:
1[2, 4, 6]
Inspecting Gradients
You can inspect gradients during debugging.
Example:
1import torch 2 3x = torch.tensor(3.0, requires_grad=True) 4 5y = x ** 2 6 7print("Before backward:", x.grad) 8 9y.backward() 10 11print("After backward:", x.grad)
Output:
1Before backward: None 2After backward: tensor(6.)
Leaf and Non-Leaf Gradients
Recall from Part 1:
1x = torch.tensor(2.0, requires_grad=True) 2 3y = x * 3
Here:
1x → leaf tensor 2y → non-leaf tensor
By default, .grad is retained for leaf tensors, but not for non-leaf tensors.
Example:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x * 3 6 7loss = y ** 2 8 9loss.backward() 10 11print("x.grad:", x.grad) 12print("y.grad:", y.grad)
The first contains the accumulated gradient, while y.grad is normally None.
If you need the intermediate gradient:
1y.retain_grad() 2 3loss.backward()
But remember that the graph must still exist when retain_grad() is called, and repeated backward passes may require retain_graph=True.
Using retain_grad()
Example:
1import torch 2 3x = torch.tensor(2.0, requires_grad=True) 4 5y = x * 3 6 7y.retain_grad() 8 9loss = y ** 2 10 11loss.backward() 12 13print("x gradient:", x.grad) 14 15print("y gradient:", y.grad)
Now PyTorch retains the gradient for the non-leaf tensor y.
Practical Example: Linear Regression Gradient
Consider a simple model:
1y = wx + b
where:
1w = weight 2b = bias
PyTorch:
1import torch 2 3x = torch.tensor(2.0) 4 5target = torch.tensor(10.0) 6 7w = torch.tensor( 8 3.0, 9 requires_grad=True 10) 11 12b = torch.tensor( 13 1.0, 14 requires_grad=True 15) 16 17prediction = w * x + b 18 19loss = (prediction - target) ** 2 20 21loss.backward() 22 23print("Prediction:", prediction) 24print("Loss:", loss) 25 26print("Weight gradient:", w.grad) 27print("Bias gradient:", b.grad)
The model prediction is:
13 × 2 + 1 = 7
The target is:
110
Therefore the model has an error.
Autograd calculates how w and b should change to reduce the loss.
This is the foundation of gradient-based optimization.
Practical Example: Manual Gradient Descent
You can implement a very small gradient descent example without using an optimizer.
1import torch 2 3x = torch.tensor(2.0) 4 5target = torch.tensor(10.0) 6 7w = torch.tensor( 8 1.0, 9 requires_grad=True 10) 11 12b = torch.tensor( 13 0.0, 14 requires_grad=True 15) 16 17learning_rate = 0.01 18 19for step in range(10): 20 21 prediction = w * x + b 22 23 loss = (prediction - target) ** 2 24 25 loss.backward() 26 27 with torch.no_grad(): 28 w -= learning_rate * w.grad 29 b -= learning_rate * b.grad 30 31 w.grad.zero_() 32 b.grad.zero_() 33 34 print( 35 f"Step {step + 1}: " 36 f"Loss={loss.item():.4f}, " 37 f"w={w.item():.4f}, " 38 f"b={b.item():.4f}" 39 )
This example demonstrates the complete mathematical idea:
1Forward Pass 2 ↓ 3Calculate Loss 4 ↓ 5Calculate Gradients 6 ↓ 7Update Parameters 8 ↓ 9Clear Gradients 10 ↓ 11Repeat
Modern PyTorch code normally uses an optimizer instead:
1optimizer = torch.optim.SGD( 2 [w, b], 3 lr=0.01 4)
Then:
1optimizer.zero_grad() 2 3loss.backward() 4 5optimizer.step()
Real-World Neural Network Workflow
A neural network training iteration typically looks like:
1Training Data 2 │ 3 ▼ 4Input Tensor 5 │ 6 ▼ 7Neural Network 8 │ 9 ▼ 10Predictions 11 │ 12 ▼ 13Loss Function 14 │ 15 ▼ 16loss.backward() 17 │ 18 ▼ 19Gradients 20 │ 21 ▼ 22optimizer.step() 23 │ 24 ▼ 25Updated Weights
For a large model, the same process occurs across millions or billions of parameters.
Backpropagation in Transformers
Transformers also rely on the same fundamental process.
A simplified Transformer training workflow is:
1Input Tokens 2 ↓ 3Embedding 4 ↓ 5Self-Attention 6 ↓ 7Feed-Forward Network 8 ↓ 9Logits 10 ↓ 11Cross-Entropy Loss 12 ↓ 13Backward Pass 14 ↓ 15Parameter Gradients 16 ↓ 17Optimizer 18 ↓ 19Updated Transformer Weights
Autograd calculates gradients through operations such as:
- Matrix multiplication
- Attention score calculation
- Softmax
- Linear layers
- Activation functions
- Layer normalization
- Embedding operations
This is one reason understanding Autograd is essential before studying Transformer training and large language models.
Common Mistakes
Calling backward() Twice on the Same Graph
This can fail:
1y.backward() 2y.backward()
The graph is normally released after the first backward pass.
If the same graph genuinely needs to be reused:
1y.backward(retain_graph=True)
Use this carefully because retaining graphs consumes additional memory.
Forgetting to Clear Gradients
Incorrect:
1for data, target in dataloader: 2 3 output = model(data) 4 5 loss = loss_function(output, target) 6 7 loss.backward() 8 9 optimizer.step()
Gradients accumulate across iterations.
Correct:
1for data, target in dataloader: 2 3 optimizer.zero_grad() 4 5 output = model(data) 6 7 loss = loss_function(output, target) 8 9 loss.backward() 10 11 optimizer.step()
Calling .backward() on a Non-Scalar Output
This may fail:
1x = torch.tensor( 2 [1.0, 2.0, 3.0], 3 requires_grad=True 4) 5 6y = x ** 2 7 8y.backward()
For a non-scalar output, provide a gradient:
1y.backward(torch.ones_like(y))
Reading .grad Before Backward
Before:
1loss.backward()
the gradient may be:
1None
Check it after the backward pass.
Expecting Non-Leaf Tensors to Store .grad
Intermediate tensors normally do not retain gradients automatically.
Use:
1tensor.retain_grad()
when you need to inspect an intermediate gradient.
Using retain_graph=True Everywhere
Do not use:
1loss.backward(retain_graph=True)
by default.
Retaining the graph consumes additional memory and is unnecessary for most training iterations.
Best Practices
- Use
loss.backward()to calculate gradients for scalar losses. - Call
optimizer.zero_grad()before each normal training iteration. - Remember that PyTorch gradients accumulate by default.
- Use
optimizer.step()after the backward pass. - Use
retain_graph=Trueonly when the graph must genuinely be reused. - Provide an explicit gradient when calling
backward()on non-scalar outputs. - Inspect
.gradwhen debugging model training. - Use
retain_grad()when you need gradients of intermediate non-leaf tensors. - Avoid unnecessary gradient tracking for tensors that do not require optimization.
- Monitor loss and gradients when debugging unstable training.
Practice Exercises
Beginner Exercise 1: Basic Gradient
Create:
1x = torch.tensor( 2 5.0, 3 requires_grad=True 4) 5 6y = x ** 2
Calculate:
1dy/dx
using:
1y.backward()
Then print:
1x.grad
Beginner Exercise 2: Chain Rule
Create:
1x = torch.tensor( 2 2.0, 3 requires_grad=True 4) 5 6y = 3 * x + 4 7 8z = y ** 2
Calculate:
1z.backward()
and verify that the gradient is:
160
Intermediate Exercise 3: Multiple Variables
Create:
1x = torch.tensor( 2 4.0, 3 requires_grad=True 4) 5 6y = torch.tensor( 7 5.0, 8 requires_grad=True 9) 10 11z = x * y + x
Calculate:
1z.backward()
Print:
1x.grad 2y.grad
Verify the mathematical derivatives.
Intermediate Exercise 4: Gradient Accumulation
Create:
1x = torch.tensor( 2 3.0, 3 requires_grad=True 4)
Perform two backward passes without clearing the gradient.
Observe the result.
Then use:
1x.grad.zero_()
and verify that the gradient has been reset.
Advanced Exercise 5: Non-Scalar Gradient
Create:
1x = torch.tensor( 2 [1.0, 2.0, 3.0], 3 requires_grad=True 4) 5 6y = x ** 2
Calculate the gradient using:
1y.backward(torch.ones_like(y))
Verify:
1tensor([2., 4., 6.])
Advanced Exercise 6: Mini Training Loop
Create a linear model:
1y = wx + b
Implement a training loop that:
- Creates training data.
- Calculates predictions.
- Calculates mean squared error.
- Calls
backward(). - Updates parameters.
- Clears gradients.
- Prints the loss for every iteration.
Try to train the model so that:
1y ≈ 2x
Interview Questions
What is backpropagation?
Backpropagation is an algorithm for calculating gradients of a loss function with respect to model parameters by traversing the computational graph backward and applying the chain rule.
What does backward() do?
backward() computes derivatives of a tensor with respect to tensors that participated in its computation and require gradients.
For a scalar loss:
1loss.backward()
calculates the gradients needed for parameter updates.
Where are gradients stored?
For leaf tensors that require gradients, PyTorch stores the accumulated gradient in:
1tensor.grad
Why do gradients accumulate?
PyTorch accumulates gradients because this behavior supports workflows such as gradient accumulation across multiple mini-batches.
Therefore, gradients normally need to be cleared before another independent optimization step.
What is the difference between zero_() and optimizer.zero_grad()?
1x.grad.zero_()
clears the gradient of one tensor.
1optimizer.zero_grad()
clears gradients associated with the optimizer's parameters.
For model training, optimizer.zero_grad() is normally preferred.
When is retain_graph=True required?
It is required when you need to perform another backward computation through a graph that would otherwise have been freed after the first backward pass.
It should not be used unnecessarily because retaining the graph increases memory usage.
Why can't a vector normally call backward() without an argument?
A backward call without an explicit gradient is directly supported for scalar outputs. For a non-scalar output, PyTorch needs an upstream gradient tensor to define the vector-Jacobian product.
Example:
1y.backward(torch.ones_like(y))
What is gradient accumulation?
Gradient accumulation means new gradients are added to the existing .grad values rather than replacing them.
Example:
1First backward: 2 2Second backward: 2 3Accumulated: 4
Module Summary
In this chapter, you learned how PyTorch performs the backward pass and calculates gradients using Autograd.
You learned:
- What backpropagation is.
- How the chain rule enables gradient calculation.
- How
backward()starts reverse-mode differentiation. - How
.gradstores accumulated gradients for leaf tensors. - How gradients are calculated for multiple variables.
- Why PyTorch accumulates gradients.
- How to clear gradients using
zero_(). - Why
optimizer.zero_grad()is used in neural network training. - How
retain_graph=Trueallows a computational graph to be reused. - How to calculate gradients for non-scalar outputs.
- How to inspect and debug gradients.
- How Autograd participates in neural network and Transformer training.
The key PyTorch training pattern to remember is:
1optimizer.zero_grad() 2 3output = model(input) 4 5loss = loss_function(output, target) 6 7loss.backward() 8 9optimizer.step()
This simple sequence represents the core training cycle of many deep learning systems:
1Forward Pass 2 ↓ 3Loss 4 ↓ 5Backward Pass 6 ↓ 7Gradients 8 ↓ 9Parameter Update
Next Part: Autograd Gradient Control
In Part 3, you will learn how to control gradient tracking and computational graphs using:
detach()detach_()torch.no_grad()torch.inference_mode()- Gradient tracking control
- Freezing model parameters
- Transfer learning
- Feature extraction
- Evaluation vs training mode
- Memory optimization
- Preventing unnecessary gradient computation
- Practical inference examples
These concepts are especially important when deploying PyTorch models, performing inference, freezing pretrained networks, and optimizing GPU memory usage.