Unigram LM: Fundamentals & EM Training
- Unigram LM is a probabilistic language model that assigns independent probabilities to each token, enabling robust subword segmentation.
- It employs an Expectation–Maximization algorithm to maximize corpus likelihood by effectively estimating token distributions.
- The model supports morphology-sensitive vocabulary discovery and, via dynamic programming, offers advantages over heuristic methods like BPE.
A unigram LLM (Unigram LM) is a probabilistic model that assigns a probability to a string based on the assumption that each token (typically a subword or word segment) is generated independently from a fixed vocabulary distribution. This non-contextual model plays a central role in modern subword tokenization algorithms—most notably in the Unigram LM tokenizer—which underpins many NLP pipelines for both training and inference, especially in neural LLM pretraining. Crucially, the Unigram LM formalism, objective, and estimation procedure support morphology-sensitive vocabulary discovery and robust segmentation, setting it apart from purely heuristic algorithms such as Byte Pair Encoding (BPE) (Bostrom et al., 2020, Land et al., 14 Dec 2025).
1. Probabilistic Formulation and Objective
The Unigram LM defines a probability distribution over a vocabulary of subword types. Given a segmentation of a text, the joint probability is
where for and (Bostrom et al., 2020, Land et al., 14 Dec 2025). The model's objective is to maximize the likelihood of an unlabeled corpus :
where
and 0 denotes all valid segmentations of 1 into tokens in 2 (Land et al., 14 Dec 2025).
At decoding time, segmentation of an unseen string 3 seeks the covering 4 that maximizes 5, tractably computed via Viterbi-style dynamic programming (Bostrom et al., 2020).
2. Expectation–Maximization Training and Vocabulary Pruning
Direct maximization of the Unigram LM objective is intractable due to the exponential number of segmentations. Parameter estimation proceeds by Expectation–Maximization (EM):
- E-Step: Compute the expected count 6 for each token 7 over all segmentations weighted by their posterior probability under current 8.
- M-Step: Update $\theta_{t}{\text{new}} = C_{t} / \sum_{u \in