Chain Rule Behind Backprop

The Multivariable Chain Rule

Consider a simple computational graph:

x --► [ Layer 1: g(x) ] --► z --► [ Layer 2: f(z) ] --► Loss L

We want to find how changing input $x$ affects final Loss $L$:

dL/dx = (dL/dz) * (dz/dx)

For multivariable vector representations ($x \in \mathbb{R}^n, z \in \mathbb{R}^m$):

The chain rule multiplies Jacobian Matrices:

$$\frac{\partial L}{\partial x} = \frac{\partial L}{\partial z} \cdot J_{z, x}$$

  FORWARD PASS:   Inputs x ──► [ Layer 1 ] ──► z ──► [ Layer 2 ] ──► Loss L
                  (Compute & Cache Activations x, z in VRAM)

  BACKWARD PASS:  Inputs x ◄── [ dL/dx = dL/dz · dz/dx ] ◄── dL/dz ◄── Loss L = 1.0
                  (Propagate Loss Gradients Backward via Chain Rule!)

Why Reverse-Mode is 1,000,000x Faster

Suppose a neural network has $N = 100,000,000$ weight parameters and $1$ scalar loss output $L$.

1. Forward-Mode Autodiff: Perturbs 1 input weight at a time. Requires $100,000,000$ separate forward passes to compute all weight gradients!
2. Reverse-Mode Autodiff (Backprop): Starts at final scalar loss $L$ ($\frac{\partial L}{\partial L} = 1.0$) and sweeps backward once. Computes all $100,000,000$ weight gradients in a single backward pass!

Step-by-Step Backprop for a Linear Layer + Activation

Let $z = W x + b$, and $a = \sigma(z)$, and Loss $L = \text{MSE}(a, y)$.

Step 1: Upstream Gradient

Assume we receive upstream loss gradient $\frac{\partial L}{\partial a}$ from layer above.

Step 2: Local Activation Derivative

Pass gradient through non-linear activation $\sigma$:

$$\frac{\partial L}{\partial z} = \frac{\partial L}{\partial a} \odot \sigma'(z) \quad (\text{Hadamard Element-wise Product})$$

Step 3: Weight and Bias Gradients

Compute gradients for weights $W$ and bias $b$:

$$\frac{\partial L}{\partial W} = \frac{\partial L}{\partial z} \cdot x^T, \quad \frac{\partial L}{\partial b} = \frac{\partial L}{\partial z}$$

Step 4: Downstream Gradient

Pass gradient back to previous layer inputs $x$:

$$\frac{\partial L}{\partial x} = W^T \cdot \frac{\partial L}{\partial z}$$

This recursive 4-step sequence repeats through all layers of the neural network!

Say this out loud

Backpropagation applies the multivariable chain rule backward through a computational graph. Reverse-mode automatic differentiation computes exact gradients for all N weights in a single backward pass, making deep learning computationally feasible. Intermediate activation tensors computed during the forward pass must be cached in VRAM to evaluate local partial derivatives during the backward pass.

Follow-ups to expect

• What is Gradient Checkpointing (Activation Recomputation)? A memory optimization that discards intermediate activations during the forward pass and re-computes them on-the-fly during the backward pass, trading 20% extra compute for 60%+ VRAM savings.
• Why does loss.backward() accumulate gradients in PyTorch? PyTorch adds gradients into .grad buffers (param.grad += new_grad) by default. This enables easy gradient accumulation across mini-batches, requiring explicit optimizer.zero_grad() calls before each step.