高效填充的块注意力机制

我们提出了块注意力 (Block-attention) 一种注意力机制，旨在解决检索增强生成 (RAG) 场景中的推理延迟和成本问题。传统方法通常以自回归方式对整个上下文进行编码。相反，块注意力将检索文档划分为离散块，每个块独立计算键值 (KV) 状态，除非是最后一个块。在 RAG 场景下，将每个段落定义为一个块，块注意力使我们能够复用之前看到的段落的 KV 状态，从而在推理期间显著降低延迟和计算开销。块注意力的实现涉及块分段、位置重新编码和针对块注意力机制对大语言模型进行微调。我们在 11 个多样化基准测试（包括 RAG、ICL 和通用领域）上的实验表明，在块微调后，块注意力模型不仅在性能上与完全注意力模型相当，而且可以无缝在块注意力和完全注意力模式之间切换，且不会产生任何性能损失。值得注意的是，块注意力显著降低了第一个 token 的生成时间 (TTFT) 和浮点运算量 (FLOPs)，仅需 45 毫秒即可输出第一个 token，输入序列总长度为 32K。相对于完全注意力模型，TTFT 和相应的 FLOPs 分别降低了 98.7% 和 99.8%。此外，在附录 A 中，我们阐述了块注意力在游戏 AI 场景中的应用及其巨大的潜在收益。我们强烈建议游戏领域的研究人员不要忽视这一部分。