Tokenization

Problem

A neural network takes numbers. To feed it text we need a function from strings to sequences of integers and back, so that each integer can index a row of an embedding table and so that the model's output can be a probability over a fixed list. The list must be small enough to fit (the embedding table and the output layer are both vocabulary-sized) yet expressive enough that any text, including names, typos, code and every language, can be written with it. Sequences should also be short, because a Transformer's cost grows with sequence length.

Intuition

There are three natural ways to cut text. Characters give a tiny vocabulary and never meet an unknown symbol, but sequences are long and each token carries little meaning. Words give short sequences, but the vocabulary is enormous and still misses new words ("unfriendliest"), which become an unknown token. Subwords sit between: frequent words are one token, rare words split into frequent pieces (un · friend · li · est). A subword tokenizer is learned from data: it keeps whatever pieces let it write the training text compactly.

Mechanics

Byte-pair encoding (BPE)

BPE, adapted to word segmentation by Sennrich, Haddow and Birch, starts with every word split into characters plus an end-of-word marker (written </w> here). Then it repeats: count every adjacent pair of symbols across the corpus, weighted by word frequency; merge the most frequent pair into a new symbol everywhere; add that symbol to the vocabulary. The number of merges is the one hyperparameter, and it sets the vocabulary size. The output is an ordered list of merges. Encoding a new word replays those merges in the order they were learned, so a word the tokenizer never saw still comes out as known pieces, down to single characters in the worst case.

WordPiece

WordPiece (from Schuster and Nakajima's speech work, later used by Google's translation system and by BERT) builds its vocabulary the same way, merge by merge, but chooses the merge that most increases the likelihood of the training data under a unigram language model over the current units, not the most frequent pair. Frequent pairs of rare units gain the most. BERT marks pieces that continue a word with ## (play · ##ing) and encodes a word greedily, longest matching piece first.

The unigram language model

The unigram method (Kudo, 2018) works the other way round. It starts from a large set of candidate pieces, gives each a probability, and treats a segmentation's probability as the product of its pieces' probabilities. Encoding picks the most probable segmentation with the Viterbi algorithm. Training alternates expectation-maximization over the piece probabilities with pruning the pieces whose removal hurts the corpus likelihood least, keeping single characters so everything stays writable. Because it defines a distribution over segmentations, it can also sample one, a regularizer Kudo calls subword regularization. SentencePiece implements both BPE and unigram directly on raw text, treating the space as an ordinary symbol (shown as ▁), so no language-specific pre-splitting is needed and decoding is lossless.

Byte-level BPE

Unicode has well over a hundred thousand characters, so even a character base vocabulary is large and can still meet unseen characters. Byte-level BPE, used by GPT-2, runs BPE over the UTF-8 bytes of the text instead. The base vocabulary is exactly 256 byte values, so there is no unknown token: any string is some sequence of bytes. GPT-2's vocabulary has 50,257 entries: 256 bytes, 50,000 merges and one end-of-text token. The price is that characters outside ASCII start as several tokens and only shrink where the training data had enough of them to learn merges.

Special tokens

Vocabularies reserve IDs for symbols that are not text: beginning and end of sequence (so the model learns where to start and when to stop), padding (to make a batch rectangular), unknown (for tokenizers that have one), separators and, in chat models, role markers. Padding needs care twice: its embedding should not be trained (nn.Embedding(padding_idx=…) keeps that row fixed and out of the gradient) and its positions must not count in the loss (CrossEntropyLoss(ignore_index=-100)). The attention mask stops real tokens from attending to padding.

Vocabulary: <pad>=0, <bos>=1, <eos>=2, the=3, cat=4, sat=5, dog=6. Targets are the IDs shifted one step left, for teacher forcing.

Worked example

Implementation

from collections import Counter

def train_bpe(word_counts, num_merges):
    words = {tuple(w) + ("</w>",): c for w, c in word_counts.items()}
    merges = []
    for _ in range(num_merges):
        pairs = Counter()
        for word, c in words.items():
            for a, b in zip(word, word[1:]):
                pairs[a, b] += c
        if not pairs:
            break
        best = max(pairs, key=pairs.get)       # most frequent pair
        merges.append(best)
        words = {merge(w, best): c for w, c in words.items()}
    return merges

def merge(word, pair):
    out, i = [], 0
    while i < len(word):
        if word[i:i + 2] == pair:
            out.append(word[i] + word[i + 1]); i += 2
        else:
            out.append(word[i]); i += 1
    return tuple(out)

def encode(word, merges):
    tokens = tuple(word) + ("</w>",)
    for pair in merges:                         # replay in learned order
        tokens = merge(tokens, pair)
    return tokens

# into the model: ids [B, T] -> vectors [B, T, d]
import torch.nn as nn
emb = nn.Embedding(num_embeddings=50257, embedding_dim=768, padding_idx=0)

Production tokenizers apply merges by rank with a priority queue instead of looping over the whole merge list, and pre-split text (on spaces, punctuation or character classes) before BPE so merges never cross those boundaries. The vocabulary size and num_embeddings must match: an ID at or above it fails the lookup.

Tradeoffs

• Vocabulary size. A bigger vocabulary means shorter sequences (cheaper attention, more text per context window) but bigger embedding and output layers and rarer, less-trained tokens.
• BPE versus unigram. BPE is simple and deterministic. Unigram gives a probability for every segmentation, enabling sampling as a regularizer; WordPiece sits between, merging by likelihood.
• Bytes versus characters. Byte-level vocabularies can write anything with no unknown token, but non-Latin scripts start at 2-4 tokens per character.
• Fairness and cost. A tokenizer trained mostly on English text makes other languages longer, so they pay more and fit less in context.
• Frozen choice. The tokenizer is fixed before training; changing it means retraining or at least re-learning the embedding table. Its quirks (spaces attached to words, odd number splits) become the model's quirks.

Check yourself

Connections

References

• Sennrich, Haddow and Birch. Neural Machine Translation of Rare Words with Subword Units. ACL, 2016 (arXiv 1508.07909).
• Kudo and Richardson. SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing. EMNLP demo, 2018 (arXiv 1808.06226).
• Kudo. Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates. ACL, 2018 (arXiv 1804.10959).
• Schuster and Nakajima. Japanese and Korean Voice Search. ICASSP, 2012 (the original WordPiece).
• Wu et al. Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation. 2016 (arXiv 1609.08144).
• Devlin, Chang, Lee and Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL, 2019 (arXiv 1810.04805).
• Radford, Wu, Child, Luan, Amodei and Sutskever. Language Models are Unsupervised Multitask Learners (GPT-2, byte-level BPE). OpenAI, 2019.
• PyTorch v2.15.0: Embedding, F.embedding, CrossEntropyLoss.