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

Abstract

Large language models generate text one token at a time. At every step, the model takes a hidden vector and scores it against every row of a large vocabulary matrix to decide which token comes next. This step is expensive, even though only a few tokens 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 all parameters and a noticeable share of per-token latency. This thesis asks a simple question: can we skip the full vocabulary scan and use an approximate 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 retraining and no changes to the transformer body. We measure token-level agreement with the dense head on Alpaca, MATH-Hard, and XNLI, downstream accuracy on MMLU Pro, and latency on an NVIDIA L40s, comparing against FlashHead on the same hardware and evaluation protocol. With a fine enough cluster grid, all three variants match the dense baseline on MMLU-Pro within noise. All variants degrade most on XNLI, suggesting the IVF clusters are less well aligned with non-English hidden states. IVF-PQ is the most attractive point on the frontier, running about 1.49× faster than the dense head in isolation and about 1.04× faster end-to-end at the cluster count that matches dense accuracy within noise. The end-to-end gain is modest because the transformer body, not the head, dominates per-step cost on this hardware.