用向量索引替代稠密输出层：LLM 解码加速的新路径

导语

大语言模型每生成一个词元，都要把隐藏状态映射到整个词表，随后再从全部候选中挑选下一个词元。对于词表规模较大的多语言模型，这一步并不只是简单的矩阵乘法：输出嵌入矩阵需要被频繁读取，内存带宽可能成为自回归解码的主要瓶颈。尤其在小模型、批量为 1 的 CPU 推理场景中，计算资源未必是最紧张的部分，数据搬运反而会拖慢响应速度。

核心思路：把输出层变成检索问题

这项工作提出，输出投影后的 Top-k 选择可以重新表述为一次最大内积搜索。给定模型的隐藏状态，系统不再对整个词表执行完整的稠密投影，而是把词元嵌入视作向量库，通过 HNSW（Hierarchical Navigable Small World）索引寻找内积较高的候选词元。

具体而言，方法包含几个关键环节：

• 建立词元向量索引：将输出嵌入中的词元表示组织成 HNSW 索引，以支持近似向量检索。
• 缩小候选范围：每一步解码只返回少量高分词元，避免读取和计算完整词表。
• 接入现有解码流程：检索到的候选及其 logits 被散射回稀疏的全词表张量，因此可以继续使用现有的采样或 Top-k 解码接口。
• 以近似换取吞吐：HNSW 并不保证穷举得到精确的全局排序，但目标是以足够小的候选集合保留实际生成所需的高分词元。

实验结果与边界

作者在 Gemma 3、Llama 3.2 和 Qwen 3 模型上进行了 CPU 推理测试。结果显示，向量索引方案能够明显加快输出投影；在端到端、批量大小为 1 的解码中，Gemma 3 270M 的吞吐提升最高达到 82%。同时，研究使用 AlpacaEval 检查生成表现，报告称在该评测下质量得以保持。

这些结果并不意味着稠密输出层在所有场景中都应被替换。方法的价值主要集中在词表较大、批量较小、输出投影占比较高的推理任务。对于大批量服务，矩阵乘法可能更容易利用硬件并行能力；而近似索引还会引入索引构建、内存布局和候选召回率等工程因素。本文摘要也没有提供不同索引规模、召回率与延迟之间的完整权衡，因此生产部署仍需结合具体模型和硬件验证。

意义与影响

这项工作提供了一个值得关注的系统优化方向：LLM 输出头不一定只能被视为固定的稠密矩阵乘法，也可以被视为面向词元向量库的检索组件。若后续实现能进一步降低索引开销，并系统评估近似检索对长文本、不同采样策略和更多模型的影响，那么它有望成为 CPU 端、小模型端侧部署及低并发服务中的补充方案。其核心启示是，在解码延迟受内存访问限制时，减少被访问的词表规模，可能比单纯优化矩阵乘法更直接。

来源：arXiv