Vocabulary Pruning via Submodular Optimization

Large language models generate text by scoring a distribution over the full vocabulary at each decoding step, making vocabulary size a significant latency and memory component, particularly in speculative decoding, early-exit inference, and edge deployment. Vocabulary pruning addresses this by retaining only a fixed subset of $K$ tokens, but existing approaches rely on frequency-based heuristics that optimize average token coverage and can leave minority languages or specialized domains catastrophically uncovered. Here, we formalize vocabulary pruning as a combinatorial optimization problem and propose maximizing the expected log retained probability mass, an objective equivalent to minimizing the expected reverse KL divergence to the teacher distribution. The objective is monotone submodular, yielding a greedy algorithm with a $(1-1/e)$-approximation guarantee. Frequency pruning is its first-order approximation: it is optimal under a sufficient common-ranking condition and can be suboptimal when the condition fails. Our experiments show that vocabulary pruning reduces memory and latency, while the proposed support selection improves accepted length and throughput in speculative decoding at the same vocabulary budget.