Kiru Lab  /  The Mathematics of Learning  /  The Calculus of Change

The Chain Rule as Plumbing

Backpropagation is the chain rule applied to a graph, cached so nothing is recomputed.

Concept  ·  about 30 minutes

If `y = g(f(x))`, the chain rule says `dy/dx = g'(f(x)) * f'(x)`. Sensitivity multiplies along a path. A deep network is a long composition, so the derivative of the loss with respect to an early weight is a product of many local derivatives.

dL/dw_1 = dL/da_n * da_n/da_{n-1} * ... * da_2/da_1 * da_1/dw_1

Why gradients vanish and explode

Look at that product. If each local derivative is slightly less than one, the product shrinks geometrically with depth and early layers stop learning. If each is slightly greater than one, it grows without bound and training diverges. Vanishing and exploding gradients are not mysterious pathologies; they are what multiplication does over many terms.

  • Residual connections add an identity path so the product includes a term that is exactly one.
  • Normalization layers keep the scale of activations — and therefore local derivatives — near one.
  • Careful initialization sets the initial product close to one at every depth.
  • Gradient clipping caps the product's magnitude after the fact, bluntly but effectively.

Every one of those four techniques is a response to the same equation. Once you see that, architecture choices stop looking like folklore.

Vanishing gradients in four lines
import numpy as np
for local_derivative in (0.9, 1.0, 1.1):
    for depth in (10, 50, 100):
        print(local_derivative, depth, local_derivative ** depth)
# 0.9 ** 100 = 2.6e-5   (vanished)
# 1.1 ** 100 = 1.4e+4   (exploded)
Reverse mode

Backpropagation is the chain rule evaluated right-to-left, from loss backward. That ordering matters: with one scalar loss and millions of parameters, reverse mode computes every derivative in roughly the cost of one forward pass. Forward mode would cost one pass per parameter.


Hold on to

  • Sensitivity multiplies along a computational path
  • Vanishing and exploding gradients are geometric consequences of depth
  • Reverse-mode ordering is what makes training large models affordable

Work through

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

  1. Derive by hand the gradient of a two-layer network with a sigmoid, then check it numerically.
    Hint

    Work backward from the loss: dL/dW2 first, then propagate to dL/dW1.

    Solution

    For a sigmoid hidden layer the local derivative is s(1-s), which peaks at 0.25. That is the vanishing-gradient problem in one number: even in the best case each sigmoid layer multiplies the gradient by at most a quarter, so ten layers gives at most 0.25**10, about 1e-6. Verify your derivation numerically before trusting it.

    import numpy as np
    
    sigmoid = lambda z: 1 / (1 + np.exp(-z))
    
    def forward_backward(x, y, W1, W2):
        z1 = x @ W1
        a1 = sigmoid(z1)
        y_hat = a1 @ W2
        loss = float(((y_hat - y) ** 2).mean())
    
        dy = 2 * (y_hat - y) / y.size
        dW2 = a1.T @ dy
        da1 = dy @ W2.T
        dz1 = da1 * a1 * (1 - a1)          # sigmoid derivative, max 0.25
        dW1 = x.T @ dz1
        return loss, dW1, dW2
  2. Simulate the product of 100 random local derivatives drawn near 1.0 and plot the distribution of the result on a log scale.
    Hint

    Draw 100 values from a distribution centred on 1.0 and take the product.

    Solution

    The distribution of the product is roughly log-normal and spans many orders of magnitude — the median sits below 1 even when the mean of the factors is exactly 1, because the log of the product is a sum of logs and log is concave. This is why depth is unstable by default and why residual connections, which guarantee a path with multiplier exactly 1, changed what depths were trainable.

    import numpy as np, matplotlib.pyplot as plt
    
    rng = np.random.default_rng(0)
    products = [np.prod(rng.normal(1.0, 0.1, size=100)) for _ in range(10_000)]
    
    plt.hist(np.log10(np.abs(products)), bins=60)
    plt.xlabel("log10 |product of 100 factors|"); plt.show()
    print("median:", np.median(products))    # well below 1

Sign in to track your progress through the lab.