摘要
arXiv:2409.15355v5 宣告类型: replace-cross
摘要:我们介绍了Block-attention,一种旨在解决检索增强生成(RAG)场景中增加的推理延迟和成本的注意力机制。传统的做法通常以自回归的方式编码整个上下文。相反,Block-attention将检索到的文档分割成离散的块,除了最后一个块外,每个块独立计算键值(KV)状态。在RAG场景中,通过将每个段落定义为一个块,Block-attention使我们能够重用之前见过的段落的键值状态,从而大大减少了推理过程中的延迟和计算开销。Block-attention的实现涉及块分割、位置重新编码以及对LLM进行微调以适应Block-attention机制。在包括RAG、ICL和一般领域的11个不同基准上进行的实验表明,在块微调后,Block-attention模型不仅实现了与全注意力模型相当的性能,而且可以在不损失性能的情况下无缝地在块注意力和全注意力模式之间切换。值得注意的是,Block-attention显著降低了第一个词生成时间(TTFT)和浮点操作(FLOPs)到非常低的水平。对于总长度为32K的输入序列,产出第一个词仅需45毫秒。与全注意力模型相比,TTFT和相应的FLOPs分别减少了98.7%和99.8%。此外,在附录A中,我们详细介绍了Block-attention在游戏AI场景中的应用及其带来的显著潜在益处。我们强烈建议游戏领域的研究人员不要忽视这一部分。