WordPiece › The WordPiece Training Algorithm

WordPiece training begins with a character-level vocabulary and builds up from there. The procedure has three stages: initialization, the merge loop, and final vocabulary construction. The overall structure is nearly identical to BPE; the critical difference is only in how merges are scored.

Understanding the training algorithm matters because it explains both what the vocabulary contains and why it looks the way it does. The tokens in a WordPiece vocabulary are not arbitrary: each one was selected because, at the moment it was created, it represented the most statistically informative unit to add, given everything that had been added before. The history of merges shapes the final vocabulary in ways that are not always obvious from looking at the tokens alone.

Initialization. The corpus is segmented into individual characters. Every word is split into its constituent characters, with the first character left as-is and every subsequent character prefixed with ##. For example, "unbelievable" becomes:

The initial vocabulary consists of all unique characters that appear in the corpus, plus a few special tokens (typically [UNK], [CLS], [SEP], [PAD], [MASK] for BERT). This initialization guarantees that every possible input can be tokenized: in the worst case, each character becomes its own token.

The ## prefix on continuation characters is applied from the very beginning, not added later. This means the character e and the continuation character ##e are different entries in the initial vocabulary, and they remain different throughout training. A merge that joins e with something after it will create a word-initial sequence, while a merge that joins ##e with something after it will create a word-internal sequence. The distinction is preserved at every level of the vocabulary.

Merge loop. At each iteration, the algorithm computes the score $P(AB) / (P(A) \cdot P(B))$ for every adjacent pair of tokens that appears in the corpus. It selects the pair with the highest score and merges it into a new token, updating all affected words in the corpus. This process repeats until the vocabulary reaches the target size.

After each merge, all token probabilities must be recomputed, because the merge changes which tokens exist and how often each appears. This is what makes WordPiece training more expensive than BPE: the full probability recalculation happens at each step. In practice, efficient implementations batch these updates and avoid recomputing from scratch by maintaining running counts that are updated locally around the affected pairs.

Final vocabulary. The result is a vocabulary that contains whole words (those frequent enough to be merged all the way), common subword pieces (frequent morphological components), and individual characters as a fallback. In BERT's case, after training on BookCorpus and Wikipedia, the vocabulary of 30,000 merges produces roughly 30,522 tokens including the special tokens.

The key difference from BPE at this algorithmic level is only in the scoring: BPE counts raw co-occurrences; WordPiece computes the PMI-based score. The overall loop structure is nearly identical. This means that if you already understand BPE's training loop, you understand WordPiece: you just swap out one scoring formula for another, and the rest follows.