Lesson 1.4: A first model: the bigram language model
Train a bigram language model by gradient descent, check that it learns the probabilities a table of counts gives, and see what a one-word memory cannot do.
By the end of this lesson you will be able to turn words into integer ids, write a bigram language model as a table of learnable numbers, compute its loss and its gradient with array operations, train it with gradient descent, check that it reaches the loss of the count-based table, sample text from it, and explain from its output why a one-word memory is not enough. You will show that an untrained model scores 3.54 nats against a blind-guess value of 3.53, that gradient descent brings it to 1.162 against the count table's 1.159, that the learned probabilities differ from the counted ones by at most 0.0017, that one learning rate reaches a loss of 1.32 in 100 steps where another is still at 3.36, and that the sentences it writes are fluent word to word and wrong as sentences. You should have read Lessons 1.1 to 1.3. The mathematics is the softmax and its gradient from Lesson 1.3. This course is education about how language models work, not a guide to any commercial product.
Lesson 1.1 built a bigram model out of counts. Lesson 1.3 built the engine that finds numbers by gradient descent. This lesson joins them. The result is the first model in the course that learns instead of counting. It is also the simplest member of the family that includes every large language model: a function from a context to a distribution over the next token, with adjustable parameters, trained by following a gradient. Everything that follows in the course changes the function in the middle, and nothing about how we train or score it.
Words as numbers
A computer cannot take the logarithm of the word "cat". The first step is to give every token in the vocabulary an integer id. We sort the 34 distinct tokens, number them from 0, and keep two lookups: stoi from a string to an id, and itos back again.
python
import numpy as np
from corpus import makecorpus
text = makecorpus()
words = text.split()
vocab = sorted(set(words))
V = len(vocab) # 34
stoi = {w: i for i, w in enumerate(vocab)}
itos = vocab
ids = np.array([stoi[w] for w in words])
cut = int(0.9 len(ids))
trainids, valids = ids[:cut], ids[cut:]
Each position t has the input token ids[t] and the target ids[t+1].
xtrain, ytrain = trainids[:-1], trainids[1:] # 26,546 pairs
xval, yval = valids[:-1], valids[1:] # 2,949 pairs
The task is now precise. We have 26,546 training pairs. In each, the input x is a token id and the target y is the id of the token that followed it. The model sees x and must put as much probability as it can on y. The 2,949 validation pairs, from the last tenth of the text, are used only to score.
The model: a table of scores
A bigram model needs, for each of the 34 possible previous tokens, a distribution over the 34 possible next tokens. We store it as a matrix W with 34 rows and 34 columns. Row i holds the scores (logits) for what follows token i. To predict after token x, we take row x of W and turn it into probabilities with a softmax:
p(next | x) = softmax(W[x])
That is the whole model. It has 34 × 34 = 1,156 parameters, the same number of cells as the count table. The difference is that nobody fills them in from counts. We begin with small random numbers and let gradient descent adjust them.
There is a second way to say the same thing that will matter later. Represent the token x as a one-hot vector: a vector of 34 numbers, all zero except a one in position x. Multiplying this vector by W picks out row x:
python
onehot = np.eye(V)[xtrain[:5]] # shape (5, 34)
print(np.allclose(onehot @ W, W[xtrain[:5]])) # True
So the lookup W[x] is the same as a matrix multiplication of a one-hot input. A bigram model is a single linear layer followed by a softmax. In Module 2 the lookup will be called an embedding, and in Module 4 the single layer will become a stack of many.
The loss and its gradient
For one pair, the loss is −ln p(y | x), the surprise at the correct answer. For the whole training set it is the average. In the array form:
python
def softmax(z):
z = z - z.max(axis=1, keepdims=True) # stable, as in Lesson 1.2
e = np.exp(z)
return e / e.sum(axis=1, keepdims=True)
def lossandprobs(W, x, y):
P = softmax(W[x]) # (N, V): a distribution for every input
loss = -np.log(P[np.arange(len(y)), y]).mean()
return loss, P
P[np.arange(len(y)), y] picks, for every pair, the probability given to the correct token. Its negative logarithm, averaged, is the cross-entropy.
Lesson 1.3 gave us the gradient of the loss with respect to the logits of one example: the probabilities minus the one-hot target. For the average over N examples it is the same thing divided by N. That is the gradient with respect to W[x], one row per example. Because several examples share the same previous token, and so the same row of W, their gradients must be added into that row. That is exactly the rule for a number used in several places, and NumPy provides np.add.at to do it:
python
def gradient(W, x, y):
loss, P = lossandprobs(W, x, y)
dlogits = P.copy()
dlogits[np.arange(len(y)), y] -= 1 # probabilities minus one-hot target
dlogits /= len(y) # average over examples
dW = np.zeroslike(W)
np.add.at(dW, x, dlogits) # add each example's gradient to its row
return loss, dW
We should check this before trusting it, in the way Lesson 1.3 taught. A numerical estimate for a single cell of W takes four lines:
python
rng = np.random.defaultrng(0)
W = rng.standardnormal((V, V)) 0.1
loss, dW = gradient(W, xtrain, ytrain)
i, j, h = 5, 7, 1e-5
Wp, Wm = W.copy(), W.copy()
Wp[i, j] += h
Wm[i, j] -= h
numeric = (lossandprobs(Wp, xtrain, ytrain)[0] - lossandprobs(Wm, xtrain, ytrain)[0]) / (2 h)
print(dW[i, j], numeric) # the two agree to many decimal places
If they agree, the gradient is right. If not, the difference points to the bug before any training run is wasted.
Training
With a loss and a gradient, training is the loop from Lesson 1.3: compute the gradient, take a step against it, repeat. Because the whole training set is small, each step uses all of it, which is called full-batch gradient descent.
python
def train(lr, steps, seed=0):
rng = np.random.defaultrng(seed)
W = rng.standardnormal((V, V)) 0.1
for step in range(steps + 1):
loss, dW = gradient(W, xtrain, ytrain)
if step in (0, 1, 10, 50, 100, 200, 500, 1000, 3000):
valloss, = lossandprobs(W, xval, yval)
print(step, round(loss, 4), round(valloss, 4))
W -= lr dW
return W
W = train(lr=10.0, steps=3000)
The first number to check
At step 0 the loss is 3.542. This is the single most useful number to check in any model you train. A model that knows nothing should give nearly equal probability to every token, so its loss should be close to ln V = ln 34 = 3.526. Ours is 3.542, a little higher because the random starting scores are not exactly equal, which is as expected.
If the starting loss were far from ln V, something would be wrong. For instance, starting W with large random numbers, with a standard deviation of 3 instead of 0.1, gives a starting loss of 7.16. The model is then confidently wrong from the start, and training has to undo that before it can learn anything. A starting loss of ln V is a sign of a well-initialised model, and we will use this check on every model from here on.
The loss curve
For a learning rate of 10:
Step · Training loss · Validation loss0 · 3.542 · 3.543
1 · 3.361 · 3.362
10 · 2.407 · 2.407
50 · 1.482 · 1.485
100 · 1.317 · 1.321
200 · 1.229 · 1.232
500 · 1.182 · 1.183
1,000 · 1.169 · 1.170
3,000 · 1.161 · 1.162
Read three things from the table. First, the model learns. It goes from a blind guess to a bigram model by following a gradient. Second, the training and validation losses stay almost equal, because the model has only 1,156 parameters and 26,546 pairs to fit them, so it cannot memorise: there is little room for overfitting. Third, the loss falls fast at the start and slowly at the end, as in Lesson 1.3, because the gradient shrinks as the model improves.
It learns what the counts say
Is the model any better than the count table? It cannot be: the count table is the best possible bigram model on the training data, since each row of counts divided by its total is exactly the distribution that minimises the training loss. We can verify that the learned probabilities agree with the counted ones:
python
counts = np.zeros((V, V))
np.add.at(counts, (xtrain, ytrain), 1)
Pcount = (counts + 0.01) / (counts + 0.01).sum(axis=1, keepdims=True)
Plearned = softmax(W)
print(np.abs(Pcount - Plearned).max()) # 0.0017
The largest difference between any probability in the two tables is 0.0017. After 3,000 steps the learned model has found the table we built by counting, to within a few thousandths. Their validation losses are 1.162 and 1.159. The remaining gap closes as the number of steps increases, because the optimum for the unsmoothed model has some probabilities that are exactly zero, which gradient descent can only approach.
So what did training gain us? For this model, nothing: it did the same job as counting more slowly. That is the point of the exercise. We have shown that the learned model and the counted model are the same thing, so that when we replace the table by something that cannot be counted, we know what the learning is trying to reproduce. In Module 3 a model with a longer memory will do what no table of counts can: share what it learns between similar contexts.
Reading the learned probabilities
The learned rows say what the grammar of the toy language is. Here are the five most probable next words for four inputs, from the trained model:
Previous token · Most probable next tokens (probability). · the (0.626), a (0.246), my (0.128)
the · ball (0.152), old (0.148), red (0.145), dog (0.114), child (0.113)
red · bus (0.999)
chased · a (0.517), the (0.480)
After a full stop, a new sentence begins. In the generator, five of the eight subjects start with "the", two with "a" and one with "my", so the true probabilities are 0.625, 0.25 and 0.125. The model learned 0.626, 0.246 and 0.128, from counting alone. After "red", only "bus" is possible, and the model has learned that to be 0.999 certain. After "chased", the words "a" and "the" are about equally likely, which is correct.
Now look at the row for "the". It gives similar probability to "ball", "old", "red", "dog" and "child": a mixture of the things that can follow "the" at the start of a sentence (subjects such as "dog" and "child") and in the middle of one (objects such as "ball" and "old"). The model cannot tell these cases apart, because it sees only the previous word. That is the whole weakness of a bigram model, and the next section shows what it costs.
Experiment with the loss and its gradient on a single prediction, to see how the surprise and the push on the correct score change with the probability the model gives it.
Writing text
To generate text, we start from a token, sample the next one from its row, and repeat. Each new token becomes the context for the next.
python
def sample(W, ntokens, seed, start="."):
rng = np.random.defaultrng(seed)
i = stoi[start]
out = []
for in range(ntokens):
p = softmax(W[i:i+1])[0]
i = rng.choice(V, p=p)
out.append(itos[i])
return " ".join(out)
print(sample(W, 30, seed=1))
print(sample(W, 30, seed=2))
Two samples from the trained model:
> the red bus today . the old tree at night . the cat saw a fox chased a fish . the cat found the red bus at night . the
> my teacher saw a loud drum . a fish quickly . a fox saw a loud drum quickly . a fish again . the cat watched a loud drum .
Read them closely. Every adjacent pair of words is perfectly valid: "red bus", "cat saw", "loud drum", "drum quickly". Locally the text looks right. But the whole is wrong in ways the corpus never is. "The red bus today ." and "a fish quickly ." have no subject and no verb. "The cat saw a fox chased a fish" has two verbs in a row. The model that has just seen "bus" cannot know whether a subject came earlier, because it has already forgotten it.
This is a precise example of what a limited context does. The bigram model knows what follows each word and nothing about how the sentence got there. A fluent-looking local surface hides a lack of structure. Real language models fail in the same way, only more subtly, when their context cannot hold what they need.
The trigram table from Lesson 1.1 had a loss of 0.992, against the bigram's 1.159, and the difference is the part of the grammar visible only from two words back: after "the red" comes "bus", after "my teacher" comes a verb. No bigram model, however well trained, can score better than 1.159 on this data, since the limit is in its memory, and not in its training.
The learning rate, once more
The learning rate of 10 looks large next to Lesson 1.3's 0.05, but the loss is an average of 26,546 terms and its gradient is correspondingly small, so the right step is bigger. After 100 steps, from the same start:
Learning rate · Validation loss after 100 steps0.1 · 3.365
1 · 2.415
10 · 1.321
50 · 1.205
At a rate of 0.1 the model has hardly moved after 100 steps, and at 50 it has nearly finished. A good learning rate is the largest one that does not make the loss rise, found by trying values separated by a factor of ten and watching the loss. Because each step of this problem is a safe, smooth descent, a very large step is fine here; in a deep model with many interacting layers the stable range is much narrower.
The pull towards uniform
One more small experiment prepares the ground for later lessons. If we add to the gradient a term proportional to W itself, dW += wd W, the update shrinks every score slightly towards zero at each step. A model whose scores are all zero gives all tokens the same probability, so this penalty, called weight decay, pulls the model towards "know nothing". In 1,500 steps at a learning rate of 10, a weight decay of 0 gives a validation loss of 1.166, 0.001 gives 1.469, and 0.01 gives 2.500.
A small amount works like the smoothing in the count table: it keeps the model from betting everything on rare events. A large amount stops the model from learning at all. The same trade-off, between fitting the data and keeping the model plain, runs through the whole course, and Lesson 5.2 describes how real training uses weight decay.
What we have, and what is missing
We now have a complete small system: data turned into integers, a model with parameters, a loss that measures it, a gradient that is checked against a numerical estimate, a training loop that reduces the loss, a held-out score, and a generator. Each of those has an equivalent in a model with billions of parameters, and none of them will change.
What is missing is memory. The bigram sees one token. To see more, we need a function that takes a sequence and mixes the information of the tokens in it. The first step is to make the representation of each token richer than a row of 34 scores, and to choose a better way to split text into tokens than words. That is Module 2. After that, Module 3 introduces the mechanism that lets each token look back at the ones before it.
Key takeaways
Give each token an integer id. A bigram model is a matrix W of shape (V, V); row x holds the logits for what follows token x. It has 34 × 34 = 1,156 parameters, and looking up a row equals multiplying a one-hot vector by W.
The loss is the average of −ln p(correct token). The gradient with respect to the logits is the probabilities minus the one-hot target, divided by the number of examples, added into the row for each input token with np.add.at.
A well-initialised model starts at a loss near ln V. Ours started at 3.542 against ln 34 = 3.526; a badly scaled start gave 7.16.
Gradient descent took the validation loss from 3.543 to 1.162, within 0.004 of the count table's 1.159, with learned probabilities within 0.0017 of the counted ones.
The model learned the generator's structure: 0.626, 0.246 and 0.128 for sentence-starting "the", "a" and "my" against true values of 0.625, 0.25 and 0.125; and 0.999 for "bus" after "red".
Its sentences are locally right and globally wrong because it sees one word. No bigram can beat 1.159 on this data; the trigram reached 0.992.
The learning rate sets the speed. After 100 steps, a rate of 0.1 had reached 3.37, a rate of 10 reached 1.32, and a rate of 50 reached 1.21.
Common mistakes
Skipping the starting-loss check. A model that does not start near ln V has a bug in its initialisation or its loss. Find it before training.
Assigning gradients instead of accumulating them. Several examples share a row, so use np.add.at. Writing dW[x] = dlogits keeps only the last example for each token.
Forgetting to divide by the number of examples. The loss is an average, so its gradient is too. A missing division scales the learning rate by N.
Evaluating on training data. Always report the held-out loss as well. Here the two agree because the model is tiny; with a bigger model they will not.
Mixing up inputs and targets. The target at position t is the token at position t + 1. Shifting by the wrong amount trains the model to copy its input.
Reading fluent text as understanding. A bigram model writes plausible word pairs, and no structure.
Check yourself
1. Why does a model that knows nothing have a loss of about ln V?
It gives every token probability 1/V, so the loss is −ln(1/V) = ln V for every example.
2. The model has V = 34 and a context of one token. How many parameters does it have, and what are the shapes of W[x] for a batch of 100 inputs and of dW?
1,156 parameters. W[x] is (100, 34); dW is (34, 34), the same shape as W.
3. Two training pairs have the same input token, 5. Their gradients with respect to the logits are g₁ and g₂. What is the gradient for row 5 of W?
g₁ + g₂, because row 5 is used by both and the contributions add.
4. Why does the trained bigram model never score better than 1.159 on the validation set?
It sees only one previous token. Information that sits two tokens back, such as the "red" in "the red bus", is unavailable, and the missing information costs about 0.17 nats.
5. A learning rate of 50 reached a loss of 1.205 in 100 steps and 0.1 reached 3.365. Does that mean 50 is the best rate in general?
No. It is the best of those four on this problem. A rate that is too large makes the loss rise or diverge, as in Lesson 1.3, so the choice must be tested.
Try this
Set aside ninety minutes. You need Python and NumPy, and the files from the earlier lessons.
1. Run the code and reproduce the table of losses. Confirm the starting loss and the gradient check.
2. Change the initial scale from 0.1 to 3. What is the starting loss, and how many steps does it take to reach 1.5?
3. Print the row for "chased" and the row for "saw". Are they close to each other? Explain with reference to the generator.
4. Make the model score its training data after training for 3,000 steps with a weight decay of 0.001. Does the validation loss fall or rise, and by how much?
5. Generate 20 sentences, and count how many are grammatical according to the generator's rule (subject, verb, object, optional adverb, full stop). What fraction do you get?
6. Compute the bigram loss at the first word of each sentence (inputs equal to the full stop), and at all the other words, separately. Which is higher, and why?
A good result is a training run that you can reproduce exactly, a loss curve that approaches the count table from above, and a clear sense of the limit imposed by a one-word memory. The next module changes how text is turned into numbers: first the tokens themselves, and then the vectors that represent them.
Read more
Andrej Karpathy's lecture "The spelled-out intro to language modeling: building makemore" builds a bigram character model and shows both the count and the gradient versions, which is the same lesson on a different dataset. Daniel Jurafsky and James Martin's Speech and Language Processing, chapter 3, covers n-gram models and their evaluation. Yoshua Bengio and colleagues' 2003 paper "A Neural Probabilistic Language Model" introduced the idea of replacing the table with a learned function.
Tags: Build an LLM Bigram model Softmax Cross-entropy Gradient descent