You Only Use Reactive Attention Slice For Long Context Retrieval

📄 arXiv: 2409.13695v1 📥 PDF

作者: Yun Joon Soh, Hanxian Huang, Yuandong Tian, Jishen Zhao

分类: cs.CL, cs.AI, cs.IR

发布日期: 2024-09-03

💡 一句话要点

提出基于反应式注意力切片的YOURA长文本检索方法，提升LLM推理效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本检索 注意力机制 大型语言模型 检索增强生成 反应式注意力 长上下文 LLM推理加速

📋 核心要点

  1. 现有检索增强生成（RAG）方法在长文本上下文中表现不足，主要依赖于基于嵌入的检索。
  2. YOURA通过计算token级别注意力对查询的“反应”程度，选择性地检索最相关的句子，提升检索效率。
  3. 实验表明，YOURA在长文本问答任务中，能够显著提升LLM的推理吞吐量，同时保持与现有方法相当的性能。

📝 摘要（中文）

为了支持大型语言模型（LLM）处理更长的上下文，本文提出了一种基于注意力的检索技术，名为You Only Use Reactive Attention slice (YOURA)。YOURA利用一种新颖的检索启发式方法，称为反应得分，来评估输入上下文中每个句子与查询句子的相关性。直观地说，该方法衡量每个token的注意力得分对查询的“反应”程度，并贪婪地检索最具反应性的句子。YOURA在内部为整个输入上下文生成一个token索引的向量（称为反应向量）。为了将每个句子映射到token索引的向量，我们提出了一种与嵌入无关的句子生成方法（Embedding-Agnostic Sentence Yield，EASY），这是一种尽力而为的token微调算法。在六个LongBench QA数据集上，对三个开源预训练LLM模型进行了评估。结果表明，该技术在服务长上下文查询时，vLLM推理吞吐量提高了30%，同时质量得分与简单但有效的truncate-middle方法几乎相同。

🔬 方法详解

问题定义：现有基于嵌入的检索方法在处理长文本上下文时存在局限性，无法有效捕捉长距离依赖关系，导致检索结果不准确。这限制了LLM在需要长上下文信息的任务中的应用。

核心思路：YOURA的核心思路是利用LLM自身计算的注意力权重来判断输入文本中哪些部分与查询最相关。通过分析每个token的注意力得分对查询的“反应”，可以更精确地定位到关键信息，从而实现更有效的检索。这种方法避免了对额外嵌入模型的依赖，直接利用了LLM的内部知识。

技术框架：YOURA主要包含以下几个阶段：1) 反应向量生成：对于输入上下文中的每个token，计算其对查询的反应得分，生成token索引的反应向量。2) 句子映射：利用Embedding-Agnostic Sentence Yield (EASY)算法，将每个句子映射到token索引的反应向量。EASY算法旨在找到句子中每个词对应的token在反应向量中的最佳位置。3) 句子排序与检索：根据句子的反应得分（例如，句子中所有token反应得分的平均值），对句子进行排序，并选择得分最高的句子作为检索结果。

关键创新：YOURA的关键创新在于提出了一种基于注意力“反应”的检索启发式方法。与传统的基于嵌入的检索方法不同，YOURA直接利用LLM的注意力机制来评估文本的相关性，无需额外的嵌入模型。此外，EASY算法提供了一种将句子映射到token索引向量的有效方法，使得可以基于token级别的注意力得分进行句子级别的检索。

关键设计：反应得分的计算方式是YOURA的关键设计之一。具体如何定义“反应”以及如何量化注意力得分的变化是需要仔细考虑的。EASY算法的具体实现，例如如何处理token与词之间的对应关系，以及如何优化映射过程，也是重要的技术细节。此外，句子反应得分的计算方式（例如，平均值、最大值等）也会影响检索效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，YOURA在LongBench QA数据集上，能够显著提升vLLM的推理吞吐量，最高可达30%，同时保持与truncate-middle方法几乎相同的质量得分。这表明YOURA能够在不牺牲性能的前提下，有效提高LLM处理长文本的效率。该结果验证了基于注意力反应的检索方法的有效性。

🎯 应用场景

YOURA技术可应用于各种需要长文本上下文理解的场景，例如长篇文档问答、信息检索、法律文本分析、医疗记录处理等。通过提高LLM处理长文本的效率和准确性，可以提升这些应用的用户体验和智能化水平。该技术还有助于降低长文本推理的计算成本，促进LLM在资源受限环境中的部署。

📄 摘要（原文）

Supporting longer context for Large Language Models (LLM) is a promising direction to advance LLMs. As training a model for a longer context window is computationally expensive, many alternative solutions, such as Retrieval Augmented Generation (RAG), have been used. However, most existing RAG methods adopt embedding-based retrieval that falls short on long contexts. To address such challenges, we propose an attention-based retrieval technique, You Only Use Reactive Attention slice (YOURA). YOURA leverages a novel retrieval heuristic called reaction score to rank the relevance of each sentence in the input context with the query sentence. Intuitively, we measure how the per-token attention score "reacts" to the query and greedily retrieves the most reactive sentences. Internally, YOURA generates a token-indexed vector (called reaction vector) for the whole input context. To map each sentence to the token-indexed vector, we propose an Embedding-Agnostic Sentence Yield (EASY), a best-effort token wiggling algorithm. We evaluate our retrieval technique on three open-source pre-trained LLM models across six LongBench QA datasets. Our technique achieves up to 30% vLLM inference throughput improvement for serving long-context queries with a nearly identical quality score to the simple yet effective truncate-middle approach.