A Distribution-Aware Approximate LM Head for Reduced-Memory Decoding

The output head of a large language model scores the entire vocabulary at every decoding step while most decoding approaches use only a small set of candidates, but this operation can be taken as a maximum inner-product search problem in which the rows of the LM head are fixed items and each hidden state is a query. We study when approximate vocabulary search can replace the dense head without changing decoding behavior. Standard search systems however optimize top-$k$ recall, a poor target here because it does not take probabilities into consideration. We propose the \MethodFull (\method), which searches a compact graph index over the rows of the LM head, optionally reranks the returned candidates, normalizing only those. At a fixed search budget, \method's time and transient memory do not depend on vocabulary size, and three operating regimes optimize our approach for low latency, low memory, or a balance of the two. On four 7–8B LLMs, rescoring the 64 returned candidates with exact logits keeps \method within $0.1$ percentage points of dense multiple-choice accuracy on every language-matched benchmark, whereas decoding from its own approximate scores loses at most $7.1$ points but can be recovered for about $6\%$ of serving throughput with a low-margin rescoring fallback while reducing the head's share of resident GPU memory by at least $2.7\times$ in vLLM, where the dense head holds up to $5.5\%$ of resident GPU memory.