The chain rule

Problem

The previous page differentiated functions you could write in one line. A network's loss is nothing like that. It is loss(softmax(linear(relu(linear(x))))), dozens or hundreds of layers deep, with every weight buried somewhere in the middle. Writing out ∂loss/∂w for each weight as one formula is hopeless, and estimating it with finite differences would take a forward pass per weight.

We need a rule that builds the derivative of a composition out of the derivatives of its parts, and a way to apply it that costs about as much as computing the loss once.

Intuition

Rates of change multiply. If u changes 3 times as fast as x, and f changes 8 times as fast as u, then f changes 3 × 8 = 24 times as fast as x. A chain of functions is a chain of gears, and the overall ratio is the product of the ratios.

The order you multiply in doesn't change the answer, but it changes the work. Starting from the loss and walking back toward the inputs, you carry a single vector, "how much does the loss care about this value", and every layer updates it with one cheap product. That backward walk is reverse mode, and backpropagation is reverse mode applied to a network.

Forward (green) computes values; backward (amber) starts from ∂f/∂f = 1 and multiplies by each operation's local derivative on the way back. Each local derivative only needs values the forward pass already computed.

Mechanics

The scalar chain rule

f(x) = g(h(x))  ⇒  f′(x) = g′(h(x)) · h′(x),   or   df/dx = df/du · du/dx

For f(x) = (3x + 1)², let u = 3x + 1. Then du/dx = 3 and df/du = 2u, so df/dx = 2(3x + 1) · 3. At x = 1, u = 4 and df/dx = 8 · 3 = 24. Expanding (3x + 1)² = 9x² + 6x + 1 and differentiating directly gives 18x + 6 = 24 too; the chain rule just never needs the expansion.

When a value feeds into several places, its derivative is the sum over the paths: if f depends on x through both u and v, then df/dx = ∂f/∂u · du/dx + ∂f/∂v · dv/dx. This is why .grad accumulates: autograd adds up the contribution of each use.

The vector form: multiplying Jacobians

For vector-valued steps the derivatives are Jacobians and the chain rule becomes matrix multiplication. If x → u = h(x) → y = g(u), then

J_y←x = J_y←u · J_u←x    [m, n] = [m, k] · [k, n]

A network with L layers has L Jacobians to multiply. When the final output is a scalar loss, its Jacobian is a row vector, the gradient, and

∇_xL = ∇_yL · J_L · J_L−1 ⋯ J₁

Forward mode versus reverse mode

Matrix products are associative, so that chain can be evaluated from either end.

• Forward mode goes right to left in the formula, in the same order as the computation: push a tangent vector v through as J₁v, then J₂(J₁v), and so on. Each pass is a Jacobian-vector product (JVP) and gives the derivative of every output along one input direction. A full gradient needs one pass per input.
• Reverse mode starts from the output and goes backwards: ∇_yL · J_L, then that times J_L−1, and so on. Each pass is a vector-Jacobian product (VJP) and gives the derivative of one output with respect to every input. A full gradient needs one pass per output.

Training has millions of inputs (the parameters) and one output (the loss), so reverse mode wins by a factor of the parameter count. The price is memory: the backward pass needs the forward pass's intermediate values, so they are all kept until backward runs.

Forward mode pays once per input, reverse mode once per output. With one loss and a billion parameters the gap is a billion to one, which is why every training framework is built on reverse mode.

Local gradients per operation

Reverse mode never forms a Jacobian. Each operation only needs a rule that turns the gradient of its output into gradients of its inputs. For the operations in a two-layer network, with ḡ meaning ∂L/∂(that value):

Notice the transpose in x̄ = Wᵀȳ. Multiplying by a Jacobian's transpose is exactly a VJP, and for a linear layer the Jacobian is W itself.

How PyTorch packages a local rule

Every differentiable operation in PyTorch pairs a forward with a backward, the VJP rule. You can write your own by subclassing torch.autograd.Function: forward computes the output and stashes what backward needs with ctx.save_for_backward; backward receives the output's gradient and returns one gradient per input. The source notes that vjp = backward: they are the same thing. An optional jvp supplies the forward-mode rule. Calling loss.backward() runs torch.autograd.backward, which walks the recorded graph from the loss back to the leaves, calling each node's VJP.

Worked example

Implementation

The same network in PyTorch; autograd reproduces every number in the table:

import torch

x  = torch.tensor([2.0, 1.0], requires_grad=True)
W1 = torch.tensor([[1.0, -1.0], [-1.0, 1.0]], requires_grad=True)
b1 = torch.tensor([1.0, 0.0], requires_grad=True)
w2 = torch.tensor([2.0, 3.0], requires_grad=True)
b2 = torch.tensor(0.0, requires_grad=True)

h = torch.relu(W1 @ x + b1)        # [2, 0]
y = w2 @ h + b2                    # 4
loss = 0.5 * (y - 3.0) ** 2        # 0.5
loss.backward()

W1.grad   # tensor([[4., 2.], [0., 0.]])
w2.grad   # tensor([2., 0.])
x.grad    # tensor([ 2., -2.])

A custom operation with a hand-written VJP, the way PyTorch's own operations are built:

class Square(torch.autograd.Function):
    @staticmethod
    def forward(ctx, u):
        ctx.save_for_backward(u)
        return u * u

    @staticmethod
    def backward(ctx, grad_output):    # grad_output = ∂L/∂(u²)
        (u,) = ctx.saved_tensors
        return grad_output * 2 * u     # ∂L/∂u, by the chain rule

x = torch.tensor(1.0, requires_grad=True)
Square.apply(3 * x + 1).backward()
x.grad                                  # tensor(24.)
torch.autograd.gradcheck(Square.apply, torch.randn(3, dtype=torch.double, requires_grad=True))

Tradeoffs

• Reverse mode trades memory for speed. One backward pass costs about two forward passes, but every intermediate the VJPs need stays in memory until it runs. Activation memory often limits batch size more than the weights do. Gradient checkpointing recomputes some activations during backward to save memory.
• Forward mode still has uses. With few inputs and many outputs, such as the sensitivity of a whole output to one hyperparameter, or Jacobian-vector products inside second-order methods, forward mode is cheaper and needs no stored graph.
• Long products shrink or blow up. A gradient that passes through 100 layers is a product of 100 Jacobians. If their typical scale is below 1 it vanishes; above 1 it explodes. Initialization, normalization and residual connections exist largely to keep that product near 1; the backpropagation page covers it.
• Hand-written backward functions are a bug source. A wrong VJP trains silently worse instead of crashing. Always check a custom backward with gradcheck, which compares it to finite differences in double precision.

Check yourself

Connections

References

• Goodfellow, Bengio and Courville. Deep Learning, section 6.5 (back-propagation and other differentiation algorithms). MIT Press, 2016.
• Rumelhart, Hinton and Williams. Learning representations by back-propagating errors. Nature 323, 1986.
• Baydin et al. Automatic Differentiation in Machine Learning: a Survey. JMLR 2018. arXiv:1502.05767
• Griewank and Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd ed. SIAM, 2008.
• PyTorch v2.15.0: torch.autograd.Function and torch.autograd.backward.