Research proposes replacing dense output projections with vector indexing to accelerate LLM inference

A new arXiv preprint proposes replacing dense output embedding matrices with an HNSW-based vector index to mitigate a memory bandwidth bottleneck during autoregressive decoding. The authors reformulate the standard output projection followed by top-k token selection as a maximum inner product search over token embeddings, retrieving only a small candidate set of high-scoring tokens.

The proposed output head integrates into existing decoding pipelines by scattering retrieved logits into a sparse full-vocabulary tensor. On CPU inference with compact models including Gemma 3, Llama 3.2, and Qwen 3, the method substantially accelerates the output projection step and improves end-to-end batch-size-one decoding throughput by up to 82% for Gemma 3 270M.

The authors report that generation quality is preserved under AlpacaEval evaluation, suggesting approximate retrieval can serve as a practical alternative to dense output projections in latency-sensitive small-batch decoding scenarios.