Accelerating LLM Inference via Vector Index Based Output Embeddings

通过基于向量索引的输出嵌入加速大语言模型推理

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

摘要： 在自回归解码过程中，庞大的输出嵌入矩阵会造成严重的内存带宽瓶颈，对于拥有大规模多语言词表的小型大语言模型（LLM）而言尤为如此。

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.

我们将“输出投影后接 Top-k 词元选择”的过程重新表述为针对词元嵌入的最大内积搜索，并使用基于 HNSW（分层可导航小世界图）的向量索引取代了稠密的词表投影。

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.

由此产生的输出头仅检索一小部分高分词元候选集，并通过将检索到的 Logits 散布到稀疏的全词表张量中，从而集成到现有的解码流水线中。

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.

在 Gemma 3、Llama 3.2 和 Qwen 3 模型的 CPU 推理测试中，该方法显著加速了输出投影过程，并将 Gemma 3 270M 模型在 Batch Size 为 1 时的端到端解码吞吐量提升了高达 82%，同时在 AlpacaEval 评估中保持了原有的生成质量。

These results suggest approximate retrieval is a practical alternative to dense output projections in latency-sensitive small-batch decoding.

这些结果表明，在对延迟敏感的小批量解码场景中，近似检索是替代稠密输出投影的一种实用方案。