Step 3 — Autodiff, and a Neural $n$-gram Model¶

Solution notebook, and the densest of the eight. Part 1 builds reverse-mode automatic differentiation from nothing — about 60 lines implementing Lecture 3's Theorem 3.1 — and validates it three ways, the third being Step 2's diamond network differentiated by hand, by the engine, by finite differences, and finally by PyTorch. Part 2 uses PyTorch's autograd to train the Bengio (2003) model, and ends with a measured result that motivates the entire second half of the course.

Runtime: the engine is instant; the $k$-sweep is ~20 s per value of $k$ single-threaded.

Part 1 — A scalar autograd engine¶

Each Value holds a number, a gradient slot, its parents, and a closure saying how to push its own gradient to those parents. That closure is the local partial $\partial\varphi_i/\partial v_j$ of Theorem 3.1; backward() is the reverse topological sweep.

Two details carry the whole proof, and both are where bugs live: reverse topological order (a node must not be processed before all its children) and += not = (a node used twice receives a contribution along each path).

In [ ]:
class Value:
    def __init__(self, data, _parents=()):
        self.data = float(data)
        self.grad = 0.0
        self._parents = _parents
        self._backward = lambda: None

    def __repr__(self):
        return f'Value({self.data:.4f}, grad={self.grad:.4f})'

    def _wrap(self, other):
        return other if isinstance(other, Value) else Value(other)

    def __add__(self, other):
        other = self._wrap(other)
        out = Value(self.data + other.data, (self, other))
        def _backward():
            self.grad  += out.grad          # d(a+b)/da = 1
            other.grad += out.grad
        out._backward = _backward
        return out

    def __mul__(self, other):
        other = self._wrap(other)
        out = Value(self.data * other.data, (self, other))
        def _backward():
            self.grad  += other.data * out.grad   # d(ab)/da = b
            other.grad += self.data  * out.grad
        out._backward = _backward
        return out

    def __pow__(self, k):
        assert isinstance(k, (int, float))
        out = Value(self.data ** k, (self,))
        def _backward():
            self.grad += k * self.data ** (k - 1) * out.grad
        out._backward = _backward
        return out

    def tanh(self):
        import math
        t = math.tanh(self.data)
        out = Value(t, (self,))
        def _backward():
            self.grad += (1 - t * t) * out.grad   # d tanh = 1 - tanh^2
        out._backward = _backward
        return out

    def exp(self):
        import math
        e = math.exp(self.data)
        out = Value(e, (self,))
        def _backward():
            self.grad += e * out.grad
        out._backward = _backward
        return out

    def log(self):
        import math
        out = Value(math.log(self.data), (self,))
        def _backward():
            self.grad += (1.0 / self.data) * out.grad
        out._backward = _backward
        return out

    def relu(self):
        out = Value(max(self.data, 0.0), (self,))
        def _backward():
            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad   # subgradient 0 at the corner
        out._backward = _backward
        return out

    # conveniences
    def __neg__(self):       return self * -1
    def __sub__(self, o):    return self + (-self._wrap(o))
    def __radd__(self, o):   return self + o
    def __rsub__(self, o):   return (-self) + o
    def __rmul__(self, o):   return self * o
    def __truediv__(self, o): return self * (self._wrap(o) ** -1)

    def backward(self):
        order, visited = [], set()
        def visit(v):
            if id(v) not in visited:
                visited.add(id(v))
                for p in v._parents:
                    visit(p)
                order.append(v)      # parents appended before children
        visit(self)
        self.grad = 1.0              # dL/dL = 1
        for v in reversed(order):    # reverse topological order
            v._backward()

Validation 1 — fan-out, against hand computation¶

$\mathcal L = (ab + a)\tanh b$ at $a=2$, $b=-1$. Both variables are used twice, so this is precisely the test that = instead of += fails. By hand, with $t = \tanh(-1)$: $\partial\mathcal L/\partial a = (b+1)t$ and $\partial\mathcal L/\partial b = a\,t + (ab+a)(1-t^2)$.

In [ ]:
import math
a, b = Value(2.0), Value(-1.0)
L = (a * b + a) * b.tanh()
L.backward()

t = math.tanh(-1.0)
da_hand = ((-1.0) + 1) * t
db_hand = 2.0 * t + (2.0 * -1.0 + 2.0) * (1 - t * t)
print(f'engine: dL/da = {a.grad:.10f}   hand: {da_hand:.10f}')
print(f'engine: dL/db = {b.grad:.10f}   hand: {db_hand:.10f}')
assert abs(a.grad - da_hand) < 1e-12 and abs(b.grad - db_hand) < 1e-12
print('fan-out test passed')

Validation 2 — finite differences¶

In double precision (Python floats are float64), so the resolution problem that bites float32 finite differences does not arise here. The ReLU test is placed away from the corners: at a corner the derivative doesn't exist, and a central difference straddling it returns $\tfrac12$ while the engine's subgradient returns $0$ or $1$ — both defensible, neither "wrong".

In [ ]:
def fd_check(f, xs, h=1e-6):
    vals = [Value(x) for x in xs]
    out = f(vals)
    out.backward()
    errs = []
    for i in range(len(xs)):
        up = list(xs); up[i] += h
        dn = list(xs); dn[i] -= h
        num = (f([Value(v) for v in up]).data
               - f([Value(v) for v in dn]).data) / (2 * h)
        errs.append(abs(num - vals[i].grad) / max(abs(vals[i].grad), 1e-12))
    return max(errs)

tests = {
    '(ab+a)tanh(b)': (lambda v: (v[0]*v[1] + v[0]) * v[1].tanh(), [2.0, -1.0]),
    'exp(ab)/(1+a^2)': (lambda v: (v[0]*v[1]).exp() / (1 + v[0]**2), [0.7, -0.4]),
    'log(a^2+b^2+1)': (lambda v: (v[0]**2 + v[1]**2 + 1).log(), [1.3, 0.5]),
    'relu(a)+relu(-b)': (lambda v: v[0].relu() + (-v[1]).relu(), [0.9, -0.6]),   # away from the corners
}
for name, (f, xs) in tests.items():
    e = fd_check(f, xs)
    print(f'{name:<18} max relative error {e:.2e}')
    assert e < 1e-6

Validation 3 — the full circle: Step 2's diamond network¶

Rebuild one forward pass of the diamond classifier from scalar Values at $\mathbf x=(0.5,0.25)^\top$, $\gamma=4$, with target inside, and differentiate the loss with respect to the input coordinates. Both coordinates are positive there, so only two of the four ReLUs are active and $r=x_1+x_2$; the loss is

$$ \mathcal L=-\log p(\mathrm{in}\mid\mathbf x) =\log\bigl(1+e^{z_{\mathrm{out}}-z_{\mathrm{in}}}\bigr) =\log\bigl(1+\exp(2\gamma(x_1+x_2-1))\bigr), $$

whose derivative in either coordinate is $2\gamma\,\sigma(2\gamma(x_1+x_2-1))=2\gamma\,p(\mathrm{out}\mid\mathbf x) =8\,\sigma(-2)\approx0.9536$. Three independent routes to that number: the formula, the engine, central differences.

In [ ]:
def diamond_loss(x1, x2, gamma=4.0):
    # forward pass of Step 2's network on scalars (Value objects), target = inside
    h = [x1.relu(), (-x1).relu(), x2.relu(), (-x2).relu()]      # the four hidden units
    r = h[0] + h[1] + h[2] + h[3]                                # |x1| + |x2|
    z_out, z_in = (r - 1) * gamma, (1 - r) * gamma              # W2 h + b2
    Z = z_out.exp() + z_in.exp()                                 # softmax denominator
    return Z.log() - z_in                                        # -log softmax(z)[in]

x1, x2 = Value(0.5), Value(0.25)
L = diamond_loss(x1, x2)
L.backward()

p_out = 1 / (1 + math.exp(-2 * 4.0 * (0.5 + 0.25 - 1)))          # sigma(2γ(x1+x2-1))
analytic = 2 * 4.0 * p_out
h = 1e-6
fd = lambda f, a, b: (f(Value(a + h), Value(b)).data - f(Value(a - h), Value(b)).data) / (2 * h)
print(f'loss                 : {L.data:.10f}   (= log(1+e^-2) = {math.log(1 + math.exp(-2)):.10f})')
print(f'analytic  dL/dx1     : {analytic:.10f}')
print(f'engine    dL/dx1, dx2: {x1.grad:.10f}, {x2.grad:.10f}')
print(f'finite-difference    : {fd(diamond_loss, 0.5, 0.25):.10f}')
assert abs(x1.grad - analytic) < 1e-6 and abs(x2.grad - analytic) < 1e-6
assert abs(fd(diamond_loss, 0.5, 0.25) - analytic) < 1e-6
print('three routes, one gradient: 0.9536')

3.2 Switch to PyTorch autograd¶

The same check once more, with PyTorch's tensor-valued engine and Step 2's diamond_net reproduced verbatim. requires_grad=True marks the point as a leaf whose gradient we want; loss.backward() is our Value.backward() — the topological sort and reverse sweep — filling point.grad. PyTorch is your Value class with three upgrades: tensors instead of scalars, a compiled backend, and a library of primitives. Conceptually nothing else. After this cell you may call .backward() for the rest of your life with a clear conscience.

In [ ]:
import torch

def diamond_net(X, gamma=4.0):                     # Step 2's solution, unchanged
    W1 = torch.tensor([[1., 0.], [-1., 0.], [0., 1.], [0., -1.]], dtype=X.dtype)
    b1 = torch.zeros(4, dtype=X.dtype)
    W2 = gamma * torch.tensor([[1., 1., 1., 1.], [-1., -1., -1., -1.]], dtype=X.dtype)
    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)
    a1 = X @ W1.T + b1
    h = torch.relu(a1)
    logits = h @ W2.T + b2
    shifted = logits - logits.max(dim=1, keepdim=True).values
    weights = shifted.exp()
    probs = weights / weights.sum(dim=1, keepdim=True)
    return a1, h, logits, probs

point = torch.tensor([[0.5, 0.25]], dtype=torch.float64,
                     requires_grad=True)
_, _, _, probs = diamond_net(point, gamma=4.0)
loss = -probs[0, 1].log()             # target = inside
loss.backward()

expected = 2 * 4.0 * probs.detach()[0, 0]
print(point.grad)                      # both entries ≈ 0.9536
print(expected)
assert torch.allclose(point.grad, torch.full((1, 2), analytic, dtype=torch.float64), atol=1e-6)
assert abs(point.grad[0, 0].item() - x1.grad) < 1e-6
print('PyTorch agrees with the engine, the formula, and finite differences')

Part 2 — The Bengio (2003) model¶

$$x = (C[a_1],\dots,C[a_k]) \in \mathbb R^{kd},\qquad z = W_2\tanh(W_1x + b_1) + b_2,\qquad p = \operatorname{softmax}(z).$$

The embedding table $C$ is the new idea: each character becomes a learned vector, and because $C$ is shared across positions and contexts, evidence about one character informs every context containing it. That sharing is what a count table cannot do (Lecture 2 §2).

We write minibatch SGD by hand — no torch.optim until Step 6.

Setup: data and tokenizer (Step 1's solution, reproduced)¶

Every notebook in this series is self-contained: it re-creates what it needs from earlier steps in one compact cell, so you can run it top to bottom without opening the others. On Colab, uncomment the download.

In [ ]:
# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
import torch

with open('input.txt') as f:
    text = f.read()
vocab = sorted(set(text))
V = len(vocab)
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for i, ch in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)
ids = torch.tensor(encode(text), dtype=torch.long)

assert len(text) == 1_115_394 and V == 65
print(f'{len(text):,} characters, vocab {V}')
In [ ]:
import torch.nn.functional as F
torch.set_num_threads(1)

n_split = int(0.9 * len(ids))
tr, va = ids[:n_split], ids[n_split:]

def windows(data, k):
    return data.unfold(0, k, 1)[:-1], data[k:]   # (N,k) contexts, targets

def run(k, d=16, hid=128, steps=20_000, B=64, seed=0, verbose=False):
    torch.manual_seed(seed)
    Xtr, Ytr = windows(tr, k)
    Xva, Yva = windows(va, k)
    C  = (torch.randn(V, d) * 0.1).requires_grad_()
    W1 = (torch.randn(k*d, hid) * (5/3) / (k*d)**0.5).requires_grad_()
    b1 = torch.zeros(hid, requires_grad=True)
    W2 = (torch.randn(hid, V) * 0.01).requires_grad_()
    b2 = torch.zeros(V, requires_grad=True)
    params = [C, W1, b1, W2, b2]
    fwd = lambda Z: torch.tanh(C[Z].view(Z.shape[0], -1) @ W1 + b1) @ W2 + b2

    for s in range(steps):
        ix = torch.randint(0, Xtr.shape[0], (B,))
        loss = F.cross_entropy(fwd(Xtr[ix]), Ytr[ix])
        for p in params:
            p.grad = None
        loss.backward()
        lr = 0.1 if s < steps * 0.75 else 0.01     # crude two-stage decay
        with torch.no_grad():
            for p in params:
                p -= lr * p.grad

    @torch.no_grad()
    def ev(Xd, Yd, iters=30, BB=4096):
        tot = 0.0
        for _ in range(iters):
            ix = torch.randint(0, Xd.shape[0], (BB,))
            tot += F.cross_entropy(fwd(Xd[ix]), Yd[ix]).item()
        return tot / iters

    return ev(Xtr, Ytr), ev(Xva, Yva), sum(p.numel() for p in params), (C, fwd)

The context-length sweep — and a surprise¶

Predict before running: does longer context help?

In [ ]:
import math
print(f"{'k':>3} {'params':>8} {'train':>8} {'val':>8} {'bits/ch':>9} {'table would need':>18}")
results = {}
for k in (1, 3, 5, 8):
    trl, val, npar, _ = run(k)
    results[k] = val
    print(f'{k:>3} {npar:>8,} {trl:>8.4f} {val:>8.4f} {val/math.log(2):>9.4f}'
          f' {float(V**(k+1)):>18.1e}', flush=True)

Reference: 2.4805, 1.9821, 2.0018, 2.0329 for $k = 1,3,5,8$. Read that carefully — it says three things.

  1. $k{=}1$ reproduces the bigram (2.48 val against Step 1's 2.45 train-set CE), as it must: the $k{=}1$ model's hypothesis class contains the bigram table.
  2. $k{=}3$ buys 0.5 nats for 4,000 extra parameters — against the $1.8\times10^{7}$-entry table a trigram counting model would need for the same context. That ratio is the whole argument of Lecture 2.
  3. $k{=}5$ and $k{=}8$ are worse than $k{=}3$, and this is not a bug. Train loss barely improves either, so it is underfitting, not overfitting: flattening a longer context into one $kd$-vector spreads fixed capacity thinner and hands SGD a harder problem.

Point 3 is the hinge of the course. The fix is not more parameters. It is a way of reading context that does not concatenate everything into a flat vector, and that can decide per input which earlier positions matter. That is attention.

Samples, and the diagnostic failure¶

Reference output from the $k{=}3$ model (seeded; the cell below reproduces it) contains real words in runs and — strikingly — the SPEAKER: line format with its preceding blank line. But the speakers it emits, ENCHAR: and DONLY:, occur zero times in the corpus. Three characters of context cannot span a name, so the model has learned the shape of a speaker line with no ability to remember any particular one. Hold that thought until Step 6, where the transformer emits WARWICK: and Montague correctly spelled.

In [ ]:
torch.manual_seed(0)
_, _, _, (C3, fwd3) = run(3, steps=20_000)
ctx, out = [stoi['\n']] * 3, []
with torch.no_grad():
    for _ in range(300):
        p = F.softmax(fwd3(torch.tensor([ctx]))[0], dim=-1)
        i = int(torch.multinomial(p, 1))
        out.append(itos[i]); ctx = ctx[1:] + [i]
print(''.join(out))

The embedding table, visualized¶

Two principal components of $C$. Vowels cluster; digits and punctuation separate from letters; capitals separate from lowercase. Nobody told the model these categories exist — they fall out of predicting the next character. Geometry learned from prediction alone is the thesis of the whole enterprise, and here it is at 11,601 parameters.

In [ ]:
import matplotlib.pyplot as plt
Cd = C3.detach()
Cc = Cd - Cd.mean(0)
U, S, Vh = torch.linalg.svd(Cc, full_matrices=False)
xy = (Cc @ Vh[:2].T).numpy()
plt.figure(figsize=(9, 7))
for i, (x, y) in enumerate(xy):
    ch = itos[i]
    label = {'\n': '\\n', ' ': '␣'}.get(ch, ch)
    color = ('tab:red' if ch in 'aeiouAEIOU' else
             'tab:blue' if ch.isupper() else
             'tab:green' if ch.islower() else 'tab:gray')
    plt.text(x, y, label, color=color, fontsize=13, ha='center')
plt.xlim(xy[:,0].min()*1.15, xy[:,0].max()*1.15)
plt.ylim(xy[:,1].min()*1.15, xy[:,1].max()*1.15)
plt.title('character embeddings, first two principal components\n'
          'red = vowels, blue = capitals, green = lowercase, grey = other')
plt.tight_layout()

→ Continue with Step 4: one attention head, and the proofs from Lecture 4 executed as assertions.