Predicting the next character

A word is a list of characters with a marker at each end

Take one word from the training pile: baker. Five characters. A model that predicts the next character needs to know two more things that the five characters do not say. Where does the word begin, and where does it stop?

Both problems are solved with one extra symbol. Write the word as . b a k e r . and the question what comes first becomes an ordinary next-character question: what follows the dot. The question when do I stop becomes: predict a dot.

The dot is not punctuation. It is a character the model can predict like any other, and it does two jobs with one symbol because start and end never appear in the same position.

Across all 80 words this gives 671 pairs. That number is the whole training set, and it is small enough that you will feel the limits of it later.

Before anything can be counted, the characters need names the computer can use. A grid cannot be indexed by the letter k.

The vocabulary is a numbering you build once and never change

Collect every distinct character in the corpus and sort them. These 80 occupations use 24 letters: no q and no x. Put the dot at the front and you have 25 symbols, numbered 0 to 24.

The dot goes first on purpose. Index 0 is the one you will type most often, since every generation starts there.

Two dictionaries carry the numbering in both directions. stoi takes a character and gives its index, which is what you need when reading text in. itos takes an index and gives the character, which is what you need when writing text out. They are the same information twice, and building both saves you from reversing a dictionary in the middle of a loop.

One thing to watch. The vocabulary is built from the training corpus, so a character that never appeared in it has no seat. Feed this model a name with an x in it and stoi["x"] raises an error rather than doing something sensible. Small corpora make this easy to hit.

With addresses in hand, the first model needs no learning at all. It only needs counting.

Counting every pair fills a grid of 625 boxes

Make a grid with 25 rows and 25 columns, all zeros. Walk through the corpus one pair at a time. For the pair k then e, go to row k, column e, and add one. That is the entire training procedure for the first model, and it takes one pass.

A bigram is just that: two adjacent characters treated as one observation. Row means the character you are standing on, column means the character that came next.

After 671 pairs the grid holds some large numbers and a great many zeros. er appears 44 times, because a large share of these words are people who do a thing: baker, farmer, welder, printer. r followed by a dot appears 56 times, which is the same fact seen from the other side.

The zip(chs, chs[1:]) line is the whole trick for walking pairs: one copy of the list starting at the beginning, one starting one step in, zipped together so each step hands you a character and its successor.

Counts are not yet predictions, though. Row k holds ten observations and row e holds a hundred, and a raw count of 5 means something completely different in each.

A row becomes a prediction when you divide it by its own total

Take row e. It says r happened 44 times, n some number of times, and so on. Add the whole row up, then divide every entry by that total. The counts become fractions that sum to exactly one, and now the row answers a question: given an e, how likely is each of the 25 characters to come next?

That set of 25 numbers adding to one is a probability distribution. Every row of the grid becomes one, and the grid becomes 25 separate predictions, one per current character.

The division has to be per row, which is what dim=1 says: sum along the columns, leaving one total for each row. And keepdim=True keeps that total shaped as a column so the division lines up row by row rather than collapsing into a single number.

Look at the row for j, which appears in joiner, judge and jeweller. Three observations, and the resulting probabilities look every bit as confident as the row for e, which rests on more than a hundred.

There is a worse problem hiding in the same place, and it is not about confidence. Of the 625 boxes in the grid, 411 are still zero.

Adding one to every box buys insurance against words you have not seen

The pair aa never occurs in these 80 words. Its count is zero, so its probability is zero, so the model says it is impossible. Not unlikely. Impossible.

That is a strong claim to make from 671 observations, and it has a hard consequence. In a moment we will score the model by taking the logarithm of the probability it assigned to each pair that actually happened. The logarithm of zero is negative infinity. One unseen pair anywhere in the test data and the score for the whole corpus is infinite, no matter how well everything else went.

Smoothing is the cheap repair: add one to every box before dividing. Every pair now has a small probability instead of none, and pairs that did occur lose a little to pay for it.

The cost is measurable and worth seeing. Scored on the same corpus it was built from, the raw model reaches 1.9436 and the smoothed one 2.2506, so smoothing looks strictly worse. That comparison is rigged: the raw model was scored only on pairs it had already seen, which is the one situation where its zeros can never be reached. The moment a new word arrives, the raw model's score is infinite and the smoothed model's is fine.

With a probability in every box, the table can finally be asked to produce a word.

Generation is one loop: look up a row, draw from it, move

Start at the dot, which is index 0. Read row 0, which is the distribution over first characters. Draw one character using those probabilities as weights. Whatever you drew becomes the new row to read. Keep going until you draw a dot, then stop.

Everything hinges on the word draw. The function torch.multinomial picks an index at random, but not uniformly: a character with probability 0.44 is picked about 44 times in every hundred, and one with 0.01 is picked about once. The high-probability characters are favoured without the low-probability ones being excluded.

The generator with a fixed seed makes the randomness repeatable. Same seed, same word, every run. That is not a detail for tidiness; without it you cannot tell whether a change you made improved anything or you simply got a different roll.

Real samples from this table: bantontr, glmar, ksmucn, ubaryspofstr. Vowels and consonants alternate roughly the way they should, words end after a plausible number of characters, and er shows up often. None of them is a word.

That is a fair description by eye. Eyes are a poor instrument for comparing two models, so the next thing to build is a number.

One number says how surprised the model was by the truth

Here is the test. Take a real word the model should find plausible, walk its pairs, and at each pair ask what probability the model gave to the character that actually came next. For baker the six answers are 0.1143, 0.0952, 0.0308, 0.1613, 0.4412 and 0.4957.

Multiplying six probabilities gives a tiny number, and multiplying six hundred gives a number too small for a computer to hold. So take the logarithm of each one and add instead. Multiplication becomes addition, and the numbers stay in a comfortable range.

Log probabilities are negative, since probabilities are below one. Flip the sign to get a positive quantity where smaller is better, and divide by the number of pairs so that long words and short words are on the same scale. That is the negative log-likelihood, and when a model is being trained against it, the same quantity is called the loss.

The word baker scores 1.8911, which is better than the corpus average of 2.2506, and that makes sense: baker is short and full of common pairs. A single word's score is not a measurement of the model, only of the fit between that model and that word.

So we have a bar to beat: 2.2506, from a table that took one pass through the data. To beat it, something has to change about what the model is allowed to look at.

Two characters of memory would need 15,625 rows and you have 671

The table's weakness is easy to name. Standing on e, it gives the same answer whether the word so far is bak or arch. One character of memory is all it has.

The obvious repair is a bigger table: index by the last two characters instead of one. That means 25 times 25 rows, each with 25 columns. Three characters of memory means 15,625 rows. The table grows by a factor of 25 for every character of context you add.

The data does not grow at all. There are 671 observations, and there always will be. Spread across 15,625 rows, almost every row is empty, and a row built from zero or one observation predicts nothing useful no matter how you smooth it.

The gap is not subtle. With three characters of context there are 15,625 possible rows and 671 observations, so on average each row would hold 0.04 observations. Nearly all of them would be empty boxes filled in by smoothing, which is another way of saying the model would be guessing uniformly.

The counting model does not fail because counting is naive. It fails because the table grows faster than any corpus can fill it.

So the second model keeps the wider window and throws away the table.

Every position in every word becomes one training example

Fix a window width. The code uses three, which it calls the block size: the model always sees exactly the last three characters and predicts the fourth.

Start each word with the window full of dots, because at the beginning there are no previous characters and the dot is what stands for nothing-here. For baker the first example is: given . . ., the answer is b. Then slide the window one step, dropping the oldest character and appending the one just used, and the second example is: given . . b, the answer is a.

Keep sliding to the end, including the final dot as an answer so the model learns where words stop. Five letters give six examples, one per position, exactly as the pair-walk did.

Note context.copy(). Without it every row of X would point at the same list, which keeps being modified, and all 671 rows would end up identical. It is a one-word fix for a bug that produces a model that trains happily and learns nothing.

Something to keep in mind about the split that comes next. The script takes the last 15 percent of these rows as held-out data, which cuts through the middle of a word and leaves the validation set made of whichever occupations happen to be at the end of the list. Splitting by whole word instead is a small change and a much more honest test, and it is what the numbers later on this page use.

Now the model has to do something with three indices. Not compare them, not average them. They are addresses.

Each character gets eight numbers that the training is allowed to change

Build a table with 25 rows, one per character, and 8 columns. Fill it with random numbers. Row 5 is now the description of character 5, and it means nothing at all yet.

Writing C[X] takes every index in X and replaces it with that character's row. A batch of 32 examples, each three characters wide, comes out as 32 by 3 by 8 numbers. No arithmetic has happened; this is a lookup.

What makes this different from the counting table is that these numbers are parameters: training adjusts them. If treating a and o similarly makes the model's guesses better, nothing stops their rows from drifting toward each other, and then anything learned about a partly transfers to o. That transfer is how the model handles contexts it has never seen.

Be careful about how much you read into these numbers. After 750 steps on 80 words, a's nearest neighbour is o, which is pleasing and might be chance. Eight columns is a size chosen because it is small enough to print, and the columns have no individual meanings: no column is the vowel column. Only the pattern across all eight is used.

Three lookups per example leaves 24 numbers. Turning those into a guess is the network.

The network is two multiplications with a squash in between

The three embeddings for one example sit as 3 rows of 8. Lay them end to end into one row of 24, which is what emb.view(emb.shape[0], -1) does. The -1 means work out this dimension from the total, so a batch of 32 becomes 32 by 24 without you having to type 24.

Flattening keeps position information: the first eight numbers are always the oldest character in the window, the last eight always the newest. The network can treat them differently because they arrive in different slots.

Then multiply by a 24 by 64 matrix, add a bias of 64 numbers, and pass the result through tanh, which squashes every value into the range from minus one to one. Without that squash the two multiplications would collapse into a single one and the model could only draw straight-line relationships. Then multiply by a 64 by 25 matrix and add 25 more numbers. Twenty five outputs, one per character. These are the logits: raw scores, not yet probabilities.

Count the numbers being learned: 200 in the embedding table, 1,536 in W1, 64 in b1, 1,600 in W2, 25 in b2. Total 3,425. The counting table had 625 boxes, so this model is about five times larger, and it is being asked to do considerably more with a window three times as wide.

Twenty five scores, and one of them corresponds to the character that actually came next. Turning that into a single number to improve is the last piece before training.

Cross entropy is the same score as before, applied to scores instead of counts

Take the 25 logits. Exponentiate each one, which makes them all positive, then divide by their total so they sum to one. That is softmax, and it turns scores into a probability distribution exactly like a normalised row of the counting table.

Now look up the probability given to the character that actually came next, take its logarithm, flip the sign. That is the same negative log-likelihood from before. Doing both steps at once is what F.cross_entropy is, and it takes logits rather than probabilities on purpose: computing the softmax and the logarithm together avoids exponentiating a large number and then immediately taking its log, which loses precision and can overflow.

Because it is the same measurement, the two models are directly comparable. A cross entropy of 2.4 from the network means the same thing as an NLL of 2.4 from the table.

Here is the moment the wider window earns its place. Given the context b a k, the trained network puts 0.62 on e. The counting table, which can only see the k, puts 0.17 on it. Same corpus, same answer, and the difference is entirely that one model knew a b and an a came before.

That comparison is on a context the model has seen. Getting it to hold on words it has never met is what training is for, and training is where the interesting failure lives.

Training is four lines repeated: measure, blame, nudge, clear

Pick 32 rows of X at random. Run them through the network and get a loss. Call loss.backward(), which works out, for each of the 3,425 numbers, whether nudging it up would raise the loss or lower it, and by how much. That quantity is the gradient. Then move every number a small step in the direction that lowers the loss.

The step size is the learning rate, 0.1 here. Too small and training crawls. Too large and each move overshoots, and the loss bounces around instead of settling.

The gradients must be cleared before each pass, which is what setting p.grad = None does. PyTorch adds new gradients to whatever is already there, so skipping the clear means every step is being steered partly by measurements from steps you have already taken.

The 32 rows are a batch. Using a sample rather than all 562 rows makes each step much cheaper and slightly wrong, which is a trade almost always worth taking: many noisy steps beat a few exact ones.

At step 0, before any learning, the loss is 3.442. Guessing uniformly among 25 characters would score 3.219, so the random starting point is slightly worse than knowing nothing, which is what random should mean. After 100 steps it is 2.251, and after 750 it is 1.413.

Falling loss looks like progress. It is worth asking, carefully, what exactly is getting better.

The training loss keeps falling long after the model stops improving

Hold out twelve of the eighty occupations before training starts. Fit on the other sixty eight, and every so often measure the loss on the twelve the model has never seen. Two numbers now, and they do not do the same thing.

The training loss falls the whole way, from 3.442 to 0.601 after ten thousand steps. The validation loss falls to 2.410 at around step 750, and then turns around and climbs to 4.209.

That climb is overfitting, and the name is exact: the model is still getting better at the sixty eight words in front of it, by learning them specifically, and that specific knowledge is worthless on any other word. With 3,425 parameters and 562 training examples there is easily enough capacity to memorise.

The repair is early stopping: keep the version from around step 750 and throw the rest away. The script's default of ten thousand steps trains straight past the useful model and hands you the memorised one.

Two things to keep in mind. The gap between the two curves is a size signal: a large gap with a small corpus usually means the model is too big or trained too long, and both are cheap to fix. And the validation set here is twelve words, so its loss is a noisy number, which is why the curve wobbles near the bottom rather than turning at one clean point.

Stopping at 750 gives a model worth sampling from. What comes out of it depends on one more choice, made after training is over.

Temperature reshapes the distribution without touching the model

Divide the logits by a number before the softmax. That number is the temperature, and it changes nothing that was learned. Same weights, same scores, different shape to the distribution they turn into.

Divide by 0.5 and the differences between logits double, so the leader pulls away and the tail is crushed. Divide by 1.5 and the differences shrink, so the distribution flattens and unlikely characters get a real chance. At exactly 1 the logits are used as they are.

The effect on generated words is direct. Low temperature gives safe, repetitive, often memorised output. High temperature gives variety and more nonsense. There is no correct setting, only a choice about which failure you would rather have.

One thing that catches people out: temperature has no effect on the loss. The loss is measured on the distribution the model actually produces, at temperature 1. Reporting a low-temperature sample as evidence the model improved is comparing two different things.

Which leaves the question the whole project exists to answer. Two models, one corpus. Which one is actually better?

The network wins by a tenth of a nat, and not for the reason you would guess

To compare fairly, both models get the same sixty eight training words and are scored on the same twelve held-out ones. The counting table, fitted on those sixty eight, scores 2.496. The network, stopped at its best point, scores 2.410. A win of 0.086.

That is a real improvement and a small one. On a corpus of eighty words it is roughly what there is to win.

Now the part that is easy to get wrong. Run the same experiment with the window set to one character, so the network sees exactly what the counting table sees, and it scores 2.379. Better than the three-character network. The window was not what helped.

What helped is that the network's numbers are shared. Every character's eight numbers are adjusted by every example that character appears in, so a rare context still gets an answer assembled from things the model learned elsewhere. The counting table has to fill each row from that row's own observations, and on 671 observations most rows are nearly empty.

What would change the answer is more words. The same code on tens of thousands of names shows the wider window paying off clearly, because then there is enough evidence behind each pattern for the model to have something to generalise from. Nothing about the architecture has to change; only the pile of text does.

Guessing the next character is the whole mechanism. Everything else is how much you can afford to remember while you do it.

A last thing worth holding on to. Both models here are doing exactly what a large language model does: assign a probability to every possible next symbol, draw one, move along. What separates them from a real system is the size of the vocabulary, the width of the window, and the number of parameters. Not the idea.

The whole project in one screen

Count pairs. Divide by row totals. Add one so nothing is impossible. Draw from a row and move. Score with the log of the probability you gave the truth. Then, when the table cannot grow any further, swap the rows for a few thousand shared numbers, slide a window over the text, and push those numbers downhill until a held-out score stops improving.

Run it yourself with python build.py --tiny for five hundred steps, or --full for ten thousand, which will show you the overfitting rather than tell you about it.