用向量索引加速大模型输出推理,提升速度近80%。
Accelerating LLM Inference via Vector Index Based Output Embeddings

- 将输出投影转为向量搜索,用HNSW索引快速找候选词。
- 在CPU上使Gemma 3 270M推理吞吐提升最高达82%。
- 适合低延迟、小批量推理场景,保持生成质量不变。
大型输出嵌入矩阵在自回归解码过程中造成显著内存带宽瓶颈,尤其对词汇量大的紧凑多语言大模型更为明显。本文将输出投影与top-k选词重构为基于词嵌入的最大内积搜索,并用基于HNSW的向量索引替代稠密词汇投影。输出头仅检索高分候选词,通过将检索到的得分分散到稀疏全词汇张量中,即可无缝集成至现有解码流程。在CPU上对Gemma 3、Llama 3.2和Qwen 3模型的实验表明,该方法显著加速输出投影,在批大小为1时,Gemma 3 270M的端到端推理吞吐提升最高达82%,且在AlpacaEval评估下生成质量保持不变。结果表明,在低延迟的小批量解码中,近似检索是稠密输出投影的可行替代方案。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。