Chapter 2 of 16

Data and tokenisation

A public-domain corpus, BPE merges learnt live, windows and the validation split.

Slides for this chapter: Transformer Training 02, Data and Tokenisation

A language model learns from text, turned into numbers. The Decoder's model reads characters; most real models read tokens, frequent fragments learnt from the text itself by byte-pair encoding (BPE). Below, BPE learns its first merges on this part's training text, the whole of Alice's Adventures in Wonderland, in your browser.

Loading the animation…

Concept

The corpus. These chapters train on Lewis Carroll's Alice's Adventures in Wonderland (1865, Project Gutenberg eBook #11, public domain), reduced to the Decoder's 64-character alphabet: lower case, straight quotes, no italics marks. That leaves 142,041 characters. The last 10% (14,204 characters, the end of the book) is held out as validation text; the model trains on the first 127,837. A held-out split matters because a model can lower its training loss by memorising; only text it has never seen shows whether it learnt something that carries over.

Windows and teacher forcing. A training example is a window of $S + 1 = 9$ consecutive characters; $S = 8$ is the model's context length, how many tokens it can see at once (real models: thousands). The first 8 are the input; the target at each position is the character after it. So one window is 8 predictions at once, and the model always predicts from the true text so far, never from its own earlier guesses (teacher forcing). The causal mask from the Decoder's attention chapter is what makes this honest: position $i$ cannot look at the character it is predicting. Each step draws 16 windows (sequences) at random places in the training text, so the batch is 16 sequences of 8 tokens: 128 tokens, and 128 predictions, per step.

Packing documents. Real corpora are many documents of different lengths. Rather than pad short ones, trainers concatenate them into one long stream with a separator token and cut fixed-length windows, so no compute is wasted on padding. A window can then straddle two documents, and some runs mask attention across the boundary so a document never attends to the one before it; Llama 3's report describes doing this, and finds it matters most for very long sequences (Grattafiori et al. 2024, §3.2). With a single book, there is no boundary here.

Why tokens rather than characters. With characters the model must spend its context and its layers assembling "the" from three pieces every time. BPE (Sennrich, Haddow & Birch 2015) starts from single symbols and repeatedly merges the most frequent adjacent pair into a new symbol. On this book, 64 merges (a vocabulary of 128) turn 142,041 characters into 85,614 tokens, 1.66 characters per token, so the same 8-token context sees about two thirds more text. GPT-2 runs the same algorithm over bytes rather than characters, so any text can be encoded, with a vocabulary of 50,257 (Radford et al. 2019). The first merges are the language's commonest fragments: " t", "he", " a", " s", "in", and by the sixth, " the" (a space and the whole word). These chapters keep training on characters, so that the model is exactly the Decoder's; the tokeniser is here to show how real models see text.

Data at scale. Small experiments often use small, simple corpora such as TinyStories, short stories written by language models in a child's vocabulary, on which models of a few million parameters already write fluent English (Eldan & Li 2023). Real pretraining corpora are web crawls plus books, code and papers, measured in trillions of tokens. Most of the work is deciding what to keep: removing near-duplicates (duplicated text is memorised and wastes compute; Lee et al. 2021), filtering for quality, and choosing the mixture of sources (FineWeb, Penedo et al. 2024, documents one such pipeline in detail). How much data a model of a given size should see is a scaling question (chapter 14): the Chinchilla analysis found about 20 tokens per parameter compute-optimal (Hoffmann et al. 2022). This toy has 23.8 training characters per parameter, by coincidence in the same range.

Maths

For a window $w_1 \dots w_{S+1}$ the inputs are $x_i = w_i$ and the targets $y_i = w_{i+1}$ , $i = 1 \dots S$ ; the model is trained to maximise

\log p_\theta(w_2, \dots, w_{S+1} \mid w_1) = \sum_{i=1}^{S} \log p_\theta(w_{i+1} \mid w_1, \dots, w_i),

the chain rule of probability, which is why one window gives $S$ training signals. A window starts at $\lfloor u\,(N - S - 1) \rfloor$ for $u$ uniform in $[0, 1)$ from the seeded generator, $N$ the number of training characters.

BPE: with the corpus split into words $w$ with counts $n_w$ , each word a sequence of symbols, the count of a pair is $\operatorname{count}(a\,b) = \sum_w n_w \cdot \#\{i : w_i = a, w_{i+1} = b\}$ . Each merge removes exactly $\operatorname{count}(a^\ast b^\ast)$ symbols from the corpus (fewer if the pair overlaps itself, as in "a a a"). Ties go to the smaller pair of ids, so the Python reference and the browser learn identical merges.

Code

The tokeniser's training loop, src/lib/training/bpe.ts (checked merge for merge against scripts/training/data.py):

for (let step = 0; step < nMerges; step++) {
  const ranked = [...pairCounts(words).values()].sort(byRank);
  if (ranked.length === 0) break;
  const best = ranked[0]!;
  const [a, b] = best.pair;
  const id = vocab.length;
  vocab.push(vocab[a]! + vocab[b]!);
  merges.push([a, b]);
  const merged = new Map<string, Word>();
  for (const w of words) {
    const ids = mergeWord(w.ids, a, b, id);
    const k = key(ids);
    const e = merged.get(k);
    if (e) e.count += w.count;
    else merged.set(k, { ids, count: w.count });
  }
  words = [...merged.values()];
  // … (records the history the animation shows)
}

And a batch, src/lib/training/data.ts:

for (let b = 0; b < B; b++) {
  const start = Math.floor(rng() * (train.length - S - 1));
  const w = train.slice(start, start + S + 1);
  starts.push(start);
  inputs.push(w.slice(0, S));
  targets.push(w.slice(1));
}

Comments

Be the first to leave a comment on this section.