Unigram LLM Tokenization: EM Training

Byte Pair Encoding and WordPiece both work by building up a vocabulary through merges: they start with characters and repeatedly combine the best pair into a new token. The unigram language model approach flips this logic entirely. It starts with a large vocabulary and prunes it down, removing tokens that contribute least to explaining the training data. This inversion is more than an implementation detail. It reflects a fundamentally different theory of what a good tokenizer should be.

Think of the difference this way. Imagine you are assembling a LEGO set. BPE starts with individual bricks and glues them together one pair at a time, always picking the most common pair. Unigram LM starts by dumping the entire inventory of every LEGO piece ever made on the table, then systematically discards the pieces that are never needed to build the models in your collection. You end up with a curated kit where every piece earns its place.

This inversion produces a different algorithm with a useful capability: it gives every token a probability, which allows the model to produce a distribution over all possible tokenizations rather than a single split. You can sample from this distribution. This exposes later models to multiple ways the same word might be split. This technique, called subword regularization, acts as a data augmentation method that makes models more less brittle to tokenization variation. When a neural network has seen "unhappiness" split as both ["un", "happiness"] and ["unhappy", "ness"] during training, it is less dependent on any single tokenization convention, and it generalizes better to inputs that arrive with slightly different surface forms.

The unigram LM algorithm is one of the two modes supported by SentencePiece (alongside BPE), and it underlies tokenizers used in models like XLNet, ALBERT, and mBART. Understanding it means understanding both a tokenization algorithm and the probabilistic philosophy that treats text segmentation as inference rather than rule application. This chapter builds the full picture: why the probabilistic view is useful, how the EM training algorithm works, how Viterbi decoding finds the best segmentation at inference time, how sampling enables regularization, and what limitations you should keep in mind when applying this approach in practice.

To understand why the unigram approach is principled, you need to appreciate just how ambiguous tokenization really is. Given the word "unhappiness," is the right split ["un", "happiness"], ["un", "happy", "ness"], ["unhappy", "ness"], or ["u", "n", "happiness"]? All of these are valid subword segmentations. Each captures different structural information about the word. The first recognizes the prefix "un" and treats "happiness" as a semantic unit. The second decomposes the word into three meaningful morphemes. The third groups the adjectival stem with its negative prefix and separates the nominalization suffix. The fourth is arguably silly but still technically valid as a string decomposition.

BPE picks one segmentation deterministically based on which merges were applied in order. This determinism feels clean, but it hides a real ambiguity. The "best" segmentation according to BPE is best only in the sense that it follows the greedy merge sequence BPE happened to learn. Another training run on a slightly different corpus might produce a different merge sequence and therefore a different segmentation of the same word. The choice of segmentation is not linguistic ground truth; it is an artifact of the algorithm.

The unigram LM takes a more honest stance: all segmentations are possible, and we should assign probabilities to each one. The segmentation with the highest probability under our learned model is the best guess, but it is still just a guess. Other segmentations are not wrong; they are just less likely given the training data. This probabilistic framing immediately permits sampling: we can draw segmentations from the distribution instead of always taking the argmax, and that diversity turns out to be tremendously useful during model training.

The unigram language model treats each token in a segmentation as independently drawn from a probability distribution over the vocabulary. The word "unigram" in the algorithm's name refers to this independence assumption: we model each token with a unigram (single-token) probability, ignoring context. Given a vocabulary $V$ where each token $x_i \in V$ has an assigned probability $P(x_i)$, the probability of a particular segmentation $\mathbf{x} = (x_1, x_2, \ldots, x_m)$ of a word is computed by multiplying together the individual token probabilities:

where:

• $\mathbf{x}$: a specific segmentation, meaning a specific sequence of subword tokens that together reconstruct the original word
• $m$: the number of tokens in this segmentation
• $P(x_i)$: the probability of the $i$-th token in the segmentation, estimated from training data
• $\prod$: the product over all token probabilities, which reflects the independence assumption that each token contributes multiplicatively to the overall probability

Why does multiplication make sense here? Because we are computing the joint probability of a sequence of independent events. If flipping heads has probability 0.5, the probability of flipping heads three times in a row is $0.5 \times 0.5 \times 0.5 = 0.125$. Analogously, if seeing token "un" has probability 0.05 and seeing "happy" has probability 0.04, the probability of the segmentation ["un", "happy"] under the independence assumption is $0.05 \times 0.04 = 0.002$.

This "unigram" formulation assumes tokens are independent. In reality, adjacent tokens are not independent: "un" appearing next to "happy" is far more likely than "un" appearing next to "ness." But the independence simplification makes the math tractable and the algorithm practical. More sophisticated models that capture token-to-token dependencies would require exponentially more parameters and far more complex inference algorithms. The unigram approximation gives us something that is good enough to rank segmentations effectively while remaining computationally manageable.

The best segmentation is the one that maximizes this probability. Among all ways to split "unhappiness," we choose the segmentation whose constituent tokens have the highest combined probability. Intuitively, the best split uses common, well-established subwords rather than obscure fragments. If "happiness" is a common enough subword that its probability exceeds the combined probability of "happy" and "ness," then ["un", "happiness"] beats ["un", "happy", "ness"]. The model's learned probabilities encode collective evidence from the training corpus about which subword units are truly productive.

this probabilistic view converts tokenization from a rule-based string operation into an inference problem. You are asking: given the evidence encoded in training data token frequencies, what is the most probable explanation for how this word was constructed from subword units? This reframing connects tokenization to a long tradition of probabilistic modeling in NLP and makes the algorithm amenable to the toolkit of statistical inference.

Before diving into the training algorithm, it is worth pausing to understand why the subword level is the right granularity for this probabilistic model. Word-level probabilities fail for rare and out-of-vocabulary words: a model that assigns probability only to whole words seen in training cannot handle new words at all. Character-level probabilities solve the coverage problem but produce extremely long sequences where each character must be processed independently, which is both slow and fails to encode the morphological structure that makes words predictable.

The subword level hits the sweet spot. Subwords are short enough that a large fraction of them appear frequently in training data, giving reliable probability estimates. They are long enough to encode meaningful linguistic units like prefixes, suffixes, and stems. And the vocabulary is small enough to be manageable as a discrete probability distribution. The unigram LM formalizes this intuition by learning exactly which subwords deserve their own probability mass and which are better decomposed into smaller pieces.

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:

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$:

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:

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:

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:

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.

At inference time, given a new word, we want to find its most probable segmentation. This is not a trivial problem: a word of length $n$ has up to $2^{n-1}$ possible segmentations, and even a 10-character word could have hundreds of valid tokenization paths through the vocabulary. Checking all of them explicitly would be far too slow.

The solution is the Viterbi algorithm, a classic dynamic programming technique adapted here to sequence segmentation. Think of Viterbi as a smart path-finding algorithm. Imagine you are walking from one end of a word to the other, stepping through it one character at a time. At each position, you can "land" by completing a valid vocabulary token that ends at that position. Your goal is to find the sequence of landings that maximizes your total log-probability score.

Imagine writing out the characters of a word as positions $0, 1, 2, \ldots, n$ where $n$ is the word length. We want to find the path from position 0 to position $n$ that passes through subword tokens with the highest combined log-probability.

Define $\text{best}[j]$ as the log-probability of the highest-probability segmentation of the prefix $w[0:j]$. We build this table up character by character, from left to right. At each position $j$, we ask: for every possible last token that ends at position $j$, what is the best we could do if we used that token? The answer is the best score for the prefix ending before that token, plus the log-probability of the token itself:

where:

• $\text{best}[j]$: the maximum log-probability achievable for the prefix ending at position $j$
• $i$: the start of the last token in the optimal segmentation of $w[0:j]$
• $w[i:j]$: the substring from position $i$ to $j$, which must be a valid token in the vocabulary $V$
• $\log P(w[i:j])$: the log-probability of that token, which contributes to the total score
• $\text{best}[i]$: the best log-probability for the prefix $w[0:i]$, already computed in an earlier step

We initialize $\text{best}[0] = 0$ because the empty string has log-probability $\log(1) = 0$. Then we fill in $\text{best}[1], \text{best}[2], \ldots, \text{best}[n]$ using this recurrence. While filling the table, we track which starting position $i^*$ achieved the maximum at each $j$, allowing us to backtrack and recover the actual best segmentation.

Why does this formula make sense? Notice that we are computing the best possible score for prefix $w[0:j]$ by considering every valid token $w[i:j]$ that ends at position $j$. For each such token, the best total score we could achieve is the best score for the prefix before the token (already computed) plus the token's own log-probability. Taking the maximum over all valid tokens gives us the globally optimal score for $w[0:j]$, and the key insight is that this optimal score at position $j$ only depends on optimal scores at earlier positions, never on positions after $j$. This is the "optimal substructure" property that makes dynamic programming applicable.

This dynamic programming approach runs in $O(n^2)$ time in the word length (or $O(n \cdot k)$ where $k$ is the maximum token length), making it fast in practice since words are rarely longer than a few dozen characters. Even for a 20-character word with a maximum token length of 16, the algorithm only needs to examine $20 \times 16 = 320$ pairs.

After filling the $\text{best}$ table, we recover the actual segmentation by backtracking. Starting at position $n$ (the end of the word), we follow the stored pointers backward: the pointer at position $n$ tells us where the last token started, say at position $i^*$. Then we follow the pointer at $i^*$ to find where the second-to-last token started, and so on, until we reach position 0. Reversing the collected tokens gives us the optimal segmentation in left-to-right order.

This backtracking step is $O(n)$ and is the same technique used in Viterbi decoding for hidden Markov models, sequence-to-sequence alignment, and many other applications. The combination of forward dynamic programming to fill the table and backward tracing to recover the solution is one of the most versatile techniques in algorithm design.

One of the most powerful features of the unigram LM approach is that it naturally supports sampling a segmentation rather than always picking the single best one. Instead of running Viterbi to find the maximum-probability segmentation, we can sample from the distribution over all segmentations, where each segmentation's probability is proportional to the product of its token probabilities.

This capability directly enables subword regularization. During training of a downstream neural network, instead of always tokenizing "unhappiness" as ["un", "happiness"], the tokenizer might also produce ["unhappy", "ness"] or even ["un", "happy", "ness"] on different training examples. The neural network must learn to understand "unhappiness" from all of these representations, making it robust to the particular segmentation convention used at inference time.

Sampling can be done efficiently using the forward-backward algorithm. The forward pass computes, for each position $j$, the total probability mass of all segmentations of the prefix $w[0:j]$, working left to right:

This is analogous to the forward algorithm in hidden Markov models, where we sum over all paths to a state rather than taking the maximum. The forward probabilities represent the total evidence for all segmentation paths that reach position $j$.

In the backward pass, we sample a path from position $n$ back to position 0. At each position $j$, we look at all valid tokens $w[i:j]$ that could end the segmentation here, compute the probability of each as $\text{fwd}[i] \cdot P(w[i:j])$, normalize to get a probability distribution, and sample one token according to these weights. We then move to the starting position $i$ of the sampled token and repeat. Each step samples the next token in the segmentation proportional to its contribution to the total segmentation probability.

The result is an algorithm that, when tokenizing the same word multiple times, can produce different segmentations. "unhappiness" might be split as ["un", "happiness"] one time and ["unhappy", "ness"] another time, with probabilities reflecting how well each split explains the training data.

Formally, the sampling procedure uses a temperature parameter $\alpha$ to control how much randomness to introduce. When $\alpha = 1$, we sample proportional to the true probability distribution. When $\alpha \to 0$, all probability mass concentrates on the single best segmentation, equivalent to Viterbi. When $\alpha \to \infty$, all segmentations become equally likely. In practice, $\alpha$ around 0.1 to 1.0 works well, with lower values providing gentle regularization and higher values introducing more aggressive diversity.

The temperature parameter makes intuitive sense if you think of $\alpha$ as controlling how "peaked" the sampling distribution is. A very low temperature makes the highest-probability segmentation overwhelmingly likely, so you rarely see alternatives. A moderate temperature spreads probability more evenly among the top few segmentations, giving you meaningful diversity without too much noise. A very high temperature gives nearly equal probability to all segmentations, including very low-probability ones that represent linguistically bizarre splits, which can hurt rather than help.

The forward-backward algorithm for sampling has a deep connection to the E-step of the EM training algorithm. During training, the E-step uses the forward-backward algorithm to compute expected counts of each token: for every word in the corpus, it computes how much probability mass flows through each token position across all possible segmentations. During inference, we use nearly the same algorithm but sample a specific path rather than summing. This connection is not coincidental: sampling from the posterior over segmentations is the same mathematical operation as computing expected counts, just realized differently. Training uses the full sum for exact gradient computation; inference uses sampling for stochastic exploration.

The unigram LM and BPE approaches embody different philosophies about what makes a good tokenizer. Understanding both philosophies helps you reason about when to use each and what tradeoffs to expect.

BPE is constructive: it starts small and grows. The vocabulary begins as individual characters, and merges create new tokens by combining the most frequent existing pairs. The merge order is determined purely by co-occurrence frequency in the current vocabulary state. BPE is deterministic given the same corpus and hyperparameters, and it produces exactly one segmentation per word at inference time. The algorithm is simple enough to implement in a few dozen lines of Python and fast enough to train on large corpora in minutes.

The unigram LM is reductive: it starts large and shrinks. The vocabulary begins with many candidates, the EM algorithm assigns each a probability based on how much it contributes to explaining the training data, and pruning removes low-value tokens. At inference time, the model can produce multiple segmentations with associated probabilities. The algorithm requires understanding the EM framework, the Viterbi forward-backward recursion, and the log-likelihood objective, which is a higher conceptual overhead than BPE.

The practical differences between the two approaches are significant in several dimensions. BPE vocabularies tend to contain more "greedy" tokens: subwords that are very common in their raw character form but may not align well with linguistic units. For example, BPE might create the token "##ing" (a suffix) early in its merge sequence simply because the character trigram "ing" is common, even if this particular three-character token is redundant given tokens like "##running," "##playing," and "##walking" that appear frequently enough to be learned later. The unigram LM is more principled: every token that survives pruning increases the corpus likelihood, meaning it is needed to compress the data more efficiently than the remaining vocabulary could without it.

Empirically, both approaches produce similar-quality tokenizations, but the unigram LM's probabilistic foundation enables subword regularization, which provides a meaningful edge in training downstream models. On tasks where morphological variation is common, such as machine translation with inflected languages or document classification across formal and informal registers, models trained with unigram sampling tend to generalize better.

One important computational tradeoff is that unigram LM training involves running EM and evaluating all segmentations, which is more expensive than BPE's simple frequency-counting approach. For large corpora with hundreds of gigabytes of text, this cost can be substantial. In practice, tokenizer training on large corpora is done once and the result is reused, so the higher training cost is typically acceptable. The cost is paid during tokenizer training, not during downstream model training or inference.

WordPiece, the third major subword algorithm, falls between BPE and unigram LM. Like BPE, it is constructive (merge-based), but it selects merges based on a likelihood criterion rather than raw frequency. WordPiece does not naturally support sampling, but its likelihood-based training gives it some of the principled quality of the unigram LM approach without the full EM machinery.

Let's trace through the core ideas with a small, concrete example. Working through a numerical example makes the abstract formulas tangible and reveals how the algorithm's decisions arise from the learned probabilities. Suppose our vocabulary contains these tokens with their probabilities:

These probabilities reflect a hypothetical trained model. "un" is the most probable token (0.05). This reflects its frequent use as a prefix across many words. "happiness" has the lowest probability (0.015) because it is a longer, more specific token that appears less frequently in a typical corpus.

To segment "unhappiness," we consider the main candidate segmentations and compute the probability of each:

Segmentation A: ["un", "happy", "ness"]

We compute the probability as the product of the three token probabilities:

In log space: $\log_{10} P_A = -1.301 + (-1.398) + (-1.523) = -4.222$

Segmentation B: ["unhappy", "ness"]

In log space: $\log_{10} P_B = -1.699 + (-1.523) = -3.222$

Segmentation C: ["un", "happiness"]

In log space: $\log_{10} P_C = -1.301 + (-1.824) = -3.125$

Comparing these values: $P_C = 0.000750 > P_B = 0.000600 > P_A = 0.000060$. The Viterbi algorithm chooses segmentation C, ["un", "happiness"], as the most probable.

This result deserves careful interpretation. Notice how the model preferred the two-token split using the compound token "happiness" over the three-token split ["un", "happy", "ness"]. Even though "happiness" has a lower individual probability (0.015) than "ness" (0.03) or "happy" (0.04), the two-token segmentation wins because multiplying two moderately small numbers gives a larger result than multiplying three small numbers. Each additional token in a segmentation multiplies in another factor less than one, so shorter segmentations tend to win unless the longer segmentation's tokens are individually much more probable.

This is the unigram LM's built-in Occam's razor: all else being equal, simpler (fewer tokens) segmentations are favored. The model will only use more tokens when the increased specificity of a longer segmentation brings enough probability gain to overcome the cost of the extra multiplication.

Now let's trace through the Viterbi recurrence for "unhappiness" (of length 11). We build the best table position by position:

• $\text{best}[0] = 0$ (empty string)
• $\text{best}[2] = \text{best}[0] + \log P(\text{un}) = 0 + (-1.301) = -1.301$ (token "un" from positions 0-2)
• At position 9 ("unhappines"), we find "happiness" runs from positions 2-11, but we are only at position 9 here, so we must continue.
• $\text{best}[11]$: we check token "happiness" from positions 2-11. $\text{best}[2] + \log P(\text{happiness}) = -1.301 + (-1.824) = -3.125$. We also check "ness" from positions 7-11: $\text{best}[7] + \log P(\text{ness})$, where $\text{best}[7]$ was computed as $\text{best}[0] + \log P(\text{unhappy}) = -1.699$ (if "unhappy" runs from 0 to 7). So this path gives $-1.699 + (-1.523) = -3.222$.

The maximum at position 11 is $-3.125$ (from the path via "un" then "happiness"), confirming our earlier calculation. The backtracking pointer at position 11 points to position 2, and the pointer at position 2 points to position 0, recovering the segmentation ["un", "happiness"].

Let's build a simplified unigram LM tokenizer from scratch to see all the moving parts. We'll implement the core E-step probability calculations, Viterbi decoding, and segmentation sampling. Building these pieces ourselves, rather than just calling a library, illuminates how the theoretical framework maps to concrete computations.

We start with our imports and dependencies. The implementation uses only standard Python libraries plus NumPy, which is readily available in any data science environment.

We begin by creating a small training corpus and extracting all substrings as vocabulary candidates. Each substring receives an initial probability proportional to its frequency. This mirrors the initialization step in the real algorithm, where all substrings up to some maximum length are added to the candidate vocabulary before EM begins.

The training corpus we use here is deliberately small and focused on English morphology, with words sharing prefixes, suffixes, and stems. This makes it easy to see how the algorithm learns recurring subwords like "ing," "ed," "un," and "er."

Our initial vocabulary contains substrings of all lengths, giving us many candidate tokens before pruning begins. The highest-probability tokens tend to be short, common substrings that appear across many words. Single characters dominate the top of the probability list because they appear in every word that contains them. As EM training proceeds and vocabulary pruning is applied, the algorithm will learn to give higher relative probability to the longer subwords that encode meaningful morphological units.

Next, we implement the Viterbi decoding algorithm that finds the most probable segmentation of any word given the current vocabulary and probabilities. The implementation faithfully follows the recurrence relation: we fill a table best_log_prob from left to right, at each position $j$ checking all valid tokens that end at position $j$ and keeping track of which starting position gave the best score.

The Viterbi decoder finds the segmentation that maximizes the product of token probabilities. Longer, more specific tokens tend to win when their probability is high enough to beat the combination of shorter fragments. Notice that "playfully" likely gets split using whatever substrings of "playfully" appear in our corpus, even if the result looks unusual by linguistic standards. This is expected: our toy corpus does not contain "playfully" as a training word, so the algorithm must piece it together from familiar parts.

With the initial vocabulary (before EM training and pruning), the segmentations tend to use longer tokens because many substrings have been seen and have non-negligible initial probabilities. After training, the vocabulary will be more selective and the segmentations will stabilize around the most productive subwords.

Now we implement the sampling procedure that makes subword regularization possible. Instead of picking the single best path through the segmentation lattice, we sample proportionally to the probability of each path. The implementation uses a forward pass to compute the total probability flowing to each position, then samples backward from the end of the word.

The key technical challenge is numerical stability. Probabilities can become extremely small for longer words because each segmentation probability is a product of many small numbers. We work in log-space throughout and use the log-sum-exp trick to avoid underflow when summing probabilities.

When we sample multiple times, we see different segmentations appear with frequencies proportional to their probabilities. The most probable segmentation dominates, but less probable alternatives also appear. This diversity is precisely what gives subword regularization its power during model training: the downstream model must learn to handle each word in multiple representations, forcing it to build more robust semantic understanding that does not depend on a single tokenization convention.

Notice that the sampling distribution is not uniform over all segmentations. Some splits appear dozens of times while others may appear only once or twice in 200 samples. This non-uniform distribution reflects the trained token probabilities, and it ensures that the most linguistically sensible segmentations still dominate the training signal even when regularization is applied.

Let's implement a simplified EM training loop that demonstrates the vocabulary pruning process. This implementation uses Viterbi EM (sometimes called "hard EM"), where each E-step assigns all probability mass to the single best segmentation rather than distributing it across all segmentations. Hard EM is computationally simpler than full EM and often works nearly as well in practice, since the best segmentation typically carries most of the probability mass anyway.

After EM training and pruning, the vocabulary has converged to tokens that improve the corpus likelihood. Common morphemes like "ing," "ed," "er," and "un" survive because they appear in many words and provide efficient coverage of the training data. Longer tokens that were idiosyncratic to only one or two words tend to be pruned away, replaced by combinations of shorter, more general tokens.

This output illustrates the algorithm's core principle in action: a token survives if and only if its presence in the vocabulary allows the corpus to be segmented more efficiently (higher log-probability) than the vocabulary without it. Every surviving token earns its place by improving compression of the training data.

A key hyperparameter in subword regularization is $\alpha$, the temperature controlling how sharply the sampling distribution peaks around the best segmentation. Understanding how $\alpha$ affects segmentation diversity is important for applying the technique effectively in practice.

The three panels show how alpha controls the tradeoff between consistency and diversity. In this implementation, alpha multiplies each token's log-probability, so it acts as a concentration parameter rather than an ordinary temperature. At alpha = 0.1, the distribution is broad and most of the probability mass lies outside the five most common splits. At alpha = 0.7, a handful of plausible segmentations account for most samples. At alpha = 1.5, the highest-probability splits dominate even more strongly.

Values in the 0.1 to 0.7 range are common starting points for subword regularization. Within that range, a lower alpha introduces more variation, while a higher alpha stays closer to the model's preferred segmentations. The appropriate value depends on how much tokenization noise the downstream task can tolerate, so it should be validated rather than treated as a universal default.

Let's visualize how the corpus log-likelihood evolves across EM iterations. Watching the log-likelihood improve confirms that the algorithm is making progress, and the pattern of improvement reveals how pruning affects learning.

The first hard-EM update substantially improves the Viterbi objective by reallocating probability mass to the tokens used by the best segmentations. The curve then plateaus: on this tiny corpus, those best paths stabilize after one update, so later hard-EM steps repeat the same assignments. The pruning stages remove candidates with negligible probability that are not used by those paths, which is why they do not create visible drops here. A larger corpus or full soft EM would generally produce a more gradual trajectory; this compact example instead makes the rapid convergence of hard assignments explicit.

In practice, you'll use SentencePiece rather than implementing the algorithm from scratch. SentencePiece provides a battle-tested, highly optimized implementation of the unigram LM algorithm that handles production-scale corpora, multilingual text, and edge cases that would take significant effort to handle in a from-scratch implementation.

SentencePiece's sample_encode method directly supports the sampling needed for subword regularization. During training a downstream model, you call this method at each epoch instead of the deterministic encode, exposing the model to varied tokenizations of the same text. The interface is intentionally simple: you replace encode with sample_encode_as_pieces and set the alpha hyperparameter, and the regularization happens automatically.

This simplicity belies the mathematical depth underneath. Every call to sample_encode_as_pieces runs the forward-backward algorithm, computes the full distribution over segmentations with temperature scaling, and samples a path from it. The implementation is highly optimized in C++ with Python bindings, making it fast enough for use in real training pipelines without becoming a bottleneck.

The key parameters for SentencePiece unigram training are:

• vocab_size: The target vocabulary size after pruning. Typical values range from 8,000 to 64,000. Larger vocabularies handle more words as single tokens but increase model size. For multilingual models, larger vocabularies (32,000 to 250,000) are common to give each language adequate coverage.
• model_type: Set to "unigram" for the probabilistic approach. SentencePiece also supports "bpe", "word", and "char" modes. The unigram mode is recommended when you want sampling capability or when you are training multilingual models.
• character_coverage: The fraction of characters from the training corpus to include in the base vocabulary. Setting to 1.0 ensures no characters are dropped, which is important for multilingual text. For corpora with rare Unicode characters, you may use 0.9995 to exclude extremely uncommon symbols.
• pad_id, unk_id, bos_id, eos_id: Integer IDs reserved for special tokens (padding, unknown, beginning-of-sequence, end-of-sequence). Setting these explicitly ensures consistent token IDs across models trained on different corpora.

SentencePiece's unigram mode is used in production in several important ways that go beyond what simple demonstrations reveal. Understanding these production characteristics helps you use the library effectively in real systems.

The algorithm operates directly on raw Unicode text without requiring any pre-tokenization step (unlike BPE implementations that often require whitespace pre-splitting). This is a significant advantage for multilingual text, where the boundary between "words" may not be meaningful, such as in Chinese, Japanese, or Thai. In these languages, there are no spaces between words, so any tokenizer that requires word boundaries as input must first run a separate word segmentation step. SentencePiece treats the entire character sequence as its input, learning subword patterns that respect the linguistic structure of each language automatically.

SentencePiece introduces a special whitespace marker, the character ▁ (U+2581, a lower one-eighth block), to encode word boundaries within the token stream. A token beginning a new word is prefixed with this character, while continuation tokens within a word have no prefix. This allows the tokenizer to reconstruct the original whitespace from the token sequence without storing it separately. For example, the sentence "hello world" might tokenize as ["▁hello", "▁world"], and you can recover the original text with spaces by stripping the ▁ markers and inserting spaces before each one.

For multilingual models like mBART, mT5, and XLM-RoBERTa, SentencePiece unigram mode trains a single shared tokenizer across all languages simultaneously. By training on multilingual data with appropriate character coverage settings, the resulting vocabulary naturally allocates token capacity to each language proportional to its representation in training data. Languages with simpler morphology (like English) tend to get many whole-word tokens, while languages with complex morphology (like Finnish or Turkish) get more subword tokens that compose into the rich inflectional paradigms those languages use.

The whitespace handling and language-agnostic character coverage combine to make SentencePiece unigram mode the preferred choice for multilingual NLP. When you see a model like mBART described as using "SentencePiece with 250,000 vocabulary," that vocabulary supports over 100 languages with appropriate coverage of their alphabets and common morphological patterns.

The unigram LM approach is principled and powerful, but it comes with real limitations that you should understand before applying it to new problems.

The most significant practical limitation is computational cost during tokenizer training. The EM algorithm requires multiple passes over the full corpus, and each pass involves computing segmentation probabilities for every word. For large corpora with hundreds of gigabytes of text, this is expensive. BPE, by contrast, can be trained with simple frequency counting that scales to large corpora more easily. In practice, this means unigram LM tokenizers are most often trained on a representative subset of the full training data rather than the entire corpus. The assumption is that a well-sampled subset captures the same morphological patterns as the full corpus, which is usually true for morphological tokenization but may not hold for rare vocabulary in very domain-specific text.

A subtler limitation is the independence assumption at the heart of the model. Treating each token's probability as independent of its neighbors is mathematically convenient but linguistically false. The probability of seeing the suffix "-ness" in reality depends heavily on what precedes it: it follows adjectives but not verbs, and it follows certain stems but not others. This means the model's probability estimates are inaccurate in absolute terms, even if they produce useful relative rankings between segmentations. For most applications this doesn't matter since we only care about which segmentation has the highest probability, but it means the absolute probability values assigned to segmentations should not be interpreted as true linguistic probabilities. If you use the probabilities for downstream tasks like calibration or uncertainty estimation, be aware of this systematic bias.

The vocabulary pruning strategy also introduces a failure mode worth understanding. The algorithm scores tokens by their marginal impact on the corpus log-likelihood, but this marginal impact depends on which other tokens are present in the vocabulary. Removing a token changes the optimal segmentations of words that used it, which changes the effective value of other tokens, which can cause the pruning to make suboptimal decisions in some cases. The prune-then-EM loop mitigates this by re-estimating probabilities after each pruning step, but it does not guarantee globally optimal pruning decisions. The final vocabulary is a local optimum, not necessarily the global optimum.

Despite these limitations, the unigram LM approach has had wide use on production NLP systems. Its main contribution goes beyond the training algorithm: it introduced probability-based tokenization that enables subword regularization. Research on multilingual models consistently finds that models trained with subword regularization are more less brittle, particularly on low-resource languages and text that differs from the tokenizer's training distribution. The SentencePiece library made these ideas practical and accessible, and it remains widely used in production systems years after its introduction.

The approach also changed how researchers think about tokenization itself. Before the unigram LM, tokenizers were viewed as deterministic preprocessing steps: you run the tokenizer, you get tokens, you move on. The probabilistic framing introduced by the unigram LM model encouraged thinking of tokenization as a modeling choice with uncertainty, opening the door to more flexible approaches to text representation. This philosophical shift influenced later work on learned tokenization, byte-level models, and tokenization-free architectures, all of which grapple with the same fundamental question of what the right granularity for text representation is. The unigram LM's answer, that there is no single right granularity and that maintaining a distribution over segmentations is more principled than committing to one, remains relevant in current research.

The approach also benefits for low-resource and morphologically rich languages. Languages like Turkish, Finnish, and Swahili have agglutinative morphology where a single word form can carry the meaning of an entire English phrase. A deterministic tokenizer trained on a balanced multilingual corpus may consistently make suboptimal choices for these languages, systematically over-segmenting or under-segmenting based on the dominant language's statistics. Subword regularization, by exposing the downstream model to multiple segmentations, allows it to learn to handle each language's morphological patterns even when the tokenizer was not perfectly trained for that language. This robustness property is especially valuable in production systems that must handle multilingual input at scale.

The unigram language model tokenization algorithm takes a top-down, probabilistic approach to building subword vocabularies. Rather than constructing a vocabulary through merges (as in BPE and WordPiece), it starts with a large candidate vocabulary and prunes it using the EM algorithm, retaining tokens that maximize the corpus log-likelihood. The result is a vocabulary where every token earns its place by improving the model's ability to compress the training data.

The key takeaways are:

• Probabilistic foundation: Every token in the vocabulary has an associated probability, which allows computing a probability distribution over all possible segmentations of any word. This is the central conceptual contribution that distinguishes the algorithm from BPE and WordPiece.
• EM training: The algorithm alternates between computing expected token counts (E-step) and updating probabilities (M-step), then prunes the lowest-impact tokens until reaching the target vocabulary size. Each EM iteration is guaranteed to improve the corpus log-likelihood. This provides a reliable convergence signal.
• Viterbi decoding: At inference time, dynamic programming efficiently finds the maximum-probability segmentation of any word in $O(n^2)$ time, making inference fast despite the complexity of the underlying probability model.
• Subword regularization: By sampling from the segmentation distribution instead of always taking the best, downstream models are exposed to diverse tokenizations during training, improving robustness to tokenization variation. The temperature parameter $\alpha$ controls the tradeoff between consistency and diversity.
• SentencePiece integration: The SentencePiece library provides a production-ready implementation that handles raw Unicode text directly, uses a whitespace marker to encode word boundaries, and supports multilingual vocabulary training across more than 100 languages.
• Limitations to keep in mind: Training cost is higher than BPE, the independence assumption means absolute probabilities are not linguistically accurate, and the pruning process may reach local rather than global optima. These limitations are acceptable for most applications but worth understanding when diagnosing unexpected behavior.

Building on the vocabulary-problem motivation from earlier in this part and the BPE and WordPiece algorithms we have already covered, the unigram LM completes our survey of the core subword tokenization approaches. Each algorithm represents a different philosophy: BPE is simple and fast (frequency-driven merges), WordPiece is likelihood-aware (merge selection based on probability), and unigram LM is fully probabilistic (top-down pruning with a principled objective). The next chapter covers SentencePiece in depth, exploring how it unifies BPE and unigram LM training under a single framework and handles the practical challenges of multilingual tokenization at production scale.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about the unigram language model tokenization algorithm.