Autograd and computational graphs in PyTorch

Overview

PyTorch's autograd system is the engine that powers automatic differentiation—it tracks operations on tensors and computes gradients automatically. This eliminates manual derivative calculations during backpropagation, making deep learning research fast and error-free.

Figure — Autograd and computational graphs in PyTorch

Core Concepts

How Autograd Works: Step-by-Step Derivation

Step 1: Forward Pass and Graph Construction

When you perform operations on tensors with requires_grad=True, PyTorch records them:

import torch
 
x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([3.0], requires_grad=True)
 
z = x * y + x**2  # z = 3x + x²

What happens internally:

x * y creates a node with grad_fn=<MulBackward>
x**2 creates a node with grad_fn=<PowBackward>
+ combines them with grad_fn=<AddBackward>
z now has z.grad_fn = <AddBackward> pointing to the entire graph

Why this matters: The graph structure encodes the chain of operations, which is exactly what we need for the chain rule.

Step 2: Backward Pass and Chain Rule

The chain rule states: if $z = f(g(x))$ , then $\frac{dz}{dx} = \frac{dz}{dg} \cdot \frac{dg}{dx}$ .

For our example $z = xy + x^2$ :

$\frac{dz}{dx} = \frac{\partial}{\partial x}(xy) + \frac{\partial}{\partial x}(x^2) = y + 2x$

$\frac{dz}{dy} = \frac{\partial}{\partial y}(xy) + \frac{\partial}{\partial y}(x^2) = x + 0 = x$

Derivation from first principles:

Starting with the definition of derivative: $\frac{dz}{dx} = \lim_{h \to 0} \frac{z(x+h, y) - z(x, y)}{h}$

$= \lim_{h \to 0} \frac{[(x+h)y + (x+h)^2] - [xy + x^2]}{h}$

$= \lim_{h \to 0} \frac{xy + hy + x^2 + 2xh + h^2 - xy - x^2}{h}$

$= \lim_{h \to 0} \frac{hy + 2xh + h^2}{h} = \lim_{h \to 0}(y + 2x + h) = y + 2x$

Why this step? We expand the function at $x+h$ , cancel terms, factor out $h$ , and take the limit. This is the fundamental definition of a derivative.

PyTorch automates this:

z.backward()  # Compute dz/dx and dz/dy
print(x.grad)  # tensor([8.]) = 3+ 2*2
print(y.grad)  # tensor([2.]) = 2

Step 3: Gradient Accumulation

By default, calling .backward() accumulates gradients into .grad:

$\text{tensor.grad} \leftarrow \text{tensor.grad} + \frac{\partial L}{\partial \text{tensor}}$

Why accumulation? In mini-batch training, you sum gradients across multiple samples before updating weights. But this means you must zero gradients between iterations:

x.grad.zero_()  # Clear previous gradients

Derivation of why accumulation is needed:

In batch gradient descent, the loss is: $L_{\text{batch}} = \frac{1}{N}\sum_{i=1}^N L_i$

The gradient is: $\frac{\partial L_{\text{batch}}}{\partial w} = \frac{1}{N}\sum_{i=1}^N \frac{\partial L_i}{\partial w}$

If we compute $\frac{\partial L_i}{\partial w}$ for each sample separately, we must sum them. Accumulation implements this sum automatically.

Formula Chain Rule in Computational Graphs

For a path $x \to a \to b \to z$ , the gradient is: $\frac{\partial z}{\partial x} = \frac{\partial z}{\partial b} \cdot \frac{\partial b}{\partial a} \cdot \frac{\partial a}{\partial x}$

Backward pass algorithm:

Start at output node $z$ with $\frac{\partial z}{\partial z} = 1$
For each node $n$ in reverse topological order: $\frac{\partial z}{\partial n} = \sum_{\text{child } c} \frac{\partial z}{\partial c} \cdot \frac{\partial c}{\partial n}$
Store $\frac{\partial z}{\partial n}$ in n.grad Why sum over children? A node can feed into multiple operations (e.g., $x$ used in both $x^2$ and $xy$ ). Each path contributes to the total gradient (multivariate chain rule).

Detailed Examples

Worked example Example 1: Simple Neural Network Layer

Problem: Compute gradients for a single linear layer $y = Wx + b$ followed by ReLU.

import torch
# Input and parameters
x = torch.tensor([[1.0, 2.0]], requires_grad=True)
W = torch.tensor([[0.5, -0.3], [0.2, 0.4]], requires_grad=True)
b = torch.tensor([[0.1, -0.2]], requires_grad=True)
 
# Forward pass
z = torch.matmul(x, W) + b  # z = xW + b
y = torch.relu(z)            # y = max(0, z)
loss = y.sum()               # Scalar loss
 
# Backward pass
loss.backward()
print(x.grad)  # Gradient of loss w.r.t input
print(W.grad)  # Gradient of loss w.r.t weights

Step-by-step reasoning:

Compute $z$ : $z = xW + b = [1, 2] \begin{bmatrix}0.5 & -0.3\\0.2 & 0.4\end{bmatrix} + [0.1, -0.2]$

$z_1 = 1(0.5) + 2(0.2) + 0.1 = 1.0$
$z_2 = 1(-0.3) + 2(0.4) - 0.2 = 0.3$
So $z = [1.0, 0.3]$

Why this step? Matrix multiplication and addition—standard forward pass. 2. Apply ReLU: $y = \max(0, z) = [1.0, 0.3]$ (both positive, so unchanged) Why this step? ReLU zeros out negatives; here nothing is zeroed.

Compute loss: $L = y_1 + y_2 = 1.3$ Why this step? We need a scalar to backpropagate from.
Backward through sum: $\frac{\partial L}{\partial y_1} = 1$ , $\frac{\partial L}{\partial y_2} = 1$ Why? $L = y_1 + y_2$ , so derivatives are constants.
Backward through ReLU: $\frac{\partial L}{\partial z_i} = \begin{cases}1 & \text{if } z_i > 0\ 0 & \text{if } z_i \le 0\end{cases}$ Here both $z$ elements are positive, so $\frac{\partial L}{\partial z} = [1, 1]$ . Why? ReLU derivative is 1 where input is positive, 0 otherwise.
Backward through linear layer: $\frac{\partial L}{\partial W_{ij}} = \frac{\partial L}{\partial z_j} \cdot x_i$

For example, $\frac{\partial L}{\partial W_{11}} = \frac{\partial L}{\partial z_1} \cdot x_1 = 1 \cdot 1 = 1$ Similarly: $W.\text{grad} = x^T \cdot \frac{\partial L}{\partial z} = \begin{bmatrix}1\\2\end{bmatrix} \cdot [1, 1] = \begin{bmatrix}1 & 1\\2 & 2\end{bmatrix}$

Why this step? From $z = xW$ , by chain rule: $\frac{\partial L}{\partial W} = x^T \frac{\partial L}{\partial z}$ (standard matrix calculus).

Worked example Example 2: Understanding

retain_graph=True Problem: Backpropagate through the same graph twice.

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
y.backward()  # First backward
print(x.grad)  # tensor([4.])
# Try second backward (will error!)
# y.backward()  # RuntimeError: graph has been freed
# Solution: retain graph
x.grad.zero_()
y = x ** 2
y.backward(retain_graph=True)
print(x.grad)  # tensor([4.])
y.backward()  # Works now!
print(x.grad)  # tensor([8.]) - accumulated!

Why this happens: PyTorch frees the graph after .backward() to save memory (dynamic graphs are use-once by default). Setting retain_graph=True keeps it alive.

When you need this: Computing multiple losses from the same forward pass (e.g., adversarial training, multi-task learning with shared encoder).

Derivation of accumulated gradient: After first .backward(): $x.\text{grad} = \frac{dy}{dx} = 2x = 4$ After second .backward(): $x.\text{grad} = 4 + 2x = 4 + 4 = 8$ Why accumulation? Each .backward() adds to .grad. Always zero gradients unless you explicitly want accumulation.

Worked example Example 3: Detaching Tensors to Stop Gradient Flow

Problem: Compute a loss that depends on some intermediate value, but don't backpropagate through it.

x = torch.tensor([3.0], requires_grad=True)
# Compute intermediate
hidden = x ** 2  # hidden = 9
# Use hidden as a CONSTANT (stop gradient)
output = hidden.detach() * x  # output = 9 * 3 = 27
output.backward()
print(x.grad)  # tensor([9.])

Why 9, not 15? If we hadn't detached: $\text{output} = x^2 \cdot x = x^3$ $\frac{d(\text{output})}{dx} = 3x^2 = 3(^2) = 27$ But with detach, hidden is treated as a constant $c = 9$ : $\text{output} = c \cdot x$ $\frac{d(\text{output})}{dx} = c 9$

When you need this: Reinforcement learning (treating old policy as constant), GAN discriminator freezing, or any time you want to stop gradients at a boundary.

Common Mistakes

Common mistake Mistake 1: Forgetting to Zero Gradients

Wrong approach:

for epoch in range(3):
   loss = model(x)
   loss.backward()
   optimizer.step()  # Gradients keep accumulating!

Why it feels right: You assume .backward() resets gradients like a clean slate.

The truth: PyTorch accumulates gradients (useful for gradient accumulation across mini-batches). Without zeroing, your gradients grow unbounded: Iteration 1: $\theta.\text{grad} = g_1$ Iteration 2: $\theta.\text{grad} = g_1 + g_2$ (wrong!) Iteration 3: $\theta.\text{grad} = g_1 + g_2 + g_3$ (very wrong!)

Fix:

for epoch in range(3):
   optimizer.zero_grad()  # Clear gradients
   loss = model(x)
   loss.backward()
   optimizer.step()

Rule of thumb: Call optimizer.zero_grad() (or tensor.grad.zero_()) at the START of each training iteration.

Common mistake Mistake 2: Calling

.backward() on Non-Scalar Wrong approach:

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
y.backward()  # RuntimeError!

Why it feels right: You think "backpropagation works on any tensor."

The truth: .backward() computes $\frac{\partial L}{\partial x}$ where $L$ is a scalar. If y is a vector, PyTorch doesn't know which scalar to differentiate with respect to (is it the sum? Mean? A specific element?). Mathematical reason: The gradient $\frac{\partial \mathbf{y}}{\partial x}$ is a Jacobian matrix (not a vector). You must reduce to a scalar first.

Fixes:

# Option 1: Sum to scalar
y.sum().backward()
 
# Option 2: Provide gradient argument
y.backward(torch.ones_like(y))  # Equivalent to sum
# Option 3: Specific element
y[0].backward()  # Only backprop through first element

When each is useful:

Sum: Standard for loss functions
Gradient argument: Custom weighting (e.g., importance sampling)
Specific element: Debugging individual neurons

Common mistake Mistake 3: In-Place Operations Breaking Autograd

Wrong approach:

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
y += 1  # In-place operation
y.backward()  # RuntimeError: modified by inplace operation

Why it feels right: In-place ops like += are memory-efficient.

The truth: Autograd needs the original value of y to compute gradients correctly. In-place ops destroy history.

Derivation of the problem:

Say we want $\frac{dL}{dx}$ where $L = (x^2 + 1)$ . We need: $\frac{dL}{dx} = \frac{d(x^2 + 1)}{dx} = 2x$

But after y += 1, the graph only sees y (the modified value), not the original $x^2$ . The link is broken.

Fix:

y = x ** 2
y y + 1  # Out-of-place (creates new tensor)
y.backward()

Exceptions: In-place ops on leaf tensors (user-created inputs) are forbidden. On intermediate tensors, they work if you're careful, but generally avoid them during training.

Dynamic vs Static Graphs

The `.grad_fn` Chain

Every tensor created by an operation stores a reference to the function that created it:

x = torch.tensor([1.0], requires_grad=True)
y = x * 2
z = y + 3
 
print(y.grad_fn)  # <MulBackward0>
print(z.grad_fn)  # <AddBackward0>
print(z.grad_fn.next_functions)  # Shows MulBackward0

This chain is the computational graph. Following grad_fn pointers backward reconstructs the entire computation.

Leaf tensors (created by user, not operations) have grad_fn=None and is_leaf=True:

print(x.is_leaf)  # True
print(y.is_leaf)  # False (created by operation)

Why this matters: Only leaf tensors store gradients by default. Intermediate gradients are freed unless you call retain_grad().

Advanced: `torch.no_grad()` and Inference Mode

During inference, you don't need gradients. Disabling autograd saves memory and speeds up computation:

with torch.no_grad():
    output = model(x)  # No graph built

Why this works: PyTorch skips grad_fn assignment, eliminating overhead.

Mathematical insight: In inference, we only need the forward pass: $\hat{y} = f(x; \theta)$

No need to compute $\frac{\partial L}{\partial \theta}$ , so no graph required.

Inference mode (newer, faster):

with torch.inference_mode():
    output = model(x)

Difference: inference_mode() is more aggressive—tensors created inside cannot be used in gradient computation later (even outside the context). Use for pure inference; use no_grad() if you might need gradients later.

Memory Management

Strategies:

Gradient checkpointing: Recompute intermediate activations during backward pass (trade compute for memory)
Delete graphs after use: PyTorch's default (graphs freed after .backward())
Detach unnecessary tensors: Stop gradients where not needed

Formula for memory cost: $\text{Memory} \approx O(N \cdot S)$ where $N$ = number of layers, $S$ = activation size. Gradient checkpointing reduces this to $O(\sqrt{N} \cdot S)$ at cost of $2 \times$ compute.

Active Recall Exercises

Recall Explain to a 12-Year-Old

Imagine you're building a really tall tower with LEGO blocks. Every time you place a block, you write down on a piece of paper: "I put block #5 on top of block #3." That list of notes is your computational graph—it's a record of how you built the tower.

Now, at the end, someone tells you, "The tower is 2 inches too tall." You need to figure out which blocks to remove. You start at the top and work backward: "If I remove this block, the tower gets 1 inch shorter. If I remove the one below, it gets another inch shorter." That's backpropagation—going backward through your notes to see how each piece contributed to the final height.

PyTorch's autograd is like having a robot that automatically writes those notes for you and calculates how changing each block affects the tower's height. You just build (forward pass), and the robot figures out the changes (backward pass) for you!

Connections

3.301-Introduction-to-PyTorch-tensors-and-operations - Foundation of tensors that autograd operates on
3.2.04-Backpropagation-algorithm - The mathematical algorithm autograd implements
3.4.01-Building-neural-networks-with-torch.nn.Module - Uses autograd to train models
3.5.02-Gradient-descent-variants - Optimizers that consume the gradients computed by autograd
4.2.03-Batch-normalization - Layer that requires careful handling of running stats (detach from graph)
5.3.01-Transfer-learning - Fine-tuning requires controlling which layers have requires_grad
6.4.02-Gradient-checkpointing - Advanced memory optimization technique for deep networks

#flashcards/ai-ml

What is a computational graph in PyTorch? :: A directed acyclic graph (DAG) where nodes are tensors/operations and edges are data flow. PyTorch builds it dynamically during the forward pass and uses it during backward pass to compute gradients via the chain rule.

What does requires_grad=True do?

Marks a tensor for gradient tracking. PyTorch records all operations on such tensors, building the computational graph that enables automatic differentiation during backpropagation.

Why must you call optimizer.zero_grad() before each training iteration?

PyTorch accumulates gradients by default (adds new gradients to existing .grad values). Without zeroing, gradients from previous iterations add up incorrectly, causing wrong parameter updates.

What's the difference between torch.no_grad() and torch.inference_mode()?

Both disable gradient tracking. no_grad() allows tensors to later be used in gradient computation (after exiting the context). inference_mode() is more aggressive—tensors created inside cannot participate in autograd later, but it's faster for pure inference.

Why does .backward() require a scalar input?

Backpropagation computes

\frac{\partial L}{\partial x}

where

L

is scalar. For vector outputs, the gradient would be a Jacobian matrix. You must reduce to scalar (via sum, mean, or indexing) or provide a gradient argument.

What does .detach() do?

Creates a new tensor that shares data with the original but has requires_grad=False and grad_fn=None. Gradients do not flow through detached tensors—they're treated as constants in the computational graph.

What is grad_fn?

A reference to the function that created a tensor (e.g., <MulBackward>, <AddBackward>). It forms the computational graph—following grad_fn pointers backward reconstructs the sequence of operations needed for backpropagation.

Why do in-place operations (like +=) sometimes break autograd?

Autograd needs original values to compute correct gradients. In-place ops overwrite history. If you modify a tensor that's needed later in the backward pass, the gradient computation fails. Use out-of-place operations (y = y + 1 instead of y += 1) during training.

What's the chain rule formula for a computational graph path

x \to a \to b \to z

\frac{\partial z}{\partial x} = \frac{\partial z}{\partial b} \cdot \frac{\partial b}{\partial a} \cdot \frac{\partial a}{\partial x}

. Autograd computes this by traversing the graph backward, multiplying local gradients at each node.

When would you use retain_graph=True?

When you need to call .backward() multiple times on the same graph (e.g., computing multiple losses from one forward pass in multi-task learning or adversarial training). By default, PyTorch frees the graph after first .backward() to save memory.

What's a leaf tensor?

A tensor created directly by the user (not as result of an operation). It has is_leaf=True and grad_fn=None. Only leaf tensors store gradients by default; intermediate tensors need retain_grad() to keep their gradients.

Derive $\frac{d}{dx}(x^2)$ from the limit definition :: $\frac{d}{dx}(x^2) = \lim_{h \to 0} \frac{(x+h)^2 - x^2}{h} = \lim_{h \to 0} \frac{x^2 + 2xh + h^2 - x^2}{h} = \lim_{h \to 0} \frac{2xh + h^2}{h} = \lim_{h \to 0}(2x + h) = 2x$

Why does gradient accumulation happen by default in PyTorch?

To support mini-batch training where you sum gradients across multiple samples:

\frac{\partial L_{\text{batch}}}{\partial w} = \frac{1}{N}\sum_{i=1}^N \frac{\partial L_i}{\partial w}

. Accumulation implements this sum. You must manually zero between training iterations.

What's the memory complexity of storing a computational graph?

O(N \cdot S)

where

N

is the number of operations/layers and

S

is the size of activations. Each operation stores inputs and metadata for the backward pass. Gradient checkpointing reduces this to

O(\sqrt{N} \cdot S)

by recomputing activations.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Dekho, autograd basically PyTorch ka woh magical engine hai jo automatically derivatives (gradients) nikaal deta hai — tumhe khud haath se calculus karne ki zarurat nahi padti. Jab tum neural network train karte ho, tumhe loss function ka gradient nikaalna padta hai lakhon parameters ke respect me. Ye kaam manually karna impossible hai aur galtiyaan hone ki full guarantee. Isliye PyTorch ek "tape recorder" ki tarah kaam karta hai — jab tum forward pass me tensors pe operations karte ho (jaise multiply, add, power), to woh chupchaap ek computational graph banata jaata hai jisme har operation record hota hai. Jaise hi tum .backward() call karte ho, woh is graph ko ulta traverse karke, chain rule laga ke saare gradients ek hi baar me nikaal deta hai.

Ismein main trick ye hai ki tum tensor pe requires_grad=True set karte ho — matlab "iss tensor ko track karo". Fir har naya tensor jo operations se banta hai, uske paas ek grad_fn hota hai (jaise <MulBackward>, <AddBackward>) jo batata hai ki woh kis operation se aaya. Yehi grad_fn saare tensors ko jodkar poora graph bana deta hai. Backward pass me PyTorch bas chain rule apply karta hai — agar $z = xy + x^2$ hai, to $dz/dx = y + 2x$ automatically nikal aata hai, aur tumhe answer .grad attribute me mil jaata hai. Yaani jo math tum kaagaz pe limit definition se derive karte, woh PyTorch milliseconds me kar deta hai.

Ek important baat jo dhyaan me rakhni hai — .backward() by default gradients ko accumulate karta hai, matlab purane gradient me naya add hota jaata hai. Ye mini-batch training ke liye kaam ka hai kyunki tum multiple samples ke gradients sum karte ho, par iska matlab hai ki har iteration ke baad tumhe x.grad.zero_() se gradients clear karne padte hain, warna galat purane values add hote rahenge aur tumhara training kharab ho jaayega. Ye samajhna kyun zaroori hai? Kyunki deep learning ka poora foundation isi automatic differentiation pe khada hai — bina iske tum bade networks efficiently train hi nahi kar sakte, aur ye concept tumhe har PyTorch project me kaam aayega.

Autograd and computational graphs in PyTorch

Overview

Core Concepts

How Autograd Works: Step-by-Step Derivation

Step 1: Forward Pass and Graph Construction

Step 2: Backward Pass and Chain Rule

Step 3: Gradient Accumulation

Detailed Examples

Common Mistakes

Dynamic vs Static Graphs

The `.grad_fn` Chain

Advanced: `torch.no_grad()` and Inference Mode

Memory Management

Active Recall Exercises

Connections

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Deep Learning Frameworks

Connections