Tokenization Subword Segmentation Bpe Wordpiece Sentencepiece and Unigram Lm
# Tokenization & Subword Segmentation: BPE, WordPiece, SentencePiece, and Unigram LM
## Introduction & Motivation
Every large language model begins its processing pipeline with tokenization: the conversion of raw text into a sequence of discrete integer identifiers drawn from a fixed vocabulary. This step is so foundational that it is easy to overlook, yet it fundamentally shapes what a model can and cannot represent efficiently. A poorly designed tokenizer inflates sequence lengths, wastes context window capacity, degrades performance on rare languages or domains, and can introduce subtle failure modes such as an inability to reliably count letters within a word or perform character-level arithmetic. Subword tokenization emerged as the dominant solution to a hard tradeoff: word-level vocabularies are enormous and cannot cover every inflection, misspelling, or neologism, while character-level or byte-level vocabularies are compact but produce very long sequences that are expensive to model and harder for a Transformer's fixed-size attention window to exploit statistically.
The subword approach threads this needle by learning a vocabulary of frequently occurring character sequences directly from a training corpus, ranging from single characters up to whole common words. Frequent words like "the" or "running" are represented as single tokens, while rare or novel words are decomposed into meaningful or at least statistically justified fragments, such as "un", "believ", "able" for "unbelievable". This gives the model graceful degradation: it never encounters an unknown-token wall, because any string can be expressed as some sequence of subword units, down to the level of individual bytes or characters if necessary.
Three algorithmic families dominate practice: Byte-Pair Encoding (BPE), used by GPT-2, GPT-3, GPT-4, and LLaMA; WordPiece, used by BERT and its descendants; and the Unigram Language Model tokenizer, popularized by SentencePiece and used in T5, ALBERT, and XLNet. Although all three produce a fixed subword vocabulary and a deterministic (or near-deterministic) segmentation procedure, they differ substantially in how the vocabulary is learned and, consequently, in the linguistic and statistical properties of the segmentations they produce. Understanding these differences is essential for practitioners choosing or building a tokenizer, diagnosing model behavior on rare or multilingual text, and reasoning about the true "cost" of a piece of text in tokens, which directly determines API billing, context budget, and inference latency.
## Core Concepts & Theory
At the heart of subword tokenization is a simple observation: character frequency statistics in natural language are highly non-uniform, and so are the frequencies of character sequences. If a training corpus is scanned for adjacent symbol pairs, a small number of pairs (such as "t" followed by "h", or "in" followed by "g") occur vastly more often than the rest. Byte-Pair Encoding, adapted from a 1994 data-compression algorithm, exploits this directly: starting from a vocabulary of individual characters (or bytes), it repeatedly finds the single most frequent adjacent pair of symbols in the corpus and merges them into a new symbol, adding that merged unit to the vocabulary. This process repeats for a fixed number of merge operations, chosen in advance to hit a target vocabulary size (typically 32,000 to 128,000 for modern LLMs). The result is an ordered list of merge rules; at inference time, encoding a new string means applying the learned merges in the same order they were learned, repeatedly merging the highest-priority applicable pair until no more merges apply.
WordPiece, introduced by Google for its speech recognition and later NLP systems, uses the same iterative merge-based construction but changes the merge selection criterion. Rather than choosing the pair with the highest raw co-occurrence count, WordPiece chooses the pair that maximizes the increase in the training corpus's likelihood under a unigram language model built from the current vocabulary, which in the common simplified implementation reduces to selecting the pair $(a, b)$ that maximizes the pointwise mutual-information-like score $ ext{count}(a,b) / ( ext{count}(a) \cdot ext{count}(b))$. This subtly changes the character of the resulting vocabulary: WordPiece tends to prefer merging symbols that co-occur *almost exclusively* with each other, even if their absolute frequency is modest, over symbols that are individually extremely common (like "e" or "t" in English) but whose particular adjacent pairing is not especially distinctive. The two algorithms can and do diverge on which merges they perform first, particularly in corpora containing both very frequent generic characters and rarer characters that appear in tight, predictable pairs.
The Unigram Language Model tokenizer, in contrast to both BPE and WordPiece, is fundamentally a *pruning* algorithm rather than a merge-building one. It begins with a large seed vocabulary, typically all substrings up to some maximum length that occur in the training corpus, each assigned an initial probability. It then alternates between an expectation step, in which the current vocabulary's unigram probabilities are used to find the most probable (Viterbi) segmentation of every training word, and a pruning step, in which the least useful subword units, those whose removal would most modestly harm the corpus's overall segmentation likelihood, are dropped, shrinking the vocabulary toward a target size. Because segmentation under the Unigram LM is probabilistic, and because the loss of each candidate token's removal can be estimated, this method naturally supports sampling multiple valid segmentations of the same string, a property exploited by subword regularization for tokenization-level data augmentation. SentencePiece is not itself a distinct algorithm; it is a language-agnostic implementation framework that supports both BPE and Unigram LM as a back end, and which crucially treats input as a raw stream of Unicode characters (or bytes) with no assumption of whitespace-delimited words, making it well suited to languages such as Japanese, Thai, or Chinese, where word boundaries are not marked by spaces.
## Mathematical Formulation
For Byte-Pair Encoding, given a corpus represented as a multiset of words with frequencies $f_w$, and each word represented as a sequence of symbols, the pair selection rule at each merge step is
$$ (a^*, b^*) = \arg\max_{(a,b)} \sum_{w \in ext{corpus}} f_w \cdot ext{count}_w(a, b) $$
where $ ext{count}_w(a,b)$ is the number of adjacent occurrences of symbols $a$ immediately followed by $b$ within word $w$'s current symbol sequence. After selecting $(a^*, b^*)$, every occurrence of that adjacent pair across the corpus is replaced by the merged symbol $a^*b^*$, and the vocabulary is extended by one entry.
WordPiece replaces the raw-count objective with a normalized score,
$$ ext{score}(a, b) = \frac{ ext{count}(a, b)}{ ext{count}(a) \cdot ext{count}(b)} $$
which approximates the mutual information between adjacent symbols $a$ and $b$ under an independence assumption, favoring statistically tight co-occurrences over merely frequent individual symbols.
The Unigram LM tokenizer defines a probabilistic model over segmentations. Given a vocabulary $V$ with per-token probabilities $p(x)$ summing to one, the probability of a particular segmentation $\mathbf{x} = (x_1, \ldots, x_k)$ of a string $s$ is
$$ P(\mathbf{x}) = \prod_{i=1}^{k} p(x_i) $$
and the most likely segmentation is found via Viterbi decoding,
$$ \mathbf{x}^* = \arg\max_{\mathbf{x} \in ext{Seg}(s)} \sum_{i=1}^{k} \log p(x_i) $$
where $ ext{Seg}(s)$ is the set of all ways to partition $s$ into vocabulary tokens. Training alternates an EM-style probability re-estimation, $p(x) \propto \sum_{w} f_w \cdot \mathbb{E}[ ext{count of } x ext{ in segmentations of } w]$, with a pruning step that removes the tokens contributing least to total corpus log-likelihood, evaluated as
$$ \Delta \mathcal{L}(x) = \mathcal{L}_{ ext{full}} - \mathcal{L}_{V \setminus \{x\}} $$
where $\mathcal{L}_V = \sum_w f_w \log P(\mathbf{x}^*_w \mid V)$ is the total training-corpus log-likelihood under the best segmentation available with vocabulary $V$. Tokens with the smallest $\Delta \mathcal{L}(x)$ are pruned first, since their removal costs the least likelihood.
A practically important derived quantity is fertility, the average number of subword tokens produced per input word (or per byte, for byte-level tokenizers),
$$ ext{fertility} = \frac{1}{N} \sum_{i=1}^{N} | ext{tokens}(w_i)| $$
which directly determines the effective context length consumed by a given amount of text and varies substantially across languages and domains for any fixed vocabulary.
## Advanced Theory & Extensions
Byte-level BPE, used by GPT-2 and its successors, operates on the 256 possible byte values rather than Unicode characters as its base alphabet. This guarantees that literally any input, including malformed Unicode, emoji, or arbitrary binary-adjacent text, can be tokenized without ever needing an unknown-token fallback, since the byte alphabet is already a complete covering set. The tradeoff is that a single multi-byte Unicode character (common in non-Latin scripts) may initially decompose into two or more byte tokens before any merges are learned, which is part of why tokenizers trained predominantly on English text exhibit markedly worse fertility on scripts such as Devanagari, Thai, or CJK ideographs unless the training corpus is deliberately balanced across languages.
SentencePiece's Unigram LM mode enables subword regularization, a training-time augmentation technique in which, instead of always using the single best (Viterbi) segmentation of a training example, an alternative segmentation is sampled from the distribution $P(\mathbf{x}) \propto \prod_i p(x_i)^{1/ au}$ for some temperature $ au$, using an n-best list or a specialized sampling algorithm. Exposing the model to multiple valid segmentations of the same underlying text during training has been shown to improve robustness to tokenization noise and modestly improve downstream translation and generation quality, since the model no longer over-fits to one arbitrary decomposition of each word.
A separate line of work questions whether learned merge-based subword vocabularies are necessary at all. Tokenizer-free approaches such as CANINE and ByT5 operate directly on characters or raw bytes, avoiding tokenization-induced artifacts (such as inconsistent handling of numbers, whitespace, or rare tokens) at the cost of much longer sequences and increased compute per unit of text, usually mitigated by downsampling strategies inside the model architecture itself. More recently, the practical concern of tokenizer-induced unfairness, whereby the same amount of semantic content costs dramatically more tokens (and therefore more compute and higher API cost) in some languages than others due to a training corpus dominated by English, has motivated multilingual vocabulary balancing techniques that upsample underrepresented languages or scripts during vocabulary construction, and post-hoc fertility audits across language sets before a tokenizer is finalized for a production model.
## Computational Considerations
Naive BPE training, in which the full pair-frequency table is recomputed from scratch after every single merge, costs $O(M \cdot |C|)$ where $M$ is the number of merges and $|C|$ is the total corpus size in symbols, which becomes prohibitive for vocabularies in the tens of thousands trained on web-scale corpora. Production implementations instead maintain the pair-count table incrementally: after each merge, only the counts of pairs adjacent to the merged positions need to be updated, typically using a priority queue (max-heap) keyed by pair frequency together with a linked-list representation of each word's current symbol sequence, reducing amortized per-merge cost to roughly logarithmic in vocabulary size rather than linear in corpus size.
The Unigram LM tokenizer's training cost is dominated by the EM re-estimation step, since computing the Viterbi segmentation for every word in the corpus at every EM iteration, and doing so again after every pruning round, is comparatively expensive; in practice, seed vocabularies are capped (commonly at a few hundred thousand candidate substrings, often derived from a suffix-array-based frequent-substring enumeration rather than all possible substrings) and pruning is done in large batches, removing a fixed fraction (often 10 to 20 percent) of the lowest-utility tokens per round rather than one at a time, trading a small amount of pruning-order optimality for a large reduction in the number of expensive EM passes required.
At inference time, encoding cost matters enormously, since tokenization sits on the critical path of every model call. Greedy longest-match-first encoding (used by some WordPiece implementations) is $O(n \cdot L)$ per word for word length $n$ and maximum token length $L$, whereas priority-merge BPE encoding using a heap over applicable merges within a word is $O(n \log n)$. Highly optimized tokenizer libraries such as Hugging Face's `tokenizers` (written in Rust) achieve encoding throughput of millions of tokens per second by precompiling merge rules into a trie or automaton structure and processing text in parallel across CPU cores, which matters at the scale of both large-batch training data preprocessing and low-latency production inference serving.
## Practical Implementation Strategies
When training a new tokenizer, vocabulary size is the single most consequential hyperparameter, and it should be chosen jointly with the expected training corpus composition rather than copied blindly from a prior model. Larger vocabularies reduce sequence length (lower fertility, cheaper training and inference per unit of text) but increase the embedding and output-projection matrix size, which for very large vocabularies can become a non-trivial fraction of total model parameters, and require correspondingly more training data per rare token to be well estimated. A useful diagnostic before committing to a vocabulary is to plot fertility (average tokens per word or per character) on a held-out sample across every language or domain the model is expected to serve, since a vocabulary that looks efficient in aggregate can hide severe fragmentation on underrepresented subsets of the training mix.
Reserved special tokens (such as beginning- and end-of-sequence markers, padding, and any structural markers required for chat formatting or tool-call delimiters) must be added to the vocabulary before or during training and must never collide with a learned merge; most tokenizer training libraries support explicitly protecting a list of strings from being split, ensuring that, for instance, a chat role marker like "<|assistant|>" is always emitted as exactly one token rather than being fragmented by the general-purpose BPE merges. Digit and number handling is a well-known source of downstream arithmetic errors: because BPE and WordPiece both merge based on frequency, common short numbers get single tokens while longer or less common numbers are split at arbitrary, frequency-driven boundaries rather than at place-value boundaries, which is one reason some more recent tokenizers explicitly force digit-by-digit or fixed-chunk-size splitting for numeric spans regardless of what the general merge statistics would otherwise select.
Because retraining a tokenizer invalidates every embedding a model has learned, tokenizer choices are effectively frozen for the lifetime of a model family; extending a vocabulary after the fact (to add support for a new language or a new special token) is possible but requires either careful initialization of the new embedding rows (commonly by averaging the embeddings of the new token's constituent subwords under the old vocabulary) followed by continued pretraining, or accepting a degraded cold-start period during fine-tuning. This makes tokenizer design deserving of the same upfront rigor as architecture and data mixture decisions, rather than treatment as a solved preprocessing detail.
## Benchmark Datasets & Evaluation
Tokenizer quality is evaluated along several distinct axes rather than a single metric. Compression efficiency is typically measured as bytes per token or characters per token on a held-out corpus, directly reflecting how much raw text a fixed context window can hold. Fertility, tokens per word, is the standard metric for cross-lingual comparison and is commonly reported per-language on multilingual benchmark suites such as FLORES-200, which provides parallel sentences across two hundred languages and is widely used to audit tokenizer fairness across the world's languages. The XNLI and XTREME benchmark suites, while primarily designed to evaluate downstream multilingual model performance rather than tokenizers in isolation, are frequently used indirectly to demonstrate that high-fertility languages (where the tokenizer fragments words heavily) correlate with worse downstream task accuracy at matched compute budgets, since the effective context and effective training signal per token is diluted.
Morphological segmentation benchmarks, most notably derived from the CELEX and Morpho Challenge datasets, are used to evaluate how well a learned subword vocabulary's segmentations align with linguistically motivated morpheme boundaries (prefixes, roots, suffixes), though it is worth noting that alignment with linguistic morphology is not a design goal of BPE, WordPiece, or Unigram LM, and empirically, best downstream task performance does not require the learned segmentation to match a linguist's morphological analysis. Tokenizer robustness is also assessed via targeted stress tests, checking behavior on numbers, code identifiers, URLs, emoji sequences, and adversarially rare Unicode inputs, since production systems must tokenize arbitrary user input reliably without crashing, silently truncating, or producing pathologically long token sequences for short inputs.
## Key Challenges & Limitations
Subword tokenization introduces systematic, hard-to-fix biases that persist for the entire lifetime of a trained model. The clearest is multilingual fertility inequality: a tokenizer whose vocabulary was learned predominantly from English (and other high-resource, Latin-script languages) will fragment text in low-resource or non-Latin-script languages into far more tokens per unit of meaning, which increases API cost, reduces effective context length, and has been empirically linked to worse downstream performance for those languages, even before accounting for any disparity in raw training data volume. This is a tokenizer-level fairness issue distinct from, and compounding, dataset-level representation imbalance.
Arithmetic and character-level tasks are notoriously difficult for subword-tokenized models precisely because the tokenizer's merge boundaries rarely align with place-value or letter boundaries: the number "1234567" might be split as "123", "456", "7" by one tokenizer and "1", "234567" by another with a slightly different vocabulary, and a model relying on positional token embeddings has no consistent representation of "which digit is in the hundreds place" across different numbers of similar magnitude. Similarly, asking a model to reverse a word or count a specific letter within it is fundamentally harder than it looks, because the model's input representation may never expose individual characters as distinct tokens for common words.
Tokenizer-and-model mismatch is a persistent operational hazard: fine-tuning or continuing pretraining of a model with a different tokenizer than the one it was originally trained with silently corrupts every embedding lookup, since token ID 4291 means something entirely different under two different vocabularies, yet no error is raised, only degraded and hard-to-diagnose performance. Whitespace and formatting sensitivity is another subtle failure mode: a leading space before a word frequently produces a different token than the same word without a leading space (since many BPE implementations include the space as part of the token, following GPT-2's convention), meaning that semantically identical strings that differ only in surrounding whitespace or line breaks can tokenize completely differently, which has real consequences for prompt engineering and for token-level metrics such as perplexity comparisons across differently formatted texts.
## Hyperparameter Tuning
Vocabulary size is the primary tunable hyperparameter and should be selected via a fertility-versus-parameter-count tradeoff curve rather than a single conventional value; common choices cluster around 32,000 for earlier BERT-era and single-language models, 50,000 to 100,000 for modern general-purpose multilingual LLMs, and beyond 100,000 for models emphasizing multilingual or code-heavy fertility efficiency, but the right choice depends heavily on training corpus composition and target deployment languages. The maximum candidate substring length used to seed a Unigram LM tokenizer's initial vocabulary trades off coverage of long, useful whole-word tokens against training cost, since a larger maximum length exponentially increases the number of seed candidates.
For WordPiece and Unigram LM, a character coverage threshold (commonly set between 0.9995 and 0.9999 in SentencePiece) determines what fraction of characters in the training corpus must be representable by base vocabulary entries before falling back to an unknown-character byte-fallback mechanism; setting this too low causes rare-script or rare-symbol text to produce unknown tokens, while setting it unnecessarily high wastes vocabulary slots on extremely rare characters that contribute little to overall corpus coverage. The pruning fraction per EM round in Unigram LM training, and the corresponding number of EM iterations, trade training wall-clock time against how close the final vocabulary is to a jointly (rather than greedily) optimal pruning order; aggressive single-pass pruning is faster but can remove tokens whose true marginal value only becomes clear once other, correlated tokens have already been removed.
## Real-World Applications & Case Studies
GPT-2 and its successors adopted byte-level BPE with a roughly 50,000-token vocabulary specifically to guarantee that any input string, including code, emoji, or malformed text, is representable without unknown tokens, a decision that has been carried forward largely unchanged through subsequent GPT model generations, with vocabulary size scaled up substantially in GPT-4-generation tokenizers to improve fertility on non-English languages and code. BERT's original WordPiece vocabulary of roughly 30,000 tokens, trained jointly on English Wikipedia and BooksCorpus, became a de facto standard for encoder-only models and illustrates the long shadow a tokenizer choice casts: because the vocabulary is baked into every checkpoint derived from that pretraining run, thousands of downstream fine-tuned BERT variants inherited the same tokenizer, including its known weaknesses on non-English and code text, for years afterward.
T5 and its multilingual variant mT5 adopted SentencePiece with a Unigram LM back end specifically to support the training-time subword regularization technique described earlier, and mT5 in particular used a substantially enlarged, deliberately language-balanced vocabulary (250,000 tokens) to reduce fertility inequality across its 101 supported languages, an explicit acknowledgment that naive frequency-driven vocabulary construction on an English-dominated corpus produces unacceptable fragmentation for lower-resource languages. Code-specialized models such as Codex and subsequent code-generation LLMs required dedicated attention to whitespace and indentation tokenization, since Python's syntactic reliance on indentation depth means that naive BPE merge statistics, learned predominantly from natural language, handle runs of leading spaces or tabs poorly unless the training corpus and merge process are specifically informed by a substantial proportion of source code, motivating specialized code tokenizers that treat common indentation patterns and syntactic punctuation clusters as first-class merge candidates.
## Integration with Other Methods
Tokenization interacts directly with positional encoding and context length budgeting: since fertility determines how many tokens a fixed amount of raw text consumes, a lower-fertility tokenizer effectively extends a model's usable context window for a given nominal token limit, which is one of the underappreciated reasons that comparing "context length" figures across model families with different tokenizers can be misleading without normalizing for fertility on representative text. Retrieval-augmented generation systems are similarly sensitive to tokenizer choice, since the token budget consumed by retrieved passages competes directly with the budget available for the query and generated answer, making tokenizer efficiency an indirect but real lever on how many retrieved documents can be included in a prompt.
Parameter-efficient fine-tuning methods that add new special tokens (for instance, task-specific control tokens or newly introduced language markers) must coordinate embedding-table resizing with the base tokenizer's existing vocabulary, typically freezing the original embedding rows and only training the newly added ones, or applying a small adapter exclusively to the new token embeddings. Speculative decoding and other inference-acceleration techniques that operate at the token level are also sensitive to tokenizer granularity, since a higher-fertility tokenizer offers more, shorter tokens for a draft model to speculate over, changing the achievable speedup profile relative to a lower-fertility tokenizer covering the same text with fewer, longer tokens.
## Future Research Directions
Learned, end-to-end differentiable tokenization, in which the segmentation boundaries themselves are optimized jointly with the downstream model's training objective rather than fixed in advance by a separate frequency-statistics-driven algorithm, remains an active area, with approaches ranging from soft, continuous relaxations of discrete segmentation to hierarchical byte-level models that learn to dynamically pool bytes into variable-length chunks. Such approaches aim to eliminate the tokenizer-model mismatch hazard entirely and to allow segmentation granularity to adapt per-input rather than being globally fixed, at the cost of added architectural and training complexity relative to the simplicity of a fixed, precomputed vocabulary.
Fairness-aware tokenizer construction, explicitly optimizing vocabulary learning to equalize fertility (or a compute-cost-weighted variant of it) across a target set of languages rather than optimizing raw corpus compression alone, is gaining attention as multilingual deployment becomes standard practice rather than an afterthought, alongside proposals for standardized fertility-audit reporting akin to a model card, so that practitioners can compare the effective cost of serving different languages before committing to a tokenizer. Finally, as models are increasingly deployed on structured and semi-structured non-natural-language input, such as source code, chemical notation (SMILES strings), protein sequences, or tabular data, domain-specific tokenizer research continues to explore whether general-purpose subword algorithms transfer adequately to these domains or whether specialized segmentation strategies, informed by domain syntax rather than corpus frequency statistics alone, are needed to achieve competitive fertility and downstream performance.
## Summary & Key Takeaways
Subword tokenization is the foundational preprocessing step that converts raw text into the discrete token sequences every language model actually operates on, and the three dominant algorithmic families, BPE, WordPiece, and Unigram LM, differ chiefly in how they construct their vocabularies: BPE merges the most frequent adjacent symbol pair, WordPiece merges the pair maximizing a normalized likelihood-based score, and Unigram LM prunes a large seed vocabulary down to a target size based on each token's marginal contribution to corpus log-likelihood, with SentencePiece serving as a widely used, language-agnostic implementation framework for the latter two. These design choices are not merely implementation details; they directly determine fertility (tokens consumed per unit of text), which in turn governs effective context length, inference cost, and, critically, fairness across languages and domains, since tokenizers trained predominantly on English systematically fragment underrepresented languages into disproportionately more tokens. Tokenizer decisions are effectively permanent once a model is pretrained, making upfront vocabulary size selection, special-token reservation, and cross-lingual fertility auditing essential engineering practices rather than afterthoughts, and known failure modes, including inconsistent numeric tokenization, whitespace sensitivity, and tokenizer-model mismatch under naive fine-tuning, remain active sources of subtle model errors that practitioners must design around.
Keywords: tokenization, subword segmentation, byte-pair encoding, BPE, WordPiece, SentencePiece, Unigram language model, vocabulary size, fertility, byte-level tokenization, Viterbi segmentation, subword regularization, morphological segmentation, tokenizer fairness, multilingual tokenization, out-of-vocabulary, merge operations, character coverage
---
## Appendix: Practical Labs
### Lab 1: Training and Applying Byte-Pair Encoding from Scratch
This lab implements the core BPE training loop (iterative most-frequent-pair merging) and the corresponding greedy encoding procedure, then verifies that total tokenized corpus length decreases monotonically as more merge operations are learned.
import re
import collections
def get_stats(corpus):
"""Count frequency of every adjacent symbol pair across the corpus."""
pairs = collections.Counter()
for word, freq in corpus.items():
symbols = word.split()
for i in range(len(symbols) - 1):
pairs[(symbols[i], symbols[i + 1])] += freq
return pairs
def merge_vocab(pair, corpus):
"""Replace every occurrence of `pair` (adjacent symbols) with their merge."""
new_corpus = {}
bigram = re.escape(' '.join(pair))
pattern = re.compile(r'(?<!\S)' + bigram + r'(?!\S)')
for word in corpus:
new_word = pattern.sub(''.join(pair), word)
new_corpus[new_word] = corpus[word]
return new_corpus
def train_bpe(word_freqs, num_merges):
"""Learn `num_merges` BPE merge rules from a word-frequency dictionary."""
corpus = {' '.join(list(w)) + ' </w>': f for w, f in word_freqs.items()}
merges = []
for _ in range(num_merges):
pairs = get_stats(corpus)
if not pairs:
break
best = max(pairs, key=pairs.get)
corpus = merge_vocab(best, corpus)
merges.append(best)
return merges, corpus
def bpe_encode(word, merges):
"""Apply learned merges, in learned order, to tokenize a new word."""
symbols = list(word) + ['</w>']
for pair in merges:
i = 0
new_symbols = []
while i < len(symbols):
if i < len(symbols) - 1 and (symbols[i], symbols[i + 1]) == pair:
new_symbols.append(symbols[i] + symbols[i + 1])
i += 2
else:
new_symbols.append(symbols[i])
i += 1
symbols = new_symbols
return symbols
def test_bpe_merges_reduce_token_count_monotonically():
word_freqs = {"low": 5, "lower": 2, "newest": 6, "widest": 3, "new": 4, "wide": 2}
lengths = []
for n in [0, 5, 10, 15]:
merges, _ = train_bpe(word_freqs, n)
total_len = sum(len(bpe_encode(w, merges)) for w in word_freqs)
lengths.append(total_len)
print("total token count vs. num_merges [0,5,10,15]:", lengths)
assert lengths == sorted(lengths, reverse=True), (
"token count should be non-increasing as more merges are learned"
)
assert lengths[-1] < lengths[0], "more merges should strictly reduce total token count"
# sanity check: a well-covered word should encode to very few tokens
merges15, _ = train_bpe(word_freqs, 15)
encoded = bpe_encode("newest", merges15)
print("'newest' encoded with 15 merges:", encoded)
assert len(encoded) <= 3
print("BPE monotonic compression test passed.")
if __name__ == "__main__":
test_bpe_merges_reduce_token_count_monotonically()### Lab 2: BPE vs. WordPiece — Divergent Merge Criteria
This lab implements both the raw-frequency (BPE) and normalized-likelihood (WordPiece) merge-selection criteria on a crafted corpus designed so that one pair has high raw frequency but low normalized score, while another has the reverse, demonstrating that the two algorithms can choose different first merges from identical data.
import re
import collections
def get_stats(corpus):
pairs = collections.Counter()
symbol_freq = collections.Counter()
for word, freq in corpus.items():
symbols = word.split()
for s in symbols:
symbol_freq[s] += freq
for i in range(len(symbols) - 1):
pairs[(symbols[i], symbols[i + 1])] += freq
return pairs, symbol_freq
def merge_vocab(pair, corpus):
new_corpus = {}
bigram = re.escape(' '.join(pair))
pattern = re.compile(r'(?<!\S)' + bigram + r'(?!\S)')
for word in corpus:
new_word = pattern.sub(''.join(pair), word)
new_corpus[new_word] = corpus[word]
return new_corpus
def bpe_score(pairs, symbol_freq):
"""BPE: rank purely by raw co-occurrence count."""
return dict(pairs)
def wordpiece_score(pairs, symbol_freq):
"""WordPiece: rank by normalized likelihood-style score."""
return {(a, b): f / (symbol_freq[a] * symbol_freq[b]) for (a, b), f in pairs.items()}
def train(word_freqs, num_merges, scorer):
corpus = {' '.join(list(w)) + ' </w>': f for w, f in word_freqs.items()}
merges = []
for _ in range(num_merges):
pairs, symbol_freq = get_stats(corpus)
if not pairs:
break
scores = scorer(pairs, symbol_freq)
best = max(scores, key=scores.get)
corpus = merge_vocab(best, corpus)
merges.append(best)
return merges
def test_bpe_and_wordpiece_diverge_on_first_merge():
# 'a','b' are individually very frequent (long repeated run) but the (a,b)
# pair's *normalized* score is diluted by their high individual frequency.
# 'x','y' are individually rare but occur almost exclusively paired together.
word_freqs = {
"ababababab": 1,
"xy": 20,
}
bpe_merges = train(word_freqs, 3, bpe_score)
wp_merges = train(word_freqs, 3, wordpiece_score)
print("BPE first 3 merges: ", bpe_merges)
print("WordPiece first 3 merges: ", wp_merges)
assert bpe_merges[0] != wp_merges[0], (
"BPE (raw frequency) and WordPiece (normalized score) should choose "
"different first merges on this crafted corpus"
)
# BPE should follow raw pair frequency: (a,b) pair count = 5*1=5 occurrences
# within 'ababababab', vs (x,y) pair count = 1*20 = 20 -> BPE picks (x,y) first.
assert bpe_merges[0] == ('x', 'y')
# WordPiece normalizes by individual symbol frequency: a,b each occur 5 times
# (freq=1 word), giving score 5/(5*5)=0.2, versus x,y score 20/(20*20)=0.05,
# so WordPiece prefers (a,b) first despite its lower raw count.
assert wp_merges[0] == ('a', 'b')
print("BPE vs WordPiece merge-criterion divergence test passed.")
if __name__ == "__main__":
test_bpe_and_wordpiece_diverge_on_first_merge()### Lab 3: Unigram Language Model Tokenizer via EM and Pruning
This lab implements a simplified Unigram LM tokenizer: seeding a large candidate vocabulary, using Viterbi decoding to find each word's best segmentation under current token probabilities, and iteratively pruning the least-used tokens down to a target vocabulary size, while guaranteeing full coverage by always retaining single characters.
import math
import collections
def seed_vocab(words, max_len=5):
vocab = collections.Counter()
for w in words:
w = w + "_"
n = len(w)
for i in range(n):
for j in range(i + 1, min(i + max_len, n) + 1):
vocab[w[i:j]] += 1
return vocab
def viterbi_segment(word, logp, max_len=5):
n = len(word)
best = [(-math.inf, -1)] * (n + 1)
best[0] = (0.0, -1)
for i in range(1, n + 1):
for j in range(max(0, i - max_len), i):
piece = word[j:i]
if piece in logp:
score = best[j][0] + logp[piece]
if score > best[i][0]:
best[i] = (score, j)
if best[n][1] == -1 and n > 0:
return None, -math.inf
pieces = []
i = n
while i > 0:
j = best[i][1]
pieces.append(word[j:i])
i = j
pieces.reverse()
return pieces, best[n][0]
def em_probs(vocab_counts):
total = sum(vocab_counts.values())
return {p: c / total for p, c in vocab_counts.items()}
def train_unigram(words, target_vocab_size, max_len=5, prune_frac=0.2):
words_tagged = [w + "_" for w in words]
vocab = seed_vocab(words, max_len)
for w in words_tagged:
for ch in w:
vocab[ch] += 1 # single characters always retained -> guarantees coverage
while len(vocab) > target_vocab_size:
probs = em_probs(vocab)
logp = {p: math.log(pr) for p, pr in probs.items()}
usage = collections.Counter()
for w in words_tagged:
pieces, _ = viterbi_segment(w, logp, max_len)
for p in pieces:
usage[p] += 1
candidates = [p for p in vocab if len(p) > 1]
candidates.sort(key=lambda p: usage.get(p, 0))
n_remove = max(1, int(len(vocab) * prune_frac))
to_remove = candidates[:n_remove]
for p in to_remove:
if len(vocab) <= target_vocab_size:
break
del vocab[p]
probs = em_probs(vocab)
logp = {p: math.log(pr) for p, pr in probs.items()}
return vocab, logp
def test_unigram_lm_prunes_while_preserving_full_coverage():
words = ["low", "lower", "lowest", "new", "newer", "newest",
"wide", "wider", "widest", "slow", "slower"]
vocab, logp = train_unigram(words, target_vocab_size=25)
print("final pruned vocab size:", len(vocab))
assert len(vocab) <= 25
# every training word must remain fully segmentable after pruning
for w in words:
pieces, score = viterbi_segment(w + "_", logp)
assert pieces is not None, f"word '{w}' lost full coverage after pruning"
print(f" {w:8s} -> {pieces}")
# a smaller target vocabulary should yield lower (more negative) total
# log-likelihood, since fewer / shorter pieces are available
_, logp_big = train_unigram(words, target_vocab_size=40)
_, logp_small = train_unigram(words, target_vocab_size=20)
ll_big = sum(viterbi_segment(w + "_", logp_big)[1] for w in words)
ll_small = sum(viterbi_segment(w + "_", logp_small)[1] for w in words)
print("total log-likelihood: big vocab =", round(ll_big, 2),
" small vocab =", round(ll_small, 2))
assert ll_small < ll_big, "shrinking the vocabulary should reduce total log-likelihood"
print("Unigram LM pruning and coverage test passed.")
if __name__ == "__main__":
test_unigram_lm_prunes_while_preserving_full_coverage()### Lab 4: Measuring Tokenizer Fertility on In-Domain vs. Out-of-Domain Text
This lab trains a BPE tokenizer on an English-like word list and measures fertility (average tokens per word) on in-domain English words versus out-of-domain strings (DNA-like sequences, code identifiers, hexadecimal constants), quantifying the fragmentation effect a mismatched tokenizer produces on unfamiliar text.
import re
import collections
def get_stats(corpus):
pairs = collections.Counter()
for word, freq in corpus.items():
symbols = word.split()
for i in range(len(symbols) - 1):
pairs[(symbols[i], symbols[i + 1])] += freq
return pairs
def merge_vocab(pair, corpus):
new_corpus = {}
bigram = re.escape(' '.join(pair))
pattern = re.compile(r'(?<!\S)' + bigram + r'(?!\S)')
for word in corpus:
new_word = pattern.sub(''.join(pair), word)
new_corpus[new_word] = corpus[word]
return new_corpus
def train_bpe(word_freqs, num_merges):
corpus = {' '.join(list(w)) + ' </w>': f for w, f in word_freqs.items()}
merges = []
for _ in range(num_merges):
pairs = get_stats(corpus)
if not pairs:
break
best = max(pairs, key=pairs.get)
corpus = merge_vocab(best, corpus)
merges.append(best)
return merges
def bpe_encode(word, merges):
symbols = list(word) + ['</w>']
for pair in merges:
i = 0
new_symbols = []
while i < len(symbols):
if i < len(symbols) - 1 and (symbols[i], symbols[i + 1]) == pair:
new_symbols.append(symbols[i] + symbols[i + 1])
i += 2
else:
new_symbols.append(symbols[i])
i += 1
symbols = new_symbols
return symbols
def fertility(text_words, merges):
total_tokens = sum(len(bpe_encode(w, merges)) for w in text_words)
return total_tokens / len(text_words)
def test_out_of_domain_text_has_much_higher_fertility():
english_train = {
"running": 10, "jumped": 8, "walking": 9, "talked": 7, "played": 6,
"faster": 5, "slower": 5, "biggest": 4, "smallest": 4, "nation": 6,
"national": 5, "international": 3, "the": 50, "and": 40, "is": 35,
"was": 30, "function": 8, "connect": 7, "connection": 6,
}
merges = train_bpe(english_train, 35)
in_domain_words = ["running", "jumped", "national", "connection", "faster"]
out_domain_words = [
"ACGTACGTACGTACGT", # DNA-like sequence
"xk9zqw4vbnmasdf", # random alphanumeric identifier
"getElementByIdAsyncCallback", # camelCase code identifier
"0x1F3A9B2C", # hex constant
]
fert_in = fertility(in_domain_words, merges)
fert_out = fertility(out_domain_words, merges)
print(f"in-domain fertility (tokens/word): {fert_in:.2f}")
print(f"out-of-domain fertility (tokens/word): {fert_out:.2f}")
assert fert_out > fert_in * 1.5, (
"out-of-domain text should fragment into substantially more "
"subword tokens per word than in-domain text"
)
assert fert_in >= 1.0 and fert_out >= 1.0
print("Tokenizer fertility / domain-mismatch fragmentation test passed.")
if __name__ == "__main__":
test_out_of_domain_text_has_much_higher_fertility()