The Outsized Vocabulary Tax: Selective-Vocabulary Decoding for Edge LLMs

Modern small language models pair a compact backbone with a large vocabulary which incur an output projection cost of $\mathcal{O}(Vd)$ per decoded token, reaching up to 28% of total parameters and $\sim$20% of decode latency on edge hardware. Recent cluster-routed decoding reduces this to $\mathcal{O}(Cd+|\mathcal{S}|d)$ without modifying any weights, by scoring only the tokens $\mathcal{S}$ in the few of $C$ clusters whose centroids best match a given hidden state. Existing methods cluster LM-head rows directly, but we show that tokens close in weight space are often not the tokens that similar hidden states select. We introduce $\text{CHES}$ (Calibrated Hidden-state Embedding Selection), a training-free method that represents each token by a count-adaptive shrinkage estimate on the unit sphere. This estimate interpolates between the mean direction of the hidden states that select the token and a prior given by its LM-head row. We further propose $\text{AVQ}$, a closed-form anisotropic centroid update that aligns clustering with inference-time inner-product routing. On Qwen2.5 and Gemma2 models, $\text{CHES}+\text{AVQ}$ reaches 95% top-1 routing accuracy while scoring only 6-8% of the vocabulary and yields up to 3.2$\times$ LM-Head speedup on Apple M4 and Intel Xeon CPUs. Across multiple downstream benchmarks (e.g., GSM8K, MATH, TriviaQA), $\text{CHES}+\text{AVQ}$ can closely match the full LM-Head performance while utilizing only $\sim$18-20% of total vocabulary.