How Unigram Tokenizers Prune Vocabularies in AI

The Unigram language model tokenizer is a unique subword tokenization method used in artificial intelligence preprocessing that operates from the top down. Unlike other popular algorithms that build a vocabulary from scratch by merging characters, the Unigram approach begins with a massive initial vocabulary of potential subwords and iteratively removes tokens. This pruning process is guided by a probabilistic language model that evaluates the importance of each subword, ensuring that the removal of specific tokens minimizes the overall loss of information across the training corpus until a target vocabulary size is reached.

Starting with a Massive Vocabulary

Most tokenization algorithms, such as Byte-Pair Encoding (BPE) and WordPiece, are bottom-up approaches. They start with single characters and iteratively merge the most frequent pairs. The Unigram tokenizer reverses this logic.

To begin, the algorithm constructs a massive seed vocabulary from the training corpus. This initial list includes all individual characters to ensure that the model never encounters an out-of-vocabulary word. On top of these base characters, the algorithm adds the most common substrings and whole words found in the text, often utilizing the Enhanced Suffix Array algorithm to find frequent substrings quickly. This initial vocabulary is intentionally much larger than the final desired size, often containing hundreds of thousands of candidate tokens.

Calculating Probabilities and Tokenization

The core assumption of the Unigram model is that the occurrence of each subword token is independent of the tokens that come before it. Because of this, the probability of a specific tokenization sequence for a word is simply the product of the probabilities of the individual subwords that make up that sequence.

During preprocessing, a single word can be segmented into subwords in multiple ways using the current vocabulary. To find the best segmentation for a word, the algorithm uses the Viterbi algorithm. The Viterbi algorithm efficiently calculates and identifies the single most probable sequence of subwords without having to brute-force every possible combination. Meanwhile, the probabilities of the tokens themselves are estimated using the Expectation-Maximization (EM) algorithm, which maximizes the likelihood of the training corpus given the current vocabulary.

Evaluating Corpus Loss

Once the probabilities of all tokens in the massive vocabulary are established, the tokenizer must decide which subwords are redundant and can be discarded. This is done by calculating the "corpus loss."

In this context, loss is defined as the negative log-likelihood of the entire training corpus. The algorithm iterates through the current vocabulary and temporarily removes one token at a time. It then recalculates how the overall likelihood of the corpus changes without that specific token.

If a removed token can be easily replaced by smaller subwords without significantly dropping the probability of the text, the loss associated with removing that token is very low. Conversely, if removing a token forces the model to use a sequence of rare, low-probability characters to represent a common word, the loss will be very high.

Iterative Pruning to the Target Size

After calculating the potential loss for every token in the vocabulary, the Unigram algorithm ranks them. It then prunes the vocabulary by permanently deleting the top \(p\) percent (typically 10% to 20%) of tokens that cause the smallest increase in corpus loss. Fundamental single characters are protected from this pruning process so that the tokenizer can always fall back on character-level segmentation if needed.

Once the least useful tokens are removed, the remaining vocabulary is smaller, and the statistical landscape of the corpus changes. Therefore, the algorithm loops back: it re-runs the Expectation-Maximization algorithm to update the probabilities of the surviving tokens, recalculates the loss for each token, and prunes another percentage of the vocabulary. This iterative cycle of estimating and pruning continues until the vocabulary shrinks to the exact target size predefined by the AI engineer.

Advantages in AI Preprocessing

By starting large and whittling down based on probability rather than simple frequency, the Unigram tokenizer produces a highly efficient and linguistically meaningful vocabulary. Furthermore, because it is based on a probabilistic model, it enables "subword regularization" during AI training. Instead of always feeding the neural network the single best tokenization of a word, the Unigram model can sample different valid segmentations based on their probabilities, making the resulting AI model more robust to noise, typos, and morphological variations.