Tokenization & Byte-Pair Encoding

Why Tokenization Exists

A neural network operates on numbers, not characters. The simplest possible mapping — one token per character — makes the vocabulary tiny but the sequences enormous, since even a short sentence becomes dozens of tokens, and the model has to learn "spelling" before it can learn anything about meaning. The opposite extreme — one token per whole word — keeps sequences short but explodes the vocabulary to the size of the language itself, and still fails the moment it meets a word it's never seen (a typo, a name, a word in another language, a brand-new term). Byte-pair encoding is the compromise the field converged on: a vocabulary of subword pieces, built so that common words end up as a single token and rare or unseen words decompose gracefully into smaller, already-known pieces.

Where BPE Comes From

The core idea, unchanged since 2016: start with a base vocabulary (individual characters, or in the byte-level variant below, individual bytes), then repeatedly find the most frequent adjacent pair of symbols across your training corpus and merge it into one new symbol. Do this enough times and you end up with a vocabulary where common patterns — whole words, common prefixes and suffixes, frequent punctuation sequences — collapse into single tokens, while everything else still decomposes into pieces the model has actually seen before.

Byte-Level BPE: GPT-2's Contribution

Training a BPE Tokenizer: Learning Merges

"Training" a BPE tokenizer doesn't involve gradient descent — it's a deterministic, count-based algorithm run once over a text corpus, producing two artifacts: a vocabulary (a mapping from integer token IDs to byte sequences) and an ordered list of merges (which pairs got combined, and in what order). The algorithm:

From-Scratch Implementation

Here's a complete, original implementation of BPE training — readable over fast, deliberately, since the point of building it yourself is understanding the algorithm, not winning a speed benchmark (that comes later, with tiktoken):

# Original from-scratch BPE trainer for this series.
# Not copied from any CS336 assignment file or OpenAI's encoder.py --
# written independently to illustrate the algorithm clearly.

from collections import Counter

def get_pair_counts(word_freqs):
    """Count frequency of every adjacent token pair across all words."""
    pair_counts = Counter()
    for word, freq in word_freqs.items():
        for i in range(len(word) - 1):
            pair_counts[(word[i], word[i + 1])] += freq
    return pair_counts

def merge_pair(word_freqs, pair):
    """Replace every occurrence of `pair` with a single merged token."""
    merged_token = pair[0] + pair[1]
    new_word_freqs = {}
    for word, freq in word_freqs.items():
        new_word = []
        i = 0
        while i < len(word):
            if i < len(word) - 1 and (word[i], word[i + 1]) == pair:
                new_word.append(merged_token)
                i += 2
            else:
                new_word.append(word[i])
                i += 1
        new_word_freqs[tuple(new_word)] = freq
    return new_word_freqs

def train_bpe(corpus_words, vocab_size, num_base_bytes=256):
    """
    corpus_words: dict mapping each word (as bytes) -> frequency in the corpus
    vocab_size:   target vocabulary size (base bytes + merges)
    Returns:      list of merges, in the order they were learned
    """
    # Represent each word as a tuple of single-byte tokens to start
    word_freqs = {
        tuple(bytes([b]) for b in word): freq
        for word, freq in corpus_words.items()
    }

    num_merges = vocab_size - num_base_bytes
    merges = []

    for _ in range(num_merges):
        pair_counts = get_pair_counts(word_freqs)
        if not pair_counts:
            break
        best_pair = max(pair_counts, key=pair_counts.get)
        word_freqs = merge_pair(word_freqs, best_pair)
        merges.append(best_pair)

    return merges

# --- try it on a tiny toy corpus ---
corpus = {
    b"low": 5, b"lower": 2, b"newest": 6, b"widest": 3,
}
learned_merges = train_bpe(corpus, vocab_size=256 + 10)
for i, (a, b) in enumerate(learned_merges):
    print(f"merge {i}: {a!r} + {b!r} -> {a+b!r}")

Encoding and Decoding

Training gives you a vocabulary and an ordered merge list. Using that tokenizer on new text is a separate, much cheaper operation: apply the learned merges, in the exact order they were learned, to any new input.

# Original encode/decode implementation, built on the vocabulary
# and merges produced by train_bpe() above.

def encode(text, merges):
    """Turn a string into a list of byte-string tokens, applying merges in order."""
    word = tuple(bytes([b]) for b in text.encode("utf-8"))

    for pair in merges:  # merges must be applied in learned order
        new_word = []
        i = 0
        while i < len(word):
            if i < len(word) - 1 and (word[i], word[i + 1]) == pair:
                new_word.append(word[i] + word[i + 1])
                i += 2
            else:
                new_word.append(word[i])
                i += 1
        word = tuple(new_word)

    return list(word)

def decode(tokens):
    """Tokens are byte strings -- concatenate and decode back to text."""
    return b"".join(tokens).decode("utf-8", errors="replace")

# --- round-trip check ---
tokens = encode("lowest", learned_merges)
print("tokens:", tokens)
print("round-trip:", decode(tokens) == "lowest")

What assignment1-basics Actually Asks

The public assignment1-basics repository's test harness (tests/adapters.py) confirms two required functions matching exactly the two halves covered above: run_train_bpe(input_path, vocab_size, special_tokens), which returns a vocabulary and merge list from a raw text corpus, and get_tokenizer(vocab, merges, special_tokens), which returns a tokenizer object implementing encode/decode using those learned merges. The README confirms the assignment trains on real corpora, not toy data — it includes direct download commands for the TinyStories dataset (a corpus of simple, short children's stories, commonly used for small-model experiments) and a Stanford-hosted subsample of OpenWebText. The specific target vocabulary size isn't stated in the README or test file — that detail lives in a PDF assignment handout not fetchable as plain text in this research pass, so this piece won't state a specific number as fact.

tiktoken and Production Tokenizers

If you want to see a real, battle-tested reference implementation of byte-level BPE encoding logic specifically (including the byte-to-unicode lookup table trick GPT-2 uses to make raw bytes printable for debugging), OpenAI's original GPT-2 repository's encoder.py is the commonly-cited canonical source — worth reading after you've built your own version above, as a way to compare your understanding against a production reference rather than as a starting point to copy from.

⚠️ Confidence Notes and Gaps

🔗 Full Reference List