PEAR: Position-Embedding-Agnostic Attention Re-weighting Enhances Retrieval-Augmented Generation with Zero Inference Overhead

📄 arXiv: 2409.19745v2 📥 PDF

作者: Tao Tan, Yining Qian, Ang Lv, Hongzhan Lin, Songhao Wu, Yongbo Wang, Feng Wang, Jingtong Wu, Xin Lu, Rui Yan

分类: cs.CL, cs.AI

发布日期: 2024-09-29 (更新: 2024-10-07)

备注: preprint

💡 一句话要点

PEAR：零推理开销的位置嵌入无关注意力重加权增强检索增强生成

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

关键词: 检索增强生成 上下文感知 注意力机制 零推理开销 位置嵌入无关

📋 核心要点

  1. 现有RAG方法中，LLM的上下文感知能力不足，导致性能下降，而增强上下文感知的方法通常效率低，且依赖特定位置嵌入。
  2. PEAR通过在代理任务上检测并重加权抑制上下文感知的注意力头，优化LLM，从而增强上下文感知能力。
  3. PEAR在各种RAG任务中，以零推理开销优于现有基线方法，并在准确性和效率方面均有提升，且与位置嵌入算法无关。

📝 摘要（中文）

检索增强生成(RAG)增强的大型语言模型(LLM)为Web搜索引入了一种新的范式。然而，LLM有限的上下文感知能力降低了它们在RAG任务上的性能。现有的增强上下文感知能力的方法通常效率低下，在推理过程中会产生时间或内存开销，并且许多方法都是为特定的位置嵌入量身定制的。本文提出了一种位置嵌入无关的注意力重加权(PEAR)方法，它以零推理开销增强LLM的上下文感知能力。具体来说，在一个专注于上下文复制的代理任务上，我们首先检测到抑制模型上下文感知能力从而降低RAG性能的注意力头。为了削弱这些注意力头的影响，我们用可学习的系数对它们的输出进行重加权。通过调整这些系数来最小化代理任务上的损失，从而优化LLM(参数冻结)。因此，系数被优化为小于1的值，从而降低了它们抑制RAG性能的趋势。在推理过程中，优化的系数被固定以重新加权这些注意力头，而无需考虑具体的任务。与以前的方法相比，我们提出的PEAR具有两个主要优点：(1)它在内存使用或推理时间方面引入了零额外的推理开销，同时在各种RAG任务中优于竞争基线，在准确性和效率方面都表现出色。(2)它独立于位置嵌入算法，确保了更广泛的适用性。

🔬 方法详解

问题定义：论文旨在解决RAG中LLM上下文感知能力不足的问题。现有方法通常会引入额外的推理开销（时间或内存），或者依赖于特定的位置嵌入方式，限制了其通用性和效率。这些痛点使得RAG在实际应用中受到限制。

核心思路：论文的核心思路是通过识别并削弱LLM中抑制上下文信息的注意力头，从而提升模型整体的上下文感知能力。通过可学习的系数对这些注意力头的输出进行重加权，降低其对最终结果的影响。这种方法的关键在于，重加权系数的学习过程是在离线完成的，因此在推理阶段不会引入额外的开销。

技术框架：PEAR方法主要包含以下几个阶段：1) 代理任务构建：设计一个上下文复制任务，用于评估和优化LLM的上下文感知能力。2) 注意力头检测：在代理任务上，识别出那些抑制上下文信息的注意力头。3) 重加权系数优化：使用可学习的系数对这些注意力头的输出进行重加权，并通过优化代理任务的损失函数来调整这些系数。4) 推理阶段：在实际的RAG任务中，使用优化后的重加权系数对相应的注意力头进行重加权，而无需进行额外的训练或调整。

关键创新：PEAR最关键的创新在于其零推理开销的特性。通过离线优化注意力头的权重，避免了在推理阶段引入额外的计算或存储负担。此外，PEAR的设计与位置嵌入算法无关，使其具有更广泛的适用性。

关键设计：PEAR的关键设计包括：1) 代理任务的选择：上下文复制任务能够有效地评估和优化LLM的上下文感知能力。2) 注意力头检测方法：需要设计一种有效的指标来识别出抑制上下文信息的注意力头。3) 重加权系数的优化策略：需要选择合适的优化算法和损失函数，以有效地调整重加权系数，并避免过拟合。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

PEAR方法在各种RAG任务中表现出色，优于现有基线方法，同时实现了零推理开销。这意味着在提升性能的同时，不会增加额外的计算或存储负担。实验结果表明，PEAR在准确性和效率方面均有显著提升，使其成为一种极具竞争力的RAG增强方法。

🎯 应用场景

PEAR方法可以广泛应用于各种需要检索增强生成技术的场景，例如问答系统、知识库构建、文本摘要等。通过提升LLM的上下文感知能力，可以提高这些应用在处理复杂查询和生成高质量回复方面的性能。该方法尤其适用于对推理效率有较高要求的场景，例如在线客服、实时搜索等。

📄 摘要（原文）

Large language models (LLMs) enhanced with retrieval-augmented generation (RAG) have introduced a new paradigm for web search. However, the limited context awareness of LLMs degrades their performance on RAG tasks. Existing methods to enhance context awareness are often inefficient, incurring time or memory overhead during inference, and many are tailored to specific position embeddings. In this paper, we propose Position-Embedding-Agnostic attention Re-weighting (PEAR), which enhances the context awareness of LLMs with zero inference overhead. Specifically, on a proxy task focused on context copying, we first detect heads which suppress the models' context awareness thereby diminishing RAG performance. To weaken the impact of these heads, we re-weight their outputs with learnable coefficients. The LLM (with frozen parameters) is optimized by adjusting these coefficients to minimize loss on the proxy task. As a result, the coefficients are optimized to values less than one, thereby reducing their tendency to suppress RAG performance. During inference, the optimized coefficients are fixed to re-weight these heads, regardless of the specific task at hand. Our proposed PEAR offers two major advantages over previous approaches: (1) It introduces zero additional inference overhead in terms of memory usage or inference time, while outperforming competitive baselines in accuracy and efficiency across various RAG tasks. (2) It is independent of position embedding algorithms, ensuring broader applicability.