Unigram Language Model Tokenization › Code Implementation
Language AI HandbookAll chapters Unigram Language Model Tokenization › Code ImplementationOne section · Sections in this chapter7 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
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.
Building the Vocabulary and Initial Probabilities
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."
import collections
from typing import Dict, List
def get_all_substrings(
word: str, min_len: int = 1, max_len: int = 10
) -> List[str]:
"""Extract all substrings of a word up to max_len."""
substrings = []
n = len(word)
for start in range(n):
for end in range(start + min_len, min(start + max_len + 1, n + 1)):
substrings.append(word[start:end])
return substrings
# Small training corpus
corpus = [
"happy",
"unhappy",
"happiness",
"unhappiness",
"happily",
"running",
"runner",
"runs",
"run",
"ran",
"play",
"playing",
"player",
"played",
"plays",
"walk",
"walking",
"walked",
"walker",
"walks",
"learn",
"learning",
"learned",
"learner",
"learns",
]
# Count word frequencies (all appear once in this toy corpus)
word_freq = collections.Counter(corpus)
# Build initial vocabulary from all substrings in all words
substring_counts: Dict[str, int] = collections.defaultdict(int)
for word, freq in word_freq.items():
for substr in get_all_substrings(word):
substring_counts[substr] += freq
# Always include individual characters
for word in corpus:
for char in word:
substring_counts[char] = substring_counts.get(char, 0) + word_freq[word]
# Initialize probabilities from counts (normalize so they sum to 1)
total_count = sum(substring_counts.values())
vocab_probs: Dict[str, float] = {
tok: count / total_count for tok, count in substring_counts.items()
}Initial vocabulary size: 227 Sample tokens by probability: 'a': 0.05825 'n': 0.05548 'l': 0.04438 'p': 0.04161 'e': 0.03883 'r': 0.03883 'y': 0.02219 's': 0.02219 'i': 0.01942 'u': 0.01664
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.
Viterbi Segmentation
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 checking all valid tokens that end at position and keeping track of which starting position gave the best score.
import math
from typing import Dict, List, Tuple
def viterbi_segment(
word: str, vocab_probs: Dict[str, float]
) -> Tuple[List[str], float]:
"""
Find the most probable segmentation of a word using Viterbi decoding.
Returns (best_tokens, best_log_prob).
"""
n = len(word)
# best_log_prob[j] = highest log-prob segmentation of word[:j]
best_log_prob = [-math.inf] * (n + 1)
# best_start[j] = start position of the last token in the best segmentation ending at j
best_start = [-1] * (n + 1)
best_log_prob[0] = 0.0 # empty prefix has log-prob 0
for j in range(1, n + 1):
for i in range(j):
substr = word[i:j]
if substr in vocab_probs and vocab_probs[substr] > 0:
log_p = math.log(vocab_probs[substr])
candidate = best_log_prob[i] + log_p
if candidate > best_log_prob[j]:
best_log_prob[j] = candidate
best_start[j] = i
# If we couldn't segment the word (word contains OOV characters), fall back to characters
if best_log_prob[n] == -math.inf:
return list(word), -math.inf
# Backtrack to recover the segmentation
tokens = []
j = n
while j > 0:
i = best_start[j]
tokens.append(word[i:j])
j = i
tokens.reverse()
return tokens, best_log_prob[n]Viterbi segmentation results: Word Segmentation Log-prob ---------------------------------------------------------------------- 'unhappiness' -> unhappines | s -10.389 'running' -> running -6.581 'playfully' -> p | l | a | y | f | u | l | l | y -inf 'walked' -> walked -6.581 'learner' -> learner -6.581
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.
Sampling Segmentations
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.
import random
from typing import Dict, List, Optional, Tuple
def sample_segment(
word: str,
vocab_probs: Dict[str, float],
alpha: float = 1.0,
seed: Optional[int] = None,
) -> List[str]:
"""
Sample a segmentation from the distribution over all segmentations.
alpha controls concentration: values below 1 flatten the distribution,
1 samples from the model distribution, and values above 1 sharpen it.
"""
if seed is not None:
rng = random.Random(seed)
else:
rng = random.Random()
n = len(word)
# forward[j] = total probability mass reaching position j
# We work in log-space for numerical stability
forward = [-math.inf] * (n + 1)
forward[0] = 0.0
# Store all valid token spans for backtracking
spans: List[List[Tuple[int, float]]] = [[] for _ in range(n + 1)]
for j in range(1, n + 1):
for i in range(j):
substr = word[i:j]
if substr in vocab_probs and vocab_probs[substr] > 0:
log_p = alpha * math.log(vocab_probs[substr])
spans[j].append((i, log_p))
# Forward sum (log-sum-exp for stability)
if spans[j]:
log_vals = [
forward[i] + lp
for (i, lp) in spans[j]
if forward[i] > -math.inf
]
if log_vals:
max_v = max(log_vals)
forward[j] = max_v + math.log(
sum(math.exp(v - max_v) for v in log_vals)
)
if forward[n] == -math.inf:
return list(word) # fallback
# Backward sampling
tokens = []
j = n
while j > 0:
candidates = [(i, lp) for (i, lp) in spans[j] if forward[i] > -math.inf]
if not candidates:
break
# Compute unnormalized weights
log_weights = [forward[i] + lp for (i, lp) in candidates]
max_w = max(log_weights)
weights = [math.exp(lw - max_w) for lw in log_weights]
total_w = sum(weights)
# Sample
r = rng.uniform(0, total_w)
cumulative = 0.0
chosen_i = candidates[0][0]
for (i, _), w in zip(candidates, weights):
cumulative += w
if r <= cumulative:
chosen_i = i
break
tokens.append(word[chosen_i:j])
j = chosen_i
tokens.reverse()
return tokensSampling segmentations of 'unhappiness' (alpha=0.7): unhappines | s ( 31/200) ########## un | happiness ( 29/200) ######### u | nhappiness ( 29/200) ######### unhap | piness ( 20/200) ###### unhapp | iness ( 18/200) ###### unh | appiness ( 12/200) #### unha | ppiness ( 11/200) ### u | n | happiness ( 4/200) #
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.
The EM Training Loop
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.
def compute_corpus_log_likelihood(
word_freq: Dict[str, int], vocab_probs: Dict[str, float]
) -> float:
"""Compute total log-likelihood of the corpus under current vocab."""
total_ll = 0.0
for word, freq in word_freq.items():
_, log_prob = viterbi_segment(word, vocab_probs)
if log_prob > -math.inf:
total_ll += freq * log_prob
return total_ll
def em_step(
word_freq: Dict[str, int], vocab_probs: Dict[str, float]
) -> Dict[str, float]:
"""
One EM iteration: compute expected token counts (E-step) and
update probabilities (M-step).
"""
# E-step: compute expected counts using best segmentation (Viterbi EM)
expected_counts: Dict[str, float] = collections.defaultdict(float)
for word, freq in word_freq.items():
tokens, log_prob = viterbi_segment(word, vocab_probs)
if log_prob > -math.inf:
for tok in tokens:
expected_counts[tok] += freq
# M-step: re-normalize to get updated probabilities
total = sum(expected_counts.values())
if total == 0:
return vocab_probs
new_probs = {tok: count / total for tok, count in expected_counts.items()}
# Keep all original vocab tokens (with tiny probability for unseen ones)
for tok in vocab_probs:
if tok not in new_probs:
new_probs[tok] = 1e-10
return new_probs
def prune_vocabulary(
vocab_probs: Dict[str, float],
word_freq: Dict[str, int],
prune_fraction: float = 0.20,
) -> Dict[str, float]:
"""
Remove the lowest-probability tokens (except single characters).
Always preserve single characters for guaranteed coverage.
"""
single_chars = {tok for tok in vocab_probs if len(tok) == 1}
non_chars = {
tok: p for tok, p in vocab_probs.items() if tok not in single_chars
}
n_remove = int(len(non_chars) * prune_fraction)
sorted_tokens = sorted(non_chars.items(), key=lambda x: x[1])
remove_set = {tok for tok, _ in sorted_tokens[:n_remove]}
new_vocab = {
tok: p for tok, p in vocab_probs.items() if tok not in remove_set
}
return new_vocab
# Run the EM-prune cycle
target_vocab_size = 80
vocab = dict(vocab_probs)
ll_history = []
vocab_size_history = []
for iteration in range(12):
# EM step to refine probabilities
vocab = em_step(word_freq, vocab)
ll = compute_corpus_log_likelihood(word_freq, vocab)
ll_history.append(ll)
vocab_size_history.append(len(vocab))
# Prune if still above target
if len(vocab) > target_vocab_size:
vocab = prune_vocabulary(vocab, word_freq, prune_fraction=0.20)
final_vocab_size = len(vocab)Starting vocabulary size: 227 Final vocabulary size: 72 Target vocabulary size: 80 Top 15 tokens by probability after EM training: 'happy ' 0.03846 ################### 'unhappy ' 0.03846 ################### 'happiness ' 0.03846 ################### 'unhappines ' 0.03846 ################### 's ' 0.03846 ################### 'happily ' 0.03846 ################### 'running ' 0.03846 ################### 'runner ' 0.03846 ################### 'runs ' 0.03846 ################### 'run ' 0.03846 ################### 'ran ' 0.03846 ################### 'play ' 0.03846 ################### 'playing ' 0.03846 ################### 'player ' 0.03846 ################### 'played ' 0.03846 ###################
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.
Comparing Alpha Values for Sampling
A key hyperparameter in subword regularization is , the temperature controlling how sharply the sampling distribution peaks around the best segmentation. Understanding how 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.
Log-Likelihood Convergence
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.
SentencePiece Unigram Mode
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.
import subprocess
import sys
# Install sentencepiece if needed
try:
import sentencepiece as spm
except ImportError:
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "sentencepiece", "-q"]
)
import sentencepiece as spm
import os
import tempfile
# Write a small training corpus to a file
corpus_text = "\n".join(
[
"the quick brown fox jumps over the lazy dog",
"machine learning models learn from data to make predictions",
"natural language processing enables computers to understand text",
"tokenization splits text into subword units called tokens",
"the unigram language model assigns probabilities to each token",
"byte pair encoding and wordpiece are also popular subword methods",
"transformers use subword tokenization for all text processing",
"vocabulary size balances coverage and efficiency in language models",
"rare words are split into smaller subword pieces during tokenization",
"common words often become single tokens in the vocabulary",
]
* 50
) # Repeat to give enough training data
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
f.write(corpus_text)
corpus_file = f.name
model_prefix = "/tmp/unigram_demo"
# Train SentencePiece with unigram mode
# Note: vocab_size is capped by the number of unique substrings in the corpus.
# We use 100 here to stay within what this toy corpus can support.
spm.SentencePieceTrainer.train(
input=corpus_file,
model_prefix=model_prefix,
vocab_size=100,
model_type="unigram",
character_coverage=1.0,
pad_id=0,
unk_id=1,
bos_id=2,
eos_id=3,
)
# Load the trained model
sp = spm.SentencePieceProcessor()
sp.load(f"{model_prefix}.model")
os.unlink(corpus_file)# Tokenize using best segmentation (Viterbi)
test_sentences = [
"tokenization is fundamental",
"subword regularization improves robustness",
"transformers process text efficiently",
]
results_best = []
for sentence in test_sentences:
tokens = sp.encode(sentence, out_type=str)
ids = sp.encode(sentence, out_type=int)
results_best.append((sentence, tokens, ids))SentencePiece tokenization (best segmentation): Input: 'tokenization is fundamental' Tokens: ['▁tokenization', '▁', 'i', 's', '▁', 'f', 'u', 'n', 'd', 'a', 'm', 'e', 'n', 't', 'al'] IDs: [22, 4, 16, 5, 4, 17, 28, 8, 59, 15, 7, 6, 8, 27, 66] Input: 'subword regularization improves robustness' Tokens: ['▁subword', '▁', 'r', 'e', 'g', 'ular', 'iz', 'a', 'tion', '▁', 'i', 'mp', 'ro', 'v', 'e', 's', '▁', 'ro', 'b', 'u', 's', 't', 'n', 'e', 's', 's'] IDs: [9, 4, 14, 6, 18, 85, 92, 15, 87, 4, 16, 94, 48, 96, 6, 5, 4, 48, 68, 28, 5, 27, 8, 6, 5, 5] Input: 'transformers process text efficiently' Tokens: ['▁', 't', 'ra', 'ns', 'for', 'm', 'ers', '▁p', 'ro', 'ces', 's', '▁text', '▁e', 'f', 'f', 'ic', 'i', 'e', 'n', 't', 'l', 'y'] IDs: [4, 27, 31, 58, 57, 7, 25, 41, 48, 78, 5, 23, 60, 17, 17, 19, 16, 6, 8, 27, 99, 20]
# Demonstrate sampling: same sentence, different segmentations
# sample_encode_as_pieces(text, nbest_size, alpha)
# nbest_size=-1 means sample from the full distribution; alpha=0.7 is the temperature.
sample_results = []
sentence = "tokenization splits text into subwords"
for _ in range(5):
tokens = sp.sample_encode_as_pieces(sentence, nbest_size=-1, alpha=0.7)
sample_results.append(tokens)Sampling segmentations for: 'tokenization splits text into subwords' ['▁tokenization', '▁split', 's', '▁text', '▁into', '▁subword', 's']
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.
Key Parameters
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.