Speeding Up LLM Token Generation Using Approximate Neighbors - An investigation into reducing inference cost by replacing full-vocabulary scoring with approximate nearest neighbor search. The work studies the trade-off between generation speed and output quality in LLMs

Sammanfattning

Large language models generate text one token at a time. At every step, the modeltakes a hidden vector and scores it against every row of a large vocabulary matrixto decide which token comes next. This step is expensive, even though only a fewtokens actually matter for the decision. In small models with big vocabularies,like Llama-3.2-1B-Instruct, this single operation accounts for roughly 21% of allparameters and a noticeable share of per-token latency.This thesis asks a simple question: can we skip the full vocabulary scan and use anapproximate nearest neighbor (ANN) index to shortlist a small set of likely tokens,then score only those exactly? We test three IVF-family indices from FAISS (IVF,IVF-PQ, IVF-SQ) as drop-in replacements for the dense head, with no retrainingand no changes to the transformer body. We measure token-level agreement with thedense head on Alpaca, MATH-Hard, and XNLI, downstream accuracy on MMLUPro, and latency on an NVIDIA L40s, comparing against FlashHead on the samehardware and evaluation protocol.With a fine enough cluster grid, all three variants match the dense baseline onMMLU-Pro within noise. All variants degrade most on XNLI, suggesting the IVFclusters are less well aligned with non-English hidden states. IVF-PQ is the mostattractive point on the frontier, running about 1.49× faster than the dense head inisolation and about 1.04× faster end-to-end at the cluster count that matches denseaccuracy within noise. The end-to-end gain is modest because the transformer body,not the head, dominates per-step cost on this hardware.