Skip to content

45. Automatic Differentiation: Letting the Compiler Write the Backward Pass

Mountain goats in space helmets on Titania, a moon of Uranus

What you will understand: how a backward pass can be generated instead of derived by hand, and how to find out whether the generated one is right. Chapters 33 to 36 derived the gradient of an attention classifier, of layer normalisation and of a whole transformer block on paper, wrote each as Mountain Goat code, and checked it with finite differences. Every one of those derivations was a place to make a mistake. This chapter writes a program, autograd.py, that takes a Mountain Goat program computing a loss and writes another Mountain Goat program that computes the loss and the gradient with respect to chosen matrices. The method is reverse-mode automatic differentiation, in about 120 lines of Python over the compiler's own IR. It is then tested four ways: every rule against finite differences of an independent evaluation, ties, the hand-derived gradients of Chapters 33, 34 and 35 (it reproduces all 16 gradient matrices of the transformer block), and ten training steps whose loss curve equals Chapter 33's.

What you need to know first: Chapters 33 to 35 (backpropagation by hand: you will see here that the rules are the same ones, applied by a program), Chapter 41 (a Python program over the text of mgc mlir). The Mountain Goat compiler is used unchanged: it compiles and runs the generated programs.

Compile and run

cd docs/part45/code                                  # plain Python 3; needs the Chapter 44 build (../../part44/code/build.sh) for mgc
./make_examples.py                                   # the ground-truth programs, copied from Chapters 33, 34, 35 (already in examples/)
./autograd.py examples/00_worked_example.mg --wrt w,x > worked_example_generated.mg    # the transformation
./mgc run worked_example_generated.mg                # run what it wrote
./autograd.py examples/h33_attention.mg --wrt q0,wk0,wv0,wo0 --descend 10 --lr 3.0 > train.mg   # ten gradient-descent steps
./descend_compare.py > descend_out.txt               # the loss curve against Chapter 33's (about 20 seconds)
./check_autograd.py > check_autograd_out.txt         # the checks of this page (about 15 seconds)
./mutation.py > mutation_out.txt                     # 24 WRONG versions of autograd.py the checker must catch (about 6 minutes)
./cost.py > cost_out.txt                             # what the generated backward pass costs (about 3 minutes; needs valgrind)
cd ../../part15/code && ./run_lit.sh                 # the whole test suite
Every listing and output on this page comes from these commands.

Why this page shows a mutation run, and why 'caught' is the wanted result

mutation.py writes broken versions of autograd.py (one wrong sign, one missing transpose, ...) and expects the checker to fail on each. "Caught" is good; "NOT CAUGHT" would be a gap. Nothing on this page is a failing test.

The idea in one example

A loss is a chain of operations. To get its derivative with respect to an input, walk the chain backwards: start with "the loss changes by 1 when the loss changes by 1" and, at each operation, turn the sensitivity of its result into sensitivities of its inputs. Each sensitivity is called an adjoint, written g below. The rule for each operation is the calculus you did by hand in Chapters 33 to 35.

The example (examples/00_worked_example.mg): loss = sum of relu(x @ w) * 2.

# The example the chapter walks through: loss = sum of relu(x @ w) * 2.
let x = [[0.5, -1.0], [2.0, 1.0]]
let w = [[1.0, 2.0], [3.0, 4.0]]
print col_sum(row_sum(relu(x @ w) * 2.0))

autograd.py asks the real front end for the program's mg dialect text, inlines function calls so it is one straight-line list of operations, and writes this program:

# generated by autograd.py from 00_worked_example.mg: prints, in order: loss, grad w, grad x
# ---- forward pass: the program's own operations, one `let` each
let x = [[0.5, -1.0], [2.0, 1.0]]
let w = [[1.0, 2.0], [3.0, 4.0]]
let v1 = (x @ w)
let v2 = relu(v1)
let v3 = (v2 * 2.0)
let v4 = row_sum(v3)
let v5 = col_sum(v4)
# ---- backward pass: the operations in reverse; each adjoint g is d(loss)/d(the value it belongs to)
let g6 = [[1.0]]
let ones7 = [[1.0], [1.0]]
let g8 = (ones7 * g6)
let ones9 = [[1.0, 1.0], [1.0, 1.0]]
let g10 = (ones9 * g8)
let g11 = (g10 * 2.0)
let zero12 = [[0.0]]
let g13 = (g11 * ge(v1, zero12))
let g14 = (g13 @ transpose(w))
let g15 = (transpose(x) @ g13)
# ---- results: the loss, then the gradient of each requested matrix
print v5
print g15
print g14

Read it in two halves:

  • Forward (lines v1 to v5): the program's own operations, one let each. v1 = x @ w, v2 = relu(v1), v3 = v2 * 2, v4 = row_sum(v3), v5 = col_sum(v4) is the loss.
  • Backward (lines g6 to g15), in reverse order: g6 = 1 is the adjoint of the loss. col_sum and row_sum spread their adjoint over the elements they added (ones * g). * 2.0 multiplies the adjoint by 2. relu lets the adjoint through only where its input was positive (ge(v1, 0)). matmul gives g @ transpose(w) to its left input and transpose(x) @ g to its right one.

Running it (mgc run) prints the loss and the two gradients:

Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[26]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[4,   4], 
 [2,   2]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[0,   0], 
 [6,   14]]

Checking by hand: x @ w = [[-2.5, -3], [5, 8]], relu keeps [[0, 0], [5, 8]], times 2 and summed is 26. The gradient with respect to w is transpose(x) @ [[0, 0], [2, 2]] = [[4, 4], [2, 2]], and with respect to x it is [[0, 0], [2, 2]] @ transpose(w) = [[0, 0], [6, 14]]. The program printed exactly these.

The rules

autograd.py knows one rule per operation of the mg dialect. In the table g is the adjoint of the result, a and b are the inputs, r is the result:

operation contribution to the adjoint of the inputs
a + b g to each
a - b g to a, -g to b
a * b g * b to a, g * a to b
a / b g / b to a, -(g * a) / (b * b) to b
-a -g
a + k, a - k, k - a, a * k, a / k, k / a g, g, -g, g * k, g / k, -(g * k) / (a * a)
a @ b g @ transpose(b) to a, transpose(a) @ g to b
transpose(a) transpose(g)
reshape(a, m, n) reshape(g, shape of a)
exp(a), log(a), sqrt(a) g * r, g / a, g / (2 * r)
relu(a) g * ge(a, 0)
ge(a, b) nothing (a comparison is flat almost everywhere)
row_sum, col_sum g spread over the summed axis (ones * g)
row_max, col_max g to the entries that attain the maximum, shared equally between ties
a broadcast of a g summed back over the broadcast axes

Two further rules make it a whole algorithm: a value used more than once receives the sum of its contributions (x * x * x uses x three times), and the operations are visited in reverse order, so every adjoint is complete before it is used. Everything else is bookkeeping. The code, in full:

#!/usr/bin/env python3
"""Reverse-mode automatic differentiation for Mountain Goat, as a source-to-source transformation.

  autograd.py prog.mg [--loss N] --wrt NAME[,NAME...]  > grad.mg

reads a Mountain Goat program, takes the N-th `print` (default the first) as the LOSS (it must be a 1x1 matrix), and writes a NEW Mountain Goat program that computes the loss and the
gradient of the loss with respect to each named `let` matrix: the same language, so the result is compiled and run by the ordinary `mgc`, with nothing new in the compiler.

How it works (the whole method is on the page): it asks the real front end for the `mg` dialect text (`mgc mlir`), inlines every function call so the program is one straight-line list of
operations, drops what the loss does not depend on, writes the FORWARD pass out as `let` statements, then walks the operations in REVERSE order. Every operation has a rule that turns the
gradient of its result (the "adjoint") into contributions to the adjoints of its inputs; a value used twice receives the sum of its contributions. The rules are the ones the hand-derived
backward passes of Chapters 33 to 36 use. Supported: add sub mul div neg, scalar operations, matmul, transpose, reshape, exp log sqrt relu, row/col sum and max, broadcast, ge (gradient 0).
Not supported (it says so): rank other than 2, dynamic shapes, permute, contract.
"""
import argparse, ast, math, os, re, subprocess, sys
here = os.path.dirname(os.path.abspath(__file__)); MGC = os.path.join(here, "mgc")

class Unsupported(Exception): pass
def shape_of(t):
    m = re.findall(r"tensor<([\dx?]+)xf64>", t)
    if not m: raise Unsupported("no tensor type in: " + t)
    d = m[-1].split("x")
    if "?" in d: raise Unsupported("dynamic shapes are not supported")
    if len(d) != 2: raise Unsupported(f"rank {len(d)} tensors are not supported")
    return (int(d[0]), int(d[1]))
class Op:
    def __init__(s, res, op, args, attrs, shape): s.res, s.op, s.args, s.attrs, s.shape = res, op, args, attrs, shape
def parse(text):
    funcs = {}; cur = None
    for raw in text.split("\n"):
        line = raw.strip()
        m = re.match(r"func\.func @(\w+)\((.*?)\) -> .*\{$", line)
        if m: cur = dict(params=[(a, shape_of(t)) for a, t in re.findall(r"%(\w+): (tensor<[\dx?]+xf64>)", m.group(2))], ops=[], ret=None); funcs[m.group(1)] = cur; continue
        if cur is None or line in ("}", ""): continue
        m = re.match(r"%(\w+) = mg\.constant dense<(.*)> : (tensor<.*>)$", line)
        if m: cur["ops"].append(Op(m.group(1), "constant", [], {"data": ast.literal_eval(m.group(2))}, shape_of(m.group(3)))); continue
        m = re.match(r"%(\w+) = func\.call @(\w+)\((.*?)\) : .*-> (tensor<.*>)$", line)
        if m: cur["ops"].append(Op(m.group(1), "call", [a.strip()[1:] for a in m.group(3).split(",") if a.strip()], {"callee": m.group(2)}, shape_of(m.group(4)))); continue
        m = re.match(r"mg\.print %(\w+) :", line)
        if m: cur["ops"].append(Op("", "print", [m.group(1)], {}, (0, 0))); continue
        m = re.match(r"func\.return %(\w+) :", line)
        if m: cur["ret"] = m.group(1); continue
        m = re.match(r"%(\w+) = mg\.(\w+) ([^{:]*?)\s*(\{.*\})?\s*: (.*)$", line)
        if m:
            attrs = {k: v.strip().strip('"') for k, v in re.findall(r'(\w+) = ("[^"]*"|[^,}]+?)(?= :|,|\})', m.group(4) or "")}
            cur["ops"].append(Op(m.group(1), m.group(2), [a.strip()[1:] for a in m.group(3).split(",") if a.strip()], attrs, shape_of(m.group(5)))); continue
    return funcs
def flatten(funcs):
    ops = []; counter = [0]
    def inline(fname, argmap, prefix):
        env = dict(argmap)
        for o in funcs[fname]["ops"]:
            if o.op == "call":
                counter[0] += 1; env[o.res] = inline(o.attrs["callee"], {p[0]: env[a] for p, a in zip(funcs[o.attrs["callee"]]["params"], o.args)}, f"{prefix}c{counter[0]}_")
            elif o.op == "print": ops.append(Op("", "print", [env[o.args[0]]], {}, (0, 0)))
            else: name = prefix + o.res; env[o.res] = name; ops.append(Op(name, o.op, [env[a] for a in o.args], dict(o.attrs), o.shape))
        return env.get(funcs[fname]["ret"])
    inline("main", {}, ""); return ops

def lit(data):
    return "[" + ", ".join("[" + ", ".join(repr(float(x)) for x in row) + "]" for row in data) + "]"
def source_lets(path):
    """[(name, literal)] of every `let NAME = [[...]]` in the source, in file order (to give the constants of the IR their names back)."""
    out = []
    for l in open(path):
        m = re.match(r"\s*let\s+(\w+)\s*=\s*(\[\[.*\]\])\s*(#.*)?$", l)
        if m:
            try: out.append((m.group(1), ast.literal_eval(m.group(2))))
            except Exception: pass
    return out
def name_constants(ops, lets):
    """Match the `let` lines to the constants of the flattened program IN ORDER: the front end creates one constant per literal, in the order the source mentions them, so the k-th let is
    the first not-yet-used constant after the previous match that has equal data (this tells apart lets with equal values, like two all-zero bias rows)."""
    consts = [o for o in ops if o.op == "constant"]; pos = 0; found = {}
    norm = lambda d: [[float(x) for x in r] for r in d]
    for name, data in lets:
        for k in range(pos, len(consts)):
            if norm(consts[k].attrs["data"]) == norm(data): found[name] = consts[k].res; pos = k + 1; break
    return found

def grad_program(ops, loss_ssa, wrt, names, steps=0, lr=None):
    """ops: flat forward operations; loss_ssa: the 1x1 value; wrt: list of constant ssa names; names: ssa name -> let name of a constant.
    steps == 0: the program prints the loss and then the gradient of each wrt constant. steps > 0: the program performs `steps` gradient-descent steps (p = p - lr * gradient, all wrt matrices
    at once) and prints the loss before the first step and after each one. Returns the text and the list of printed labels."""
    by = {o.res: o for o in ops if o.res}
    need = set(); stack = [loss_ssa]
    while stack:                                  # keep only what the loss depends on
        n = stack.pop()
        if n in need: continue
        need.add(n); stack += by[n].args
    fwd = [o for o in ops if o.res in need]
    out = []; count = [0]
    def new(prefix="v"): count[0] += 1; return f"{prefix}{count[0]}"
    ones = {}
    def ones_of(shape):
        if shape not in ones: ones[shape] = new("ones"); out.append(f"let {ones[shape]} = {lit([[1.0] * shape[1] for _ in range(shape[0])])}")
        return ones[shape]
    zero = [None]
    def zero11():
        if zero[0] is None: zero[0] = new("zero"); out.append(f"let {zero[0]} = [[0.0]]")
        return zero[0]
    shared = {}                                   # constants that are not differentiated: written once, used by every step
    def forward(params):
        """Write the forward pass; `params` maps each wrt constant to the name of its CURRENT value. Returns ssa name -> variable name."""
        var = {}
        for o in fwd:
            if o.op == "constant":
                if o.res in params: var[o.res] = params[o.res]; continue
                if o.res not in shared: shared[o.res] = names.get(o.res) or new(); out.append(f"let {shared[o.res]} = {lit(o.attrs['data'])}")
                var[o.res] = shared[o.res]; continue
            v = new(); var[o.res] = v; a = [var[x] for x in o.args]
            if o.op in ("add", "sub", "mul", "div"): e = f"({a[0]} {dict(add='+', sub='-', mul='*', div='/')[o.op]} {a[1]})"
            elif o.op == "scalar":
                k = float(o.attrs["value"].split()[0]); sym = dict(add="+", sub="-", mul="*", div="/")[o.attrs["op"]]
                e = f"({k!r} {sym} {a[0]})" if o.attrs["reversed"] == "true" else f"({a[0]} {sym} {k!r})"
            elif o.op == "neg": e = f"(0.0 - {a[0]})"
            elif o.op == "matmul": e = f"({a[0]} @ {a[1]})"
            elif o.op == "transpose": e = f"transpose({a[0]})"
            elif o.op == "reshape": e = f"reshape({a[0]}, {o.shape[0]}, {o.shape[1]})"
            elif o.op in ("exp", "log", "sqrt", "relu"): e = f"{o.op}({a[0]})"
            elif o.op == "ge": e = f"ge({a[0]}, {a[1]})"
            elif o.op == "reduce": e = f"{'row' if int(o.attrs['axis'].split()[0]) == 1 else 'col'}_{o.attrs['kind']}({a[0]})"
            elif o.op == "broadcast": e = f"({a[0]} * {ones_of(o.shape)})"
            else: raise Unsupported(f"operation mg.{o.op} is not supported by autograd")
            out.append(f"let {v} = {e}")
        return var
    def backward(var):
        """Write the backward pass for the forward pass named by `var`; returns ssa name -> variable holding d(loss)/d(that value)."""
        adj = {}
        def add_adj(x, expr):
            if x in adj: v = new("g"); out.append(f"let {v} = ({adj[x]} + {expr})"); adj[x] = v
            else: v = new("g"); out.append(f"let {v} = {expr}"); adj[x] = v
        seed = new("g"); out.append(f"let {seed} = [[1.0]]"); adj[loss_ssa] = seed
        for o in reversed(fwd):
            if o.res not in adj or o.op == "constant": continue
            g = adj[o.res]; a = [var[x] for x in o.args]; r = var[o.res]
            if o.op == "add": add_adj(o.args[0], g); add_adj(o.args[1], g)
            elif o.op == "sub": add_adj(o.args[0], g); add_adj(o.args[1], f"(0.0 - {g})")
            elif o.op == "mul": add_adj(o.args[0], f"({g} * {a[1]})"); add_adj(o.args[1], f"({g} * {a[0]})")
            elif o.op == "div": add_adj(o.args[0], f"({g} / {a[1]})"); add_adj(o.args[1], f"(0.0 - (({g} * {a[0]}) / ({a[1]} * {a[1]})))")
            elif o.op == "neg": add_adj(o.args[0], f"(0.0 - {g})")
            elif o.op == "scalar":
                k = float(o.attrs["value"].split()[0]); sym = o.attrs["op"]; rev = o.attrs["reversed"] == "true"
                if sym == "add": add_adj(o.args[0], g)
                elif sym == "sub": add_adj(o.args[0], f"(0.0 - {g})" if rev else g)
                elif sym == "mul": add_adj(o.args[0], f"({g} * {k!r})")
                elif sym == "div": add_adj(o.args[0], f"(0.0 - (({g} * {k!r}) / ({a[0]} * {a[0]})))" if rev else f"({g} / {k!r})")
            elif o.op == "matmul": add_adj(o.args[0], f"({g} @ transpose({a[1]}))"); add_adj(o.args[1], f"(transpose({a[0]}) @ {g})")
            elif o.op == "transpose": add_adj(o.args[0], f"transpose({g})")
            elif o.op == "reshape": s = by[o.args[0]].shape; add_adj(o.args[0], f"reshape({g}, {s[0]}, {s[1]})")
            elif o.op == "exp": add_adj(o.args[0], f"({g} * {r})")
            elif o.op == "log": add_adj(o.args[0], f"({g} / {a[0]})")
            elif o.op == "sqrt": add_adj(o.args[0], f"({g} / ({r} * 2.0))")
            elif o.op == "relu": add_adj(o.args[0], f"({g} * ge({a[0]}, {zero11()}))")
            elif o.op == "ge": pass                                          # a comparison is flat almost everywhere: no gradient flows through it
            elif o.op == "broadcast":
                s_in, s_out = by[o.args[0]].shape, o.shape
                if s_in == (s_out[0], 1) and s_out[1] != 1: e = f"row_sum({g})"
                elif s_in == (1, s_out[1]) and s_out[0] != 1: e = f"col_sum({g})"
                elif s_in == (1, 1): e = f"col_sum(row_sum({g}))"
                else: e = g
                add_adj(o.args[0], e)
            elif o.op == "reduce":
                axis = int(o.attrs["axis"].split()[0]); s = by[o.args[0]].shape
                if o.attrs["kind"] == "sum": add_adj(o.args[0], f"({ones_of(s)} * {g})")
                else:   # max: the gradient goes to the positions that attain the maximum, split equally between ties (so the total over a row or column is exactly g)
                    mask = new("m"); fn = "row_sum" if axis == 1 else "col_sum"; out.append(f"let {mask} = ge({a[0]}, {r})")
                    add_adj(o.args[0], f"({mask} * ({g} / {fn}({mask})))")
            else: raise Unsupported(f"operation mg.{o.op} has no gradient rule")
        for w in wrt:
            if w not in adj: raise Unsupported("the loss does not depend on " + names.get(w, w))
        return adj
    labels = ["loss"]; params = {}
    if steps == 0:
        out.append("# ---- forward pass: the program's own operations, one `let` each"); var = forward(params)
        out.append("# ---- backward pass: the operations in reverse; each adjoint g is d(loss)/d(the value it belongs to)"); adj = backward(var)
        out.append("# ---- results: the loss, then the gradient of each requested matrix"); out.append(f"print {var[loss_ssa]}")
        for w in wrt: out.append(f"print {adj[w]}"); labels.append("grad " + names.get(w, w))
    else:
        for w in wrt: params[w] = names.get(w) or new(); out.append(f"let {params[w]} = {lit(by[w].attrs['data'])}")
        for s in range(steps + 1):
            out.append(f"# ---- step {s}: forward" + (", backward, update" if s < steps else " only (the loss after the last update)")); var = forward(params)
            out.append(f"print {var[loss_ssa]}"); labels.append(f"loss after {s} steps" if s else "loss before any step")
            if s < steps:
                adj = backward(var); upd = {}
                for w in wrt: v = new("p"); out.append(f"let {v} = ({params[w]} - ({adj[w]} * {lr!r}))"); upd[w] = v
                params = upd
        labels = labels[1:]
    return "\n".join(out) + "\n", labels

def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]); ap.add_argument("program"); ap.add_argument("--loss", type=int, default=0, help="which print is the loss (0 = the first)")
    ap.add_argument("--wrt", required=True, help="comma-separated names of `let` matrices to differentiate with respect to")
    ap.add_argument("--descend", type=int, default=0, metavar="STEPS", help="instead of printing the gradients, write a program that takes STEPS gradient-descent steps on the --wrt matrices and prints the loss after each")
    ap.add_argument("--lr", type=float, default=0.1, help="learning rate for --descend"); a = ap.parse_args(argv)
    text = subprocess.run([MGC, "mlir", a.program], capture_output=True, text=True)
    if text.returncode: sys.exit(text.stderr)
    try: ops = flatten(parse(text.stdout))
    except Unsupported as e: sys.exit("autograd: " + str(e))
    prints = [o for o in ops if o.op == "print"]
    if a.loss >= len(prints): sys.exit(f"autograd: the program has only {len(prints)} print statement(s)")
    loss = prints[a.loss].args[0]; by = {o.res: o for o in ops if o.res}
    if by[loss].shape != (1, 1): sys.exit(f"autograd: the loss must be a 1x1 matrix, this one is {by[loss].shape[0]}x{by[loss].shape[1]}")
    found = name_constants(ops, source_lets(a.program)); names = {}; wrt = []
    for w in a.wrt.split(","):
        if w not in found: sys.exit(f"autograd: no `let {w} = [[...]]` line of {a.program} matches a constant of the program")
        names[found[w]] = w; wrt.append(found[w])
    try: prog, labels = grad_program(ops, loss, wrt, names, a.descend, a.lr)
    except Unsupported as e: sys.exit("autograd: " + str(e))
    sys.stdout.write(f"# generated by autograd.py from {os.path.basename(a.program)}: prints, in order: " + ", ".join(labels) + "\n" + prog)
if __name__ == "__main__": main()

Three decisions in it are worth knowing about:

  • It is a transformation of programs, not a compiler pass. Like Chapter 41's back end, it reads the text of mgc mlir. That keeps it short, keeps the real compiler as the only thing that runs the numbers, and means its output is an ordinary Mountain Goat program that can be read, edited and compiled with any option of mgc. The cost is that it parses printed IR with regular expressions.
  • Ties in a maximum. row_max is not differentiable where two entries tie. The rule gives each tied entry an equal share, so that the total over the row is exactly g. That is a choice among valid subgradients, but it is the one that keeps softmax(x - row_max(x)) invariant to the shift the max implies. The version that gave each tied entry the full g is one of the mutants below.
  • relu at exactly zero uses ge(a, 0), gradient 1 there, which is also what Chapters 34 and 35 wrote by hand, so the two agree even at a kink.

What it generates for a transformer block

Chapter 35's block (embeddings, three layer norms, causal attention, residuals, a feed-forward network and an output layer, 16 trainable matrices) is examples/h35_transformer_step1.mg: Chapter 35's own definitions, data and step-1 forward and hand-derived backward code, ending with a print of the loss and then the 16 hand-derived gradients. make_examples.py copies those lines from Chapter 35's file (they are not re-typed). Giving autograd.py that program and the names of the 16 matrices:

./autograd.py examples/h35_transformer_step1.mg --wrt e0,g10,n10,wq0,wk0,wv0,wo0,g20,n20,w10,c10,w20,c20,gf0,nf0,wu0 > h35_grad.mg

writes a 288,339-byte program (the loss and the 16 gradients, computed by the generated backward pass) that compiles and runs with the ordinary mgc in a few seconds.

Does it give the right gradients?

check_autograd.py runs five groups of checks.

1. Every rule, against an independent reference. Ten small programs between them use every operation: m01 add, sub, mul, div; m02 all scalar operations on either side and negation; m03 broadcasting of a row and a column; m04 matmul and transpose; m05 exp, log, sqrt, relu; m06 sums, maxima and means in both directions; m07 reshape; m08 a value used three times; m09 the stable softmax and cross-entropy; m11 layer normalisation. The reference is reference.py: it evaluates the same flattened program in plain Python, with its own code for every operation and none of autograd.py's rules, and differentiates the loss by central finite differences in double precision (step 10⁻⁶, so the reference is good to about 10⁻⁹). The generated program is compiled and run by mgc and prints six digits, so the comparison is made at 10⁻⁵ relative. Chapters 33 to 35 checked their gradients with finite differences on the compiled program, limited to the six printed digits and to a step of 0.01; this reference is much sharper.

2. Ties. col_sum(row_max(m)) for m = [[1, 3, 3], [2, 2, 0]] must give [[0, .5, .5], [.5, .5, 0]] (hand-worked: finite differences do not apply at a kink).

3. The hand-derived backward passes of the book. Chapter 33's attention classifier (four gradient matrices), Chapter 34's layer normalisation, and Chapter 35's transformer block (16 gradient matrices, 648 elements). For each, the program contains the chapter's own hand-derived gradient code; autograd.py differentiates the loss print; the two sets of numbers must agree.

4. Training. Ten gradient-descent steps at learning rate 3.0 on Chapter 33's classifier, using only the generated gradients (--descend 10 --lr 3.0), against the loss that Chapter 33's own trainer printed:

steps  loss with generated gradients    Chapter 33 (hand-derived)
    0                        1.60134                      1.60134
    1                        1.51356                      1.51356
    2                        1.36719          (not printed there)
    3                        1.09744          (not printed there)
    4                       0.812335          (not printed there)
    5                       0.625597                     0.625597
    6                       0.495025          (not printed there)
    7                       0.403562          (not printed there)
    8                       0.342366          (not printed there)
    9                       0.302609          (not printed there)
   10                       0.273815                     0.273815

5. Refusals. A loss that is not a 1 × 1 matrix, an unknown name, a matrix the loss does not depend on, a loss that is flat in a matrix (it goes only through ge), and a rank-3 program are each refused with a message.

-- 1. every rule, against finite differences of an independent evaluation
  ok   all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
-- 2. ties
  ok   col_sum(row_max(m)) for m = [[1,3,3],[2,2,0]]: the gradient is [[0, .5, .5], [.5, .5, 0]]
-- 3. the hand-derived backward passes of Chapters 33, 34 and 35
  ok   Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
  ok   Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
  ok   Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
-- 4. training with the generated gradients
  ok   10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.51356 0.625597 0.273815 after 0, 1, 5, 10 steps, the same as Chapter 33's hand-derived training (1.60134 1.51356 0.625597 0.273815); the loss fell at every step
-- 5. refusals
  ok   a non-1x1 loss, an unknown name, an unused matrix and a loss that is flat in a matrix are refused with a message
  ok   a program that is not rank 2 (Chapter 28's) is refused with a message: autograd: rank 3 tensors are not supported
all checks pass

Two things to read from this output. First, all 112 elements of the ten small programs agree with a reference that shares no code with the rules. Second, the hand-derived gradients of three chapters were right: an independent mechanical derivation reproduces all of them, to the six digits the programs print. That is evidence for the hand derivations as much as for the tool; it is also an equality at six printed digits, not at the last bit (the two programs add numbers in different orders).

Are the checks good enough? Twenty-four wrong versions.

NOTE: this script deliberately breaks autograd.py. 'caught' lines are EXPECTED: they show the checker can detect the mistake.
baseline (nothing broken): checker passes (expected)
mul: each operand's gradient uses its OWN value (not the other's): caught by 5 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.60644 3.39843e+64 nan a
sub: the second operand's gradient has the wrong sign: caught by 5 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.92243 3.10859e+06 1.262
div: the denominator's gradient has the wrong sign: caught by 5 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.51355 0.663782 0.347988
div: the denominator's gradient divides by b, not by b squared: caught by 5 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.51357 0.581809 0.277632
k - x: the gradient has the wrong sign: caught by 1 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
k / x: the gradient has the wrong sign: caught by 1 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
x * k: the gradient ignores k: caught by 4 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 2.92362 1.43159e+14 4.311
neg: the gradient has the wrong sign: caught by 1 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
matmul: the left gradient does not transpose b: caught (the checker did not finish: RuntimeError: mgc run /tmp/ch45c_nzpslrt5/m04_matmul.mg failed (1): /tmp/ch45c_nzpslrt5/m04_matmul.mg:20: erro)
matmul: the right gradient is g @ transpose(a) instead of transpose(a) @ g: caught (the checker did not finish: RuntimeError: mgc run /tmp/ch45c_9_xctptf/m04_matmul.mg failed (1): /tmp/ch45c_9_xctptf/m04_matmul.mg:21: erro)
transpose: the gradient is not transposed back: caught by 2 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
reshape: the gradient is not reshaped back: caught by 1 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
exp: the gradient uses the input, not the output: caught by 4 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.55456 1.54535 1.54661 a
log: the gradient multiplies by x instead of dividing: caught by 4 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 57.1972 123.37 726.182 af
sqrt: the factor 2 is forgotten: caught by 3 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
relu: no mask (the gradient flows through negative inputs too): caught by 2 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
sum: the gradient is not broadcast over the reduced axis: caught (the checker did not finish: RuntimeError: mgc run /tmp/ch45c_b8vx44z2/m07_reshape.mg failed (1): /tmp/ch45c_b8vx44z2/m07_reshape.mg:13: er)
max: tied maxima each get the FULL gradient (no sharing): caught by 1 failing check(s):
      col_sum(row_max(m)) for m = [[1,3,3],[2,2,0]]: the gradient is [[0, .5, .5], [.5, .5, 0]]
a value used twice: the second contribution replaces the first: caught by 5 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.61025 1.60162 1.53103 a
broadcast: the gradient of a column vector is summed over the wrong axis: caught by 1 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
naming: constants are matched by value only, ignoring order (equal-valued lets collide): caught by 1 failing check(s):
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
descent: the update ADDS the gradient (climbs the loss instead of descending): caught by 1 failing check(s):
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.6869 21.3963 321337 aft
descent: the learning rate is ignored (always 1): caught by 1 failing check(s):
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.5723 1.38631 0.902248 a
descent: the weights are never updated (every step uses the starting weights): caught by 1 failing check(s):
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.60134 1.60134 1.60134 a
the loss is differentiated with a seed of 2 instead of 1: caught by 6 failing check(s):
      all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
      col_sum(row_max(m)) for m = [[1,3,3],[2,2,0]]: the gradient is [[0, .5, .5], [.5, .5, 0]]
      Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
      Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
      Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
      10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.42424 0.350446 0.127162

done (autograd.py itself was never edited; the broken copy is removed)

All 24 are caught. Which check catches which is the interesting part:

  • The hand-derived programs of Chapters 33 to 35 alone would have missed six of the rules. The mutants for k - x, k / x, negation, reshape, broadcasting of a column and max ties are caught only by the small programs (or the tie test), because the chapters' programs happen not to use those combinations. A check built only from the book's own programs would have given false confidence.
  • The constant-matching mutant is caught only by Chapter 35. autograd.py gives the IR's constants their names back by matching let lines to constants in order, because Chapter 35 has several lets with equal values (three all-zero bias rows, three all-ones scale rows). Matching by value alone picks the wrong constant.
  • Four mutants (a missing transpose in matmul, both matmul sides, and the broadcast-free sum rule) are caught because the generated program does not compile (a shape mismatch), reported as "the checker did not finish"; that is a valid detection but a less specific one than a wrong number.
  • The three training mutants (the update adds the gradient, the learning rate is ignored, the weights never change) are caught only by the training check; the gradient checks cannot see them.

What it costs

The generated backward pass against the hand-derived one, for the two big cases (cost.py): the number of let statements, the size of the source, and the instructions the compiled program executes (valgrind's exact count, clang -O0, mgc's defaults of Chapter 43):

program                                        let statements  source bytes  instructions executed
Chapter 33 classifier          hand-derived                 7        12,602              8,077,504
Chapter 33 classifier          generated                  103        21,000              3,875,248
Chapter 35 transformer block   hand-derived                75       192,514             41,048,294
Chapter 35 transformer block   generated                  307       288,339             45,533,249
  • Chapter 35's transformer block: the generated program executes 11% more instructions than the hand-derived one (45.5 million against 41.0 million) and is 1.5 times the size of the source. That is the price of a mechanical derivation with no human simplification: where a person fuses the three layer-norm terms into one expression, the generated code spells out each step.
  • Chapter 33's classifier: the generated program executes half as many instructions (3.9 million against 8.1 million). That is not a property of automatic differentiation: Chapter 33 wrote each gradient as a separate function that recomputes the forward pass, and the generated program computes the forward pass once and shares it. A hand-derived backward pass written with shared intermediates, as Chapter 35's is, loses that advantage.
  • The let counts are not comparable for the first case (Chapter 33's seven lets call functions that hide hundreds of operations).

Reverse mode also keeps every forward value alive until the backward pass uses it: the generated program is a list of lets, and every forward intermediate is read again later. Nothing here saves memory by recomputing, as production systems do (checkpointing). On these sizes that did not matter; it would on a large model.

Tests

One new lit file test/autograd45/check.mlir (the suite is now 157 tests). The checker is check_autograd.py (the five groups above); its output:

-- 1. every rule, against finite differences of an independent evaluation
  ok   all 10 programs: the generated gradients match finite differences in every one of 112 matrix elements
-- 2. ties
  ok   col_sum(row_max(m)) for m = [[1,3,3],[2,2,0]]: the gradient is [[0, .5, .5], [.5, .5, 0]]
-- 3. the hand-derived backward passes of Chapters 33, 34 and 35
  ok   Chapter 33's attention classifier: loss and 4 gradient matrices (48 elements) equal the hand-derived ones
  ok   Chapter 34's layer normalisation: loss and 1 gradient matrices (4 elements) equal the hand-derived ones
  ok   Chapter 35's transformer block: loss and 16 gradient matrices (648 elements) equal the hand-derived ones
-- 4. training with the generated gradients
  ok   10 gradient-descent steps (lr 3.0) on Chapter 33's classifier with the GENERATED gradients: loss 1.60134 1.51356 0.625597 0.273815 after 0, 1, 5, 10 steps, the same as Chapter 33's hand-derived training (1.60134 1.51356 0.625597 0.273815); the loss fell at every step
-- 5. refusals
  ok   a non-1x1 loss, an unknown name, an unused matrix and a loss that is flat in a matrix are refused with a message
  ok   a program that is not rank 2 (Chapter 28's) is refused with a message: autograd: rank 3 tensors are not supported
all checks pass

Limits and what is not established

  • It is a transformation of programs, not a compiler feature. There is no mgc grad; autograd.py parses the printed IR with regular expressions, like Chapter 41's back end, and needs mgc to be the one from this chapter's directory (or any that prints the same IR).
  • The loss is the first print of the program (or the one chosen by --loss), and must be a 1 × 1 matrix. The parameters are let matrices given by name, found by matching the let lines to the constants in order. A let whose value is not a literal matrix cannot be named; a literal in an expression that equals a later let's value and comes before it in the file could be matched by mistake (the checker's three ground-truth programs and the ten small ones are not affected).
  • Rank 2 only, static shapes only. Rank-3 tensors, permute, contract and dynamic shapes are refused.
  • No gradient through comparisons. ge gives 0, so a loss that depends on a parameter only through a comparison is refused ("does not depend"), and a loss that mixes a smooth path with a comparison differentiates the smooth part only.
  • Agreement with the hand-derived gradients is at six printed digits, not bitwise; the independent finite-difference reference covers only ten small programs (112 elements), because the big ones are too slow to evaluate in plain Python.
  • No checkpointing, no sparsity, no simplification: the generated program is larger and slower than a hand-tuned one (11% more instructions on the transformer block) and keeps every intermediate alive.
  • Second derivatives: the output is an ordinary Mountain Goat program, so it can be differentiated again; Chapter 46 does exactly that (Hessian-vector products) and tests it.
  • Large programs compile slowly: the generated 200-step training programs of Chapters 35 and 36 were not attempted (at -O2 even the hand-derived 200-step program takes over 29 minutes to compile, Chapter 44).

Reproducing

cd docs/part45/code
./autograd.py examples/00_worked_example.mg --wrt w,x > worked_example_generated.mg && ./mgc run worked_example_generated.mg
./descend_compare.py > descend_out.txt && ./check_autograd.py > check_autograd_out.txt
./mutation.py > mutation_out.txt && ./cost.py > cost_out.txt
cd ../../part15/code && ./run_lit.sh          # 157 tests

Chapter summary

  • Reverse-mode automatic differentiation is a list of local rules, one per operation, applied to the program's operations in reverse order, with a value used twice receiving the sum of its contributions. autograd.py does it as a source-to-source transformation that writes a new Mountain Goat program.
  • It reproduces, to the six printed digits, the hand-derived gradients of Chapter 33's attention classifier, Chapter 34's layer normalisation and all 16 gradient matrices of Chapter 35's transformer block; and ten gradient-descent steps with the generated gradients give exactly Chapter 33's loss curve (1.60134 → 1.51356 → 0.625597 → 0.273815).
  • Every rule matches central finite differences of an independent plain-Python evaluation, in all 112 elements of ten small programs.
  • 24 wrong versions of the tool are caught; the book's own programs alone would have missed six of the rules, which is why the small programs exist.
  • The generated backward pass of the transformer block costs 11% more instructions than the hand-derived one; it keeps all forward values alive; and it covers rank-2 programs without comparisons in the gradient path.

Self-check questions

Each answer is collapsed; try the question first.

  1. Why does reverse mode walk the operations backwards, and what happens to a value that is used twice?

    Answer

    The derivative of the loss with respect to an operation's result must be complete before it can be turned into contributions to that operation's inputs; walking the operations in reverse order guarantees that every operation after it has already contributed. A value used by several operations receives the sum of their contributions (the chain rule over all the paths from it to the loss): x * x * x gives x three contributions.

  2. Derive by hand what autograd.py wrote for col_sum(row_sum(relu(x @ w) * 2.0)) and say why ge(v1, 0) appears.

    Answer

    The loss's adjoint is [[1]]; col_sum and row_sum spread it over the elements they summed (ones * g); * 2.0 multiplies by 2; relu passes the adjoint where its input was positive, which is ge(x @ w, 0) as a 0/1 mask; matmul gives g @ transpose(w) to x and transpose(x) @ g to w. With x @ w = [[-2.5, -3], [5, 8]], only the second row is positive, so the gradient with respect to w is [[4, 4], [2, 2]].

  3. Why does the checker compare against a plain-Python reference instead of only against Chapters 33 to 35?

    Answer

    The hand-derived programs of those chapters do not use every operation: six of the 24 wrong versions (k - x, k / x, negation, reshape, a column broadcast, ties in a maximum) pass all three chapters' comparisons and are caught only by the small programs and the tie test. The reference also shares no code with the rules, so a mistake in a rule cannot be hidden by the same mistake in the reference.

  4. What does a tie in a maximum do to the gradient rule, and why share the gradient?

    Answer

    row_max has no derivative where two entries are equal. Sharing the gradient equally between the tied entries keeps the total over the row equal to the incoming gradient, so a softmax built as exp(x - row_max(x)) still has a gradient that is unchanged by shifting x. Giving each tied entry the full gradient breaks that (the total becomes larger), which the tie test catches.

  5. The generated backward pass of Chapter 33 runs in half the instructions of the hand-derived one, and that of Chapter 35 in 11% more. Why the difference?

    Answer

    Chapter 33's hand-derived gradients are separate functions that each recompute the forward pass, so the program does the forward work several times; the generated program does it once. Chapter 35's hand-derived backward pass shares its intermediates, so the comparison is between two programs that both do the forward pass once, and the generated one pays for not simplifying (for example, spelling out each step of the layer-norm gradient).

  6. What does it mean that the generated gradients equal the hand-derived ones, and what does it not mean?

    Answer

    It means two independent derivations (a person's, and a mechanical one) agree on every printed element of 16 gradient matrices, so they are both very probably right. It does not mean they are identical to the last bit (the programs add in different orders and print six digits), and it does not say anything about operations the chapters never used.

  7. Name three limits of the claims on this page.

    Answer

    Any of: rank 2 and static shapes only; the loss must be a 1 × 1 first print; no gradient flows through comparisons; the independent finite-difference reference covers only the ten small programs; the transformation reads printed IR with regular expressions; the 200-step training programs were not attempted; second derivatives were not tried; the generated program keeps all forward values alive.