Week 6 · Lesson 3
Automatic differentiation
Training means nudging every weight in the direction that lowers the loss, and that direction is the gradient. %g = gradient %loss %w asks your compiler for it. Your compiler answers by building more graph: it walks backward from the loss and applies the chain rule, one op at a time.
What gradient means
For a scalar y and a value x of any shape, gradient y x has x's shape, and element i is ∂y/∂x[i]: how fast y changes when you wiggle x[i] alone. If y = x · x and x is 3, the gradient is 6.
Your compiler computes this while building the program, not while running it. The result is an ordinary graph of muls and adds, so everything else (run, render, hash consing, staging) works on it unchanged. You can even take the gradient of a gradient.
The chain rule, backward
Every op knows its local derivatives. For u = mul a b, ∂u/∂a = b and ∂u/∂b = a. The chain rule says that if you already know g = ∂y/∂u, then u sends g · b to a and g · a to b. Those are a and b's shares of the gradient.
So start at y with g = 1.0 (∂y/∂y), and walk the graph from y back toward x. At each node, add up the shares it has received, then pass on its own shares to its operands. When the walk reaches x, the sum there is the answer. This is reverse mode. One backward walk gives the gradient with respect to every weight at once, which is why it's what trains networks: one loss, millions of weights.
A node has to receive all of its shares before it passes anything on. The spec orders the walk to guarantee that: take the post-order of y's graph (depth first, operands left to right, each node once) and visit it in reverse. Every node comes before the nodes it uses.
The rules for elementwise ops
The spec lists a share for every op. These are the ones from calculus, written as ops (u is the op's own result):
| op | shares, given g |
|---|---|
add a b | g to both |
mul a b | mul g b to a, mul g a to b |
recip a | g · -(u · u), since d(1/a) = -1/a² |
max a b | all of g to the larger, half each on a tie |
where c a b | g to the side c picked, 0 to the other, nothing to c |
exp2 a | g · u · ln 2 |
log2 a | g / (a · ln 2), as mul g (recip (mul a 0.6931472)) |
sqrt a | g / (2u) |
sin a | g · sin(a + π/2): cos, written with sin |
cast, bitcast | g, but only from f32 to f32 |
Comparisons and integer ops have no rule because they don't need one: their results aren't f32, and only f32 nodes carry gradients. A path that leaves f32 is a dead end. That's why floor (a cast to i32 and back) has gradient 0 for |x| < 2^23, which is also the true derivative almost everywhere.
Adding shares up, in order
In (x · x) · x, x receives three shares. In exact arithmetic the order doesn't matter, but in f32 it does: (a + b) + c and a + (b + c) can round differently. So the spec fixes it. Shares are added in the order they arrive during the reverse walk, as a left fold: add c1 c2, then add that and c3. Because the walk runs backward, arrival order is often the reverse of the order the statements appear in the file. The test grad_accumulation_order has three shares whose sum depends on the order, and grad_accumulation_other_order checks that the program-order sum is not what you get.
Which nodes take part
Only nodes on a path from x to y. A node that doesn't depend on x can't pass anything back to x, so the walk skips it. Skipping isn't just faster. A node off the path may be an op with no rule (a reduce op=mul of some unrelated data, say), and that's fine as long as x doesn't reach it. The walk also stops at x: x's own operands don't matter, because x is what you're differentiating against.
What can be on the path is limited. Gradients don't go through call, load, store or loops, so a program like that is invalid. That's the syllabus's rule: autodiff works on pure graphs. x itself may be a load, though. The weights of a model are.
Backpropagation on a tape
A tape is a program as a list: inputs, constants, and add or mul of earlier entries. Return the derivative of the last entry with respect to every entry. Forward for values, backward for gradients.
detach: "treat this as a constant"
detach x has the value of x and passes no gradient. gradient (mul x (detach x)) x is x, not 2x: the detached copy counts as a constant. It's how you stop a gradient on purpose. Lesson 6.5 uses it on the max subtracted for a stable softmax, whose gradient cancels out anyway.
When gradient is invalid
- The loss isn't an f32 scalar, or x isn't an f32 value (a buffer, an i32).
- x is a literal.
gradient %y 2.0would mean every2.0in the program, because equal literals are one hash-consed node. - Either depends on an open range: no gradients inside hand-written loops.
- A
call,loadorstorelies on a path from x to y, or an op with no rule does (reduce op=mul).
If y doesn't depend on x at all, that's not an error. The gradient is zeros of x's shape.
f'(x) for five x at once
Slot 0 holds five values of x. For f(x) = x³ - 2x, store f'(x) for each into slot 1, using gradient. The starter stores f(x).
Check yourself
What is gradient of y = (x · x) · x at x = 2?
d/dx x³ = 3x², or follow the figure.
x receives three shares: 4 from the outer mul (g · c with c = 4), then 4 and 4 from the inner mul (2 · x each). 4 + 4 + 4 = 12 = 3 · 2².
Why reverse mode and not forward mode for training?
Forward mode pushes one input's wiggle forward through the program.
Forward mode needs one pass per input. With a million weights that's a million passes. Reverse mode needs one pass per output, and a loss is one number.
The walk visits nodes in reverse post-order. Why that order?
What would go wrong if a node passed its gradient on before receiving its last share?
Its operands would get a partial gradient. Reverse post-order is a reverse topological order, so users always come first. And because f32 addition isn't associative, a fixed walk also fixes the bits.
What is gradient (mul x (detach x)) x at x = 3?
The detached x is a constant with value 3.
d/dx (x · c) = c = 3. Without the detach it would be 2x = 6.
y = cast (cast x dtype=i32) dtype=f32. What is gradient y x?
Which dtype is the middle node?
The i32 node never receives a share, so nothing reaches x and the gradient is 0. That matches the math: rounding to an integer is flat almost everywhere.
gradient %y 2.0: why is it invalid?
How many nodes does add %a 2.0 and mul %b 2.0 make for the literal?
One. A gradient with respect to it would mix every use of 2.0 together, which is never what anyone means. If you want a variable, load it or stage it.
y uses x, and also reduce %d op=mul of some data %d that x doesn't touch. Is gradient y x valid?
Which nodes receive shares?
Only nodes that depend on x. The reduce's value is used, but it never needs a rule. The test grad_reduce_mul_off_path checks exactly this.