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

Released with Lecture 3 · ~4 hours (the long one — start early) · Starts from: your Step 1 tokenized corpus and the matrix-layer ideas of Step 2.

In what follows, there are blocks of code provided for you. It’s important that you can read the code. If you are not proficient with Python, then I propose the following workflow. For each block of code, keywords are provided for the python syntax which is being used, so you can look up python syntax and functions you don’t recognise. To help, beneath many code blocks you will find Python background foldouts: self-contained tutorials, with live Python boxes you can edit and run right on the page, for syntax that is new in this step (syntax already introduced is covered in earlier steps’ foldouts). Then I suggest that you first experiment with changing the code or printing out intermediate variables, until you have figured out what you think it does. Then point AI toward the page and step number, provide your precise explanation of what you believe each line of the code is doing, and ask for corrections. In this way you will quickly become a code reader (if not a code writer). — Kate

Goal

Two things, in order:

1. Build reverse-mode automatic differentiation yourself—a roughly 60-line scalar “autograd engine” in the style of micrograd. Validate it on a graph with fan-out, with finite differences, and on the hand-wired diamond network from Step 2. After this you may use $\,\texttt{loss.backward()}\,$ for the rest of your life with a clear conscience.
2. Use PyTorch’s autograd to train the first learned neural language model of the course: Bengio et al. (2003)—an embedding matrix plus an MLP reading a context of $k$ characters.

On paper first

1. A computation is a DAG whose nodes hold values $v_i$ computed by primitives from their parents. Reverse-mode AD computes adjoints $\bar v_i = \partial \mathcal{L}/\partial v_i$ by the chain rule, in reverse topological order: $\bar v_i \mathrel{+}= \bar v_j \,\partial v_j/\partial v_i$ summed over children $j$. Write out this recursion for $\mathcal L = (a b + a)\tanh(b)$, by hand, both forward-mode (derivatives along with values) and reverse-mode, and note what each pass costs.
2. Cheap gradient principle (Lecture 3): reverse mode computes the gradient with respect to all $n$ inputs in $O(1)$ function evaluations, vs $O(n)$ for forward mode or finite differences. For a billion-parameter model, that constant is the entire field.
3. Why must the traversal be in reverse topological order, and why does the update use += rather than =? (Both bugs will otherwise find you in Task 3.1; fan-out — a node used twice — is the test case.)

Tasks

3.1 A scalar autograd engine

Implement a class Value wrapping a float, supporting +, *, tanh, exp, log, relu, and **, that records the computation graph and backpropagates:

class Value:                              # class definition
    def __init__(self, data, _parents=()):  # constructor; self; default argument (an empty tuple)
        self.data = data                  # attribute assignment
        self.grad = 0.0
        self._parents = _parents
        self._backward = lambda: None     # lambda function; how to push my grad to my parents

    def __add__(self, other):             # dunder method: overloads the + operator
        other = other if isinstance(other, Value) else Value(other)  # conditional (ternary) expression; isinstance
        out = Value(self.data + other.data, (self, other))  # tuple
        def _backward():                  # nested function (closure)
            self.grad  += out.grad        # d(a+b)/da = 1
            other.grad += out.grad
        out._backward = _backward         # functions are values: stored in an attribute
        return out

    def __mul__(self, other):
        # you write this one, and tanh/exp/log/relu/pow in the same pattern
        ...                               # Ellipsis: placeholder for missing code

    def backward(self):
        order = []                        # list; topological sort of the graph
        visited = set()                   # set
        def visit(v):                     # recursive function
            if v not in visited:          # membership test (not in)
                visited.add(v)
                for p in v._parents:
                    visit(p)
                order.append(v)
        visit(self)
        self.grad = 1.0                   # dL/dL = 1
        for v in reversed(order):         # reversed
            v._backward()

class Counter:
    def __init__(self, start=0):   # the constructor, run by Counter(...)
        self.count = start         # an attribute on this particular object
    def bump(self):                # a method; self = the object it's called on
        self.count += 1

c = Counter()
c.bump()
c.bump()
print(c.count)
d = Counter(start=100)             # a second, independent object
d.bump()
print(d.count)
print(c.count)                     # c was not affected

class Vec:
    def __init__(self, x, y):
        self.x, self.y = x, y
    def __add__(self, other):          # defines what + means for Vec
        return Vec(self.x + other.x, self.y + other.y)
    def __repr__(self):                # defines how a Vec prints
        return "Vec(" + str(self.x) + ", " + str(self.y) + ")"

print(Vec(1, 2) + Vec(10, 20))
def todo():
    ...                                # placeholder body: runs, does nothing
print(todo())                          # a function with no return gives None

x = 7
parity = 'odd' if x % 2 == 1 else 'even'   # a value picked inline
print(parity)
print(isinstance(x, int))
print(isinstance('7', int))
v = '5'
v = v if isinstance(v, int) else int(v)    # the wrap-if-needed pattern
print(v + 1)

def make_adder(n):
    def add(x):            # defined inside: it remembers this n
        return x + n
    return add             # functions are values — return one

add5 = make_adder(5)
add100 = make_adder(100)   # a separate closure with its own n
print(add5(3))
print(add100(3))
print(add5(3))             # add5 still remembers n = 5

def countdown(n):
    if n == 0:
        return             # the base case: stop recursing
    print(n)
    countdown(n - 1)       # the function calls itself
countdown(3)
vals = [3, 1, 4]
print(3 in vals)           # membership test
print(9 not in vals)
seen = set()
seen.add('a')
print('a' in seen, 'b' in seen)   # fast membership: why visited is a set
for v in reversed(vals):
    print(v)

Validate it three ways (this is the substance of the task):

1. Fan-out test: $\mathcal L=(ab+a)\tanh(b)$ at $a=2,b=-1$; check $a.\mathrm{grad}$ and $b.\mathrm{grad}$ against your hand calculation.

2. Compare your engine with central finite differences on a few random expressions. Use float64 and several step sizes.

3. The full circle: rebuild one forward pass of Step 2’s diamond network from scalar $\texttt{Value}$ objects at $\mathbf{x}=(0.5,0.25)^\top$, $\gamma=4$, with target “inside.” Because both coordinates are positive there,

   $$\mathcal L = \log\left(1+\exp\left(2\gamma(x_1+x_2-1)\right)\right),$$

   and both input derivatives equal

   $$\frac{\partial\mathcal L}{\partial x_i} = 2\gamma\,p(\mathrm{outside}\mid\mathbf{x}) \approx0.9536.$$

   Check this analytic value, your engine, and finite differences against one another. The ReLU primitive also gives you a piecewise-linear test; do not place a finite-difference check exactly at a ReLU corner.

3.2 Switch to PyTorch autograd

Repeat the diamond check with PyTorch’s tensor-valued engine:

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)

W = torch.randn(3, 3, requires_grad=True)
loss = (W * W).sum()      # every operation on W is recorded
loss.backward()           # the reverse sweep: your Value.backward()
print(W.grad)             # d(loss)/dW = 2W, accumulated into .grad
W.grad.zero_()            # zero it before the next step — or gradients add up
frozen = W.detach()       # same numbers, no graph attached

PyTorch is your Value class with three upgrades: tensors instead of scalars, a C++/GPU backend, and a library of primitives. Nothing else.

3.3 The Bengio model

Now a model that genuinely needs autodiff. Context length $k$ (start $k=3$), embedding dimension $d=16$, hidden width $h=128$:

with parameters $C \in \mathbb{R}^{V \times d}$ (the embedding matrix — each character becomes a learned vector), $W_1 \in \mathbb{R}^{h \times kd}$, $W_2 \in \mathbb{R}^{V \times h}$.

1. Build the dataset of (context, next-char) pairs with a rolling window over ids. Split train = first 90%, val = last 10% — from now on, all reported numbers are validation numbers.
2. Implement the forward pass with tensor ops (C[X] gathers embeddings; .view(-1, k*d) flattens the context; use F.cross_entropy(logits, ys) — it is exactly Lecture 2’s stabilized log-sum-exp loss, and you have now earned it).
3. Write minibatch SGD yourself (no torch.optim until Step 6): sample a random batch of 64 contexts per step, for p in params: p -= lr * p.grad, inside torch.no_grad(); don’t forget to zero the grads (why does PyTorch accumulate instead of overwrite? — you know this from your own engine).
4. Train ~20k steps. Plot train and val loss curves. Compare full-batch loss-per-step against minibatch loss-per-wallclock-second to see why stochasticity wins.
5. Sample from the model (roll the context window as you generate).
6. Sweep $k = 1, 3, 5, 8$ and tabulate val loss and parameter count against the $65^{k+1}$ entries the corresponding count matrix would need. Two things to explain, and the second is the interesting one:
   • $k{=}1$ should land near the bigram’s ≈2.45 (why must it? — what family does the $k{=}1$ model contain?);
   • the sweep is not monotone in $k$. See the checkpoints below before you decide you have a bug.

Checkpoints

• Fan-out test: $a.\mathrm{grad}$, $b.\mathrm{grad}$ match hand computation exactly.

• Step 2 diamond-network input gradient reproduced by your engine, finite differences, and PyTorch ($\le 10^{-6}$).

• Bengio model, $d{=}16$, hidden 128, 20k steps of batch-64 SGD. Measured sweep (your numbers will vary by a few hundredths with seed):

  $k$	params	train	val	bits/char	count matrix would need
  1	11,601	2.4706	2.4805	3.579	$4.2\times10^{3}$
  3	15,697	1.9171	1.9821	2.860	$1.8\times10^{7}$
  5	19,793	1.9054	2.0018	2.888	$7.5\times10^{10}$
  8	25,937	1.9372	2.0329	2.933	$2.1\times10^{16}$

  Read the table carefully — it says three things:

  1. $k{=}1$ reproduces the bigram (2.48 val vs 2.45 train-set CE in Step 1), as it must: the $k{=}1$ model contains the bigram matrix in its hypothesis class.
  2. Going to $k{=}3$ buys a large gain — 0.5 nats — for 4,000 extra parameters, versus the $10^7$-entry matrix a trigram-counting model would need for the same context. That ratio is the entire argument of Lecture 2.
  3. But $k{=}5$ and $k{=}8$ are worse than $k{=}3$, and this is not a bug. At fixed hidden width and fixed step budget, flattening a longer context into one $kd$-dimensional vector spreads the same capacity thinner and gives SGD a harder problem; the extra context isn’t worth what it costs. Note train loss barely improves either, so this is not overfitting — it is underfitting from a bad architecture for long context.

  Point 3 is the motivation for the rest of the course. The fix is not “more parameters”; it is a way of reading context that doesn’t concatenate everything into one flat vector and that can decide, per input, which earlier positions matter. That is attention (Lecture 4).

• Samples have mostly-real short words and plausible morphology. Genuine output from the $k{=}3$ reference solution:

  Seacood whrre,
  This shally
  Beard we erest it hen kess,
  Path horecen in gut co te they livith me same. I give! hat plave vay
  he, yeerfabe fide he to heart had toyter ment:
  Whyre,
  Opere awn the as I when the is party
  Make I heatest of me with in Sear there losm.
  
  LOET:
  He halled A by lord--prustreigno

  Compare against Step 1’s bigram output. Real words now appear in runs (the as I when the is party, heart had), and — with only three characters of context — the model has picked up the SPEAKER: line format and the blank line before it. It has not picked up which speakers exist: LOET: is a plausible-looking name that is not in the corpus. That failure is diagnostic. Three characters of context cannot span a whole name, so the model has learned the shape of a speaker line without the ability to remember any particular one. Holding an identity across dozens of characters is exactly what attention will buy us.

• Your (step, bits/char) table gains a row: 2.86 bits/char at $k{=}3$ (Step 1 bigram: 3.54).

If you’re stuck

• Engine gradients wrong only when a variable is reused: your += is an =, or your topological sort visits a node before all its children.
• RuntimeError: element 0 ... does not require grad: you rebuilt W with an operation that dropped requires_grad, or sliced the graph with .detach()/.item() mid-computation.
• Loss decreasing but samples garbage: check you’re rolling the context window during generation (most common bug this week: feeding a stale context).

Going further (optional)

• Visualize the embedding matrix: PCA C to 2D and plot the characters. Vowels cluster; punctuation clusters; capital letters cluster. Nobody told the model about any of these categories. This picture — geometry learned from prediction alone — is the thesis of the whole enterprise, three lectures early.
• Add log and softmax composite nodes to your engine and derive their local gradients; compare the numerical stability of naive log(softmax) vs the fused form.
• Learning-rate finder: sweep lr exponentially over one epoch, plot loss vs lr on a log axis, and locate the cliff.

Catch-up

Run the data-loading and tokenization cells from $\texttt{solutions/step-01.ipynb}$, then copy your \texttt{diamond_net} function from Step 2 for the validation check. If you did not complete Step 2, its supplied forward-pass code is enough; the new work here is differentiating it.