Bigram Language Model: Make the First Training Step Explainable

A transition table is a model

Row t of a (V,V) table contains next-token logits.

Input [0,1,0] reuses row 0 twice.

The table has V squared parameters and remembers only the current token.

1 / 8

Make softmax numerically safe

Subtract the row maximum before exponentiation.

For [1000,1001], use shifted logits [-1,0].

Keep arithmetic in Value; the common shift changes no probabilities.

2 / 8

Probability stability is not loss stability

Logits [1000,-1000], target 1.

The target probability can underflow, yet the correct loss is about 2000.

Use log(sum(exp(z-m))) - (z_target-m) directly.

3 / 8

Predict one update

Row a starts at [0,0]. Observe a -> b, with b=1.

Loss is log(2); gradient is [0.5,-0.5].

At learning rate 0.1, what is the updated row?

4 / 8

Explain the update

New row a: [-0.05,0.05].

P(b given a) rises to about 0.524979; loss falls to about 0.644397.

Row b does not change. Target selection is different from input row selection.

5 / 8

Check the reduction

Mean cross-entropy divides by position count.

Duplicating all examples preserves the mean gradient.

Clear leaves, use a fresh graph, backward once, update each identity once.

6 / 8

Generate from the last token

Copy a nonempty seed. Read last-token logits, divide by positive finite temperature, sample, append.

Return seed plus new tokens. Lower temperature sharpens the distribution. Zero is invalid here.

7 / 8

Expose the model limit

Both abx and cby end with b before different targets.

A bigram model cannot condition on a versus c.

Next: give the classifier an ordered context window.

8 / 8
←/→ or click edges to navigate · ? help · N notes · F fullscreen