RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

这篇论文提出了RetrievalAttention，一种无需训练的方法，旨在加速注意力计算并减少GPU内存消耗，以便将LLMs扩展到更长的上下文。RetrievalAttention利用了注意力机制的动态稀疏性，建议在CPU内存中使用近似最近邻搜索（ANNS）索引来存储KV向量，并在生成过程中使用向量搜索来检索最相关的向量。然而，我们观察到现成的ANNS索引常常对注意力机制中查询向量和关键向量之间的分布差异（OOD）无效。RetrievalAttention通过设计一种注意力感知的向量搜索算法来解决OOD挑战，该算法可以适应查询向量的分布。我们的评估表明，RetrievalAttention只需要访问1-3%的数据即可保持高模型准确性。这导致长上下文LLMs的推理成本显着降低，GPU内存占用也大大降低。特别是，RetrievalAttention只需要一张单独的NVIDIA RTX4090（24GB）来为具有8B参数的LLMs提供128K标记的服务，能够在0.188秒内生成一个标记。

本文旨在解决Transformer-based Large Language Models (LLMs)中注意力计算的二次时间复杂度问题，导致长上下文推理的极慢推理延迟和高GPU内存消耗的问题。

RetrievalAttention是一种训练-free的方法，可以加速注意力计算并减少GPU内存消耗。它利用注意力机制的动态稀疏性，使用CPU内存中的近似最近邻搜索(ANNS)索引来存储KV向量，并在生成过程中使用向量搜索来检索最相关的向量。

RetrievalAttention通过设计注意力感知的向量搜索算法来解决out-of-distribution(OOD)问题，以适应查询向量的分布。实验表明，RetrievalAttention只需要访问1-3%的数据就能保持高模型精度，从而显著降低长上下文LLMs的推理成本和GPU内存占用。

与此相关的最新研究包括：1）Lin et al. (2020)提出了一种使用局部敏感哈希(LSH)来加速注意力计算的方法，2）Kitaev et al. (2020)提出了一种使用固定大小的记忆矩阵来近似注意力计算的方法。