Accelerating LLM Inference via Vector Index Based Output Embeddings

arXiv:2608.27460 · cs.CL, cs.LG · Submitted 2026-07-01 · Read on arXiv

cs.CL, cs.LG

Submitted: 2026-07-01

Updated: 2026-07-01

Comments: ICML 2026 - AdaptFM Workshop

License: http://creativecommons.org/licenses/by/4.0/

The gist: Large output embedding matrices create a significant memory bandwidth bottleneck during autoregressive decoding, especially for compact LLMs with large multilingual vocabularies.

Terminology

Abstract

Large output embedding matrices create a significant memory bandwidth bottleneck during autoregressive decoding, especially for compact LLMs with large multilingual vocabularies. We reformulate the output projection followed by top-k token selection as a maximum inner product search over token embeddings and replace the dense vocabulary projection with an HNSW-based vector index. The resulting output head retrieves only a small candidate set of high-scoring tokens and can be integrated into existing decoding pipelines by scattering retrieved logits into a sparse full-vocabulary tensor. On CPU inference with Gemma 3, Llama 3.2, and Qwen 3 models, our method substantially accelerates the output projection and improves end-to-end batch-size-one decoding throughput by up to 82% for Gemma 3 270M, while preserving generation quality under AlpacaEval evaluation. These results suggest approximate retrieval is a practical alternative to dense output projections in latency-sensitive small-batch decoding.

Related papers