通过注意力不变性,精准定位文档关键部分,加速RAG预填充计算。
SIFT: Selective-Index For Fast Compute of RAG Prefill by Exploiting Attention Invariance

- 离线提取文档中高注意力位置,运行时仅计算这些区域。
- 预填充时间提升1.71倍,精度与全重算相差不到1%。
- 仅存储位置信息,存储量比完整键值对小24000倍,适合部署。
检索增强生成(RAG)通过注入相关文档提升大模型输出质量,但会增加提示长度并延长首个词生成时间(TTFT)。由于相同文档常重复出现,传统方法虽预计算键值(KV)张量,但因现代GPU上磁盘传输延迟高,反而比全重算更慢。同时粗粒度重算损害精度。本文提出SIFT:利用注意力不变性,离线提取文档中高注意力位置,构建两个紧凑的位向量存储。发现两个关键规律:(1) 局部注意力不变性——文档内高注意力位置不随上下文变化;(2) 交叉注意力一致性——文档内高注意力键也吸引后续文档的关注。运行时仅在标记位置计算注意力,避免存储大量KV数据,彻底规避磁盘开销。实验显示,该方法将预填充阶段的TTFT提升1.71倍,且精度与全重算相比仅下降1%以内。
原文摘要 · Abstract (English)
Retrieval-Augmented Generation (RAG) injects LLM queries with relevant documents to improve response quality. This injection increases prompt length and slows time to first token (TTFT). Unlike standard queries, RAG queries have a unique property of context reuse where the same documents recur across user queries. Thus, fully recomputing documents for every RAG query does redundant compute and increases TTFT. Prior works precompute KV tensors of RAG documents offline and coarsely recompute some tokens during online prefill. However, such KV reuse is often slower than full recomputation on modern GPUs due to high-latency disk transfers. Further, such a coarse-grained recomputation degrades accuracy. To address these limitations, this paper proposes SIFT: Selective-Index For Fast Compute of RAG Prefill by Exploiting Attention Invariance. SIFT processes documents offline and extracts fine-grained locations of high attention scores for each document. Next, we identify the following attention invariance insights that enable us to exploit the extracted locations during runtime: (1) Local-Attention Invariance: The location of high attention scores within a document remain invariant to surrounding documents. This helps us predict the location of high scores where the document attends to itself. (2) Cross-Attention Consistency: Keys with high intra-document attention also attract cross-attention from subsequent documents. This helps us predict the location of high scores where the document attends to future documents. Critically, SIFT stores no KV data and only stores locations of high scores in the form of two compact bit vectors. SIFT's storage is up to 24,000x smaller than KV tensors, obviating costly disk transfers. During prefill, SIFT computes the attention only for the marked locations and improves TTFT by 1.71x while holding accuracy within 1% of full recompute.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。