Learn AI Create the Future — with Pranay
Main site

Tokeniser Internals

The unigram tokeniser

Unigram works the opposite way to BPE. It starts with too many pieces, gives each a probability, and prunes the ones the corpus can afford to lose.

Read these first

On this page 7
  1. Every piece carries a probability
  2. Why "unigram"
  3. The training loop, in plain terms
  4. The bonus: more than one answer
  5. Where you have seen this
  6. Remember this
  7. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

A unigram tokeniser starts with far too many pieces, scores how useful each one is, and throws away the least useful. Then it repeats.

Think of packing for a long trip. You do not start with an empty bag and add items one by one. You pile everything you own on the bed, then remove whatever you can manage without.

BPE packs the empty bag. Unigram empties the overfull one.

The result is different in a way that matters. Removing from a full pile lets you judge each item against everything else already there. Adding one at a time means you can never take back an early bad choice.

Every piece carries a probability

The other big difference: each piece in a unigram vocabulary has a number attached, saying how likely it is.

That gives the tokeniser something BPE cannot do. Faced with a word that can be cut several ways, it can score every cut and pick the best one.

   the word "lowest", with this vocabulary:

     ["lowest"]              probability 0.010
     ["low", "est"]          probability 0.008
     ["lo", "west"]          probability 0.001
     ["l","o","w","e","s","t"] probability 0.000000004

   pick the highest

BPE has no such choice. It replays a fixed list of glue operations and whatever comes out, comes out.

Why "unigram"

A unigram model is the simplest possible language model. It assumes each piece is independent of the others. So the probability of a whole cut is the probability of each piece multiplied together.

That assumption is plainly false about language. It is fine here, because the tokeniser is choosing between cuts of the same short word, not modelling meaning.

The training loop, in plain terms

  1. Start with a huge candidate list: every common substring in the text.
  2. Work out how much the text likes each candidate.
  3. For each candidate, ask: if I removed this, how much worse would the text get?
  4. Remove the worst ten to twenty percent.
  5. Repeat until the vocabulary is the size you asked for.

Step 3 is what makes this careful. A piece survives if the text genuinely needs it, not because it appeared early.

The bonus: more than one answer

Because every cut has a probability, the tokeniser can sample instead of always picking the best.

That turns out to be useful during training. Showing a model the same word cut two different ways teaches it that both cuts mean the same thing. This is called subword regularisation, and it measurably helps low-resource translation.

Where you have seen this

  • T5 and Flan-T5, through SentencePiece.
  • ALBERT, XLNet and mT5.
  • Llama's original tokeniser, which used SentencePiece with the BPE algorithm rather than unigram — the library and the algorithm are separate choices.
  • Most multilingual models, where unigram handles many scripts more evenly.

Remember this

  • Unigram starts with too many pieces and prunes; BPE starts with none and merges.
  • Every piece has a probability, so whole segmentations can be scored and compared.
  • That also allows sampling different cuts of the same word, which helps training.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install tokenizers

Written against tokenizers 0.22.2. The scoring and Viterbi parts need no libraries at all.

Scoring segmentations, and finding the best one efficiently

unigram.py
import math

# A unigram tokeniser is a vocabulary where every piece carries a probability.
VOCAB = {"low": 0.10, "lo": 0.05, "l": 0.02, "o": 0.03, "w": 0.04,
         "est": 0.08, "es": 0.02, "e": 0.06, "s": 0.05, "t": 0.05,
         "lowest": 0.01, "west": 0.02}
LOGP = {p: math.log(v) for p, v in VOCAB.items()}

def all_segmentations(word, pieces):
    if not word:
        yield []
        return
    for i in range(1, len(word) + 1):
        if word[:i] in pieces:
            for rest in all_segmentations(word[i:], pieces):
                yield [word[:i]] + rest

word = "lowest"
scored = sorted(((sum(LOGP[p] for p in seg), seg) for seg in all_segmentations(word, VOCAB)),
                reverse=True)
print(f"every way to cut {word!r} with this vocabulary: {len(scored)}")
for score, seg in scored:
    print(f"  log-probability {score:8.3f}   {seg}")

def viterbi(word, logp):
    """Best score for every prefix, built left to right. No enumeration needed."""
    best = [(-math.inf, None)] * (len(word) + 1)
    best[0] = (0.0, None)
    for end in range(1, len(word) + 1):
        for start in range(end):
            piece = word[start:end]
            if piece in logp and best[start][0] > -math.inf:
                cand = best[start][0] + logp[piece]
                if cand > best[end][0]:
                    best[end] = (cand, start)
    if best[-1][1] is None and len(word):
        return None, -math.inf
    out, i = [], len(word)
    while i > 0:
        j = best[i][1]
        out.append(word[j:i]); i = j
    return out[::-1], best[-1][0]

seg, score = viterbi(word, LOGP)
print(f"\nViterbi picks {seg} with score {score:.3f}")
print("same as the best of the full list:", seg == scored[0][1])
print(f"work done: {len(word)}^2 = {len(word)**2} steps, "
      f"instead of enumerating {len(scored)} segmentations")

# Unigram also gives you sampling, which BPE cannot do.
print("\nprobabilities over the segmentations (subword regularisation samples from this):")
z = sum(math.exp(s) for s, _ in scored)
for s, seg in scored[:4]:
    print(f"  {math.exp(s)/z:6.3f}  {seg}")

print("\na character missing from the vocabulary breaks the whole word:")
print("  viterbi('lowz'):", viterbi("lowz", LOGP)[0])
Output
every way to cut 'lowest' with this vocabulary: 12
  log-probability   -4.605   ['lowest']
  log-probability   -4.828   ['low', 'est']
  log-probability   -6.908   ['lo', 'west']
  log-probability   -8.740   ['lo', 'w', 'est']
  log-probability   -9.210   ['low', 'es', 't']
  log-probability  -11.107   ['low', 'e', 's', 't']
  log-probability  -11.331   ['l', 'o', 'west']
  log-probability  -13.122   ['lo', 'w', 'es', 't']
  log-probability  -13.163   ['l', 'o', 'w', 'est']
  log-probability  -15.019   ['lo', 'w', 'e', 's', 't']
  log-probability  -17.545   ['l', 'o', 'w', 'es', 't']
  log-probability  -19.442   ['l', 'o', 'w', 'e', 's', 't']

Viterbi picks ['lowest'] with score -4.605
same as the best of the full list: True
work done: 6^2 = 36 steps, instead of enumerating 12 segmentations

probabilities over the segmentations (subword regularisation samples from this):
   0.518  ['lowest']
   0.415  ['low', 'est']
   0.052  ['lo', 'west']
   0.008  ['lo', 'w', 'est']

a character missing from the vocabulary breaks the whole word:
  viterbi('lowz'): None

Reading that output

Twelve ways to cut a six-letter word. That count grows fast with word length, which is why enumerating is not an option in practice.

Fewer pieces usually wins, and not always. ['lowest'] scores -4.605 with a single piece at probability 0.01. ['low','est'] scores -4.828 with two pieces at 0.10 and 0.08. The single piece wins by a small margin. Move lowest to probability 0.005 and the two-piece cut takes over. Segmentation is a genuine optimisation, not a bias toward longer pieces.

The bottom row is the all-characters cut, at -19.442. Multiplying six small probabilities gives something astronomically unlikely. This is why unigram tokenisers avoid character-level output without any special rule against it.

Viterbi finds the same answer in 36 steps. It walks the word left to right, keeping only the best score for each prefix. The saving over enumeration grows exponentially with word length.

The probability column is what BPE cannot offer. lowest gets 0.518 and low est gets 0.415. Sampling from that distribution during training shows the model both readings of the same word.

lowz returns None. With z absent, no path reaches the end of the word. Real implementations avoid this with byte fallback, covered in the next lesson.

Training a real unigram tokeniser, offline

unigram_train.py
from tokenizers import Tokenizer, models, trainers, pre_tokenizers

CORPUS = ["low low low low low lower lower newest newest newest",
          "newest newest newest widest widest widest new new newer"] * 20

tok = Tokenizer(models.Unigram())
tok.pre_tokenizer = pre_tokenizers.Whitespace()
tok.train_from_iterator(CORPUS, trainers.UnigramTrainer(
    vocab_size=30, special_tokens=["<unk>"], unk_token="<unk>", show_progress=False))

print("vocab size:", tok.get_vocab_size())
for w in ["low", "lowest", "newest", "widen"]:
    print(f"  {w:>7} -> {tok.encode(w).tokens}")
Output
vocab size: 14
      low -> ['low']
   lowest -> ['low', 'e', 's', 't']
   newest -> ['newe', 's', 't']
    widen -> ['w', 'i', 'd', 'e', 'n']

The trainer stopped at 14 pieces despite being asked for 30, because pruning removed everything the tiny corpus could spare. widen falls back to characters rather than becoming [UNK], which is the visible difference from WordPiece on the same input.

Common mistakes

Expecting the shortest segmentation. Unigram maximises probability, not brevity. A rare long piece can lose to two common short ones.

Enabling sampling at inference. Subword regularisation belongs in training. At inference you want the deterministic Viterbi path, or your outputs stop being reproducible.

Confusing unigram with SentencePiece. SentencePiece is a library. Unigram is one of the algorithms it implements, and BPE is another. Llama used SentencePiece with BPE, T5 used SentencePiece with unigram.

Setting the vocabulary size above what the corpus supports. As the output shows, the trainer stops early. Check get_vocab_size() after training rather than assuming you got what you asked for.

Try it yourself

Change VOCAB["lowest"] from 0.01 to 0.005 and rerun. The winning segmentation flips to ['low', 'est']. Then raise VOCAB["west"] until ['lo', 'west'] wins. Watching the ranking move as probabilities change is the fastest way to understand what the algorithm is actually optimising.

What to learn next

Researcher — Mathematics and papers.

The model

Kudo (2018) defines a unigram language model over subword sequences. For a segmentation $\mathbf{x} = (x_1, \dots, x_M)$ of a sentence:

$$ P(\mathbf{x}) = \prod_{i=1}^{M} p(x_i), \qquad \sum_{x \in \mathcal{V}} p(x) = 1 $$

The best segmentation of a string $X$ is

$$ \mathbf{x}^{*} = \arg\max_{\mathbf{x} \in S(X)} P(\mathbf{x}) $$

over the set $S(X)$ of segmentations consistent with vocabulary $\mathcal{V}$. This is computed by Viterbi in $O(|X|^2)$ with a naive scan, or $O(|X| \cdot L_{\max})$ where $L_{\max}$ bounds piece length, which is the practical implementation.

Training

The vocabulary and the probabilities are estimated jointly. Since $\mathcal{V}$ is latent, the procedure alternates:

  1. Seed. Build a large candidate set, typically all substrings up to a length limit, pruned by frequency. The reference implementation uses the Enhanced Suffix Array to enumerate frequent substrings efficiently.
  2. E step. With $\mathcal{V}$ fixed, estimate $p(x)$ by EM, maximising the marginal likelihood $\sum_s \log \sum_{\mathbf{x} \in S(s)} P(\mathbf{x})$ over the corpus.
  3. Prune. For each piece $x$, compute $\mathrm{loss}(x)$: the drop in corpus log-likelihood if $x$ were removed and every sentence re-segmented without it. Drop the bottom $\eta$ percent, typically 10 to 20.
  4. Repeat 2 and 3 until $|\mathcal{V}|$ reaches the target.

Characters are always retained so the model can, in principle, segment anything in its alphabet.

Subword regularisation

Because $P(\mathbf{x})$ is a distribution, one can sample from the $l$-best segmentations with a temperature:

$$ P(\mathbf{x}_i \mid X) \approx \frac{P(\mathbf{x}i)^{\alpha}}{\sum{j=1}^{l} P(\mathbf{x}_j)^{\alpha}} $$

with $\alpha$ controlling sharpness ($\alpha \to \infty$ recovers Viterbi, $\alpha \to 0$ gives uniform over the $l$-best). Sampling is done efficiently with Forward-Filtering Backward-Sampling on the lattice, avoiding explicit $l$-best enumeration.

Kudo reports consistent BLEU gains across language pairs, largest in low-resource settings, and improved robustness to noisy input. Provilkov et al. (2020) achieve a comparable effect for BPE by dropping merges at random, and report similar gains.

Unigram versus BPE, empirically

Bostrom and Durrett (2020) hold everything else fixed and compare BPE and unigram tokenisation for pretraining. Findings:

  • Unigram segmentations align better with morphology in both English and Japanese, judged against gold morphological analyses.
  • Downstream results on GLUE, SQuAD and other tasks favour unigram, modestly but consistently.
  • The advantage is attributed to the pruning objective evaluating pieces against the whole vocabulary, where BPE's greedy merging cannot revise early decisions.

Despite that, byte-level BPE dominates decoder-only production models. The reasons are practical: it composes with the byte-level guarantee that eliminates unknown tokens, encoding is faster, and the ecosystem standardised on it.

Complexity

  • Training: dominated by repeated EM over the lattice. Each iteration is $O(N \cdot L_{\max})$ over corpus size $N$, with $O(\log(|\mathcal{V}_{\text{seed}}| / |\mathcal{V}|))$ pruning rounds at 20 percent per round. Substantially more expensive than BPE training.
  • Encoding: $O(|X| \cdot L_{\max})$ per word by Viterbi, against $O(m \log m)$ for BPE with a rank heap. Comparable in practice.
  • Storage: one float per piece in addition to the string, so slightly larger than a BPE merge table.

Papers

What to learn next

Related terms

What to learn next

These follow on from what you just read.

  • Tokeniser Internals

    SentencePiece

    SentencePiece treats the space as an ordinary character, so it works on languages that do not use spaces and can rebuild the exact original text.

  • Tokeniser Internals

    Byte-level BPE

    Convert text to raw bytes first, then run BPE. With 256 starting symbols the tokeniser can represent literally anything, and unknown tokens become impossible.

  • Tokeniser Internals

    Training a tokeniser on your own text

    A tokeniser has five separate stages, and each one changes the result. Here is the whole pipeline built from scratch, saved, and loaded into transformers.