Which Tokens Can a Device Afford to Score? Vocabulary Pruning by Log-Mass Maximization

A language model on an edge device runs under tight compute and memory constraints, and its output vocabulary incurs two distinct costs. The embedding matrix can hold up to a third of the weights at sub-billion parameter scale. On top of that, every decoding step has to score the full vocabulary. Restricting inference to K tokens shrinks that layer by up to 256X in memory and 92X in the time it contributes per decoded token. However, which K tokens to keep is an open question. The standard practice keeps the most frequent tokens, which optimizes average coverage and can strip the tokens minority languages depend on. We instead choose the tokens by greedily maximizing an explicit objective: the expected log retained probability mass under the deployment distribution. Maximizing this objective is equivalent to minimizing reverse KL to the teacher, and it is monotone submodular, which is what lets a standard greedy algorithm inherit a (1-1/e) guarantee. On a ten-language multilingual dataset, this recovers the quality that pruning otherwise takes away from minority scripts: Arabic reverse KL falls by 0.53 nats and Russian generation gains +1.4 chrF, at negligible cost to English. The same selection of tokens decodes faster, raising speculative-decoding acceptance by +5.4pp on Arabic and measured throughput by up to 6.1%, and it shrinks memory further when combined with quantization, reaching a 3.5X reduction while keeping 95% of full-precision quality. Because the tokens are chosen offline, a device pays nothing for this at inference time.