Vocabulary trimming architecture — embedding budgets, keep sets, merge closure, byte fallback

Why architecture matters here

The first reason this matters architecturally is that it attacks the one cost that does not shrink when you shrink the model. Distillation, pruning, and quantisation all reduce the layers. None of them touch V. Quantisation reduces the embedding table's bytes-per-parameter but keeps every row; a 4-bit 128k-token table is still 128k rows wide. Vocabulary trimming is the only technique in the compression toolbox that removes rows, and it therefore composes with all of them: a trimmed, quantised model is smaller than either alone, multiplicatively.

The second is that the embedding table has an unusually bad memory access profile, which makes its size hurt more than the parameter count suggests. It is accessed by random index — one row per token, scattered — so it has no locality worth speaking of and every lookup is a potential cache miss. On a device with a small cache and a memory-bandwidth-bound decoder, a table that does not fit is not just occupying space; it is a source of stalls in the hot path. Shrinking it can improve latency in ways the parameter count does not predict.

The third is the output side, which is often forgotten and is frequently the larger win. The LM head projects the hidden state to V logits, and unlike the embedding lookup — which touches one row — this is a full d x V matrix multiply executed on every single generated token. On a small model with a large vocabulary the head can be a substantial fraction of per-token decode FLOPs, and the softmax over 128k entries is a real cost on weak hardware. Halving V halves that work directly, every token, forever.

The fourth is that this is a decision about product scope expressed in the artefact, and that is what makes it risky as well as valuable. Trimming to English encodes 'this product is English-only' into the weights. That is fine when it is true and it is a quiet catastrophe when it stops being true — a locale expansion that would have been a configuration change becomes a re-quantisation, re-validation, and re-release of every model artefact. Understanding the technique means understanding that you are trading a memory saving today for a coupling between the model artefact and the market it serves.

The architecture: every piece explained

The corpus survey is the empirical foundation. You take a body of text representative of what the deployment will actually process — real prompts, real documents, real outputs, in the real language mix — tokenize all of it with the original tokenizer, and count which IDs fire. What comes back is invariably a brutal power law: a few thousand tokens carry almost all the mass, a long tail fires occasionally, and a large fraction of the vocabulary never appears at all. That never-appears set is the prize, and its size is why the technique is worth the trouble.

The keep set is the survey plus judgement, and the judgement is where the safety lives. It is the union of observed tokens; every special token (BOS, EOS, PAD, and any chat-template markers, which are load-bearing and whose absence breaks the model in confusing ways); the complete byte-level fallback alphabet; and a deliberate margin of tail tokens the survey did not see but plausibly would. That margin is cheap insurance — it costs a few thousand rows and it is what stops an unlucky corpus from becoming a production quality regression.

The byte fallback set deserves its own paragraph because it is the structural guarantee that makes the whole technique safe rather than merely usually-fine. Byte-level BPE tokenizers can represent any byte sequence as a sequence of single-byte tokens. If you keep all 256 of those, then no string is unencodable, no matter how aggressively you trimmed everything else. Unexpected Cyrillic does not throw; it tokenizes into bytes. It tokenizes badly — many tokens per character, burning context and producing text the model handles poorly — but it degrades instead of failing. Keep the bytes.

The ID remap is the mechanical core. Keeping 32,000 of 128,000 tokens means the survivors need contiguous IDs 0..31999, because IDs are tensor row indices and a sparse index set is not a tensor. So you build an old-to-new mapping, and then everything that referred to a token by ID must be rewritten through it: the embedding rows, the head rows, the head bias if there is one, the tokenizer's own vocabulary file, the merge table, the special token IDs in the config, and any generation parameters that name IDs — eos_token_id, suppressed tokens, logit bias maps. Miss one and you have a model that runs and is subtly, silently wrong.

The merge table repair is the subtlety that separates a working trim from a broken one, and it is where naive implementations fail. A BPE tokenizer is not a lookup table; it is a program. It encodes by starting from bytes and applying ranked merge rules, and a rule like ('to', 'ken') -> 'token' is only reachable if 'to' and 'ken' both exist. Delete an intermediate token that no survey ever saw as a final token, and you have silently broken every merge path that passed through it. Frequent tokens stop forming. The tokenizer still works — it just produces different, longer output than it used to. Merge rules whose inputs or output you dropped must be deleted, and any token whose only role is as a stepping stone to a kept token must itself be kept.

End-to-end flow

A trim run begins with the corpus, and the corpus is the thing that determines whether the artefact is any good. You assemble text that matches production in language mix, domain, and register — including the model's own outputs, not just its inputs, since the head is being trimmed too and the model must remain able to emit every token it needs. Then you tokenize the lot with the original tokenizer and accumulate a frequency table over IDs. This is an embarrassingly parallel pass over text and is fast even on a large corpus.

Next you construct the keep set: observed IDs, union all specials, union the 256 byte tokens, union a margin from the frequency tail. Then you close it under the merge graph — walk the merge rules and, for every rule producing a kept token, ensure both inputs are kept, recursively. This closure step typically pulls in a handful of tokens the survey never saw, which is exactly correct: they are structural, not statistical. The result is a keep set that is both empirically justified and internally consistent.

Now the remap. Sort the keep set — usually by original ID, which keeps the low-numbered specials where tooling expects them — and assign new contiguous IDs. Materialise this as an explicit array: new_to_old[new_id] = old_id. That single array is the entire transformation, and it is worth serialising alongside the model, because it is what lets you debug a suspicious token six months later by mapping it back to what it used to be.

The tensor surgery is then almost anticlimactic. E_new = E[new_to_old] — one fancy-index operation produces the new embedding table with rows in the right order. W_new = W[new_to_old] for the head, and the bias likewise if present. If the model ties embeddings and head weights, you slice once and re-tie, and you must verify the tie survived the operation, because a framework that silently untied them has just doubled the size of the artefact you were trying to shrink.

Then the tokenizer, which is where the care is required. Rewrite the vocab mapping to the new IDs. Filter the merge table to rules whose inputs and output all survive, preserving rank order, because BPE applies merges by rank and reordering them changes the tokenization. Rewrite every special token ID in every config file. Rewrite generation config. At this point the artefact is complete and it is also completely unvalidated.

Validation is the step nobody budgets for and everybody needs. Round-trip a held-out corpus: encode with the trimmed tokenizer, decode, confirm you get the original text back byte for byte — this catches merge damage immediately. Measure fertility, the average tokens per word, against the original tokenizer on the same text; a rise means merge paths broke and every subsequent cost — context, latency, quality — rose with it. Then run the model's real evaluations, because the only claim that matters is that quality held. And run a deliberate out-of-scope check: feed it a language you trimmed away and confirm it produces byte-fallback garbage rather than an exception, because that is the difference between a model with a narrow scope and a model with a crash.

Failure modes and mitigations

• Unrepresentative survey corpus. The keep set is inferred from data, so a corpus that misses a register — user typos, product names, the code blocks in support tickets — deletes tokens production needs. The failure appears in production as quiet quality loss on the exact inputs you did not think to sample.
• Broken merge paths. Dropping an intermediate token that the survey never saw as final silently disables every merge through it. The tokenizer does not error; it just produces longer sequences. Fertility measurement is the only thing that catches this, and only if you run it.
• Missing byte fallback. Without the full byte alphabet, an unexpected input is unencodable and you get an exception or a silent replacement character in a production request. The 256 rows this costs are the cheapest insurance in the entire pipeline.
• Stale special token IDs. The config still says eos_token_id: 128001 and that row is gone or now means something else. The model generates without ever stopping, or stops instantly. The symptom looks nothing like a tokenizer bug.
• Untied weights after slicing. Models that tie embeddings and head can lose the tie during surgery. The artefact is twice the expected size and, worse, the two copies can be updated independently by any subsequent fine-tune. Assert the tie explicitly after slicing.
• Locale expansion after trimming. The product adds a language and the model physically cannot represent it efficiently. This is not a bug, it is the trade you made, and it is only a crisis if nobody wrote it down. Record the trim scope in the artefact's manifest.
• Trimming a model you will later fine-tune. Fine-tuning data may contain tokens the trim removed. The training run either errors on unknown IDs or quietly byte-falls-back and learns from a degraded representation. Decide the order deliberately: trim last.

Operational playbook

• Version the survey corpus with the artefact. The keep set is a function of the corpus, so the corpus is an input to the model. Hash it, store it, and record which corpus produced which trim. Re-trimming without it is guesswork.
• Ship the new_to_old map alongside the model. It is small and it is the only thing that makes a trimmed model debuggable. Without it, a suspicious token ID in a production log is unanswerable.
• Gate on fertility, not just on evaluations. Compare average tokens-per-word against the untrimmed tokenizer on held-out text and fail the build on a regression beyond a threshold. Fertility is the earliest and clearest signal that merge repair went wrong.
• Monitor byte-fallback rate in production. The fraction of requests that hit byte tokens is the live measure of how wrong your keep set is. A rising rate means the input distribution is moving away from what you trimmed for — an early warning with weeks of lead time.
• Keep a generous tail margin. The parameters are cheap and the failure they prevent is expensive. Trimming to exactly the observed set is optimising the wrong variable; a few thousand extra rows costs megabytes and buys robustness against a corpus you sampled imperfectly.
• Round-trip test on every build. Encode-decode over a held-out corpus with byte-exact comparison is fast, deterministic, and catches the entire class of merge-table damage. There is no reason for it not to be a unit test.
• Record the trim scope in the model card. 'English and code, trimmed to 32k' is a product-visible property. It belongs where the next team will find it, not in the pipeline script that produced it.