A bigram model

Generating text

Think of each row as a loaded die with one face per character. The faces aren't equally likely. Characters that often came next get bigger faces.

To write text, start on the newline, the break between two descriptions. Roll the newline's die to pick the first character. Then roll that character's die to pick the next one. Keep going until you roll a newline or reach 80 characters. Choosing at random by the chances like this is called sampling.

import numpy as np
import requests

text = requests.get("https://api.gradient.test/v1/corpus").json()["text"]
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
ids = np.array([stoi[c] for c in text])
counts = np.ones((len(chars), len(chars)))
np.add.at(counts, (ids[:-1], ids[1:]), 1)
probs = counts / counts.sum(axis=1, keepdims=True)

rng = np.random.default_rng(0)
for _ in range(3):
    current, out = stoi["\n"], ""
    for _ in range(80):
        current = rng.choice(len(chars), p=probs[current])
        if chars[current] == "\n":
            break
        out += chars[current]
    print(repr(out))

This cell counts over the whole catalog, because it only writes text and doesn't measure a loss. rng.choice is the die roll: it picks a character ID using the chances in the current row. The random generator starts from a fixed seed (0), so you get the same three lines every run.

The output has the texture of the catalog. Common letter pairs appear, and spaces land in sensible places. But there are no real words and no meaning. The model remembers one character only. After an "a", it has no idea which word it's in. So it can't spell anything longer than two letters on purpose.

Smoothing: never say never

Some pairs never show up in the training text. Their count is zero, so their chance is zero. The model is saying "this can never happen."

Then a new description uses that pair. The loss for that step is $-\log 0$, which is infinite. The loss is an average, so one infinite step makes the whole average infinite.

Think of a weather forecaster who says "0% chance of rain", and then it rains. By this score, the forecaster was infinitely wrong. A careful forecaster never says exactly 0%.

The fix is called smoothing. Before you divide, add a small extra count to every cell of the table, even the empty ones. Adding 1 is called add-one smoothing. Adding some other amount is called add-α smoothing. The Greek letter α (alpha) stands for that extra amount.

The top is the old count plus α. On the bottom, $V$ is the number of different characters. Every row has $V$ cells, and each one got an extra α. So the row total grows by α times $V$, and each row still adds up to 1.

In plain words: pretend you saw every possible pair a few extra times. Then count and divide as before. Pairs you never saw get a small chance instead of zero.

Back to "abracadabra". It uses $V = 5$ letters: a, b, c, d and r. Row "a" has the counts 0, 2, 1, 1 and 0. With α = 1, add 1 to each: 1, 3, 2, 2 and 1. The new total is 9, which is the old total of 4 plus 1 × 5. So the chance of "b" after "a" is 3/9, about 0.333. The pair "ar" never appeared, yet it now gets 1/9, about 0.111.

from collections import Counter

text = "abracadabra"
letters = sorted(set(text))                   # a, b, c, d, r: V = 5
pairs = Counter(zip(text, text[1:]))
alpha = 1
row = [pairs[("a", b)] + alpha for b in letters]
total = sum(row)
for b, n in zip(letters, row):
    print(f"P({b} | a) = {n}/{total} = {n / total:.3f}")

The five chances it prints add up to 1.

The amount you add matters. Too little, and unseen pairs still get tiny chances. Too much, and the extra counts drown out the real ones. Every row drifts toward equal chances, and the model forgets what it learned. The exercise tries several amounts on the catalog.

Neural models get this protection for free. Their last step is a softmax, and a softmax never outputs exactly zero.

Scoring a tiny model

Here's the loss again, written out for a bigram model. Number the test text's characters $x_0, x_1, x_2$ and so on. Call the last one $x_n$, so the text has $n$ neighboring pairs.

The symbol Σ (sigma) means "add up". Here it adds over every pair, from the first (i = 1) to the last (i = n). The log is the natural log, the same as NumPy's np.log.

In plain words: walk through the test text one pair at a time. For each pair, look up the chance the model gave the real next character. Take minus the log of it. Then average.

Train on "abracadabra" with α = 1, and test on "radar". It has four pairs: ra, ad, da and ar. The model gives them chances of 3/7, 2/9, 2/6 and 1/9. Minus the logs are about 0.847, 1.504, 1.099 and 2.197. Their average is about 1.412. Without smoothing, "ar" has a chance of 0, and the loss is infinite.

import math
from collections import Counter

train, test = "abracadabra", "radar"
letters = sorted(set(train))
V = len(letters)
pairs = Counter(zip(train, train[1:]))

def prob(a, b, alpha):
    row_total = sum(pairs[(a, c)] for c in letters)
    return (pairs[(a, b)] + alpha) / (row_total + alpha * V)

for alpha in [0, 1]:
    losses = []
    line = f"alpha {alpha}:"
    for a, b in zip(test, test[1:]):
        p = prob(a, b, alpha)
        losses.append(-math.log(p) if p > 0 else math.inf)
        line += f"  {a}->{b} {p:.3f}"
    print(line, f"  loss {sum(losses) / len(losses):.3f}")

The same model, learned by gradient descent

So far you built the table by counting. Now build it a different way: by training, as in Phase 4. This is the key step toward neural language models.

Make a table $W$ with one row per character and one column per possible next character. It has $V$ rows and $V$ columns. Row a holds a raw score for every character that might follow a. These raw scores are called logits. A softmax turns a row of logits into chances that add up to 1.

Fetching row a is exactly what an embedding layer does. As you saw in Phase 4, that's the same as multiplying a one-hot vector by the table. So the whole model is one layer.

Training works as usual. The scores start out random. The model sees each pair from the training text. Cross-entropy measures how surprised it was by the real next character. Then gradient descent nudges the scores to cut that surprise. The optimizer here is Adam.

import numpy as np
import requests
import torch
import torch.nn as nn

text = requests.get("https://api.gradient.test/v1/corpus").json()["text"]
docs = text.split("\n")
train, val = "\n".join(docs[:270]), "\n".join(docs[270:])
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
tr = torch.tensor([stoi[c] for c in train])
va = torch.tensor([stoi[c] for c in val])

torch.manual_seed(0)
W = nn.Embedding(len(chars), len(chars))         # row a = logits for the character after a
opt = torch.optim.Adam(W.parameters(), lr=0.1)
for step in range(150):
    loss = nn.functional.cross_entropy(W(tr[:-1]), tr[1:])
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 50 == 0 or step == 149:
        with torch.no_grad():
            val_loss = nn.functional.cross_entropy(W(va[:-1]), va[1:]).item()
        print(f"step {step:>3}  training loss {loss.item():.3f}  held-out loss {val_loss:.3f}")

Both losses fall as training goes on. The trained layer ends at a held-out loss of about 2.17. That's the same as the counting model with light smoothing. (The exercise compares smoothing amounts.) It's the same model, learned two ways. (torch here is the course's in-browser PyTorch subset.)

The neural version is worth the extra effort because it can grow, and a counting table can't. For two characters of context, a table would need a row for every pair of characters. It also can't share what it learns between similar contexts. The neural version can take more context and more layers, as the next lessons show.