Tokenization and pretraining

1. Intuition

Tokens. Word-level vocabularies can't handle new words, typos, or code; character-level sequences get very long. Subword tokenization is the compromise: common words become one token, rare words split into pieces ("unbelievable" might be "un", "believ", "able"). Byte-level BPE starts from raw bytes, so any string, in any language, can be encoded.

The pipeline. Text → tokens → integer IDs → embedding vectors → transformer blocks → a vector per position → logits over the vocabulary → softmax probabilities for the next token.

The context window is the maximum number of tokens the model can attend to at once (prompt plus generated output). Anything outside it does not exist for the model.

Pretraining is self-supervised: every position in every document is a training example whose label is the next token. The loss is cross-entropy, the negative log of the probability assigned to the true next token. The result is a base model that continues text fluently; it is not yet an assistant that follows instructions (that comes from adaptation).

2. Visualization

the model learns to predict the next token from the previous tokens
training the model on more data makes the predictions better
a transformer model uses attention to weigh the tokens in its context
the learning rate controls how fast the model learns during training
tokenization splits text into tokens and the model predicts tokens
the model is trained to predict the next token in the training data

Try it (predict first, then check):

1. Start at 0 merges and predict how many tokens your sentence has. Then step merges up and watch which pairs join first.
2. Type a word that never appears in the training corpus. Predict how it will be split, then check.
3. Press Fill the context and shrink the window. What has to give when text no longer fits?

3. The math

Symbols

Autoregressive objective

Thanks to the causal mask, one forward pass over a sequence computes the loss for all $T$ positions in parallel.

Perplexity

Perplexity is the effective number of tokens the model is "choosing between." A model guessing uniformly over 50,000 tokens has loss $\ln 50000 = 10.82$ and perplexity 50,000.

Byte-pair encoding

Start with characters (or bytes). Repeatedly count every adjacent pair across the corpus and merge the most frequent pair into a new token, until the vocabulary reaches the target size.

Worked example 1: BPE merges

Toy corpus word counts: "low" ×5, "lower" ×2, "newest" ×6, "widest" ×3 (with an end-of-word marker </w>).

1. Most frequent pair: e s (6 in "newest" + 3 in "widest" = 9). Merge into es.
2. Next: es t (9). Merge into est.
3. Next: est </w> (9). Merge into est</w>.
4. Next: l o (5 + 2 = 7). Merge into lo.
5. Next: lo w (7). Merge into low.

After five merges, "newest" is n e w est</w> and "low" is low </w>. Frequent fragments become single tokens.

Worked example 2: cross-entropy and perplexity

A model assigns these probabilities to the actual next tokens in a 4-token sequence: $0.6, 0.25, 0.9, 0.05$.

Mean loss $= (0.511 + 1.386 + 0.105 + 2.996)/4 = 1.250$. Perplexity $= e^{1.250} = 3.49$. One surprising token (probability 0.05) contributes more than half the loss.

4. Implementation

The full script is code/15-llms/bpe_pretrain_scratch.py (stdlib and NumPy only, runs in about a second on a laptop CPU). It builds the whole path from text to a trained model: a byte-level BPE tokenizer, a sliding-window dataset, and a pretraining loop with hand-written gradients.

Build it from scratch

Step 1: split into chunks, count pairs. Real tokenizers first cut text into words and whitespace runs so a merge never glues the end of one word to the start of the next. Each distinct chunk is stored once with its count, so counting pairs costs one pass over the distinct chunks, not the whole text.

def pretokenize(text):
    return re.findall(r"\S+|\s+", text)          # merges never cross a word/space boundary

def count_pairs(units):                           # units: {tuple of ids: count}
    pairs = Counter()
    for sym, c in units.items():
        for a, b in zip(sym, sym[1:]):
            pairs[(a, b)] += c
    return pairs

Predict: in the toy corpus, "newest" appears 6 times. How much does it add to the count of the pair (e, s)?

Step 2: the training loop. Start from the 256 byte values, so any string in any language is representable and there is no unknown token. Each round merges the most frequent pair into a new ID. Ties go to the smallest pair, so two runs on the same text learn the same merges.

def train_bpe(text, vocab_size):
    vocab = {i: bytes([i]) for i in range(256)}   # ids 0..255 are raw bytes: nothing is ever unknown
    merges = {}                                   # (a, b) -> new id, in the order learned
    units = Counter(tuple(w.encode("utf-8")) for w in pretokenize(text))
    while len(vocab) < vocab_size:
        pairs = count_pairs(units)
        if not pairs:
            break
        best = min(pairs, key=lambda p: (-pairs[p], p))   # most frequent; ties -> smallest pair
        if pairs[best] < 2:
            break
        new_id = len(vocab)
        vocab[new_id] = vocab[best[0]] + vocab[best[1]]
        merges[best] = new_id
        units = Counter({apply_merge(s, best, new_id): c for s, c in units.items()})
    return vocab, merges

apply_merge (in the script) walks one chunk left to right and replaces each (a, b) with the new ID. On the toy corpus this learns es, est, lo, low, then ew. That differs from worked example 1 at the third merge: byte-level BPE has no </w> end-of-word marker, so est </w> is not a pair; the space after a word is its own chunk instead.

Predict: why stop when the best pair has count 1?

Step 3: encode by replaying merges in learned order. Encoding a new string does not count anything. It repeatedly finds the adjacent pair with the earliest merge rank and applies it, until no learned merge applies. Decoding is just concatenating each ID's bytes.

def encode(text, merges):
    rank = {p: r for r, p in enumerate(merges)}   # earlier merge = higher priority
    ids = []
    for w in pretokenize(text):
        sym = tuple(w.encode("utf-8"))
        while len(sym) > 1:
            pair = min(zip(sym, sym[1:]), key=lambda p: rank.get(p, math.inf))
            if pair not in rank:
                break                             # no learned merge applies any more
            sym = apply_merge(sym, pair, merges[pair])
        ids.extend(sym)
    return ids

Predict: "café" contains é, which is two bytes in UTF-8. If the training text never had an é, how many tokens does "é" become?

Step 4: turn one long ID stream into training examples. Pretraining data is a single stream of token IDs. Cut it into windows of $T + 1$ IDs; the input is the first $T$, the target is the same window shifted left by one. Every position in a window is a training example, which is why one forward pass gives $T$ losses.

def windows(ids, T, stride):
    ids = np.asarray(ids)
    starts = range(0, len(ids) - T, stride)       # each window needs T + 1 ids
    X = np.stack([ids[s:s + T] for s in starts])
    Y = np.stack([ids[s + 1:s + T + 1] for s in starts])
    return X, Y                                   # both [n_windows, T]

Predict: with stride smaller than $T$, windows overlap. What goes wrong if you split train and validation after making overlapping windows?

Step 5: the smallest possible language model. A real LLM puts transformer blocks between the embedding and the output projection. To keep the backward pass short, this model skips them: it looks up the current token's embedding and maps it straight to logits, so it can only use one token of context (a neural bigram model). The loss is the same cross-entropy as section 3.

def forward(params, X, Y):
    h = params["E"][X]                            # [B, T, d]   token embeddings
    logits = h @ params["W"]                      # [B, T, V]   scores for the next token
    logits -= logits.max(-1, keepdims=True)
    p = np.exp(logits); p /= p.sum(-1, keepdims=True)
    loss = -np.log(np.take_along_axis(p, Y[..., None], -1)).mean()
    return loss, h, p

Predict: before any training the weights are small random numbers. What loss should step 0 print, roughly? (Hint: what is the loss of a uniform guess over $V$ tokens?)

Step 6: backward pass and update. The gradient of softmax plus cross-entropy with respect to the logits is $p - \text{onehot}(y)$, divided by the number of positions because the loss is a mean. The rest is the chain rule through two matrix products. np.add.at matters: if a token appears twice in a batch, its embedding row must receive both gradients.

def train_step(params, X, Y, lr):
    loss, h, p = forward(params, X, Y)
    d_logits = p.copy()
    d_logits[np.arange(Y.shape[0])[:, None], np.arange(Y.shape[1]), Y] -= 1   # softmax + CE gradient: p - onehot
    d_logits /= Y.size                                                          # mean over all positions
    dW = h.reshape(-1, h.shape[-1]).T @ d_logits.reshape(-1, d_logits.shape[-1])
    dE = np.zeros_like(params["E"])
    np.add.at(dE, X, d_logits @ params["W"].T)    # scatter-add: a token used twice gets both grads
    params["W"] -= lr * dW
    params["E"] -= lr * dE
    return loss

Predict: if you wrote dE[X] += ... instead of np.add.at, which tokens would get too little gradient?

The main block samples random batches of 16 windows, takes 300 steps, and prints validation loss next to two baselines: a uniform guess and a unigram model (token frequencies only).

Use the library

In practice you load a trained tokenizer and write the loop in PyTorch, where autograd replaces step 6:

import tiktoken                                      # pip install tiktoken
enc = tiktoken.get_encoding("cl100k_base")
ids = enc.encode("Tokenization splits text into subword units; unbelievable!")
print(len(ids), [enc.decode([i]) for i in ids])

import torch.nn.functional as F
# logits: [B, T, V] from the model; tokens: [B, T+1] one window per row
logits = model(tokens[:, :-1])                       # inputs: positions 0..T-1
targets = tokens[:, 1:]                              # targets: the same window shifted by one
loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
loss.backward(); optimizer.step(); optimizer.zero_grad()
perplexity = loss.exp()

What the libraries add: tiktoken and Hugging Face tokenizers run the same merge replay in Rust with a smarter pre-tokenization regex (splitting digits, punctuation, and contractions), and handle special tokens like <|endoftext|>. A real training loop adds AdamW, a learning rate that warms up and then decays on a cosine curve, gradient clipping, mixed precision, and checkpoints. The objective does not change.

Run it

uv run code/15-llms/bpe_pretrain_scratch.py

Key lines of output (synthetic training paragraph):

first 5 merges: ['es', 'est', 'lo', 'low', 'ew']
tokens: ['the', ' ', 'token', 'i', 'z', 'er', ' ', 'pr', 'ed', 'i', 'c', 't', 's', ' ', 'r', 'ar', 'e', ' ', 'tokens', '.', ' ', 'n', 'a', '<c3>', '<af>', 'v', 'e', ' ', 'c', 'a', 'f', '<c3>', '<a9>', '!']
49 bytes -> 34 tokens; round trip ok: True
423 tokens -> X (52, 8), Y (52, 8)
baselines on validation: uniform ln V = 5.756, unigram = 3.771
step   0  train loss 5.764  val loss 5.737  val perplexity  310.2
step 200  train loss 2.012  val loss 3.747  val perplexity   42.4
step 300  train loss 1.613  val loss 3.942  val perplexity   51.5
greedy continuation of 'the': ' the the the the'

Read it as a pretraining run in miniature. Step 0 sits at $\ln V$, a uniform guess. Words seen in training ("token", "tokens") become single tokens, while "é", never seen, falls back to its two raw bytes. Training loss keeps falling, but validation loss bottoms out near step 200, barely below the unigram baseline, and then rises: with 423 tokens the model memorizes rather than generalizes, which is why real pretraining uses trillions of tokens and watches held-out loss. The greedy continuation loops because a bigram model only sees the last token; decoding covers why greedy decoding repeats even in large models.

5. Engineering

Token economics. Cost and latency scale with tokens. The same content can take very different token counts across languages and tokenizers (non-English text and code often take more). Measure with the actual tokenizer of the model you deploy.

Context management. Long prompts cost more and attention over long contexts is expensive; models can also attend poorly to material buried in the middle of very long contexts. Retrieve and include only what is relevant (RAG).

BPE is one of three tokenizer algorithms. WordPiece (BERT) also merges pairs, but picks the pair that most raises corpus likelihood, roughly count(ab) / (count(a) · count(b)), instead of the raw count. Unigram (T5, many SentencePiece models) works the other way: start with a large candidate vocabulary and prune the tokens whose removal hurts corpus likelihood least. SentencePiece is a library that trains BPE or Unigram on raw text and writes spaces as ▁; tiktoken only encodes with a pre-built vocabulary. For an LLM you will almost always reuse the base model's tokenizer; you train one only when you pretrain.

Pretraining data quality (deduplication, filtering, mixing) matters as much as quantity. Training saves checkpoints: snapshots of the weights you can resume from or evaluate.

What pretraining does not give you: instruction following, refusal behavior, or calibrated truthfulness. The model learned to produce likely text, which is not the same as true text.

Key distinction

An LLM produces a probability distribution over the next token. It is not a database of facts. High probability means "this is the kind of thing that tends to come next," not "this is correct."

Watch out

Failure modes. Tokenization quirks (arithmetic on numbers split into odd chunks, trailing-space differences, words split differently with and without capitalization); prompts silently truncated at the context limit; comparing perplexities across models with different tokenizers (not comparable).

Common mistakes

• Estimating cost from characters or words instead of tokens.
• Forgetting to shift targets by one position when computing the language-modeling loss.
• Treating a base model's fluent continuation as an answer to a question.

6. Knowledge check

Practice exercise. A model assigns probabilities 0.5, 0.5, 0.125 to three actual next tokens. Compute the mean loss and perplexity.

Implementation challenge. The model in section 4 sees one token of context. Give it two: concatenate the embeddings of the previous two tokens (h becomes [B, T, 2d], W becomes [2d, V]), update the backward pass, and compare validation loss with the bigram model on the same windows. Does more context help on 423 tokens of training text?

Summary

• Subword tokenization (BPE) balances vocabulary size and sequence length; byte-level BPE can encode any text.
• Tokens become IDs, then embeddings, then logits over the vocabulary; the context window caps how many tokens the model sees.
• Pretraining minimizes next-token cross-entropy over massive text; perplexity $= e^{\text{loss}}$.
• A pretrained base model predicts likely text, not guaranteed truth.

Next: Decoding

Related: The transformer architecture · Learning paradigms · Model card: large language model

Go deeper, code-first (AI Engineering from Scratch): Text Generation Before Transformers — N-gram Language Models · Multilingual NLP · GPT — Causal Language Modeling · Tokenizers: BPE, WordPiece, SentencePiece · Building a Tokenizer from Scratch · Data Pipelines for Pre-Training · Pre-Training a Mini GPT (124M Parameters)

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): BPE Tokenizer From Scratch, Tokenized Dataset with Sliding Window, Training Loop and Evaluation, Subword Tokenization — BPE, WordPiece, Unigram, SentencePiece.

Interview angle