Guides / Matrix calculus

Jacobians and the chain rule: outer Jacobian on the left, and backpropagate with the transpose

The chain rule for vectors is the scalar chain rule with matrices in place of numbers, and matrices don't commute. In numerator layout the Jacobian of a composition is the outer map's Jacobian times the inner map's, in that order. A backward pass runs the same product the other way: for a scalar loss, the gradient with respect to a layer's input is the transposed Jacobian times the gradient with respect to its output. The mistakes come from multiplying in the order the forward pass runs, from evaluating a derivative at the wrong point, and from forgetting the transpose when the matrix is square and the shapes can't warn you.

Why it works

Each map is approximately linear near the point: du ≈ J_f dx, then dy ≈ J_g du. Substituting gives dy ≈ J_gJ_f dx, so the Jacobian of the composition is the product with the outer map on the left: the matrix that acts last is written first, as in y = B(Ax) = (BA)x.

For a scalar loss, transpose the whole product: ∇_xL = J_f^TJ_g^T∇_yL. Evaluated right to left, every step is a transposed Jacobian times a vector, so no Jacobian has to be multiplied by another Jacobian. That is backpropagation, and it is why a network's gradient costs about as much as its forward pass. An elementwise layer's Jacobian is diag(φ′(z)), and multiplying a vector by a diagonal matrix is an elementwise product, so that step never builds the matrix at all.

Example 1: multiply the Jacobians, outer on the left

Let u = (x₁ + x₂, x₁x₂) and y = (u₁², u₁ + u₂). Find ∂y/∂x at x = (1, 2).

1. u = (3, 2)The inner map at the point. The outer Jacobian is evaluated here.
2. ∂u/∂x = 11x₂x₁Row i holds the partials of u_i; at (1, 2) the second row is (2, 1).
3. ∂y/∂u = 2u₁011At u = (3, 2) the first row is (6, 0).
4. 60111121Outer Jacobian on the left.
5. ∂y/∂x = 6632Check directly: y₁ = (x₁ + x₂)² has partials 2(x₁ + x₂) = 6, and y₂ = x₁ + x₂ + x₁x₂ has (1 + x₂, 1 + x₁) = (3, 2).

Example 2: a backward pass through a tanh layer

A layer computes z = Wx, h = tanh(z), with W = 1−120 and x = (1, 1). The loss's gradient arriving at h is g = (1, 2). Find ∇_xL.

1. z = (0, 2)The forward pass: 1 − 1 and 2 + 0.
2. h ≈ (0, 0.9640)tanh 0 = 0 and tanh 2 ≈ 0.9640.
3. tanh′(z) = 1 − h²The derivative at the input z, computed from the saved output h.
4. ≈ (1, 0.07065)1 − 0.9640².
5. ∇_zL ≈ (1, 0.1413)tanh′(z) ⊙ g: the diagonal Jacobian acts entry by entry.
6. ∇_xL = W^T∇_zLThe transposed Jacobian of z = Wx.
7. ≈ (1.2826, −1)Rows of W^T are (1, 2) and (−1, 0). A finite-difference check of L = g^Ttanh(Wx) gives the same numbers.

Example 3: the gradient for the weights

For the layer of Example 2, find ∇_WL.

1. dL = (∇_zL)^TdzL depends on W only through z.
2. dz = (dW)xx is fixed while W changes.
3. dL = tr((∇_zL x^T)^TdW)Write the number as a trace and cycle the product.
4. ∇_WL = ∇_zL x^TAn outer product: 2 × 1 times 1 × 2, the shape of W.
5. ≈ 110.14130.1413Row i is (∇_zL)_i times x^T = (1, 1). Each weight's gradient is the error at its output times the input it multiplies.

Example 4: why backprop multiplies from the loss end

A loss L(h₃) sits on three layers of width n = 1000. Its gradient is the product J₃J₂J₁ of a 1 × n row and two n × n Jacobians. Count the multiplications for each order.

1. (J₃J₂)J₁From the loss end: each step is a row times a matrix.
2. n² + n² = 2 × 10⁶A 1 × n row times an n × n matrix costs n² multiplications, twice.
3. J₃(J₂J₁)From the input end: a matrix times a matrix first.
4. n³ + n² ≈ 10⁹An n × n times n × n product costs n³.
5. about 500 times cheaperThe product is the same either way, since matrix multiplication is associative. Only the cost differs, and starting from the scalar loss keeps every intermediate a vector.

Common mistakes, and the exact line they happen on

1. Multiplying in the order the forward pass runs

The composition of Example 1.

1. ∂u/∂x and ∂y/∂u correctRight so far: both Jacobians are right at the point.
2. (∂u/∂x)(∂y/∂u)This is the mistake. f runs first, so its Jacobian looks like it goes first. Both are 2 × 2, so the product exists and the shapes give no warning, but it is 71131, not 6632. The outer Jacobian goes on the left.

2. Evaluating the derivative at the output

The tanh layer of Example 2.

1. h ≈ (0, 0.9640)Right so far.
2. tanh′(h) ≈ (1, 0.4434)This is the mistake. The derivative belongs at the point tanh was applied to, z = (0, 2). tanh′(h) = 1 − tanh²(h) applies tanh twice. The right factor is 1 − h² ≈ (1, 0.0707); the wrong one is about six times too big in the second entry and sends (2.7735, −1) back to x.

3. Multiplying by W instead of W^T

The last step of Example 2, with ∇_zL ≈ (1, 0.1413).

1. ∇_zL ≈ (1, 0.1413)Right so far.
2. W∇_zL ≈ (0.8587, 2)This is the mistake. The forward pass multiplies by W, and the backward pass multiplies by its transpose. W is square, so the wrong product has the right shape and nothing complains. The gradient is W^T∇_zL ≈ (1.2826, −1); a finite-difference check catches the swap at once.

Practice

1. With A a 5 × 3 matrix and B a 2 × 5 matrix, what is ∂y/∂x for y = B(Ax), and its shape?

   Show answer

   BA, which is 2 × 3.

2. Find the Jacobian of h = ReLU(z) at z = (2, −1, 3).

   Show answer

   diag(1, 0, 1): ReLU′ is 1 where z > 0 and 0 where z < 0.

3. For L = g^T(Wx) with W = 1234 and g = (1, −1), find ∇_xL.

   Show answer

   W^Tg = (1 − 3, 2 − 4) = (−2, −2).

4. Find ∇ sin(a^Tx) at x = (2, −1) for a = (1, 2).

   Show answer

   cos(a^Tx) a, and a^Tx = 0, so (1, 2).

5. For the layer of Example 2, what is ∇_WL if the input is x = (2, 0) instead and ∇_zL is unchanged?

   Show answer

   ∇_zL x^T: first column (2, 0.2826), second column 0. A weight that multiplies a zero input gets no gradient.