LLM Vocabulary Pruning
Adapted from
antoinelouis/mtem-pruner — a HuggingFace Space that prunes embedding model vocabularies to a single language.
This version extends the concept to causal (decoder-only) language models, using Qwen/Qwen3.5-2B as the target.
Result
A pruned model that keeps only English + Portuguese tokens:
| Metric | Original | Pruned | Reduction |
|---|
| Vocab size | 248,044 | 117,752 | 52.5% |
| Model size (disk) | 4.55 GB | 3.24 GB | 28.8% (‑1.31 GB) |
GGUF is included for 4bit. It weighs on 4bit the same as unsloth's UD-Q2_K_X but with higher bpw.
How it works
1. Frequency estimation
Each language's 1M‑sentence dataset is tokenized with the original model. Tokens that appear at least once are marked as frequent for that language.
2. BPE merge‑graph traversal
The original Space targets WordPiece‑based embedding models (e.g., multilingual-e5), which have no merge dependencies — tokens can be freely added or removed.
BPE tokenizers are different: nearly every token is either an input to a merge rule, a result of one, or both. Removing a token that a merge produces would leave the merge dangling; removing one that a merge consumes produces the same error.
So the script:
- Scans all merge rules and collects every token that appears as a merge input (left or right side).
- Sets
keep = merge_inputs ∪ frequent_tokens.
- Filters the merge list to only the rules whose result (left+right concatenation) is in
keep.
- Builds a new BPE model from the filtered vocab and filtered merges.
Because merge_inputs is a subset of keep, every kept merge has its inputs available — the graph remains self‑consistent.
3. Weight pruning
For a causal LM, two components are pruned:
- Input embedding layer:
vocab_size × hidden_size → 117752 × hidden_size. Missing rows are copied from the original by the new‑ID → original‑ID mapping built during tokenizer pruning.
- LM head (output projection):
hidden_size × vocab_size → similarly reduced. If tied to the embedding matrix (as in Qwen3.5‑2B), it's shared — only one matrix changes.
Usage
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3tokenizer = AutoTokenizer.from_pretrained("Qwen3.5-2B-en-pt", trust_remote_code=True)
4model = AutoModelForCausalLM.from_pretrained("Qwen3.5-2B-en-pt", trust_remote_code=True)
Limitations
This is an experiment. The datasets used (Leipzig Corpora news, 1M sentences per language) are small and domain‑specific. A production‑grade version would require larger, more diverse corpora for accurate frequency estimation.
- BPE merge density caps pruning at ~50%. The 65,178 merge‑input tokens are always required, so the minimum vocab size is floor(65,178 / 248,044) ≈ 26% even for a single language.
- Only embedding + LM‑head weights are reduced. The transformer layers (~70% of parameters) are untouched. A 52% vocab cut translates to ~29% disk savings.
- The model has not been fine‑tuned after pruning — it may need a short adaption step to recover full fluency.
