Beyond Frequency: Output-Vocabulary Pruning by Log-Mass Maximization

Abstract

A tokenizer defines the vocabulary a language model scores at every decoding step. Output-vocabulary pruning modifies this vocabulary after training: it keeps the model and tokenizer fixed and restricts the set of output tokens scored at inference, compressing the vocabulary-sized output layer. The dominant approach, frequency-based pruning, optimizes average token coverage, which may lead to stripping the tokens that minority languages and specialized domains depend on. We address this by giving the selection problem an explicit objective: the expected log retained probability mass under the deployment distribution, which we show is equivalent to minimizing reverse KL to the teacher and is monotone submodular. We therefore select the retained tokens by greedy maximization of this objective, which inherits a (1 - 1/e)-approximation guarantee. Frequency pruning is its linear surrogate, optimal only when contexts agree on a common token ranking. Across synthetic and multilingual deployments, greedy outperforms both frequency- and magnitude-based pruning, with the largest gains where that agreement breaks down. Under English-dominant calibration, greedy cuts Arabic reverse KL by 0.53 nats and improves Russian generation by +1.4 chrF, at negligible cost to English; the same support raises speculative-decoding draft acceptance by +5.4pp on Arabic. Our main experiments use Qwen3.5-0.8B, and a cross-family check on Gemma-3-1B confirms the qualitative pattern.

Cite

@inproceedings{krasheninnikov2026beyond,
  title={Beyond Frequency: Output-Vocabulary Pruning by Log-Mass Maximization},
  author={Krasheninnikov, E. and Ivanov, S. and Kozodoi, N. and Tarasova, N. and Zhang, A.},
  booktitle={{COLM} 2026 Workshop on Tokenization},
  year={2026}
}