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:

$$ \mathcal{L}(\theta) = \frac{1}{N} \sum_{i=1}^{N} \ell(f(x^{(i)}; \theta), y^{(i)}) $$

where $\ell$ is a differentiable loss function. Gradient descent updates:

$$ \theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L}(\theta_t) $$

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:

$$ \frac{\partial y}{\partial x_i} = \sum_{j=1}^{m} \frac{\partial y}{\partial g_j} \frac{\partial g_j}{\partial x_i} $$

In matrix form, using the Jacobian $J_g \in \mathbb{R}^{m \times n}$:

$$ \nabla_x y = J_g^T \nabla_g y $$

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:

$$ \begin{aligned} f_1(x_1, x_2) &= x_1^2 + x_2 \\ f_2(x_1, x_2) &= 2x_1 x_2 \\ f_3(x_1, x_2) &= x_2^2 \end{aligned} $$

The Jacobian is a 3×2 matrix containing all the partial derivatives:

$$ J = \begin{bmatrix} \frac{\partial f_1}{\partial x_1} & \frac{\partial f_1}{\partial x_2} \\ \frac{\partial f_2}{\partial x_1} & \frac{\partial f_2}{\partial x_2} \\ \frac{\partial f_3}{\partial x_1} & \frac{\partial f_3}{\partial x_2} \end{bmatrix} = \begin{bmatrix} 2x_1 & 1 \\ 2x_2 & 2x_1 \\ 0 & 2x_2 \end{bmatrix} $$

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:

$$ \begin{aligned} z_1 &= W_1 x + b_1 \\ a_1 &= \sigma(z_1) \\ z_2 &= W_2 a_1 + b_2 \\ \hat{y} &= \sigma(z_2) \\ L &= \frac{1}{2}\|\hat{y} - y\|^2 \end{aligned} $$

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):

$$ \frac{\partial L}{\partial v} = \sum_{i=1}^{k} \frac{\partial L}{\partial c_i} \frac{\partial c_i}{\partial v} $$

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:

Computational Graph

Complete Derivation for Feedforward Networks

Network Architecture

Consider an $L$-layer fully-connected network:

$$ \begin{aligned} a^{[0]} &= x \\ z^{[l]} &= W^{[l]} a^{[l-1]} + b^{[l]}, \quad l = 1, \ldots, L \\ a^{[l]} &= \sigma^{[l]}(z^{[l]}), \quad l = 1, \ldots, L-1 \\ \hat{y} &= a^{[L]} = \sigma^{[L]}(z^{[L]}) \end{aligned} $$

Loss Functions

Mean Squared Error (Regression):

$$ L_{\text{MSE}} = \frac{1}{2}\|\hat{y} - y\|^2, \quad \frac{\partial L}{\partial \hat{y}} = \hat{y} - y $$

Cross-Entropy (Classification):

$$ L_{\text{CE}} = -\sum_{i} y_i \log(\hat{y}_i), \quad \frac{\partial L}{\partial \hat{y}_i} = -\frac{y_i}{\hat{y}_i} $$

Softmax + Cross-Entropy (Combined):

$$ \frac{\partial L}{\partial z^{[L]}} = \hat{y} - y \quad \text{(simplified gradient)} $$

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:

$$ \frac{\partial L}{\partial z^{[L]}} = \hat{y} - y $$

(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:

$$ \frac{\partial L}{\partial z^{[l]}} = (W^{[l+1]})^T \frac{\partial L}{\partial z^{[l+1]}} \odot \sigma'(z^{[l]}) $$

(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:

$$ \begin{aligned} \frac{\partial L}{\partial W^{[l]}} &= \frac{\partial L}{\partial z^{[l]}} (a^{[l-1]})^T \\ \frac{\partial L}{\partial b^{[l]}} &= \frac{\partial L}{\partial z^{[l]}} \end{aligned} $$

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 Function object with forward() and backward() methods
  • The computational graph is built dynamically during forward pass
  • loss.backward() traverses the graph in reverse, calling backward() on each node
  • Gradients accumulate in .grad attributes

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

← Back to Blog