arXiv:2409.10516cs.LGcs.CL2024-09NeurIPS被引 135

用向量检索加速长文本生成,降低显存占用。

RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

  • 通过向量搜索只调用最相关的键值对,跳过冗余计算。
  • 仅需1--3%的键值数据即可接近全注意力精度。
  • 适合需要处理超长文本的部署场景,如文档问答。

基于Transformer的大语言模型日益重要,但其注意力机制存在二次时间复杂度问题,导致长上下文推理速度极慢且显存消耗高。本文提出RetrievalAttention,一种无需训练的方法,通过利用注意力机制的动态稀疏性,在CPU内存中为键值(KV)向量构建近似最近邻索引,并在生成时通过向量搜索检索最相关部分。由于查询向量与键向量分布差异大,现有索引效果不佳。为此,RetrievalAttention设计了关注注意力分布的自适应向量搜索算法。实验表明,该方法在仅访问1--3%数据的情况下仍能保持接近全注意力的精度,显著降低长上下文大模型的推理开销。特别地,仅需一张NVIDIA RTX4090(24GB)即可支持80亿参数模型处理128K tokens输入,每生成一个词耗时0.188秒。

原文摘要 · Abstract (English)

Transformer-based Large Language Models (LLMs) have become increasingly important. However, due to the quadratic time complexity of attention computation, scaling LLMs to longer contexts incurs extremely slow inference speed and high GPU memory consumption for caching key-value (KV) vectors. This paper proposes RetrievalAttention, a training-free approach to both accelerate attention computation and reduce GPU memory consumption. By leveraging the dynamic sparsity of attention mechanism, RetrievalAttention proposes to build approximate nearest neighbor search (ANNS) indexes for KV vectors in CPU memory and retrieve the most relevant ones through vector search during generation. Unfortunately, we observe that the off-the-shelf ANNS indexes are often ineffective for such retrieval tasks due to the out-of-distribution (OOD) between query vectors and key vectors in the attention mechanism. RetrievalAttention addresses the OOD challenge by designing an attention-aware vector search algorithm that can adapt to the distribution of query vectors. Our evaluation demonstrates that RetrievalAttention achieves near full attention accuracy while only requiring access to 1--3% of the data. This leads to a significant reduction in the inference cost of long-context LLMs, with a much lower GPU memory footprint. In particular, RetrievalAttention only needs a single NVIDIA RTX4090 (24GB) to serve 128K tokens for LLMs with 8B parameters, which is capable of generating one token in 0.188 seconds.

长文本生成注意力优化显存节省

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。