Unigram LM Tokenizer

The generative story: a bag of independently drawn pieces

Fix a vocabulary V of subword pieces and give each piece s a probability p(s) > 0 with Σ_{s ∈ V} p(s) = 1. The generative story is deliberately naive: to produce a sentence, draw pieces one at a time, independently, from that single distribution, and concatenate them. No context, no conditioning on what came before — that is the unigram assumption, and it is the whole model.

So a segmentation x = (x_1, …, x_M) has probability P(x) = ∏_{i=1..M} p(x_i), or in log space log P(x) = Σ_i log p(x_i). Be honest about what this is: it is not a language model over text. It has no idea that able un is nonsense and un able is not. It is a segmentation prior — a scoring function whose only job is to rank the ways one observed string could have been cut. That is a much smaller ambition than modelling language, and it is exactly enough. BPE, by contrast, has no probability to offer at all; its output is whatever its merge sequence produces.

The string is observed; the segmentation is not

You never see the pieces. You see the character string X, and the segmentation that produced it is a latent variable. The probability the model assigns to the string is a marginal — a sum over every way X can be cut:

P(X) = Σ_{x ∈ S(X)} ∏_{i} p(x_i), where S(X) is the set of all valid segmentations.

S(X) is enormous. Picture the string as a graph: the nodes are the n+1 character boundaries 0 … n, and there is a directed edge from t to t+|s| for every piece s ∈ V that matches the characters starting at position t. A segmentation is a path from node 0 to node n. If every substring were in the vocabulary the path count would be the number of compositions of n, namely 2^(n-1) — over half a million for a 20-character word. Enumeration is hopeless, but the sum factorizes over prefixes, which makes it a linear-time dynamic program.

One recurrence, two semirings

Let a[t] be the total probability of all segmentations of the prefix X[1..t]. Any such segmentation ends in exactly one piece, so conditioning on that last piece gives the recurrence. Swap the sum for a max and the identical recurrence returns the single best path instead of the total — marginalization and Viterbi are the same algorithm over two different semirings.

lattice: nodes 0..n = character boundaries
         edge t → t+|s| for every s ∈ V with X[t+1 .. t+|s|] = s

forward (Σ, ×):   a[0] = 1
                  a[t] = Σ_s  a[t-|s|] · p(s)              P(X) = a[n]

Viterbi (max, +):  d[0] = 0
                  d[t] = max_s ( d[t-|s|] + log p(s) )   bp[t] = argmax_s

both range over s ∈ V such that X[t-|s|+1 .. t] = s
cost: O(n · L) edges, L = longest piece (SentencePiece caps L, ~16)

Work in logs: products of thousands of small probabilities underflow float64 fast. Backpointers walk the winning path back from n to 0, and a trie over V finds the matching pieces at each position.

A lattice worked by hand

Take X = "unable" and a vocabulary fragment. The remaining 0.305 of probability mass sits on pieces that do not match this string, so it never enters the arithmetic.

p:  un .20   able .10   a .08   ab .02   b .01   u .04
    le .06   l    .03   e .10   n  .05   unable .005

a[0] = 1
a[1] = .04                                    ("u")
a[2] = .20 + .04×.05        = .202           ("un" | "u","n")
a[3] = .202 × .08          = .01616
a[4] = .202×.02 + .01616×.01 = .0042016
a[5] = .0042016 × .03      = 1.26048e-4
a[6] = .202×.10 + .0042016×.06 + 1.26048e-4×.10 + .005
     = .0202 + 2.52096e-4 + 1.26048e-5 + .005 = .0254647008 = P(X)

There are 11 paths through this lattice; enumerating all 11 and summing their products reproduces .0254647008 exactly, which is the cheapest possible test of an implementation. Viterbi takes maxima instead: un|able = .20 × .10 = .0200 wins, beating the single-piece unable = .0050 and un|ab|le = 2.4e-4.

EM: soft counts over segmentations

With the vocabulary fixed, training means choosing the p(s) that maximize Σ_X log P(X) over the corpus. The segmentation is hidden, so this is textbook EM: the E-step computes the expected number of times each piece was used, the M-step renormalizes those counts. The expected counts need a backward pass to match the forward one.

backward:  b[n] = 1     b[t] = Σ_s p(s) · b[t+|s|]   (s = X[t+1 .. t+|s|])

E-step   posterior of the edge for piece s spanning t → t+|s|:
         γ(s,t) = a[t] · p(s) · b[t+|s|] / a[n]
         c(s) = Σ_corpus Σ_t γ(s,t)        fractional, not integer

M-step   p(s) ← c(s) / Σ_{s′} c(s′)

on "unable":  c(un) = .7957   c(u) = .0080   c(unable) = .1963

Those three numbers sum to 1.0000, and they must: every path starts with exactly one of un, u, or unable, so their posteriors partition the mass. Note that unable banks .196 of a count despite losing Viterbi. That is what soft counts buy — a piece survives on its share of the probability mass, not on whether it won an argmax.

Pruning: which piece can we afford to lose?

EM tunes probabilities but never changes V. The vocabulary is chosen by deletion. Seed a large candidate set — on the order of a million frequent substrings, harvested with a suffix array — then repeatedly ask of every piece: how much corpus log-likelihood would we lose if it vanished? The exact criterion is Δ(s) = L(V) − L(V \ {s}) ≥ 0, and computing it exactly means re-running inference once per candidate, which at |V| ≈ 10^6 is out of reach.

approximation: only sentences whose Viterbi path uses s are affected;
for each, re-segment without s and sum the log-likelihood drop.

Δ(able):  best path without "able" is "unable" = .005
            Δ = log(.0200 / .0050) = log 4 = 1.386 nats

Δ(ab):    Viterbi never used "ab"        → approximation says 0
            exact marginal: P(X) .0254647 → .0252102
            Δ = log(1.010096)             = 0.010 nats

The gap is instructive. The approximation is wrong in absolute terms — ab does carry a little mass, through four losing paths — but it agrees that ab is cheap and able is not. It is trustworthy for ranking pieces, which is all pruning needs. Each round keeps roughly the top three-quarters by loss, re-runs EM on the survivors, and repeats until |V| hits the target.

The structural opposite of BPE

Lay the two side by side and almost every axis inverts. BPE starts with an alphabet and grows: V = V_0 + k after k merges, one new token each. Unigram starts with a million candidates and shrinks toward the target. BPE's trained artifact is an ordered merge list, and encoding means replaying it in rank order; unigram's is an unordered set of (piece, log p) pairs, and encoding is a search over the lattice.

The criteria differ in scope, too. A BPE merge is chosen for the local frequency of an adjacent pair and is never revisited, so a token exists because it happened to be common early. A unigram piece survives because removing it would cost the corpus likelihood as a whole — a global objective, re-evaluated every round. And the difference that matters most downstream: a merge sequence yields exactly one answer, with nothing to sample. Unigram hands you a distribution.

Sampling the posterior instead of taking the argmax

The posterior over segmentations is already sitting in the forward table: P(x | X) = P(x) / P(X) = P(x) / a[n]. For "unable" that reads un|able at .0200/.02546 = 78.5%, unable at 19.6%, un|ab|le at 0.9%. Drawing from this distribution instead of taking the Viterbi one turns the tokenizer into a stochastic data augmenter.

You can sample exactly, with no n-best truncation, by forward-filtering backward-sampling: run the forward pass once, then walk back from node n, choosing the incoming edge s with probability proportional to a[t − |s|] · p(s). That draws from the full lattice in one backward sweep. A temperature α reshapes it by sampling proportional to P(x)^α: as α → ∞ the distribution collapses onto Viterbi, at α = 1 it is the true posterior (78.5 / 19.6 / 0.9), at α = 0.2 it flattens to roughly 46 / 35 / 19, and as α → 0 it approaches uniform over paths. Ambiguity is the point: the model practices many cuts of the same string, so no representation hinges on one arbitrary boundary.

Costs and the ways it breaks

Encoding is cheap. A 100-character sentence with L = 16 relaxes at most 1,600 lattice edges, each an add-and-compare in log space — microseconds, invisible next to a forward pass. Training is where the money goes: EM over the whole corpus, several iterations per pruning round, and a memory-hungry seed set. That cost is offline and paid once, so for a CPU-bound SLM the runtime story is a wash against BPE.

Four things bite. Single characters and byte pieces must be unprunable — drop one and some string has no path at all, making P(X) = 0 and log P(X) = −∞; SentencePiece pins the character set and byte fallback precisely for this. EM finds a local optimum, so the seed set genuinely matters. The pruning loss holds every other probability fixed, which is only defensible because EM re-runs after each round. And sample at training time, decode with Viterbi at inference — a stochastic tokenizer in production makes outputs irreproducible for no benefit.