Unigram Language Model Tokenization › Viterbi Decoding: Finding the Best Segmentation

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:

$$\text{best}[j] = \max_{i < j, \, w[i:j] \in V} \left( \text{best}[i] + \log P(w[i:j]) \right)$$

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.