Chapter 3: Automatic Differentiation (Autograd)

1. What is Automatic Differentiation

Automatic Differentiation (AD) is a technique for computing the derivative (gradient) of a function defined in a program accurately and efficiently by leveraging the structure of its computation. Training a neural network requires differentiating the loss function with respect to millions of parameters, and computing this by hand is not realistic. Autograd automates this enormous amount of differentiation.

Three Approaches to Computing Derivatives

There are broadly three approaches to computing the derivative of a function.

PyTorch's autograd adopts automatic differentiation, and specifically a method called Reverse-mode Automatic Differentiation. Reverse mode is extremely well suited to the structure of deep learning, where there are many input variables and few outputs (typically millions of parameters and a single scalar output, the loss).

Let's first compare the results of numerical differentiation and autograd to confirm that autograd returns an exact gradient.

import torch

def f(x):
    return x ** 2

# Numerical differentiation (approximation via finite differences)
def numerical_grad(f, x, eps=1e-5):
    return (f(x + eps) - f(x - eps)) / (2 * eps)

x_value = 3.0
approx_grad = numerical_grad(f, x_value)
print(f"Approximate gradient from numerical differentiation: {approx_grad:.6f}")

# Automatic differentiation (PyTorch autograd)
x = torch.tensor(x_value, requires_grad=True)
y = f(x)
y.backward()
print(f"Exact gradient from automatic differentiation: {x.grad.item():.6f}")
# The derivative of f(x) = x^2 is 2x, so the correct answer at x=3 is 6.0

Numerical differentiation returns an approximate value (extremely close to 6.000000, but with a small error), while autograd returns the analytically exact value 6.0 based on the chain rule. This difference becomes increasingly important for training stability as models grow larger.

2. Computation Graphs and Gradient Computation

At the heart of autograd is the Computational Graph. A computation graph is a Directed Acyclic Graph (DAG) that represents operations as nodes and the flow of data as edges. In PyTorch, this computation graph is dynamically built at runtime every time an operation is performed on a tensor (this approach is called Define-by-Run).

Forward Pass and Backward Pass

There are two directions of processing in a computation graph.

• Forward Pass: Executes the computation from input to output, and simultaneously builds the computation graph
• Backward Pass / Backpropagation: Applies the Chain Rule from output to input, computing the gradient of each variable

The chain rule is the mathematical property that the derivative of a composite function can be computed as "the product of the local derivatives at each stage." When $y = g(f(x))$, the following holds.

$$\frac{dy}{dx} = \frac{dy}{du} \cdot \frac{du}{dx} \quad (u = f(x))$$

Autograd traces the graph built during the forward pass and mechanically applies this chain rule at each node. Let's confirm this with actual code.

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x ** 2          # y = x^2
z = y * 3            # z = 3y = 3x^2
w = z + 1             # w = z + 1 = 3x^2 + 1

# Each tensor holds a reference (grad_fn) to the operation that created it
print(f"y.grad_fn: {y.grad_fn}")
print(f"z.grad_fn: {z.grad_fn}")
print(f"w.grad_fn: {w.grad_fn}")
# Example output:
# y.grad_fn: <PowBackward0 object at 0x7f8b1c0a3400>
# z.grad_fn: <MulBackward0 object at 0x7f8b1c0a3460>
# w.grad_fn: <AddBackward0 object at 0x7f8b1c0a34c0>

w.backward()
print(f"dw/dx = {x.grad.item()}")
# w = 3x^2 + 1, so dw/dx = 6x = 6*2 = 12.0

The grad_fn attribute held by each tensor provides the clue for tracing the edges of the computation graph. When w.backward() is called, PyTorch traces the path w → z → y → x in reverse, multiplying the local derivatives at each stage to obtain the final gradient. This process is illustrated in the diagram below.

Solid lines represent the forward computation path, and dashed lines represent the path along which gradients flow during the backward pass. Thanks to the computation graph, no matter how complex the chain of operations becomes, an exact gradient can be obtained simply by mechanically applying the chain rule.

3. torch.Tensor and requires_grad

Whether a PyTorch tensor tracks gradient computation is controlled by the requires_grad attribute. Every operation on a tensor for which this is True is recorded in the computation graph.

How to Set requires_grad

There are mainly two ways to set it: specifying it at tensor creation, and enabling it after creation.

import torch

# Scalar Tensor (0-dimensional)
x_scalar = torch.tensor(3.0, requires_grad=True)
print(f"x_scalar: {x_scalar}, requires_grad: {x_scalar.requires_grad}, is_leaf: {x_scalar.is_leaf}")

# Vector Tensor (1-dimensional)
x_vector = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
print(f"x_vector: {x_vector}, requires_grad: {x_vector.requires_grad}, is_leaf: {x_vector.is_leaf}")

# Enable requires_grad after creation
y = torch.randn(3)
print(f"requires_grad immediately after creation: {y.requires_grad}")  # False

y.requires_grad_(True)  # The trailing underscore indicates an in-place operation
print(f"requires_grad after setting: {y.requires_grad}")  # True

Let's also touch on the concept of a Leaf Tensor introduced here. A tensor with requires_grad=True that the user creates directly (such as x_scalar or x_vector in the example above) is a "leaf," and its is_leaf is True. On the other hand, a tensor newly generated by an operation becomes a "non-leaf" intermediate tensor.

import torch

try:
    int_tensor = torch.tensor([1, 2, 3], requires_grad=True)
except RuntimeError as e:
    print(f"Error: {e}")
# Example output:
# Error: Only Tensors of floating point and complex dtype can require gradients

# Solution: specify a floating-point dtype such as float32
float_tensor = torch.tensor([1, 2, 3], dtype=torch.float32, requires_grad=True)
print(f"float_tensor: {float_tensor}, requires_grad: {float_tensor.requires_grad}")

Gradient Differences Between Leaf and Intermediate Tensors

By default, only leaf tensors have a value stored in .grad after backward() is executed. The .grad of intermediate tensors remains None to save memory. If you also want to inspect the gradient of an intermediate tensor, you need to call retain_grad() beforehand.

import torch

x = torch.tensor(2.0, requires_grad=True)  # leaf tensor
y = x ** 2                                  # intermediate tensor (non-leaf)
z = y * 3                                   # intermediate tensor (non-leaf)

z.backward()
print(f"x.grad (leaf): {x.grad}")       # 12.0 -> value is stored
print(f"y.grad (non-leaf): {y.grad}")   # None -> not retained by default

# Use retain_grad() to keep the gradient of an intermediate tensor
x2 = torch.tensor(2.0, requires_grad=True)
y2 = x2 ** 2
y2.retain_grad()  # instruct PyTorch to retain y2's gradient
z2 = y2 * 3
z2.backward()
print(f"y2.grad (retained): {y2.grad}")  # dz2/dy2 = 3.0 -> retained and populated

4. How to Use the backward() Method

The backward() method executes reverse-mode automatic differentiation, computing the gradient for every leaf tensor in the computation graph (those with requires_grad=True) and storing it in each tensor's .grad attribute.

backward() for Scalar Output

When the output is a scalar (a single element), you can compute the gradient simply by calling backward() with no arguments. The same applies when there are multiple variables.

import torch

a = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(3.0, requires_grad=True)

c = a ** 2 + b ** 3  # c = a^2 + b^3 = 4 + 27 = 31
c.backward()

print(f"dc/da = {a.grad}")  # dc/da = 2a = 4.0
print(f"dc/db = {b.grad}")  # dc/db = 3b^2 = 27.0

backward() for Vector Output

When the output is not a scalar (i.e., a vector or matrix), you must pass a gradient argument to backward() with the same shape as the output. This represents an "upstream gradient" vector flowing in from the output side, and mathematically corresponds to computing a Jacobian-Vector Product.

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2  # Vector output: y = [1, 4, 9]

# Since y is not a scalar, backward() requires a gradient vector to be passed
gradient = torch.tensor([1.0, 1.0, 1.0])
y.backward(gradient=gradient)

print(f"x.grad: {x.grad}")  # dy/dx = 2x = [2, 4, 6]

Passing a vector of all ones as gradient is equivalent to computing the gradient with respect to the sum of the outputs. In practice, when computing a loss function, it's common to convert the output to a scalar with sum() or mean() before calling backward(), as shown below.

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2
loss = y.sum()  # Convert vector to scalar
loss.backward()

print(f"x.grad: {x.grad}")  # [2, 4, 6] -> same result as the previous example

The retain_graph Option

By default, once backward() is executed, the intermediate buffers of the computation graph are freed to save memory. Consequently, calling backward() twice on the same graph results in an error. If you want to run the backward pass multiple times while keeping the graph, specify retain_graph=True.

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x ** 3  # y = x^3

y.backward(retain_graph=True)  # Keep the graph instead of freeing it
print(f"1st gradient: {x.grad}")  # dy/dx = 3x^2 = 12.0

# Since retain_graph=True was specified, backward() can be called again
x.grad.zero_()
y.backward()
print(f"2nd gradient: {x.grad}")  # dy/dx = 3x^2 = 12.0 (recomputed on the same graph)

5. Gradient Accumulation and Initialization

PyTorch's autograd has an important behavior that often trips up beginners: .grad accumulates (adds up) by default. Each time backward() is called, the new gradient is not overwritten but added to the existing value of .grad.

import torch

x = torch.tensor(2.0, requires_grad=True)

# First backward pass
y1 = x ** 2
y1.backward()
print(f"1st gradient: {x.grad}")  # dy1/dx = 2x = 4.0

# Second backward pass (accumulates because the gradient wasn't reset)
y2 = x ** 4
y2.backward()
print(f"2nd gradient (after accumulation): {x.grad}")  # 4.0 + dy2/dx(=4x^3=32.0) = 36.0

# Reset the gradient to zero
x.grad.zero_()
y3 = x ** 4
y3.backward()
print(f"Gradient after reset: {x.grad}")  # 32.0 only (no accumulation)

In an actual training loop, it's standard to use the optimizer.zero_grad() method provided by a torch.optim optimizer, instead of calling tensor.grad.zero_() one tensor at a time. Below is an example of a minimal loop that trains a linear regression model using gradient descent.

import torch

torch.manual_seed(0)
w = torch.randn(1, requires_grad=True)
b = torch.randn(1, requires_grad=True)
optimizer = torch.optim.SGD([w, b], lr=0.01)

X = torch.randn(20, 1)
y_true = 2 * X + 1

for step in range(3):
    y_pred = w * X + b
    loss = ((y_pred - y_true) ** 2).mean()

    optimizer.zero_grad()  # Reset gradients at the start of each step
    loss.backward()
    optimizer.step()       # Update parameters using the gradients

    print(f"Step {step+1}: loss={loss.item():.4f}")

Running this loop, you can observe that loss decreases with each step, because optimizer.zero_grad() resets the gradient at every step. If you forget to call optimizer.zero_grad(), gradients will accumulate without bound, and training will not proceed correctly.

6. Techniques for Stopping Differentiation

It's not always necessary to track the gradients of every operation. At inference time, or in transfer learning scenarios where you want to freeze some parameters, you can deliberately stop gradient tracking to reduce memory usage and computation time. PyTorch offers three main approaches.

Method 1: torch.no_grad()

torch.no_grad() is a context manager that stops the construction of the computation graph itself for operations performed within its block. It is most commonly used during inference.

import torch

x = torch.tensor(2.0, requires_grad=True)

# Normal computation - a computation graph is built
y = x ** 2
print(f"y.requires_grad: {y.requires_grad}")  # True

# Inside torch.no_grad(), no computation graph is built
with torch.no_grad():
    z = x ** 2
    print(f"z.requires_grad: {z.requires_grad}")  # False

# A common pattern used at inference time
def predict(model_w, model_b, x_input):
    with torch.no_grad():
        return model_w * x_input + model_b

w = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
prediction = predict(w, b, torch.tensor(5.0))
print(f"Prediction: {prediction}, requires_grad: {prediction.requires_grad}")
# Prediction: 11.0, requires_grad: False

Method 2: detach()

detach() returns a new tensor that has the same value as the original but is detached from the computation graph. Note that it shares memory (data) with the original tensor.

import torch

x = torch.tensor(3.0, requires_grad=True)
y = x ** 2

# Create a new Tensor detached from the graph with detach()
y_detached = y.detach()
print(f"y.requires_grad: {y.requires_grad}")            # True
print(f"y_detached.requires_grad: {y_detached.requires_grad}")  # False

# y and y_detached share memory, so an in-place change to one affects the other
y_detached.mul_(2)
print(f"y after modification: {y}")  # y's value also changes to 18.0

Method 3: requires_grad_(False)

You can also permanently switch a tensor's own requires_grad attribute to False. This is often used in transfer learning to "freeze" some layers of a pretrained model so their parameters are not updated.

import torch

# A common pattern in transfer learning: freeze some parameters
w1 = torch.randn(3, requires_grad=True)
w2 = torch.randn(3, requires_grad=True)

# Freeze w1 (exclude it from gradient computation)
w1.requires_grad_(False)

x = torch.randn(3)
y = (w1 * x).sum() + (w2 * x).sum()
y.backward()

print(f"w1.grad: {w1.grad}")  # None (frozen, so no gradient is computed)
print(f"w2.grad is None: {w2.grad is None}")  # False (w2's gradient is computed as usual)

Each of the three methods has situations where it fits best. The table below summarizes guidelines for choosing among them.

Exercises

# Write your code here
import torch

x = torch.tensor(2.0, requires_grad=True)
# Compute f(x) = 3x^3 - 2x^2 + 5 and call backward()

import torch

x = torch.tensor(2.0, requires_grad=True)
f = 3 * x ** 3 - 2 * x ** 2 + 5
f.backward()
print(f"df/dx = {x.grad.item()}")
# df/dx = 9x^2 - 4x = 9*4 - 4*2 = 36 - 8 = 28.0

import torch

x = torch.tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)

y = x ** 2
y.sum().backward()
print(f"1st (dy/dx = 2x): {x.grad}")
# tensor([2., 4., 6., 8.])

# Continue without resetting, computing the gradient of z = x^3
z = x ** 3
z.sum().backward()
print(f"2nd (after accumulation): {x.grad}")
# dz/dx = 3x^2 = [3, 12, 27, 48] is added,
# resulting in tensor([ 5., 16., 33., 56.])

import torch

w = torch.tensor(2.5, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
x = torch.tensor(4.0)

with torch.no_grad():
    y_pred = w * x + b
    print(f"Prediction: {y_pred.item()}")           # 2.5*4 + 1.0 = 11.0
    print(f"requires_grad: {y_pred.requires_grad}")  # False

Review of Learning Objectives

Let's review the learning objectives set out at the beginning of this chapter.

• ✅ Explain how automatic differentiation works: We learned the differences from numerical and symbolic differentiation, and how autograd uses the chain rule to compute exact gradients
• ✅ Understand the relationship between computation graphs and gradient computation: We confirmed how the computation graph is dynamically built during the forward pass, and how the chain rule is applied during the backward pass
• ✅ Master requires_grad and backward(): We practiced setting requires_grad on both scalar and vector tensors, and how to call backward() (including the gradient argument)
• ✅ Correctly handle gradient accumulation and initialization: We confirmed that .grad accumulates by default, and the importance of resetting it with zero_() and optimizer.zero_grad()
• ✅ Appropriately stop differentiation with torch.no_grad() and other tools: We learned the three techniques torch.no_grad(), detach(), and requires_grad_(False), and when to use each

Summary

In this chapter, we learned how PyTorch's autograd (automatic differentiation) works.

• ✅ Unlike numerical or symbolic differentiation, autograd computes gradients accurately and efficiently using the computation graph and the chain rule
• ✅ PyTorch dynamically builds the computation graph with every operation (Define-by-Run), and each tensor records the operation that created it via grad_fn
• ✅ A tensor with requires_grad=True becomes a "leaf tensor," and after backward(), its gradient is stored in .grad (intermediate tensors require retain_grad())
• ✅ Scalar outputs allow backward() to be called without arguments, but vector outputs require a gradient argument
• ✅ Since .grad accumulates by default, training loops must reset it with zero_grad() at every step
• ✅ By choosing among torch.no_grad(), detach(), and requires_grad_(False) as appropriate, you can stop unnecessary gradient tracking and save memory and computation cost

Reference Resources

• PyTorch Official Documentation: Autograd mechanics
• PyTorch Official Tutorial: A Gentle Introduction to torch.autograd
• PyTorch Official Notes: Autograd mechanics (detailed)