Computational Graphs

Learning Objectives

By the end of this section, you will be able to:

1. Understand computational graphs as the fundamental data structure for representing neural network computations
2. Visualize forward and backward passes as information flowing through the graph in opposite directions
3. Apply the chain rule on graphs to systematically compute gradients for any computation
4. Explain how autograd systems work by recording operations and replaying them in reverse
5. Analyze gradient flow patterns including addition, multiplication, and branching operations
6. Understand the memory vs computation trade-offs in modern deep learning frameworks

Why This Matters: Computational graphs are the secret weapon of modern deep learning. They transform the complex problem of computing gradients through millions of parameters into a simple, systematic algorithm. Every major deep learning framework—PyTorch, TensorFlow, JAX—is built on this foundation.

The Big Picture

Imagine you need to find the derivative of a complex function involving dozens of operations and millions of parameters. Doing this by hand would be impossibly tedious and error-prone. Computational graphs provide an elegant solution: represent the computation as a directed graph, then compute gradients automatically using the chain rule.

The Origin Story

The idea of using graphs to represent computations dates back to the 1960s, but the breakthrough for neural networks came in the 1980s with the development of backpropagation. Researchers realized that by structuring computations as directed acyclic graphs (DAGs), they could:

1. Compute outputs efficiently in a single forward pass
2. Compute all gradients efficiently in a single backward pass
3. Reuse intermediate values, avoiding redundant computation

This insight is what makes training neural networks with billions of parameters feasible. Without computational graphs and efficient backpropagation, modern AI would not exist.

Two Key Insights

What Is a Computational Graph?

A computational graph is a directed graph where:

• Nodes represent values (inputs, intermediate results, outputs)
• Edges represent operations or data flow between values

A Simple Example

Consider the function $L = (wx + b - y)^2$, which is the squared error loss for linear regression. We can break this into elementary operations:

Each operation becomes a node in our graph, with edges showing the flow of data from inputs to outputs. This decomposition is key: each elementary operation has a simple, known gradient.

Why Graphs?

Using a graph structure provides several advantages:

• Automatic differentiation: Gradients computed systematically, not derived by hand
• Efficiency: Intermediate values cached during forward pass, reused in backward pass
• Generality: Works for any computation that can be expressed as a DAG
• Parallelization: Independent operations can be computed in parallel

Anatomy of a Computational Graph

Let's examine the components of a computational graph in detail:

Node Types

Edge Information

Each edge in the graph carries two pieces of information:

1. Forward direction: The value flowing from one node to the next
2. Backward direction: The local gradient $\partial \text{output} / \partial \text{input}$

The Gradient Function

Every operation in the graph has an associated gradient function (or grad_fn in PyTorch terminology). This function computes how the operation's output changes with respect to each input.

Interactive: Forward and Backward Pass

Explore the computational graph for $L = (wx + b - t)^2$. In the forward pass, values flow from inputs to outputs. In the backward pass, gradients flow from outputs back to inputs.

The Chain Rule on Graphs

The chain rule is the mathematical foundation of backpropagation. On computational graphs, it takes a particularly elegant form.

Single Path

For a simple chain $x \to f \to g \to L$, the chain rule gives:

This is just multiplying local gradients along the path from $L$ to $x$.

Multiple Paths

When a node has multiple outgoing edges (its value is used by multiple downstream operations), we sum the gradients from all paths:

This is the multivariate chain rule. If $x$ influences $L$ through multiple intermediate nodes $y_1, y_2, \ldots$:

The Backward Pass Algorithm

The backward pass computes gradients efficiently by:

1. Starting at the output with $\partial L / \partial L = 1$
2. Processing nodes in reverse topological order
3. At each node, multiplying incoming gradient by local gradient
4. Accumulating gradients when paths merge (multiple outgoing edges)

Interactive: Chain Rule Exploration

Explore how the chain rule applies to different computational graph structures. Click on different examples to see single-path chains, multi-path graphs, and neural network layers.

Key Gradient Flow Patterns

Understanding how gradients flow through common operations is essential for debugging neural networks and designing architectures.

Addition Gate

For $z = x + y$:

Pattern: Addition acts as a gradient distributor. The incoming gradient is passed unchanged to both inputs.

Multiplication Gate

For $z = x \cdot y$:

Pattern: Multiplication acts as a gradient swapper. Each input receives the gradient multiplied by the other input's value.

Max Gate

For $z = \max(x, y)$:

Pattern: Max acts as a gradient router. The full gradient flows only to the larger input; the smaller input receives zero gradient.

Copy/Fan-out

When a value is used multiple times:

Pattern: Copying acts as a gradient accumulator. Gradients from all uses are summed.

Autograd: Automatic Differentiation

Modern deep learning frameworks implement automatic differentiation (autograd) to compute gradients without manual derivation. Understanding how this works demystifies framework internals.

The Two Phases

1. Forward Pass (Tape Recording): As operations execute, the framework records each operation, its inputs, and its output onto a "tape" (also called a "trace" or "computational graph").
2. Backward Pass (Tape Playback): Starting from the output, the framework walks the tape in reverse, applying the chain rule at each step to compute gradients.

What Gets Stored?

For each operation, the tape stores:

• The operation type (add, multiply, matmul, etc.)
• References to input tensors
• The output tensor
• A gradient function (grad_fn) that knows how to compute local gradients
• Any intermediate values needed for the backward pass

Define-by-Run vs Define-and-Run

PyTorch's define-by-run approach means the graph is constructed fresh for each forward pass, enabling dynamic computation (different graphs for different inputs) and easier debugging.

Interactive: How Autograd Works

Step through PyTorch code and see how the computational tape is built during the forward pass and traversed during the backward pass.

x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

When requires_grad=True, PyTorch will record operations on these tensors.

Implementation in PyTorch

Let's see how computational graphs work in PyTorch, from basic tensor operations to understanding the grad_fn chain.

z.grad_fn = <MulBackward0>, y.grad_fn = <AddBackward0>

∂L/∂w = 2y × x = 2(7)(2) = 28

1import torch
2
3# Create input tensors with gradient tracking
4x = torch.tensor(2.0, requires_grad=True)
5w = torch.tensor(3.0, requires_grad=True)
6b = torch.tensor(1.0, requires_grad=True)
7
8# Forward pass - operations are recorded
9z = w * x        # z = 6.0
10y = z + b        # y = 7.0
11loss = y ** 2    # loss = 49.0
12
13# Examine the computational graph
14print(f"loss.grad_fn: {loss.grad_fn}")
15print(f"  -> next: {loss.grad_fn.next_functions}")
16
17# Backward pass - traverse graph in reverse
18loss.backward()
19
20# Gradients are now computed
21print(f"\n∂L/∂w = {w.grad}")  # 2*y*x = 2*7*2 = 28
22print(f"∂L/∂b = {b.grad}")    # 2*y*1 = 2*7 = 14
23print(f"∂L/∂x = {x.grad}")    # 2*y*w = 2*7*3 = 42

Gradient Accumulation

An important detail: gradients in PyTorch are accumulated, not replaced. If you call backward() multiple times without zeroing gradients, they add up:

1import torch
2
3x = torch.tensor(3.0, requires_grad=True)
4
5# First forward-backward
6y1 = x ** 2
7y1.backward()
8print(f"After first backward: x.grad = {x.grad}")  # 6.0
9
10# Second forward-backward (without zeroing)
11y2 = x ** 3
12y2.backward()
13print(f"After second backward: x.grad = {x.grad}")  # 6.0 + 27.0 = 33.0
14
15# The right way: zero gradients before each backward
16x.grad.zero_()
17y3 = x ** 2
18y3.backward()
19print(f"After zeroing and backward: x.grad = {x.grad}")  # 6.0

Common Operations and Their Gradients

Here are the gradient rules for common neural network operations. Understanding these helps debug gradient flow issues.

Memory vs Computation Trade-offs

Building computational graphs comes with memory costs. Understanding these trade-offs is crucial for training large models.

What Gets Stored?

During the forward pass, activations (intermediate values) must be saved for the backward pass. For a network with $L$ layers:

• Activations: All intermediate outputs (grows with batch size and model width)
• Graph structure: Links between operations (relatively small)
• Intermediate values: Values needed for gradients (e.g., for BatchNorm)

Memory-Saving Techniques

1# Disable graph building for inference
2with torch.no_grad():
3    output = model(input)  # No grad_fn, saves memory
4
5# Gradient checkpointing for memory-efficient training
6from torch.utils.checkpoint import checkpoint
7output = checkpoint(expensive_layer, input)

Knowledge Check

Test your understanding of computational graphs with this quiz:

Summary

Computational graphs are the foundation of modern deep learning, enabling efficient and automatic gradient computation.

Key Concepts

Key Takeaways

1. Computational graphs decompose complex functions into simple, differentiable operations
2. The forward pass computes values while building the graph; backward pass computes gradients
3. The chain rule on graphs: multiply local gradients along paths, sum across multiple paths
4. Autograd systems record operations to a "tape" and replay in reverse for gradients
5. Understanding gradient flow patterns (add distributes, multiply swaps, max routes) helps debugging
6. Memory-computation trade-offs are important for training large models

Looking Ahead

Now that you understand computational graphs, you're ready to see the full backpropagation algorithm in the next section. We'll derive the exact equations for gradient computation in multi-layer networks and implement backprop from scratch.

Exercises

Conceptual Questions

1. Draw the computational graph for $L = \log(\sigma(wx + b))$ where $\sigma$ is the sigmoid function. Label all intermediate nodes.
2. For the graph in question 1, write out the local gradient at each edge. Then compute $\partial L / \partial w$ using the chain rule.
3. Explain why the gradient at a node that splits into multiple branches must be the sum of gradients from all branches, not the product or average.
4. If a ReLU neuron always outputs 0 (its input is always negative), what happens to its gradient? Why is this called a "dead neuron"?

Coding Exercises

1. Manual Gradient Check: Create a simple computation $y = 3x^2 + 2x + 1$ in PyTorch. Compute $\partial y / \partial x$ using autograd, then verify by computing the gradient manually.
2. Graph Exploration: Build a small neural network (2 layers) in PyTorch. After a forward pass, explore the grad_fn chain by following output.grad_fn.next_functions. Print the operation names.
3. Memory Analysis: Create tensors of increasing size with and without requires_grad=True. Measure memory usage to understand the overhead of gradient tracking.

Challenge Exercise

Implement a Mini Autograd: Create a simple Variable class that:

1. Stores a value and tracks whether it requires gradients
2. Implements __add__, __mul__, and __pow__ that return new Variables and record the operation
3. Has a backward() method that computes gradients using the chain rule

Start with just scalar values before attempting vectors/matrices.

In the next section, we'll use computational graphs to derive the backpropagation algorithm for multi-layer neural networks—the core algorithm that makes deep learning possible.