Fortify the Shortest Stave in Attention: Enhancing Context Awareness of Large Language Models for Effective Tool Use

📄 arXiv: 2312.04455v4 📥 PDF

作者: Yuhan Chen, Ang Lv, Ting-En Lin, Changyu Chen, Yuchuan Wu, Fei Huang, Yongbin Li, Rui Yan

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

发布日期: 2023-12-07 (更新: 2024-06-04)

备注: ACL 2024 main

💡 一句话要点

提出Attention Buckets方法，增强LLM上下文感知能力，显著提升工具使用性能。

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

关键词: 大型语言模型 上下文感知 注意力机制 工具使用 旋转位置编码

📋 核心要点

  1. 现有LLM在工具使用等任务中，由于注意力机制的波形模式，容易忽略关键上下文信息，导致性能下降。
  2. Attention Buckets通过多个并行进程处理输入，每个进程使用不同的旋转位置嵌入基角，生成不同的注意力波形。
  3. 实验表明，Attention Buckets能显著提升LLM在工具使用、RAG等任务中的性能，7B模型可达到与GPT-4相当的水平。

📝 摘要（中文）

本文揭示了大型语言模型（LLM）注意力分配中固有的波形模式会显著影响其在需要高度上下文感知任务中的性能，例如利用LLM进行工具使用。具体而言，当关键信息位于注意力波形的波谷区域时，模型可能会忽略这些信息，从而导致性能下降。为了解决这个问题，我们提出了一种名为Attention Buckets的新型推理方法。它允许多个并行进程处理LLM的输入。每个进程使用不同的旋转位置嵌入基角，从而创建独特的注意力波形。通过用另一个进程的注意力峰值补偿特定进程的注意力波谷，我们的方法增强了LLM对各种上下文位置的感知，从而降低了忽略关键信息的风险。在最大的工具使用基准测试中，我们的方法将一个7B模型提升到最先进的性能，与GPT-4相当。在其他基准测试和一些RAG任务中，Attention Buckets也表现出显著的性能提升，这些任务也需要对上下文内容有透彻的理解。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）在处理需要高度上下文感知的任务，特别是工具使用时，由于其注意力机制的固有波形模式导致的关键信息被忽略的问题。现有方法未能有效解决注意力分配不均，导致模型性能受限。

核心思路：核心思路是通过创建多个具有不同注意力波形的并行处理分支（即Attention Buckets），利用不同分支的注意力峰值来补偿其他分支的注意力波谷，从而增强模型对整个上下文的感知能力，降低忽略关键信息的风险。

技术框架：Attention Buckets方法的核心在于并行处理流程。首先，将输入数据复制到多个分支中。每个分支使用不同的旋转位置嵌入（RoPE）的基角。RoPE是一种常用的位置编码方法，通过旋转矩阵来表示位置信息。不同的基角会导致不同的注意力波形。然后，每个分支独立地进行前向推理。最后，将各个分支的输出进行某种形式的聚合（例如，平均或加权平均）以得到最终的预测结果。

关键创新：关键创新在于通过调整旋转位置嵌入的基角，创造出多个具有不同注意力分布的“buckets”，从而弥补了单一注意力机制可能存在的盲点。这种并行处理和注意力互补的机制，使得模型能够更全面地捕捉上下文信息。与现有方法相比，Attention Buckets不需要修改模型结构或训练过程，可以直接应用于现有的LLM。

关键设计：关键设计包括旋转位置嵌入基角的选择。论文可能探讨了如何选择合适的基角，以确保不同的注意力波形能够有效地互补。此外，分支数量的选择也是一个重要的参数。过多的分支会增加计算成本，而过少的分支可能无法充分覆盖整个上下文空间。论文可能还涉及了如何聚合不同分支的输出，例如，使用简单的平均或更复杂的加权平均方法。

📊 实验亮点

Attention Buckets方法在工具使用基准测试中取得了显著成果，将一个7B模型提升到与GPT-4相当的性能水平，达到了state-of-the-art。此外，在其他需要上下文理解的任务（如RAG）中也观察到了明显的性能提升，验证了该方法的有效性和泛化能力。

🎯 应用场景

该研究成果可广泛应用于需要高度上下文理解的自然语言处理任务，例如智能助手、代码生成、文档摘要、问答系统等。通过提升LLM的上下文感知能力，可以提高这些应用在复杂场景下的准确性和可靠性，具有重要的实际应用价值和商业潜力。

📄 摘要（原文）

In this paper, we demonstrate that an inherent waveform pattern in the attention allocation of large language models (LLMs) significantly affects their performance in tasks demanding a high degree of context awareness, such as utilizing LLMs for tool-use. Specifically, the crucial information in the context will be potentially overlooked by model when it is positioned in the trough zone of the attention waveform, leading to decreased performance. To address this issue, we propose a novel inference method named Attention Buckets. It allows LLMs to process their input through multiple parallel processes. Each process utilizes a distinct base angle for the rotary position embedding, thereby creating a unique attention waveform. By compensating an attention trough of a particular process with an attention peak of another process, our approach enhances LLM's awareness to various contextual positions, thus mitigating the risk of overlooking crucial information. In the largest tool-use benchmark, our method elevates a 7B model to achieve state-of-the-art performance, comparable to that of GPT-4. On other benchmarks and some RAG tasks, which also demand a thorough understanding of contextual content, Attention Buckets also exhibited notable enhancements in performance.