Introduction
Backpropagation is the workhorse algorithm of modern deep learning. First formalized by Rumelhart, Hinton, and Williams in 1986, it enables efficient gradient computation in neural networks through systematic application of the chain rule. Despite its age, backpropagation remains the foundation of training algorithms in frameworks like PyTorch and TensorFlow.
This post provides a practical guide to backpropagation, covering:
- Core mathematical foundations using the chain rule
- Computational graph perspective - the key to understanding backpropagation
- Implementation walkthrough of my neural network from scratch
- How modern frameworks like PyTorch automate this process
Mathematical Foundations
The Supervised Learning Setup
Given a dataset $\mathcal{D} = \{(x^{(i)}, y^{(i)})\}_{i=1}^{N}$, we seek to learn a function $f: \mathcal{X} \rightarrow \mathcal{Y}$ parameterized by $\theta \in \mathbb{R}^d$. The empirical risk is:
where $\ell$ is a differentiable loss function. Gradient descent updates:
The challenge: computing $\nabla_\theta \mathcal{L}$ efficiently for complex, compositional functions.
The Chain Rule: Multivariate Form
For a composition $y = f(g(x))$ where $g: \mathbb{R}^n \rightarrow \mathbb{R}^m$ and $f: \mathbb{R}^m \rightarrow \mathbb{R}$, the gradient is:
In matrix form, using the Jacobian $J_g \in \mathbb{R}^{m \times n}$:
This is the foundation of backpropagation: propagate gradients backward through Jacobian transposes.
What is a Jacobian? A Concrete Example
The Jacobian is simply a matrix of all partial derivatives. For example, if we have a function that takes 2 inputs and produces 3 outputs:
The Jacobian is a 3×2 matrix containing all the partial derivatives:
In neural networks, the Jacobian tells us how each output neuron changes with respect to each input neuron.
Computational Graph Perspective
The computational graph is the key abstraction that makes backpropagation both elegant and efficient. A neural network can be represented as a directed acyclic graph (DAG) where:
- Nodes represent intermediate values (tensors, scalars)
- Edges represent dependencies and operations
- Each operation knows how to compute its local gradient
Why Computational Graphs Matter
The computational graph perspective provides several critical insights:
- Modularity: Each operation is independent and composable
- Automatic differentiation: The graph structure enables automatic gradient computation
- Dynamic computation: Modern frameworks build graphs on-the-fly during forward pass
Example: Two-Layer Network
Consider the computation:
The computational graph has nodes: $\{x, W_1, b_1, z_1, a_1, W_2, b_2, z_2, \hat{y}, y, L\}$. Each node stores its value during the forward pass and accumulates gradients during the backward pass.
Forward Pass: Building the Graph
During the forward pass, we evaluate nodes in topological order (parents before children). This builds the computational graph dynamically:
for node in topological_order(graph):
node.value = node.operation(*[parent.value for parent in node.parents])
# Store intermediate values for backward pass
node.cache_for_backward()
Backward Pass: Traversing in Reverse
The backward pass traverses the graph in reverse topological order, accumulating gradients using the chain rule:
output.grad = 1.0 # Initialize: dL/dL = 1
for node in reversed(topological_order(graph)):
for parent in node.parents:
# Chain rule: accumulate gradient contributions
local_grad = node.compute_local_gradient(parent)
parent.grad += node.grad * local_grad
The Gradient Accumulation Rule
For a node $v$ with multiple children $\{c_1, c_2, \ldots, c_k\}$ (i.e., its value is used in multiple operations):
This is the multivariate chain rule in action. When a value is used in multiple
places,
gradients from all downstream paths must be summed. This is why we accumulate gradients
with +=
rather than overwriting them.
Local Gradients: The Building Blocks
Each operation only needs to know its local gradient—how its output changes with respect to its inputs. The chain rule handles the rest:
- Matrix multiplication: $\frac{\partial (Wx)}{\partial W} = x^T$, $\frac{\partial (Wx)}{\partial x} = W^T$
- Element-wise activation: $\frac{\partial \sigma(z)}{\partial z} = \sigma'(z)$
- Addition: $\frac{\partial (a+b)}{\partial a} = 1$, $\frac{\partial (a+b)}{\partial b} = 1$
This modularity is what makes automatic differentiation possible: we can compose arbitrary operations, and as long as each knows its local gradient, the full gradient emerges automatically.
Visualizing the Computational Graph
Here's what a computational graph looks like for a simple two-layer network. The forward pass (solid arrows) computes values from left to right, while the backward pass (dashed arrows) propagates gradients from right to left:
Complete Derivation for Feedforward Networks
Network Architecture
Consider an $L$-layer fully-connected network:
Loss Functions
Mean Squared Error (Regression):
Cross-Entropy (Classification):
Softmax + Cross-Entropy (Combined):
Backward Pass: Computing Gradients
Backpropagation computes gradients in three simple steps:
Step 1: Start at the output
Compute how the loss changes with respect to the output layer:
(This is the gradient at the output layer - how wrong we were)
Step 2: Propagate backward through layers
For each layer going backward, multiply by the weight matrix and the activation derivative:
(The $\odot$ means element-wise multiplication)
Step 3: Compute parameter gradients
Once we have gradients for each layer, compute how to update weights and biases:
These gradients tell us how to adjust each weight and bias to reduce the loss.
Implementation from Scratch
Let's implement a modular neural network architecture that demonstrates the computational graph principles. This implementation uses a layer-based design where each layer is responsible for its own forward and backward computations:
Modular Neural Network Architecture
This is the core architecture from my universal-approximator repository. The key insight is to separate concerns: the network manages the computational graph, layers handle their local computations, and the optimizer updates parameters.
import numpy as np
class NeuralNetwork:
def __init__(self):
self.layers = []
self.loss_function = None
self.optimizer = None
def add(self, layer):
"""Add a layer to the network (builds the computational graph)"""
self.layers.append(layer)
def set_loss(self, loss_function):
"""Set the loss function for training"""
self.loss_function = loss_function
def set_optimizer(self, optimizer):
"""Set the optimizer and initialize it with network parameters"""
self.optimizer = optimizer
self.optimizer.initialize(self.layers)
def forward(self, x):
"""Forward pass: traverse the computational graph"""
for layer in self.layers:
x = layer.forward(x)
return x
def backward(self, grad_output):
"""Backward pass: traverse the graph in reverse"""
for layer in reversed(self.layers):
grad_output = layer.backward(grad_output)
def train_step(self, x, y):
"""Single training step: forward, loss, backward, update"""
# Forward pass through network
output = self.forward(x)
# Compute loss
loss = self.loss_function.forward(output, y)
# Backward pass: start with loss gradient
grad = self.loss_function.backward(output, y)
self.backward(grad)
# Update parameters using optimizer
self.optimizer.step(self.layers)
return loss
def predict(self, x):
"""Inference mode: just forward pass"""
return self.forward(x)
Understanding the Architecture
This design embodies the computational graph principles:
- Layers as nodes: Each layer is a node in the computational graph
- Forward builds the graph: The forward pass constructs the graph dynamically
- Backward traverses in reverse: Gradients flow backward through the graph
- Local gradients: Each layer only computes its local gradient
- Separation of concerns: Network structure, optimization, and loss are decoupled
Example Layer Implementation
Here's how a dense (fully-connected) layer would implement the forward and backward methods:
class DenseLayer:
def __init__(self, input_size, output_size):
# He initialization for weights
self.W = np.random.randn(output_size, input_size) * np.sqrt(2.0 / input_size)
self.b = np.zeros((output_size, 1))
self.cache = {}
def forward(self, x):
"""Compute z = Wx + b and cache inputs for backward pass"""
self.cache['x'] = x
return self.W @ x + self.b
def backward(self, grad_output):
"""Compute local gradients using cached values"""
x = self.cache['x']
m = x.shape[1] # batch size
# Parameter gradients (for optimizer)
self.dW = (1/m) * (grad_output @ x.T)
self.db = (1/m) * np.sum(grad_output, axis=1, keepdims=True)
# Gradient w.r.t. input (for previous layer)
return self.W.T @ grad_output
Usage Example
# Build the network (construct computational graph)
model = NeuralNetwork()
model.add(DenseLayer(784, 128))
model.add(ReLULayer())
model.add(DenseLayer(128, 64))
model.add(ReLULayer())
model.add(DenseLayer(64, 10))
# Configure training
model.set_loss(CrossEntropyLoss())
model.set_optimizer(SGD(learning_rate=0.01))
# Training loop
for epoch in range(num_epochs):
for x_batch, y_batch in data_loader:
loss = model.train_step(x_batch, y_batch)
print(f'Epoch {epoch}, Loss: {loss:.4f}')
This architecture makes the computational graph explicit: each layer is a node, and the forward/backward methods define the edges. The network class manages the graph traversal, while layers handle their local computations—exactly as the theory prescribes.
Modern Automatic Differentiation
PyTorch Example
Modern frameworks automate backpropagation through dynamic computational graphs:
import torch
import torch.nn as nn
# Define network
model = nn.Sequential(
nn.Linear(784, 128),
nn.ReLU(),
nn.Linear(128, 64),
nn.ReLU(),
nn.Linear(64, 10)
)
# Forward pass (builds computational graph)
x = torch.randn(32, 784, requires_grad=True)
y_pred = model(x)
# Compute loss
loss = nn.functional.cross_entropy(y_pred, y_true)
# Backward pass (automatic!)
loss.backward()
# Gradients are now in model.parameters()[i].grad
How PyTorch Implements Backpropagation
PyTorch uses reverse-mode automatic differentiation:
- Each operation creates a
Functionobject withforward()andbackward()methods - The computational graph is built dynamically during forward pass
loss.backward()traverses the graph in reverse, callingbackward()on each node- Gradients accumulate in
.gradattributes
Custom Autograd Function
class MyReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return input.clamp(min=0)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input
# Usage
my_relu = MyReLU.apply
Common Challenges
Vanishing and Exploding Gradients
When training deep networks, gradients can become problematically small (vanishing) or large (exploding) as they propagate backward through many layers.
Why does this happen? During backpropagation, we multiply many gradients together. If these values are consistently less than 1, the product shrinks exponentially. If they're greater than 1, it explodes.
Simple solutions:
- ReLU activations: Unlike sigmoid, ReLU doesn't squash gradients (derivative is 1 for positive values)
- Proper initialization: Start with weights that keep gradients stable (He initialization for ReLU)
- Batch normalization: Normalize activations to prevent extreme values
- Gradient clipping: Cap gradient values to prevent explosion
Conclusion
Backpropagation is the algorithm that makes deep learning possible. By systematically applying the chain rule through a computational graph, we can efficiently compute gradients for millions of parameters. The beauty is in its simplicity: each layer only needs to know its local gradient, and the full gradient emerges automatically.
Key Takeaways:
- Backpropagation uses the chain rule to compute gradients efficiently
- Computational graphs make the process visual and modular
- Each operation only needs to compute its local gradient
- Modern frameworks like PyTorch automate this, but the core algorithm remains the same
- Understanding backpropagation helps you debug training issues and design better architectures
References
Check out my universal-approximator repository to see a complete implementation from scratch!
Recommended Learning Resource:
Andrej Karpathy's "Neural Networks: Zero to Hero" (YouTube)
- An excellent video series on building neural networks from scratch