Source-linked AI summary
Accelerating LLM Inference via Vector Index Based Output Embeddings
Martin Loretz, Sepp Hochreiter
TL;DR
Large output embedding matrices create a memory-bandwidth bottleneck during autoregressive decoding for compact LLMs with large multilingual vocabularies. The paper reformulates output projection and top-k selection as MIPS and replaces the dense projection with an HNSW vector index. The method achieves up to 82% end-to-end throughput improvement for single-batch inference while preserving generation quality near the exact baseline.
Problem
Large multilingual vocabularies make dense output projection a memory-bandwidth bottleneck during batch-size-one decoding.
Method
The method uses an HNSW-based vector index to retrieve high-scoring token embeddings through MIPS instead of evaluating the full vocabulary projection.
Results
Up to 82% end-to-end throughput improvement is reported for single-batch inference, while AlpacaEval win rates against the exact baseline range from 45.9% to 49.1%.
Takeaways & Limitations
Approximate retrieval is presented as a practical alternative to dense output projections for latency-sensitive small-batch decoding.
Takeaways & Limitations
The approach currently relies on CPU execution, and its relative benefit may decrease for heavily quantized dense baselines or large batch sizes.
Abstract
from arXiv · showhide
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.
1. Introduction
Compact LLMs with large multilingual vocabularies face a memory-bandwidth bottleneck during batch-size-one decoding. The paper replaces dense output projection and top-k selection with approximate vector retrieval integrated into existing models.
- Motivation: Compact edge-deployed LLMs often retain 100k–250k-token vocabularies despite compact transformer backbones.Each decoding step must stream the large output embedding matrix from memory.
- Approach: The paper reformulates output projection followed by top-k sampling as Maximum Inner Product Search over token embeddings.The hidden state serves as the query, and retrieval targets only high-scoring tokens.
- Scope: The method targets latency-sensitive small-batch inference rather than high-throughput server inference dominated by efficient batched GEMM.It is designed as a post-training, drop-in modification across Gemma 3, Llama 3.2, and Qwen 3 models.
- Approach: An HNSW-based output head retrieves candidate logits without evaluating the full vocabulary projection.Retrieved logits can be scattered into a sparse full-vocabulary tensor for integration with existing decoding pipelines.
2. Background
Autoregressive token generation commonly truncates the vocabulary distribution, while exact selection of the largest logits remains expensive. The paper frames this selection as maximum inner product search and motivates HNSW for approximate retrieval.
- Token Sampling: Top-k sampling restricts generation to the highest-scoring tokens to avoid incoherent outputs from the probability distribution’s long tail.Truncation methods differ in how they determine the distribution cutoff.
- MIPS: Finding the k tokens maximizing inner products with hidden state h is a Maximum Inner Product Search problem.The exhaustive computation over Wout has O(|V| · d) complexity.
- HNSW: HNSW constructs a multilayer proximity graph and retrieves vectors by greedy traversal toward the query.Sparse upper layers provide long-range connections for rapid coarse-grained search.
3. Vector Index Embedding
The vector-index output head treats output embeddings as a static MIPS database, retrieves only the logits needed for truncated decoding, and reconstructs a sparse vocabulary tensor for compatibility. HNSW balances retrieval quality and latency, reducing projection complexity while adding graph-memory overhead.
- Dense Projection: The dense projection maps hidden state h through Wout to a vocabulary logit vector z ∈ R^|V|.Wout contains one embedding vector per vocabulary token.
- MIPS Reformulation: Selecting the largest logits is exactly MIPS because the rows of Wout serve as a static database of token vectors queried by h.This avoids computing the full vector when truncated decoding needs only top-scoring tokens.
- Index Selection: HNSW was the only evaluated candidate providing sufficient retrieval accuracy to preserve generative quality.It offered the best trade-off between retrieval quality and search latency among the considered approximate methods.
- Pipeline Integration: Retrieved logits are scattered into their vocabulary positions while unretrieved entries are set to −∞.This reconstructs a full tensor compatible with standard processors such as top-p and min-p.
- Memory Trade-off: The vector index adds a 5–10% output-layer memory overhead for graph storage under 32-bit representations.Memory grows from |V| · 4d bytes to |V| · (4d + 8M) bytes.
- Complexity: The final projection’s asymptotic complexity decreases from O(d · |V|) to approximately O(d · log |V|).Graph traversal has higher constants from random memory accesses, but the gains dominate at small batch sizes.
4. Experiments
Experiments show that HNSW-based output retrieval accelerates projection and batch-size-one decoding on CPU while maintaining high retrieval recall and near-baseline generation quality, with benefits constrained at larger batches.
- Experimental setup: CPU evaluations replace standard output projections in Gemma 3, Llama 3.2, and Qwen 3 models, using ef = 200 for primary experiments.The evaluation uses single-precision CPU inference, and ef = 200 was selected because no subjective quality degradation was observed.
- 4.1. Output Embedding Benchmarks: 12× speedup is exceeded for Gemma models at batch size 1 in output-projection throughput versus exact matrix multiplication.The advantage is strongest at low batch sizes, where memory bandwidth limits projection performance.
- 4.2. Full Model Benchmarks: Large batches can negate the vector index advantage when many unique retrieved tokens require loading a substantial fraction of the embedding matrix.Beyond this crossover, exact GEMM benefits from cache locality and contiguous memory access.
- 4.3. Index Accuracy: 99%+ Recall@2 at ef = 200 virtually eliminates divergence from the exact baseline, while ef = 400 reaches up to 99.9% recall.Recall improves monotonically with search depth; ef = 100 exceeds 96% recall but slightly degrades generation quality.
- 4.4. Generation Quality: 45.9%–49.1% length-controlled win rates on AlpacaEval indicate negligible generation-quality impact relative to the exact baseline.The evaluation compares vector-index generation with exact-baseline generation at ef = 200 across model families.
5. Conclusion
The paper concludes that vector-index output heads can substantially accelerate CPU inference for small LLMs with large vocabularies while preserving core generation quality. Its main scope boundary is reliance on CPU execution, with GPU acceleration and quantized support left for future work.
- 5. Conclusion: Replacing the dense output projection with a vector-index mechanism reduces compute and memory-bandwidth requirements for the output operation.The approach eliminates evaluation of the full projection while preserving the model’s core generative capabilities.
- 5. Conclusion: 82% end-to-end inference-throughput improvement is reported during single-batch inference for small models with large vocabularies.The conclusion highlights relevance to on-device deployment in mobile, automotive, and industrial environments.
- 5. Conclusion: CPU reliance is the primary limitation because sequential graph-based index traversal presents a bottleneck for GPU architectures.Future work targets GPU-accelerated vector indices and quantized data types; benefits may also decrease for heavily quantized baselines or large batches.
A. Profiler Traces
Figure 3 compares per-token profiler traces for standard dense output projection and HNSW-based vector index embedding in Gemma 3 270M at batch size one. The vector index replaces the dominant matrix multiplications with a short search and enables approximately 2× faster decoding.
- Profiler Traces: Approximately 2× faster decoding is achieved with the vector index embedding than with the standard output projection.The comparison uses Gemma 3 270M at batch size 1 during the final tokens of a 64-token generation.
- Profiler Traces: The baseline trace shows per-token latency dominated by large output-projection matrix multiplications.
- Profiler Traces: The vector index replaces the dense projection with an HNSW index search that occupies only a fraction of each decoding step.The search appears as a narrow blue segment at the end of the decoding step.