Unigram Language Model Tokenization › Sampling Segmentations: Subword Regularization

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.