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

TL;DR

This paper proposes a novel inference method named Attention Buckets, which allows LLMs to process their input through multiple parallel processes and enhances LLM's awareness to various contextual positions, thus mitigating the risk of overlooking crucial information.

Abstract

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.

Figures (4)

• Figure 1: (a) Task illustration: Presented with multiple key-value pairs and a target key (highlighted in bold), the model is required to accurately retrieve and generate the value associated with this key from an extensive context. (b) We illustrate the position-related fluctuation in accuracy of Llama-2-7B on this in-context retrieval task. (c) The pattern of the attention score exhibits fluctuations, which we term the "attention waveform". Our study reveals a connection between the position-related fluctuations in LLMs' performance and this attention waveform.
• Figure 2: (a) The overview of how a typical Transformer-based Large Language Model (LLM) generates the next token based on context $\mathcal{C}$. This LLM comprises $M$ layers, though for simplicity, only the inner workings of a single layer are shown. (b) The overview of our proposed Attention Buckets augmenting the context awareness of LLMs: Upon receiving context $\mathcal{C}$, it creates $N$ (specifically 3 in this example) parallel copies for processing. Each parallel stream employs a distinct RoPE base. The resulting output distributions $p_j$ are weighted and summed based on the prediction confidence $\alpha_j$, culminating in the final predicted distribution $\hat{p}$ used for decoding the next token.
• Figure 3: (a) RoPE Base Searching: we measure the distance from the candidate bases' attention peaks (or troughs) to the attention troughs (or peaks) corresponding to items in $\mathcal{B}_{c}$. For demonstration clarity, we illustrate only a partial context that corresponds to one waveform period. (b) Attention waveform corresponding to $\mathcal{B}_{c}$ searched by hyperparameters detailed in $\S$\ref{['sec:search']}.
• Figure 4: Fair and convincing experimental operations. (a) We apply padding prior to the "Query" to accurately locate the final token of the target key-value pair at a desired position, which corresponds to an attention peak or trough. (b) We use paddings to maintain consistency in prompt length across various rounds.