Week 6 · Lesson 4
Gradients through movement ops and reduce
A model is mostly reshapes, broadcasts, matmuls and sums, so the gradient pass has to get through all of them. Each one has a partner: the op that sends a gradient back to where the numbers came from. Learn the pairs and you never have to derive the gradient of a matmul.
Moving numbers, moving them back
A movement op doesn't compute anything. Each output element is some input element, or a zero from padding. So the share an input element receives is the gradient of the output position it went to. For the pure rearrangements, the gradient just goes back through the opposite move:
| op | its gradient, sent back to the operand |
|---|---|
reshape | reshape g back to the operand's shape. Row-major order never changed, so neither does anything else. |
permute order=p | permute g by the inverse order: inv[p[k]] = k |
pad pads=[l0 r0 ...] | shrink g l0 l1 ...: the padding zeros were constants, so their gradient is dropped |
shrink x o0 ... | pad g with zeros around the window: elements outside it didn't affect y |
stack a0 ... an | index g j to aj |
index x j | g padded back to x's shape, at position j |
stage x | stage g: the gradient of a staged value is staged too |
Shrink and index need literal offsets for their gradients, because the pad that undoes them needs constant pad widths. A shrink at a computed offset is fine in a program, but you can't differentiate through it.
expand and reduce are partners
expand copies one value into many places. Each copy affects y, so the original receives the gradient of every copy, added up. The gradient of an expand is a sum over the axes it grew. The spec makes that concrete: permute the grown axes to the front (in increasing order, the rest after), reduce op=add pop= how many there were, then reshape to the operand's shape.
It works the other way too. reduce op=add makes one number out of many, and each of the many contributed with weight 1. So every input gets the full gradient of the sum: the gradient of a reduce is an expand. Reshape g to put size-1 axes where the popped axes were, then expand.
reduce max picks a winner
The max of several numbers depends on only one of them, the largest. So reduce op=max sends the gradient to the position that holds the max and 0 to the rest. When several positions tie, the spec splits the gradient evenly among them, like max a b splits a tie in half. In ops: compare every input to the result expanded back, cast the hits to 1.0, count them with a reduce, and multiply the hits by g / count.
reduce op=mul has no rule. Its derivative is "the product of all the others", and the cheap way to compute that, u / x[i], divides by zero when an input is zero. It's not on any path the course needs, so the spec makes it invalid on a gradient path.
matmul needs no rule
matmul is sugar for reshape, permute, expand, mul and reduce, and each of those has a rule. Run them backward and the textbook result appears by itself. Say A is M×K and B is K×N. For C = A @ B with G = ∂y/∂C:
- The reduce over k expands
Gback along the k axis. - The mul sends A's copy
G · B, and B's copyG · A, elementwise on theK×M×Nblock. - A's expand sums over the N axis it grew, and its permute and reshape move the result back into place.
The result is ∂y/∂A = G Bᵀ and ∂y/∂B = Aᵀ G, the two matmuls every deep learning library hard-codes. Your compiler derives them.
Undoing permute and expand
Write inverse_order, the order that undoes a permute, and summed_axes, the axes an expand grew (its gradient sums over them).
The gradient of a matmul
Slot 0 holds A (2x3) and slot 1 holds B (3x2). For L = the sum of every element of A @ B, store dL/dA into slot 2. The starter takes the gradient with respect to the wrong matrix.
Staging the backward pass
Week 4's fusion recomputes an unstaged value at every use. That's cheap for a forward pass, but a gradient graph reuses things a lot. Each element of the gradient with respect to a matmul's output is used once per input feature, 784 times for MNIST. Recomputing it at every use multiplies that work by 784. The spec's rule for stage helps: a staged value's gradient is staged too. Stage the logits of a layer, and their gradient is computed once into a buffer, exactly where the backward pass needs it.
How you write a gradient's movement ops doesn't change any bits: they only move numbers. You can skip a reshape to the same shape, or merge two permutes. What has to match the spec exactly is the arithmetic: each share's ops, the order shares are added, and the fold order of the reduces in expand's gradient.
Check yourself
What is the gradient of permute x order=[2 0 1], given g?
inv[p[k]] = k.
p = [2 0 1]: inv[2] = 0, inv[0] = 1, inv[1] = 2, so [1 2 0]. Permuting by p and then by inv is the identity. A reshape would put the numbers back in the wrong places.
x has shape (3, 1) and is expanded to (3, 4). What shape does the gradient pass reduce, and over what?
Which axis grew?
Axis 1 grew from 1 to 4. Each of x's 3 values was copied 4 times, so each receives the sum of 4 shares. reduce pops leading axes, so the grown axis is permuted to the front first.
reduce [1 5 5 2] op=max is 5. With g = 1, what share does the element at index 1 get?
How many positions hold the max?
Two positions tie at 5, so each gets 1 / 2 = 0.5, and the others get 0.
The gradient of pad x pads=[2 1] for a 3-element x?
Where in the 6-element output did x's elements go?
Positions 2, 3 and 4. The padded zeros (positions 0, 1 and 5) were constants, so their gradient is dropped.
L is the sum of every element of C = A @ B. What is dL/dA?
G = dL/dC is all ones.
(G Bᵀ)[i][k] = Σⱼ G[i][j] B[k][j] = Σⱼ B[k][j]: the k-th row sum of B, the same for every row i. That is the exercise above.
Why does the stage rule matter for training speed?
How many times is each element of dL/dlogits used in dL/dW?
Once per input feature, 784 times for MNIST. Staged, it's computed once and read 784 times. The value is identical either way.