Working Notes on Large Language Models

3. From transformer to language model

Dr. Ali Al Bataineh, Higher Colleges of Technology, United Arab Emirates

The models in chapters 1 and 2 read twelve symbols. A language model reads text, and the gap between those two is mostly three things: a tokenizer that turns text into a sequence of integers and back, a model wide and deep enough to hold something about the language, and a sampler that turns the model's distribution over the next token into an actual next token. This chapter builds all three on one book, Alice's Adventures in Wonderland, 144,016 characters, and trains a model with 924,416 parameters on it. The result reads like English from a distance and falls apart up close, and the chapter spends as much time on the falling apart as on the reading, because the failure is the lesson: a model this size memorises a book this short in about eleven minutes of training on a laptop, and the held-out loss tells you exactly when it starts.

The numbers come from 03-language-models.ipynb as executed on the laptop described in chapter 1. Training took 672 seconds.

3.1 Tokens

A model needs a fixed vocabulary of symbols, and English does not come with one. Characters are too few: every word is a long sequence and the model spends its capacity on spelling. Words are too many: the book has thousands, a real corpus has millions counting typos and inflections, and anything unseen at training time cannot be represented at all. Byte-pair encoding sits between. Sennrich et al. (2016) adapted it from data compression to splitting words, starting from characters; the byte-level version used here follows GPT-2 (Radford et al., 2019). Start with 256 symbols, one per byte, so that every possible input is representable, and then learn larger symbols from the text itself.

The learning rule is one line. Count every adjacent pair of symbols in the training text; take the most frequent pair; replace every occurrence of it with a new symbol; repeat. Each merge adds one symbol to the vocabulary and shortens the text. Stop when the vocabulary is the size you wanted.

Done on the book with a vocabulary of 1,024, the 768 merges take 5 seconds, and the compression at each stage is in the table. The first 64 merges mostly build common two- and three-letter fragments and take an average token from about one character to 1.55. The returns then shrink, and the last 256 merges together add only a quarter of a character.

vocabulary characters per token
320 1.55
512 2.18
768 2.57
1,024 2.81

Figure 3.1 shows the whole curve. The first merge joins a space to t, the start of the and to, which are everywhere in English. By merge 768 the merges are building fragments of single names, such as inah from Dinah, the cat Alice keeps mentioning. Each later merge buys less because it serves fewer words.

Characters per token as the vocabulary grows, one point per merge

Figure 3.1. Average characters per token over the whole book as the vocabulary grows from 257 to 1,024, one point per merge. The dots are the four sizes in the table; the labels show the token made by a few of the merges.

Here is the sentence The Queen turned crimson with fury. after encoding, token by token:

['The', ' Queen', ' tur', 'ned', ' c', 'ri', 'm', 's', 'on', ' with', ' f', 'ury', '.']

Two things to notice. Common words are one token and carry their leading space with them, so the model never has to decide whether a space goes before a word. Rare words are spelled out from fragments: crimson appears twice in the book and gets five tokens. A tokenizer trained on a hundred gigabytes of text would have crimson whole and Queen too, but the mechanism is identical, and the quirks you will meet in large models (numbers split oddly, words that tokenize differently with and without a space in front) are all visible here at small scale.

Encoding a new string is the merges replayed in the order they were learned, which the notebook does in encode_chunk. The order matters: a later merge can only apply to symbols that earlier merges have already produced.

The whole book is 51,281 tokens. The last 10% of them, 5,129 tokens, are held out and never trained on.

3.2 The model

The transformer is chapter 2's, with rotary positions and the cache, made larger: four blocks, a 128-dimensional embedding, four heads of 32 dimensions, a context of 128 tokens. One addition: the output layer that maps the final vector to vocabulary logits shares its weight matrix with the input embedding. The embedding says what vector each token has; the output layer asks how much the final vector looks like each token's vector; using the same matrix for both is a reasonable thing to do and saves 1,024 by 128 parameters. Press and Wolf (2017) showed that tying the two lowers perplexity, and small models commonly do it.

Total: 924,416 parameters, of which 131,072 are the tied embedding, 256 are the final layer norm, and the rest are the four blocks.

3.3 Training, and when to stop

AdamW at a constant learning rate of 0.001, 1,500 steps, each a batch of 32 windows of 128 tokens drawn at random from the training text. Every 150 steps the loss is measured on 20 fixed batches from the held-out tail of the book, so the curve is comparable point to point. Figure 3.2 shows both curves.

Training and held-out loss for the book model

Figure 3.2. Loss in nats per token, every 150 steps. The grey line is $\ln 1024 = 6.93$, the loss of a uniform guess.

The training loss goes from above the uniform baseline to 0.246 and would go lower. The held-out loss goes to 4.807 at step 300 and then rises, to 10.944 by step 1,500, which is worse than guessing uniformly. The model has not stopped learning; it has stopped learning the language and started learning the book. Each training token has been seen more than a hundred times by the end (1,500 steps of 32 windows of 128 tokens is 6.1 million tokens drawn from a pool of 46,152), and a model with 924,416 parameters can store a 46,152-token text with room to spare.

The right model is therefore not the last one. The notebook keeps the weights from the step with the lowest held-out loss, step 300, and that is the model the rest of this chapter samples from and the one saved to checkpoints/03-gpt.pt for chapters 4 to 6 and 8. Its held-out perplexity is $e^{4.807} = 122$: on average, the model's next-token distribution on unseen text is as spread out as a uniform choice among 122 tokens. Large models on ordinary English are far lower, although perplexities are only comparable between models that share a tokenizer and a test text. Most of the gap is the difference between a book and the internet.

3.4 Memorisation, measured

Overfitting is a word; here is the thing itself. Take 20 passages of 32 tokens from the training text, give them to the model, and let it continue greedily for 40 tokens. Compare with the 40 tokens the book actually continues with.

The model at the end of training reproduces 86% of those tokens exactly. The model at step 300 reproduces 2%. The first number says the final model is a lossy copy of the book that can be read back by prompting it with any passage. The second says the step-300 model is not; it has the statistics of the author's English and little of the text itself.

This is the mechanism behind two things you will hear about large models: that they can emit training data verbatim, and that evaluation on data the model trained on tells you nothing. Chapter 8 returns to the second as "contamination".

3.5 Sampling

The model gives a probability for each of the 1,024 tokens. A sampler decides what to do with that. The notebook runs one function with five settings on the same prompt, Alice looked at the, with the same random seed, so the differences are the sampler's.

Greedy takes the most probable token every time. It is deterministic and it loops: once a phrase is the most likely continuation of itself, the model repeats it.

Temperature divides the logits by a number $T$ before the softmax. Below 1 it sharpens the distribution toward the top tokens; above 1 it flattens it toward uniform. At 0.8 the text stays closer to common words and phrases; at 1.5 it is mostly fragments, because the long tail of the vocabulary gets real probability and the tail is where the nonsense is.

Top-k keeps only the $k$ most probable tokens and renormalises. It cuts the tail off at a fixed count, which is too many tokens when the model is confident and too few when it is not.

Top-p, or nucleus sampling (Holtzman et al., 2020), keeps the smallest set of tokens whose probabilities add up to at least $p$, so the cut moves with the shape of the distribution: a few tokens when the model is sure, many when it is not. It is widely used in practice, often together with a temperature slightly below 1.

if top_p is not None:
    probs, order = torch.softmax(logits, dim=-1).sort(descending=True)
    keep = probs.cumsum(0) - probs < top_p          # the smallest set reaching p
    logits[order[~keep]] = float("-inf")

Figure 3.3 shows what the samplers are working with. After Alice looked at the, the single most likely next token is a line break, at 30%, because the text file breaks its lines at a fixed width and the model has learned where they fall; the next most likely, right, has 2.5%. At a temperature of 1, the 40 tokens that top-k keeps hold 64.5% of the probability, and top-p with $p = 0.9$ keeps 174 tokens. At 0.8 top-p keeps 77, and at 1.5 it keeps 350. That is the sense in which top-p adapts and top-k does not.

The model's next-token distribution at three temperatures, with the top-k and top-p cut-offs

Figure 3.3. The step-300 model's distribution over the next token after Alice looked at the, at temperatures 0.8, 1.0 and 1.5. Left: the probability of every token, sorted, on log scales; the vertical line, repeated on the right, is the top-k cut at 40. Right: cumulative probability; the horizontal line is $p = 0.9$, and the legend gives how many tokens top-p keeps at each temperature.

The samples themselves are in the notebook. Read them with the perplexity of 122 in mind: the local grammar is mostly right, names and quotation marks appear in the right places, and nothing holds together for more than a clause. That is what a million parameters and one book can learn. It is also, at a vastly larger scale of parameters and data, what a frontier model is doing.

Where this goes wrong

Trusting the training loss. It went down the whole time. If the notebook had not measured held-out loss, nothing would have said that the model was getting worse after step 300. Always hold data out, always measure on it, and keep the checkpoint that was best on it, not the last one.

A tokenizer trained on the wrong text. The merges come from the training text. Tokenize something from a different distribution (code, another language, text with numbers) and it falls back to bytes, the sequence gets long, and the model, which never saw long byte runs, does badly. The tokenizer is part of the model.

Decoding bytes that are not a character. Byte-level tokens can split a multi-byte character, and a sampler can emit half of one. decode passes errors="replace", which is why a stray replacement character shows up in some samples. Real systems do the same; look for it in their output.

Exercises

  1. Retrain the tokenizer with a vocabulary of 512 and of 4,096, and retrain the model on each. Compare held-out loss per token and per character. Which comparison is fair?
  2. The held-out loss stops improving before the training loss does. Find the step where the gap starts to open and explain, in terms of the corpus size, why a model this size overfits a book this short.
  3. Remove weight tying (give head its own weight matrix). Report the parameter count and the held-out loss after 1,500 steps.
  4. Implement repetition penalty: divide the logit of any token already in the last 64 tokens by 1.3 before sampling. Show a greedy sample before and after.
  5. The sampler truncates the context to the last 128 tokens. Add the chapter 2 cache so that generation of 120 tokens does not recompute the prompt each time, and check that the greedy output is unchanged.

Notebook for this chapter: 03-language-models.ipynb, 672 seconds of training and 695 seconds in total. It reads data/alice.txt from the repository and downloads nothing.

Further reading