BPE Training
GPT-2
Hard
Byte Pair Encoding training algorithm
Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.
Overview
Byte Pair Encoding (BPE) is a compression-style vocabulary learner: start from a small alphabet and repeatedly glue the most common adjacent pair into a new symbol. Sennrich et al. (2015) used that loop to build subword vocabularies for translation. GPT-2 adopts the same greedy merge process but changes the units and the legal pairs, because a language model that must score any Unicode string cannot afford a 130k-code-point base alphabet or a UNK token.
Section 2.2 is explicit about three design choices and thin on the training loop itself. Specified: operate on UTF-8 bytes (base size 256); do not let BPE merge across character categories, with an exception for spaces; target vocabulary 50,257. Inferred from Sennrich plus the released vocab.bpe: count pair frequencies on a pre-tokenized corpus, merge the global argmax, repeat. The official GitHub tree ships the finished encoder.json / vocab.bpe and the encode loop; it does not contain a trainer.
The space exception is the important GPT-2 twist. Naive byte-level BPE spends merges on copies of the same word with different trailing punctuation (dog., dog!, dog?). Blocking merges that cross category boundaries (letter↔︎punctuation, and so on) keeps those variants sharing a stem, and still lets a leading space attach to a word so " dog" can be one token.
How it works
Standard BPE (Sennrich et al.). Treat each training word as a sequence of symbols from a base alphabet, often characters plus an end-of-word marker. Let (c(x,y)) be the number of times symbols (x) and (y) appear adjacent in the current segmentation of the corpus. Each iteration adds one merge:
(x^\star,y^\star)=\arg\max_{(x,y)}\,c(x,y),\qquad v \leftarrow xy^\star
and rewrites every occurrence of that pair as the new symbol (v). After (M) merges the vocabulary is (\lvert\Sigma_0\rvert + M) (plus any reserved specials). Encoding later replays merges in the order they were learned, not by re-counting a new corpus.
GPT-2’s byte-level variant. The paper replaces (\Sigma_0) with the 256 UTF-8 bytes. That is enough to represent every Unicode string. Reference BPE tools that run on Unicode code points would need on the order of (10^5) base symbols before any merge, which the authors call prohibitively large relative to the 32k–64k vocabularies common at the time.
Directly merging on raw bytes, however, is a “greedy frequency heuristic” that fragments the budget. The paper’s fix: prevent BPE from merging across character categories for any byte sequence, with a space exception so a word and its preceding space can still fuse. Character category here is the usual Unicode idea (letter, number, punctuation, …) applied after bytes have been interpreted as text. The paper does not give a precise predicate or pseudocode; implementations typically refuse a pair when the two sides belong to different categories, unless one side is a space.
What 50,257 is. Table 2 / §2.3: “The vocabulary is expanded to 50,257.” That figure is the final size used by the Transformer, not the merge count alone. In the released artifacts it is 256 byte symbols plus learned merges plus at least the special <|endoftext|> token (and the byte-remap alphabet used at encode time). The paper does not itemize the 50,257 slots.
Specified vs inferred. Specified by Radford et al.: byte base, category-boundary constraint, space exception, final vocab size, and the motivation (no UNK, less waste on punctuated copies of the same word). Inferred: pair-count / argmax-merge loop, use of a pre-tokenizer before counting, and the exact category test. The encode-time regex in src/encoder.py is a consumer of the trained table; whether the trainer used that same regex is not stated.
After training, the merge list is an ordered sequence ((x_1,y_1),\ldots,(x_M,y_M)). Rank (i) is more binding than rank (j>i) at encode time. That ranking is what Encoder.bpe consults; frequency on a new string is irrelevant.
Official code
There is no trainer in openai/gpt-2. src/encoder.py only applies a finished table: get_encoder reads vocab.bpe (skipping the header line) into bpe_merges and zips them to ranks. To study the product of training, read that file and encoder.json next to a released checkpoint. Do not invent a train_bpe.py path in this repository.
Watch-outs
- Counting pairs on raw characters (Unicode scalars) is not GPT-2 BPE. You will either explode the base vocab or introduce UNKs the paper rejected.
- Forgetting the category lock lets punctuation glue onto words and burns merge slots on
dog./dog!variants — the failure mode §2.2 calls out. - Forgetting the space exception over-fragments: common words split from their leading space and compression suffers.
- The merge list is ordered. Shuffling
vocab.bpeor sorting by a new corpus’s frequencies trains a different tokenizer than the one the 1.5B weights expect.
Sources
- Paper: Language Models are Unsupervised Multitask Learners (Radford et al., 2019), §2.2–2.3
- Training algorithm: Sennrich, Haddow, and Birch, “Neural Machine Translation of Rare Words with Subword Units,” 2015
- Code (encode-side only): openai/gpt-2
src/encoder.py