The Chain Rule and Gradients

How the Chain Rule routes error signals backward through deep computational graphs and how Gradient Vectors guide multi-dimensional optimization.

26 minBeginnerCode Examples

The Core Thesis: A deep neural network is not a single formula—it is a nested chain of数十 or hundreds of layers, where the output of Layer 1 becomes the input to Layer 2. The Chain Rule is the mathematical relay race that passes error signals backward from the final Loss all the way to the very first layer, while the Gradient Vector bundles every parameter's sensitivity into a single compass pointing toward lower error.

  1. Why Do We Need the Chain Rule? (Composite Functions)

In basic algebra, a function maps an input xx directly to an output y=f(x)y = f(x). In Deep Learning, however, operations are stacked inside one another as Composite Functions:

L(w)=f(g(h(w)))L(w) = f\big(g(h(w))\big)

Think of this as a gear train or a chain of dominoes:

Step 1: Linear Layer
z=wx+bz = wx + b

Weight ww directly controls the pre-activation logit zz.

Step 2: Activation
a=σ(z)a = \sigma(z)

Logit zz controls the non-linear neuron activation aa.

Step 3: Loss Function
L=(a−y)2L = (a - y)^2

Activation aa controls the final prediction error LL.

Notice the problem: the weight ww does not appear directly inside the Loss formula L=(a−y)2L = (a - y)^2! To know how nudging ww changes the Loss LL, we must multiply the sensitivities across every link in the chain.

  1. The Single-Variable Chain Rule (Multiplying Local Slopes)

If variable ww affects zz, and zz affects aa, and aa affects LL, the Chain Rule states that the overall rate of change dLdw\dfrac{dL}{dw} is the product of each intermediate step's rate of change:

dLdw=dLda⋅dadz⋅dzdw\dfrac{dL}{dw} = \dfrac{dL}{da} \cdot \dfrac{da}{dz} \cdot \dfrac{dz}{dw}

The Gear Ratio Analogy: Suppose Gear A turns 3×3\times faster than your hand (dzdw=3\dfrac{dz}{dw} = 3), Gear B turns 2×2\times faster than Gear A (dadz=2\dfrac{da}{dz} = 2), and Gear C turns 5×5\times faster than Gear B (dLda=5\dfrac{dL}{da} = 5). How fast does Gear C turn when you move your hand? Exactly 5×2×3=30×5 \times 2 \times 3 = 30\times faster!

Step-by-Step Worked Example (A 1-Neuron Network):

Suppose input x=2x = 2, weight w=3w = 3, bias b=1b = 1, and true target y=10y = 10. Let us use a simple linear activation a=za = z with Squared Loss L=(a−y)2L = (a - y)^2:
1. Forward Pass (Compute Values)
z=wx+b=(3)(2)+1=7z = wx + b = (3)(2) + 1 = 7
L=(z−y)2=(7−10)2=(−3)2=9L = (z - y)^2 = (7 - 10)^2 = (-3)^2 = 9
2. Backward Pass (Chain Rule)
∂L∂z=2(z−y)=2(−3)=−6\dfrac{\partial L}{\partial z} = 2(z - y) = 2(-3) = -6
∂z∂w=x=2  ⟹  ∂L∂w=(−6)(2)=−12\dfrac{\partial z}{\partial w} = x = 2 \implies \dfrac{\partial L}{\partial w} = (-6)(2) = -12
⚡ Knowledge Check

During backpropagation, a neuron computes z=wx+bz = wx + b with input x=4x = 4. It receives an upstream gradient from the next layer of ∂L∂z=−2.5\dfrac{\partial L}{\partial z} = -2.5. What is the weight gradient ∂L∂w\dfrac{\partial L}{\partial w} and the bias gradient ∂L∂b\dfrac{\partial L}{\partial b}?

A) dL/dw = -10.0 and dL/db = -2.5▼
✓ Correct!Since ∂z∂w=x=4\dfrac{\partial z}{\partial w} = x = 4 and ∂z∂b=1\dfrac{\partial z}{\partial b} = 1, multiplying by the upstream gradient −2.5-2.5 gives ∂L∂w=(−2.5)(4)=−10.0\dfrac{\partial L}{\partial w} = (-2.5)(4) = -10.0 and ∂L∂b=(−2.5)(1)=−2.5\dfrac{\partial L}{\partial b} = (-2.5)(1) = -2.5.
B) dL/dw = 1.5 and dL/db = -2.5▼
✕ Incorrect.The Chain Rule multiplies the upstream gradient by the local derivative (−2.5×4=−10.0-2.5 \times 4 = -10.0), rather than adding them (−2.5+4=1.5-2.5 + 4 = 1.5).

  1. The Multivariable Chain Rule (When Branches Split and Merge)

What happens when a single parameter or activation feeds into multiple downstream pathways? For example, in a Transformer or ResNet, an embedding vector x\mathbf{x} splits into both a Skip Connection (Residual Path) and a Sub-Layer Path, which later merge back together.

The Multivariable Chain Rule (Total Derivative Rule) states that when a variable xx influences the Loss LL through multiple parallel paths u1(x)u_1(x) and u2(x)u_2(x), you SUM the gradients from all paths:

∂L∂x=∂L∂u1∂u1∂x+∂L∂u2∂u2∂x=∑k∂L∂uk∂uk∂x\dfrac{\partial L}{\partial x} = \dfrac{\partial L}{\partial u_1}\dfrac{\partial u_1}{\partial x} + \dfrac{\partial L}{\partial u_2}\dfrac{\partial u_2}{\partial x} = \sum_{k} \dfrac{\partial L}{\partial u_k}\dfrac{\partial u_k}{\partial x}
Rule 1: Gradients in Series MULTIPLY

When operations happen one after another in a single line (w→z→a→Lw \to z \to a \to L), their local derivatives multiply together.

Rule 2: Gradients in Parallel ADD

When a variable fans out into multiple branches that all affect the Loss, the gradients flowing back from those branches add together at the split point!

Why Residual Connections (y=F(x)+x\mathbf{y} = F(\mathbf{x}) + \mathbf{x}) Saved Deep Learning: By the multivariable chain rule, the gradient of a Residual block is ∂L∂x=∂L∂y⋅(∂F∂x+1)\dfrac{\partial L}{\partial \mathbf{x}} = \dfrac{\partial L}{\partial \mathbf{y}} \cdot \left(\dfrac{\partial F}{\partial \mathbf{x}} + 1\right). Even if the layer gradient ∂F∂x\dfrac{\partial F}{\partial \mathbf{x}} shrinks to 00, the +1+1 from the skip branch guarantees the upstream gradient ∂L∂y\dfrac{\partial L}{\partial \mathbf{y}} flows backward unharmed!

  1. Computational Graphs and Local Gate Patterns

Deep learning frameworks like PyTorch and JAX break every complex formula down into a Computational Graph made of tiny elementary gates. Every gate only needs to know its own local inputs and the Upstream Gradient arriving from the right:

Gate TypeForward FormulaBackward Rule (Given Upstream g=∂L∂zg = \dfrac{\partial L}{\partial z})Intuitive Behavior
Add Gate (+)z=x+yz = x + y∂L∂x=1⋅g,∂L∂y=1⋅g\dfrac{\partial L}{\partial x} = 1 \cdot g, \quad \dfrac{\partial L}{\partial y} = 1 \cdot gGradient Distributor: Copies upstream gradient equally to both branches.
Multiply Gate (*)z=x⋅yz = x \cdot y∂L∂x=y⋅g,∂L∂y=x⋅g\dfrac{\partial L}{\partial x} = y \cdot g, \quad \dfrac{\partial L}{\partial y} = x \cdot gSwap-Multiplier: Multiplies upstream gradient by the other input's forward value.
Max Gate (ReLU)z=max⁡(x,y)z = \max(x, y)gg to larger input, 00 to smallerGradient Router: Routes 100%100\% of gradient to the winning input and 00 to the loser.
⚡ Knowledge Check

A Multiply Gate computes z=x⋅yz = x \cdot y where x=3x = 3 and y=−4y = -4 (so z=−12z = -12). During backpropagation, the upstream gradient arriving at zz is ∂L∂z=5\dfrac{\partial L}{\partial z} = 5. What are the gradients ∂L∂x\dfrac{\partial L}{\partial x} and ∂L∂y\dfrac{\partial L}{\partial y}?

A) dL/dx = -20 and dL/dy = 15▼
✓ Correct!A Multiply Gate acts as a swap-multiplier: ∂L∂x=y⋅∂L∂z=(−4)(5)=−20\dfrac{\partial L}{\partial x} = y \cdot \dfrac{\partial L}{\partial z} = (-4)(5) = -20, and ∂L∂y=x⋅∂L∂z=(3)(5)=15\dfrac{\partial L}{\partial y} = x \cdot \dfrac{\partial L}{\partial z} = (3)(5) = 15.
B) dL/dx = 15 and dL/dy = -20▼
✕ Incorrect.Remember that ∂(xy)∂x=y\dfrac{\partial (xy)}{\partial x} = y (not xx). You must multiply the upstream gradient by the other input variable!

  1. The Gradient Vector, Contour Lines, and Directional Derivatives

Once the Chain Rule has computed the partial derivative ∂L∂wi\dfrac{\partial L}{\partial w_i} for every parameter in our network, we assemble them into the Gradient Vector ∇wL\nabla_{\mathbf{w}} L:

∇wL=[∂L∂w1amp;∂L∂w2amp;…amp;∂L∂wd]T∈Rd\nabla_{\mathbf{w}} L = \begin{bmatrix} \dfrac{\partial L}{\partial w_1} & \dfrac{\partial L}{\partial w_2} & \dots & \dfrac{\partial L}{\partial w_d} \end{bmatrix}^T \in \mathbb{R}^d

Why does the Gradient Vector always point in the direction of steepest ascent? Because of the Directional Derivative formula! If you take a small step in any unit direction u\mathbf{u} (∥u∥2=1\Vert{}\mathbf{u}\Vert{}_2 = 1), the rate of change of the Loss in that direction is the dot product of the gradient with u\mathbf{u}:

DuL=∇L⋅u=∥∇L∥2∥u∥2cos⁡(θ)=∥∇L∥2cos⁡(θ)D_{\mathbf{u}} L = \nabla L \cdot \mathbf{u} = \|\nabla L\|_2 \|\mathbf{u}\|_2 \cos(\theta) = \|\nabla L\|_2 \cos(\theta)
θ=0∘\theta = 0^\circ (cos⁡θ=+1\cos\theta = +1)

Stepping along +∇L+\nabla L maximizes the dot product, giving the steepest uphill increase in error.

θ=90∘\theta = 90^\circ (cos⁡θ=0\cos\theta = 0)

Stepping perpendicular to ∇L\nabla L causes zero change in Loss! This means ∇L\nabla L is always orthogonal to contour lines.

θ=180∘\theta = 180^\circ (cos⁡θ=−1\cos\theta = -1)

Stepping along −∇L-\nabla L makes cos⁡θ=−1\cos\theta = -1, giving the steepest downhill decrease in error!

  1. Vanishing vs. Exploding Gradients (When the Chain Breaks)

Because a deep network with KK layers multiplies KK local derivatives together via the Chain Rule, repeated multiplication creates two famous failure modes:

∂L∂h1=∂L∂hK⋅∏j=2K∂hj∂hj−1\dfrac{\partial L}{\partial \mathbf{h}_1} = \dfrac{\partial L}{\partial \mathbf{h}_K} \cdot \prod_{j=2}^{K} \dfrac{\partial \mathbf{h}_j}{\partial \mathbf{h}_{j-1}}
1. Vanishing Gradients (∣∂hj∂hj−1∣<1\left\vert{}\dfrac{\partial h_j}{\partial h_{j-1}}\right\vert{} \lt 1)

If each layer's local derivative is around 0.250.25 (like Sigmoid), then across 2020 layers the gradient shrinks to 0.2520≈10−12≈00.25^{20} \approx 10^{-12} \approx 0. Early layers receive zero signal and freeze completely!

Modern Fixes: ReLU / GELU activations, Residual Skip Connections, and LayerNorm.

2. Exploding Gradients (∣∂hj∂hj−1∣>1\left\vert{}\dfrac{\partial h_j}{\partial h_{j-1}}\right\vert{} \gt 1)

If each layer's local derivative is around 2.02.0, then across 2020 layers the gradient balloons to 220≈1,048,5762^{20} \approx 1{,}048{,}576, causing weights to overflow into NaN in a single step!

Modern Fixes: Gradient Clipping (clip_grad_norm_), Kaiming/Xavier initialization, and AdamW.

⚡ Knowledge Check

Why does a 50-layer neural network built with Sigmoid activations fail to train its early layers, whereas a 50-layer network with ReLU and Residual Connections (y=F(x)+x\mathbf{y} = F(\mathbf{x}) + \mathbf{x}) trains smoothly?

A) Sigmoid derivatives (≤ 0.25) multiply to ~0; Skip connections add a +1 gradient highway▼
✓ Correct!By the Chain Rule, multiplying fifty numbers smaller than 0.250.25 crushes the gradient to zero. ReLU has a derivative of 11 for positive inputs, and Residual blocks add +1+1 to the local derivative via the parallel-branch sum rule.
B) Sigmoid cannot output positive numbers▼
✕ Incorrect.Sigmoid outputs strictly positive numbers between 00 and 11. Its flaw in deep hidden layers is that its derivative σ′(z)\sigma^\prime(z) is at most 0.250.25, causing repeated Chain Rule multiplications to vanish.

  1. Visual Explanation: The Chain Rule Computational Graph

Look at how a 2-layer neural network executes the Forward Pass from left to right to cache intermediate activations, and then executes the Backward Pass (Reverse-Mode Autodiff) from right to left by multiplying local derivatives:

  1. Python Implementation: 2-Layer Backpropagation From Scratch

Here is a complete, runnable NumPy implementation of the Matrix Chain Rule across a 2-layer neural network (X→W1→ReLU→W2→Loss\mathbf{X} \to \mathbf{W}_1 \to \text{ReLU} \to \mathbf{W}_2 \to \text{Loss}), showing how shapes and transposes align perfectly during backpropagation:

chain_rule_backprop.pyPython 3 · NumPy
import numpy as np # 1. Setup Input Batch X (4 samples, 3 features) and Targets Y (4, 1)np.random.seed(42)X = np.random.randn(4, 3)Y = np.array([[1.0], [0.0], [1.0], [0.0]]) # Initialize Layer 1 (3 -> 5 neurons) and Layer 2 (5 -> 1 output)W1 = np.random.randn(3, 5) * 0.5W2 = np.random.randn(5, 1) * 0.5 # 2. FORWARD PASS (Cache intermediates for the Chain Rule)Z1 = X @ W1                                 # Shape: (4, 5)H1 = np.maximum(0, Z1)                      # ReLU Activation: (4, 5)YHat = H1 @ W2                              # Final Prediction: (4, 1)loss = np.mean((YHat - Y) ** 2)             # Mean Squared Error # 3. BACKWARD PASS (Matrix Chain Rule from Output to Input)dYHat = (2.0 / len(X)) * (YHat - Y)         # dL/dYHat -> Shape: (4, 1)dW2 = H1.T @ dYHat                          # dL/dW2 = H1^T @ dYHat -> (5, 1) dH1 = dYHat @ W2.T                          # Pass gradient to Layer 1 -> (4, 5)dZ1 = dH1 * (Z1 > 0)                        # Multiply by ReLU local derivativedW1 = X.T @ dZ1                             # dL/dW1 = X^T @ dZ1 -> (3, 5) print("Initial Loss:", round(float(loss), 4))print("Gradient Shapes -> dW1:", dW1.shape, "| dW2:", dW2.shape)

Pro Tip (The Golden Shape Invariant of Gradients): Notice that W1 has shape (3, 5) and its gradient dW1 also has shape (3, 5)! A parameter's gradient tensor must always have the exact same shape as the parameter itself. If you ever forget whether to transpose X.T @ dZ1 or dZ1 @ X.T, just match the outer shapes!

Key Points

✓The Chain Rule computes the derivative of nested composite functions by multiplying local derivatives along the path (∂L∂w=∂L∂a∂a∂z∂z∂w\dfrac{\partial L}{\partial w} = \dfrac{\partial L}{\partial a}\dfrac{\partial a}{\partial z}\dfrac{\partial z}{\partial w}).
✓When a variable splits into multiple parallel paths that later merge (such as Residual skip connections), the multivariable Chain Rule sums the gradients from all branches.
✓In computational graphs, Add gates distribute upstream gradients equally, Multiply gates swap and multiply inputs, and Max/ReLU gates route gradients to the winning input.
✓The Gradient Vector ∇L\nabla L is always orthogonal (perpendicular) to loss contour lines and points in the direction of steepest ascent; −∇L-\nabla L points toward steepest descent.
✓Repeatedly multiplying local derivatives smaller than 11 across deep layers causes Vanishing Gradients, while multiplying derivatives larger than 11 causes Exploding Gradients.
✓For any weight matrix W\mathbf{W} of shape (m×n)(m \times n), its gradient ∂L∂W\dfrac{\partial L}{\partial \mathbf{W}} always has the exact same shape (m×n)(m \times n).

Common Mistakes

✕ Discarding forward-pass activations before running the backward pass.

Because the local derivative of a linear layer Z=XW\mathbf{Z} = \mathbf{X}\mathbf{W} with respect to W\mathbf{W} is XT\mathbf{X}^T, backpropagation requires the forward activations X\mathbf{X} and H1\mathbf{H}_1 to be stored in GPU memory until gradients are computed.

✕ Forgetting to sum gradients when a tensor is used in two places.

Whenever a tensor branches into two operations (like Query/Key/Value projections or skip connections), overwriting the gradient instead of accumulating (+=) violates the multivariable Chain Rule.

✕ Forgetting to zero out accumulated gradients between training steps in PyTorch.

Because PyTorch automatically sums gradients on .backward() to support branching graphs and gradient accumulation, failing to call optimizer.zero_grad() mixes old gradients from previous batches into the current step.

✕ Mismatching matrix transpose order during manual backpropagation.

For Z=XW\mathbf{Z} = \mathbf{X}\mathbf{W}, the weight gradient is always XT⋅∂L∂Z\mathbf{X}^T \cdot \dfrac{\partial L}{\partial \mathbf{Z}} and the input gradient is ∂L∂Z⋅WT\dfrac{\partial L}{\partial \mathbf{Z}} \cdot \mathbf{W}^T. Always verify that the output shape matches the target tensor's shape.

The Big Picture

Isolated Single-Layer Calculus

Can Only Train 1 Layer → No Credit Assignment for Hidden Neurons → Shallow Models Only

Chain Rule + Reverse-Mode Autodiff (Backpropagation)

Final Loss → Multiply Local Slopes Backward Layer-by-Layer → Full Gradient Vector ∇L → Deep Learning

The important conceptual shift is realizing that the Chain Rule solves the Credit Assignment Problem: even if a weight is buried 100 layers deep inside a Transformer, multiplying local derivatives backward along the computational graph tells that weight its exact responsibility for the final prediction error.

Remember: Every time you call loss.backward() in PyTorch, you are executing the Multivariable Chain Rule from right to left across a computational graph to populate the Gradient Vector ∇L\nabla L.