Unigram Language Model Tokenization › Training: The EM Algorithm

How do we learn the token probabilities? The challenge is circular: we need token probabilities to determine which segmentations are best, but we need to know the best segmentations to estimate token probabilities. This circularity is not unusual in machine learning. Mixture models, hidden Markov models, and latent Dirichlet allocation all face the same chicken-and-egg problem. The standard solution is the Expectation-Maximization (EM) algorithm, which iterates between two steps until convergence.

Think of EM as a negotiation process. In the first step, you make your best guess about which segmentations are most probable given the current token probabilities. In the second step, you update the token probabilities based on which tokens appeared in those best-guess segmentations. Then you repeat, each iteration producing probabilities that are more consistent with the segmentations and segmentations that are more consistent with the probabilities. The two steps reinforce each other until they stop changing, at which point you have reached a local optimum.

The EM algorithm is guaranteed to never decrease the log-likelihood of the training data. Each iteration produces token probabilities that make the corpus at least as probable as the previous iteration's probabilities. This monotone improvement property means you can always tell progress is being made, even if convergence to the global optimum is not guaranteed.

Before EM begins, we construct a large initial vocabulary. A common strategy is to include all characters in the training corpus (to guarantee any word can be tokenized), all substrings up to some length, and possibly all words that appear frequently. This initial vocabulary might contain tens of thousands of candidates. Call it $V_0$.

The choice of initial vocabulary matters more than it might seem. If a token that would be useful is absent from $V_0$, it can never be added: the algorithm only prunes, it never introduces new tokens. This is a fundamental asymmetry with BPE, which creates new tokens via merges. To ensure no useful token is missed, practitioners typically include all substrings up to length 16 or even longer, accepting the computational cost of starting with a very large vocabulary.

We also need initial probability estimates for each candidate token. A natural choice is to initialize $P(x)$ proportional to the token's frequency in the training corpus, normalized so all probabilities sum to one. Frequency-proportional initialization ensures that the EM algorithm starts from a reasonable guess: common substrings receive higher initial probabilities, giving the E-step something sensible to work with in the first iteration.

Given the current probability estimates, for every word in the training corpus, we compute the probability of every possible segmentation of that word. Then we calculate expected counts for each token by summing, across all words and all segmentations of those words, the probability that a given token appears in that word's segmentation, weighted by how many times the word itself appears.

The word "expected" here has a precise statistical meaning. We are not just counting how often a token appears in the single best segmentation. We are computing a weighted average count across all segmentations, where each segmentation is weighted by its probability. If token "un" appears in five different segmentations of "unhappiness," and those segmentations have probabilities 0.5, 0.2, 0.15, 0.1, and 0.05, then "un" receives an expected count contribution of $0.5 + 0.2 + 0.15 + 0.1 + 0.05 = 1.0$ from the word "unhappiness" (assuming it appears once in the corpus). This soft counting, rather than hard assignment to a single segmentation, is what allows the EM algorithm to explore the full space of possible segmentations rather than committing prematurely to one.

More formally, let $w$ denote a word and $S(w)$ denote the set of all possible segmentations of $w$. The expected count of token $x$ in the corpus is computed by summing over every word in the corpus, and for each word, over every possible segmentation of that word:

$$c(x) = \sum_{w \in \text{corpus}} \text{freq}(w) \cdot \sum_{\mathbf{x} \in S(w)} P(\mathbf{x} \mid w) \cdot \mathbb{1}[x \in \mathbf{x}]$$

where:

• $c(x)$: the expected count of token $x$ across the entire corpus
• $\text{freq}(w)$: the number of times word $w$ appears in the training corpus
• $P(\mathbf{x} \mid w)$: the probability of segmentation $\mathbf{x}$ given word $w$, computed from the current probability estimates and normalized over all segmentations of $w$
• $\mathbb{1}[x \in \mathbf{x}]$: an indicator that equals 1 if token $x$ appears somewhere in segmentation $\mathbf{x}$, and 0 otherwise

Why does this formula make sense? Notice that $\sum_{\mathbf{x} \in S(w)} P(\mathbf{x} \mid w) \cdot \mathbb{1}[x \in \mathbf{x}]$ is exactly the posterior probability that token $x$ appears in the segmentation of word $w$. Multiplying by $\text{freq}(w)$ and summing over all words gives the total expected number of times we observe token $x$ when we segment the full corpus. This is precisely the soft count that maximum likelihood estimation needs.

In practice, computing this sum over all segmentations is done efficiently using the Viterbi forward-backward algorithm, which we discuss below. The number of possible segmentations grows exponentially with word length (a word of length $n$ has at most $2^{n-1}$ segmentations), so direct enumeration is infeasible for any word longer than a few characters. The forward-backward algorithm computes the same result in polynomial time by reusing intermediate computations through dynamic programming.

Given the expected counts from the E-step, we update the probability of each token using maximum likelihood estimation. Maximum likelihood estimation for a categorical distribution (a distribution over a finite set of outcomes, like our vocabulary) has a beautifully simple closed-form solution: the probability of each outcome is proportional to its expected count.

To see why, imagine you are estimating the probability of each face on a six-sided die given some observed rolls. If you observed face 1 appearing 30 times, face 2 appearing 20 times, face 3 appearing 10 times, and so on, the maximum likelihood estimate is simply the fraction of rolls that showed each face. The same logic applies here: the maximum likelihood probability of token $x$ is the fraction of all expected token occurrences that belong to $x$:

$$P(x) = \frac{c(x)}{\sum_{x' \in V} c(x')}$$

where:

• $P(x)$: the updated probability of token $x$
• $c(x)$: the expected count of token $x$ from the E-step
• $\sum_{x' \in V} c(x')$: the total expected count across all tokens in the vocabulary, serving as a normalizing constant that ensures all probabilities sum to one

Why does this formula make sense? Notice that the denominator is just the sum of all expected counts, which equals the expected total number of tokens in the corpus when segmented optimally. Dividing each token's expected count by this total gives its relative frequency, which is the maximum likelihood estimate of its probability under the unigram independence assumption.

This update simply says: the probability of token $x$ is proportional to how often we expect to see it when we segment the corpus according to the current probability estimates. Tokens that appear in high-probability segmentations of common words receive higher probabilities. Tokens that are rarely needed, either because they appear in low-probability segmentations or because the words containing them are rare, get lower probabilities.

After each EM cycle converges, we remove a fraction, typically 10-20%, of the tokens with the lowest probability scores from the vocabulary. But raw probability is not the right criterion for pruning. A token can have a low probability simply because it covers rare linguistic patterns, yet still be essential for tokenizing those patterns. If we remove it, words containing those patterns may become unrepresentable or degrade severely.

The key insight is that we should measure a token's impact on the training data by computing how much the corpus log-likelihood would drop if we removed that token from the vocabulary. Specifically, for each token $x$ we consider a modified vocabulary $V \setminus \{x\}$ and compute the change in log-likelihood:

$$\Delta\mathcal{L}(x) = \mathcal{L}(V) - \mathcal{L}(V \setminus \{x\})$$

A token with $\Delta\mathcal{L}(x) \approx 0$ can be pruned with minimal cost: the remaining tokens can cover the same words nearly as well. A token with large $\Delta\mathcal{L}(x)$ is irreplaceable and should be kept. Tokens that can easily be replaced by combinations of other tokens have low impact and are pruned first. Single characters are always preserved, since they guarantee that any word can be tokenized as a fallback.

After pruning, we run EM again with the reduced vocabulary. We repeat this prune-then-EM cycle until we reach the target vocabulary size. The result is a compact vocabulary where every token earns its place by maximizing the probability of the training data. The final vocabulary is an optimized compression rather than a frequency list of the training corpus's morphological structure.

The overall objective that guides this entire process is to maximize the log-likelihood of the training corpus. Log-likelihood is easier to work with than raw likelihood because products of small probabilities become sums of log-probabilities, which avoids numerical underflow.

To compute the log-likelihood, we need the probability of each word in the corpus. Because any word can be segmented in multiple ways, the probability of the word is the sum of the probabilities of all its segmentations:

$$P^*(w) = \sum_{\mathbf{x} \in S(w)} P(\mathbf{x}) = \sum_{\mathbf{x} \in S(w)} \prod_{i=1}^{m} P(x_i)$$

where $P^*(w)$ is the marginal probability of word $w$, obtained by marginalizing over all possible segmentations. This marginal probability treats the segmentation as a latent variable: we do not observe which segmentation was used, so we sum over all possibilities.

The total corpus log-likelihood is then:

$$\mathcal{L} = \sum_{w \in \text{corpus}} \text{freq}(w) \cdot \log P^*(w)$$

where:

• $\mathcal{L}$: the total log-likelihood of the corpus, which the algorithm tries to maximize
• $\text{freq}(w)$: how many times word $w$ appears in the training corpus
• $\log P^*(w)$: the log of the marginal probability of word $w$, summed over all its segmentations

Why does this formula make sense? Notice that we are giving more weight to words that appear more frequently. "The" appearing 10,000 times contributes 10,000 times as much to the objective as a word appearing once. This makes intuitive sense: a good tokenizer should above all do well on common words, and the EM algorithm's objective function reflects this priority.

Maximizing this log-likelihood means finding token probabilities that make the training corpus as probable as possible. Tokens that appear in many high-probability segmentations of common words naturally receive high probabilities. Tokens that can easily be decomposed into other tokens contribute less to the objective and are candidates for pruning.