skip to content

In backprop through a matrix product C = A B, what are the gradients with respect to A and B?

level: middleimportance: must knowfreq 70%

answer

  1. one local rule per operand
  2. the other operand is the one that transposes
  3. the upstream gradient itself is never transposed
  4. shapes decide which side it sits on

basics

~20 s

With G as the upstream gradient of the loss with respect to C, dL/dA = G B^T and dL/dB = A^T G. Matmul has two separate local rules, one per operand, and each carries exactly one transpose.

solid answer

~50 s

A matrix product has two operands, so it has two local rules. Write `C = A B` with `A` of shape (m, k), `B` of shape (k, n), and the upstream gradient `G = dL/dC` of shape (m, n). Then `dL/dA = G B^T` and `dL/dB = A^T G`. Each rule multiplies the upstream gradient by the *other* operand, transposed, and the side it sits on is fixed: `G` stays on the left for `dL/dA` and on the right for `dL/dB`. The fastest sanity check is shapes -- `dL/dA` must be (m, k), and (m, n) times (n, k) is the only way to get there; `dL/dB` must be (k, n), and (k, m) times (m, n) is the only arrangement that conforms. Matrix products do not commute, so swapping the sides is a real error, not a cosmetic one.

code

python · 16 lines
python
A = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]     # 2x3
B = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]   # 3x2
G = [[1.0, 2.0], [3.0, 4.0]]               # 2x2 upstream dL/dC

def mm(X, Y):
    return [[sum(X[i][k] * Y[k][j] for k in range(len(Y)))
             for j in range(len(Y[0]))] for i in range(len(X))]

def t(X):
    return [list(r) for r in zip(*X)]

dA = mm(G, t(B))   # (2x2)(2x3) -> 2x3, same shape as A
dB = mm(t(A), G)   # (3x2)(2x2) -> 3x2, same shape as B

print(dA)   # [[1.0, 2.0, 3.0], [3.0, 4.0, 7.0]]
print(dB)   # [[13.0, 18.0], [17.0, 24.0], [21.0, 30.0]]

go deeper

for a junior

Be able to state that a matrix product node produces two gradients, one for each operand, and that each involves a transpose. Knowing that the result must come back in the same shape as the operand is already most of the battle.

for a middle

Expect to derive both rules on the spot from the index form of the product and to justify the transposes with a shape check. Be ready to say why the two rules are not mirror images of each other.

for a senior

Show that you re-derive rather than recall: given an unusual forward convention or a batched product with a shared operand, produce the correct backward and name the extra summation. Being able to explain why the Jacobian is never materialised signals you understand reverse mode, not just a formula.

for a principal

Own the cost argument. Each matmul backward is roughly one forward's worth of arithmetic per operand, which sets the two-thirds-of-compute-in-backward rule of thumb and drives decisions about recomputation, which activations to retain, and where the memory-versus-compute line sits for a given model.

## The setup A *matrix product* (matmul) node takes two inputs and produces one output: `C = A B`, where `A` has shape (m, k), `B` has shape (k, n), and `C` has shape (m, n). The shared dimension `k` is the one that gets summed over. Backpropagation gives every node an *upstream gradient*: the derivative of the scalar loss `L` with respect to that node's output. Call it `G = dL/dC`, which has the same shape as `C`, namely (m, n). The node's job is to turn `G` into gradients for its inputs -- one the same shape as `A`, one the same shape as `B`. ## Deriving it from indices The definition of the product, entry by entry, is ``` C[i][j] = sum over p of A[i][p] * B[p][j] ``` Apply the chain rule to a single entry of `A`. The entry `A[i][p]` influences `C[i][j]` for every column `j`, and no other row of `C` at all, so ``` dL/dA[i][p] = sum over j of G[i][j] * B[p][j] ``` That sum contracts the second index of `G` with the second index of `B`, which is exactly the (i, p) entry of `G B^T`. Hence `dL/dA = G B^T`. The same argument on `B`: the entry `B[p][j]` influences `C[i][j]` for every row `i`, so ``` dL/dB[p][j] = sum over i of G[i][j] * A[i][p] ``` which contracts the first index of `A` with the first index of `G`, that is, the (p, j) entry of `A^T G`. Hence `dL/dB = A^T G`. ## Why the transposes and the sides are not interchangeable Two separate things are fixed by the derivation, and candidates usually remember one and improvise the other. 1. **Which operand transposes.** The gradient for one input is built from the upstream gradient and the *other* input, transposed. `A` never appears in its own gradient. 2. **Which side the upstream gradient sits on.** Matrix multiplication is not commutative, so `G B^T` and `B^T G` are different objects -- usually one of them is not even a legal product. Both facts fall straight out of shape checking, which is why shape checking is the practical way to recover the rules under interview pressure. `dL/dA` must be (m, k). `G` is (m, n) and `B^T` is (n, k), so `G B^T` is (m, k) -- it fits, and no other arrangement of those two factors does. `dL/dB` must be (k, n). `A^T` is (k, m) and `G` is (m, n), so `A^T G` is (k, n). If you can multiply the pieces at all, you have almost certainly found the right rule. ## The Jacobian you never build Formally the derivative of a matmul is a Jacobian: a linear map from the m*k entries of `A` to the m*n entries of `C`. Written out, that is an (m*n) by (m*k) object, which for realistic layer widths is astronomically large and almost entirely zero. Reverse-mode autodiff never materialises it. It only ever needs the *vector-Jacobian product* -- the action of the transposed linear map on one upstream gradient -- and for a matmul that action is itself just a matmul. This is the general pattern: the backward of a linear op is the adjoint (transpose) of the same linear map, which is why transposes appear at all. ## Consequences worth knowing - **Both operands are needed at backward time.** `dL/dA` uses `B` and `dL/dB` uses `A`, so the node's inputs must still be available when the reverse sweep reaches it. - **The forward convention leaks into the backward.** If a layer is written as `C = A W^T` rather than `C = A W`, the rules re-derive to `dL/dA = G W` and `dL/dW = G^T A`. Do not memorise a layer formula; memorise the plain product rule and re-derive. - **Batched products.** When a stack of matrices is multiplied, the rules apply per batch element. If one operand is shared across the batch rather than stacked, its gradient additionally sums over the batch dimension, because sharing a value across a dimension is a broadcast and the backward of a broadcast is a sum. - **Cost symmetry.** Each backward rule is one matmul of roughly the same arithmetic cost as the forward, which is where the usual rule of thumb -- backward costs about twice the forward -- comes from for a network dominated by matrix products.

  • Why does reverse-mode autodiff never build the full Jacobian of a matmul?
    For `C = A B` with `A` of shape (m, k) and `C` of shape (m, n), the Jacobian with respect to `A` is an (m*n) by (m*k) array -- enormous and mostly zero. Reverse mode only needs its action on one upstream gradient, the vector-Jacobian product, and for a matmul that action is itself a single matmul. Building the Jacobian would cost orders of magnitude more memory and time for the same result.
  • A layer is written as C = A W^T instead of C = A W. How do the two rules change?
    Re-derive rather than patch. With `C = A W^T`, treat the second operand as `W^T`: its gradient is `A^T G`, and transposing back gives `dL/dW = G^T A`. The first operand's gradient becomes `dL/dA = G W`, since the other operand transposed is now `(W^T)^T = W`. The lesson is that the transposes belong to the plain product rule, not to the layer.
  • In a batched product where one operand is shared across all batch elements, what extra step does its gradient need?
    The shared operand's gradient is computed per batch element and then summed across the batch dimension. Reusing one value in many products is a broadcast along that dimension, and the backward of a broadcast is a sum over it, so the per-element contributions accumulate into a single gradient of the operand's own shape.

saying these in an interview costs you the question

  • Gives one rule and assumes the other is symmetric
  • Writes B^T G instead of G B^T, ignoring non-commutativity
  • Transposes the upstream gradient instead of the other operand
  • Claims the Jacobian of a matmul is explicitly formed and stored
  • Says only the weight gradient is needed, not the input gradient

context

resolving links…