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

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.

## introduction recent works that augmenting large language models (llms, e.g., gpt series @xcite @xcite ) with tools have achieved advancements in various fields, such as humancomputer interactions (qin et al., 2023b; schick * equal contribution. † corresponding authors: yongbin li and rui yan (shuide.lyb@alibaba-inc.com, ruiyan@ruc.edu.cn). 1 we release our code at https://github.com/ fiorina1212/attention-buckets . et al., 2023), automating multi-modal tasks @xcite , and enhancing the overall efficiency of language-related applications @xcite . in this paradigm, upon receiving a user's intent, a large language model accesses multiple tools, typically in the form of apis. it then selects the most suitable one by referring to the relevant tool documentation and provides an accurate and suitable response. considering the integration of extensive information into the context, tool-use tasks demand a high level of context understanding and awareness from llms. despite the achievements made by current llmbased tool-use frameworks, in our practical experience, we observed that llms exhibit varying levels of awareness concerning different positions within the context. for instance, llms may overlook certain tools within the context, resulting in a failed call; however, by altering the position of these tools, the task can be successfully executed. such variations significantly affect the performance of llms in tool-use. this observation is consistent with the findings from a previous study @xcite that investigated a simple in-context retrieval task. when llms are presented with multiple key-value pairs and instructed to retrieve the value associated with a specific key, the index of the queried target key results in significant fluctuations in accuracy. figure 1 (a) provides a visual representation of the instructions for this task. figure 1(b) shows this fluctuation we replicated using the llama-2-7b @xcite . in our study, we go beyond the superficial fluctuations previously observed and identify that these position-related performance differences are closely associated with the model's fluctuating attention allocation. specifically, we observed a waveform pattern in the attention "intensity" (referred to as the attention waveform in this paper) when llms retrieve the same ## attention waves impact on context awareness in this section, we demonstrate that position-related performance fluctuations of llms are influenced by the underlying attention waveform. ## enhancing context awareness via interleaving attention waveform based on the insights presented above, we introduce a novel approach to sidestepping the above challenge, with the goal of improving the llm performance in tool-use by enhancing its context awareness. our method focuses on the inference stage of llms and does not require training. we first provide preliminary definitions, followed by a detailed introduction to our approach. ## exploring applications for retrieval-augmented generation considering our proposed attention buckets enhances the model's context awareness, it should be effective in other tasks demanding high contextual information utilization. this section explores the effectiveness and generality of attention buckets through a representative rag task: open-domain question answering (odqa). many current open-domain qa methods employ a retrieval-augmented generation (rag) paradigm @xcite @xcite , where a retriever @xcite ) gathers relevant documents, followed by a generative reader @xcite to find answers, demanding high context awareness. we conduct experiments on two popular benchmarks naturalquestion (nq, @xcite and webqa @xcite , with 3,610 and 2,032 test samples, respectively. we assess the models based on their accuracy in providing method nq webqa fid-xl (3b) 50.1 50.8 llama-2 (7b) 48.5 51.7 + asort 48.9 52.1 + usc 47.6 51.7 + attention buckets 50.3 (1.8↑) 53.1(1.4↑) we use llama-2-7b @xcite ) as our backbone model and compare it with fid-xl @xcite , the exclusive odqa model trained on multiple odqa benchmarks, including nq and webq. the results are shown in moreover, as the evaluation process for odqa is both objective and convenient, it enables us to investigate whether our identified b c approaches optimality through the examination of numerous permutations of base values. conducting this study on toolbench is too costly for us, as it needs the prohibitive use of gpt-4 and chatgpt for api calls. we leave this ablation study to the appendix f. the llm processes these documents to choose the suitable tool, inputs the necessary arguments, and relays the tool's output back to the users. many studies @xcite @xcite have found that large language models, including gpt-4 (openai, 2023), often exhibit hallucinations, such as inventing non-existent functions and arguments, or failing to adapt to changes in interactive environments. these highlight the critical need for enhancing these models' context awareness. current research heavily focuses on integrating multiple reasoning pathways to address errors caused by insufficient contextual comprehension. these techniques include react @xcite , dfs tree search @xcite , selfconsistency @xcite , etc. in contrast, our approach focuses on a fundamental solution: enhancing contextual awareness. it is both orthogonal to and stackable with those reasoning methods. ## conclusion in this paper, we delved into the waveform patterns observed in attention scores and found that waveform of the attention score could potentially affect the model's context awareness, particularly in relation to the position of crucial information within the context. we propose attention buckets, an inference augmentation method designed to enhance the model's context awareness. this method combines various attention patterns, which are controlled by different rope bases. our approach has achieved state-of-the-art (sota) performance on the current largest tool-use benchmark while showing the applicability to a wider range of rag tasks.