Step 7 — Tokenizer and Sampler

Released with Lecture 7 · ~3 hours · Starts from: Step 6 (or solutions/step-06.ipynb).

In what follows, there are blocks of code provided for you. It’s important that you can read the code. If you are not proficient with Python, then I propose the following workflow. For each block of code, keywords are provided for the python syntax which is being used, so you can look up python syntax and functions you don’t recognise. To help, beneath many code blocks you will find Python background foldouts: self-contained tutorials, with live Python boxes you can edit and run right on the page, for syntax that is new in this step (syntax already introduced is covered in earlier steps’ foldouts). Then I suggest that you first experiment with changing the code or printing out intermediate variables, until you have figured out what you think it does. Then point AI toward the page and step number, provide your precise explanation of what you believe each line of the code is doing, and ask for corrections. In this way you will quickly become a code reader (if not a code writer). — Kate

Goal

Replace the two “toy” ends of your pipeline with the real thing: a byte-pair-encoding tokenizer you train yourself (entering the model), and a proper sampler with temperature, top- $k$ , and nucleus decoding (leaving it). Plus: pool the class’s Step 6 experiments into a homemade scaling plot.

On paper first

BPE as greedy compression. The algorithm: start with the byte/char alphabet; repeatedly find the most frequent adjacent pair of tokens in the corpus and merge it into a new token; stop after $M$ merges. Argue each merge is the greedy step reducing corpus length (in tokens) the most among all single merges. What quantity per merge would a Huffman-style optimal scheme control instead, and why is BPE’s greediness fine in practice?
Tokenization changes the units of perplexity. Your model reports loss in nats per token. Show that the honest, tokenizer-independent metric is bits per character: $\text{bpc} = \dfrac{(\text{nats/token})}{\ln 2} \cdot \dfrac{\#\text{tokens}}{\#\text{chars}}$ , and explain the fallacy of comparing per-token perplexities across different vocabularies (a tokenizer that halves the token count must roughly double the per-token difficulty just to break even).
Sampling as tilting. Temperature $\tau$ replaces $p_i \propto e^{z_i}$ with $p_i \propto e^{z_i/\tau}$ . Identify the limits $\tau \to 0$ and $\tau \to \infty$ , and show sampling at $\tau$ is sampling from the distribution minimizing $\mathbb E[\text{-logit}] + \tau^{-1}\cdot(\text{entropy penalty})$ … i.e. the Gibbs variational principle, one more time.

Tasks

7.1 BPE trainer

On the raw Shakespeare text, with the 65-character vocabulary as the base alphabet, implement:

def get_pair_counts(ids):            # ids: list of ints
    ...                              # Ellipsis: you write this — count adjacent pairs

def merge(ids, pair, new_id):        # replace every occurrence of `pair` by `new_id`
    ...

merges = {}                          # empty dictionary; (id, id) -> new_id, in training order
ids = encode(text)                   # your Step 1 char-level encode
for m in range(NUM_MERGES):          # start with NUM_MERGES = 256
    pair = max(get_pair_counts(ids), key=...)   # max with a key= function
    merges[pair] = V + m             # dictionary assignment; the key is a tuple
    ids = merge(ids, pair, V + m)

Python background: max with key=

max(items) compares the items themselves; max(items, key=f) compares f(item) instead and returns the item (not the score) that scores highest. Iterating over a dictionary yields its keys, so with a count dictionary, max(counts, key=counts.get) means “the key with the largest count” — precisely the most-frequent-pair step of BPE. (min and sorted accept key= the same way.)

words = ['to', 'be', 'or', 'not']
print(max(words))                    # alphabetical winner
print(max(words, key=len))           # the longest word wins instead
counts = {'th': 8, 'he': 5, 'e ': 11}
print(max(counts))                   # max over KEYS: alphabetical again
print(max(counts, key=counts.get))   # the key with the biggest count

Python background: tuples as dictionary keys

Dictionary keys can be any immutable value — strings, numbers, and tuples all work (lists do not: they can change, so they cannot be looked up reliably). A tuple key like ('t', 'h') is exactly what the merge table needs: the pair itself indexes its new token id.

merges = {}                    # an empty dictionary
merges[('t', 'h')] = 65        # a tuple as the key
merges[('h', 'e')] = 66
print(merges)
print(merges[('t', 'h')])
pair = ('h', 'e')
print(pair in merges)          # membership test works on keys

Log the first 30 merges with their string forms. You’ll watch English assemble itself out of pure frequency counting — the first ten on Shakespeare are

 0 'e '   1 'th'   2 't '   3 's '   4 'd '
 5 ', '   6 'ou'   7 'er'   8 'in'   9 'y '

then an, :\n, or, o , en, \n\n, … Three things to notice and write a sentence about: (i) word-final patterns like e , t , s dominate the very top — the space is the most predictable character in English; (ii) :\n and \n\n are the model discovering the play format (SPEAKER: then newline, blank line between speeches) with no notion of what a play is; (iii) by merge 28 you get whole words (and ), i.e. BPE crosses from morphology into vocabulary on its own.

Report the compression ratio chars/token at $M \in \{64, 256, 512, 1024\}$ merges.

7.2 Encoder / decoder

bpe_decode: each token id → its string; concatenate. (Store the vocabulary as id → string as you build merges.)
bpe_encode: start from characters, apply the merges in training order (this order-dependence matters — construct a short input where applying available merges in the wrong order gives a different tokenization).
Round-trip test on text the tokenizer never saw (grab a sonnet not in input.txt): bpe_decode(bpe_encode(s)) == s, exactly.

7.3 Retrain on BPE tokens

Retokenize the corpus with $M = 256$ merges ( $V' = 321$ , 568,210 tokens), and retrain your Step 6 model (same config otherwise; the embedding and unembedding grow to $V'$ ). Two things to confront:

Per-token loss will be higher than Step 6’s. Convert both runs to bits/char (paper 2) for the honest comparison — did BPE help at equal training steps? (Each BPE token also sees more text in the same T_max window — effective context in characters roughly doubles. Discuss in two sentences.)
Samples: same generate loop, new decoder. Word-level errors should drop; new failure mode appears at token boundaries. Find an example.

7.4 The sampler

One function, replacing plain multinomial sampling:

def sample_next(logits, temperature=1.0, top_k=None, top_p=None):  # None default = optional argument
    ...

Python background: None as an optional-argument default

None is Python’s “no value here” object. Defaulting an argument to None is the standard way to mark it optional: inside the function you test if top_k is not None: and apply the feature only when the caller actually supplied a value. (Use is None / is not None for this test — None is a unique object, and identity is the idiomatic check.)

def describe(name, nickname=None):     # nickname is optional
    if nickname is None:
        return name
    return name + ', known as ' + nickname
print(describe('Katherine'))
print(describe('Katherine', nickname='Kate'))
x = None
print(x is None)
print(x == 0, x == '')     # None is not 0, not empty string — it's its own thing

Temperature scales logits before softmax.
Top- $k$ keeps the $k$ largest logits, sets the rest to $-\infty$ , renormalizes.
Nucleus (top- $p$ ): keep the smallest set of tokens with cumulative probability ≥ $p$ (sort, cumsum, cut, renormalize).

Then a decoding gallery from your trained model — same prompt ("ROMEO:"), generations at: $\tau \in \{0.3, 0.8, 1.0, 1.5\}$ ; greedy ( $\tau \to 0$ ); top- $k{=}5$ ; top- $p{=}0.9$ . Annotate: where does greedy fall into a loop? (It will — Holtzman et al.’s degeneration, on your own model.) Where does $\tau{=}1.5$ lose English? Which setting reads best, and what does that say about the mismatch between “highest likelihood” and “best text”?

7.5 Homemade scaling law

Pool the class’s Step 6.5 tables (or run your own sweep: $d \in \{32, 64, 128, 256\}$ at $L=4$ ). Plot best val loss vs parameter count on log–log axes and fit a line through the non-saturated points. Report your exponent; Kaplan et al. got $\approx 0.076$ for loss vs params (very different data and scale — the point is the straightness, not the number). One paragraph: where must your line break down, at each end, on 1MB of Shakespeare? (Irreducible entropy on one end; the Lecture 7 data-bottleneck discussion on the other.)

Checkpoints

Round-trip exact on unseen text. Measured compression on the 1,115,394 characters of Shakespeare:

merges $M$ vocab tokens chars/token
64 129 743,756 1.500
256 321 568,210 1.963
512 577 487,961 2.286
1024 1089 414,322 2.692

So $M{=}256$ gives ≈ 1.96 chars/token — note the returns are strongly diminishing in $M$ (quadrupling merges from 256 to 1024 buys only 37% more compression), which is worth one sentence: what does that say about where the entropy of English actually lives?
First merges are recognizably English digraphs and function words, per the list above.
BPE model within ~0.1 bits/char of the char model at equal steps (either direction — the interesting part is your unit conversion and discussion, not the winner at this tiny scale).
Decoding gallery produced, with the greedy-loop specimen captured.
A log–log plot with a credible straight segment and a fitted slope.

merges $M$	vocab	tokens	chars/token
64	129	743,756	1.500
256	321	568,210	1.963
512	577	487,961	2.286
1024	1089	414,322	2.692

If you’re stuck

BPE training slow: recount pairs from the current ids each merge — but merge should be one pass, and 256 merges on 1.1M chars runs in ~2–4 minutes of pure Python (1024 merges took ~8 minutes on our reference machine). Past that, profile before optimizing.
Round-trip fails on unseen text only: your bpe_encode applies merges by corpus frequency instead of training order, or your base alphabet can’t cover a character (fine for Shakespeare-only; the real fix — 256 raw bytes as the alphabet — is the “going further”).
Nucleus keeps the whole vocabulary: remember to sort descending before the cumulative sum, and to include the token that crosses $p$ .

Going further (optional)

Byte-level BPE (GPT-2’s choice): base alphabet = 256 bytes, so any Unicode string round-trips with zero out-of-vocabulary cases. Retrain on the same corpus; check round-trip on “Zeilenumbrüche 改行 🎭”.
Paste your favourite sonnet into tiktokenizer.vercel.app and compare GPT-4’s segmentation of it with your tokenizer’s.
Beam search: implement width- $B$ beam decoding and exhibit the classic result that pure likelihood-maximization gives worse text than sampling — tying back to your paper 3 discussion.

Catch-up

solutions/step-06.ipynb gives you a trained checkpoint (load with torch.load); 7.1–7.2 and 7.4 need no training at all, so this step is another good re-entry point.