Abstract
Large output embedding matrices create a significant memory bandwidth bottleneck during autoregressive decoding, particularly for compact large language models (LLMs) with extensive multilingual vocabularies. This paper reformulates the output projection followed by top-k token selection as a maximum inner product search (MIPS) over token embeddings, replacing the dense vocabulary projection with a Hierarchical Navigable Small World (HNSW)-based vector index. Our approach, called the vector-index output head, retrieves only a small candidate set of high-scoring tokens and integrates seamlessly into existing decoding pipelines by scattering retrieved logits into a sparse full-vocabulary tensor.
In CPU inference experiments 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 as measured by AlpacaEval. These results demonstrate that approximate retrieval offers a practical alternative to dense output projections in latency-sensitive small-batch decoding scenarios, which remain critical for edge and on-device deployments in 2026.
Introduction
The autoregressive decoding process in LLMs presents a fundamental performance challenge: each generated token requires a full projection through the output embedding matrix, which often contains millions of parameters for models with large vocabularies. This dense matrix multiplication creates a memory bandwidth bottleneck, especially for compact models designed for edge deployment, where memory access patterns dominate computational costs.
As of 2026, the trend toward increasingly multilingual and specialized vocabularies—combined with the proliferation of small, efficient models for on-device inference—has amplified this bottleneck. While prior work has explored various approximations to output projections, the direct application of approximate nearest neighbor search to token embeddings remains largely unexplored in production settings.
Method
We reframe the output projection as a maximum inner product search (MIPS) problem. Rather than computing the full logit vector (the projection of the hidden state onto all token embeddings), we construct an HNSW-based index over token embeddings, enabling approximate retrieval of the most likely tokens.
The core pipeline operates as follows: during decoding, the hidden state is used to query the HNSW index, retrieving a small candidate set of high-scoring tokens. These retrieved logits are then scattered into a sparse full-vocabulary tensor, allowing direct integration with existing sampling and generation routines that expect logits over the complete vocabulary. This design ensures that the method can be dropped into current implementations without requiring architectural changes to the underlying transformer.
Key design considerations include:
- Index construction: Token embeddings are indexed offline, with no modification to the trained model weights.
- Candidate set size: The number of retrieved tokens is configurable, trading off approximation quality against speed.
- Sparse scattering: Retrieved logits are placed in a pre-allocated sparse tensor to maintain compatibility with standard APIs.
Experiments
We evaluate our method on three model families—Gemma 3, Llama 3.2, and Qwen 3—across a range of model sizes, focusing on CPU inference with batch size one, which represents the most common deployment pattern for edge and interactive applications.
Speedup Results
Our method delivers consistent speedups on the output projection phase, with the most dramatic gains observed for models with larger vocabularies. For Gemma 3 270M, we achieve up to 82% improvement in end-to-end decoding throughput, while Llama 3.2 and Qwen 3 models show significant but more modest gains due to their smaller vocabulary sizes.
Quality Preservation
Using AlpacaEval as our evaluation benchmark, we confirm that generation quality remains intact compared to dense baseline models. The approximation introduced by the vector index does not measurably degrade output quality across diverse instruction-following tasks, suggesting that the candidate sets contain the tokens that matter for high-quality generation.
Discussion
The success of approximate retrieval in this setting highlights an important shift in LLM inference optimization. As vocabularies continue to grow in size and linguistic diversity, the dense output projection becomes proportionally more expensive. Our results indicate that the latent space of hidden states exhibits sufficient structure that a vector index over token embeddings can effectively approximate the true top-k selections.
Limitations and Future Work
While our method excels in latency-sensitive small-batch decoding on CPU, its benefits may diminish in GPU settings with highly optimized matrix multiplication kernels or with very large batch sizes where memory bandwidth is less constrained. Additionally, while we demonstrate quality preservation on AlpacaEval, further evaluation on specialized domains (e.g., code generation, mathematical reasoning) would strengthen confidence in generalizability.
Future research directions include:
- Adaptive candidate set sizes based on confidence thresholds
- Learning-to-index approaches that refine embedding spaces for better MIPS performance
- Extension to beam search and speculative decoding pipelines
Conclusion
We present a vector-index-based approach to output projection that substantially accelerates LLM inference in latency-sensitive settings. With up to 82% end-to-end throughput improvement on CPU decoding without quality degradation, our method establishes approximate retrieval as a viable and practical alternative to dense output projections for small-batch decoding—a finding that grows increasingly relevant as vocabulary sizes expand and edge deployment of LLMs becomes more pervasive.
via ArXiv CL+LG