Fortify the Shortest Stave in Attention: Enhancing Context Awareness of Large Language Models for Effective Tool Use
Yuhan Chen, Ang Lv, Ting-En Lin, Changyu Chen, Yuchuan Wu, Fei Huang, Yongbin Li, Rui Yan
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 their 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. 1
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 4db0038e-82c2-478d-8c2a-9ec56ba97829Cited by top-tier papers12
- Found in the Middle: How Language Models Use Long Contexts Better via Plug-and-Play Positional EncodingZhenyu Zhang, Runjin Chen, Shiwei Liu, Zhewei Yao et al.NeurIPS 2024 · 99 citations
- Attributive Reasoning for Hallucination Diagnosis of Large Language ModelsYuyan Chen, Zehao Li, Shuangjie You, Zhengyu Chen et al.AAAI 2025 · 33 citations
- Mixture of In-Context Experts Enhance LLMs' Long Context AwarenessHongzhan Lin, Ang Lv, Yuhan Chen, Chen Zhu et al.NeurIPS 2024 · 25 citations
- POSITION BIAS MITIGATES POSITION BIAS: Mitigate Position Bias Through Inter-Position Knowledge DistillationYifei Wang, Feng Xiong, Yong Wang, Linjing Li et al.EMNLP 2025 · 16 citations
- StreamingDialogue: Prolonged Dialogue Learning via Long Context Compression with Minimal LossesJianan Li, Quan Tu, Cunli Mao, Zhengtao Yu et al.NeurIPS 2024 · 13 citations
Builds on16
- Retrieval-Augmented Generation for Knowledge-Intensive NLP TasksPatrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni et al.NeurIPS 2020 · 19,162 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- Toolformer: Language Models Can Teach Themselves to Use ToolsTimo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu et al.NeurIPS 2023 · 5,989 citations
- Reflexion: language agents with verbal reinforcement learningNoah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan et al.NeurIPS 2023 · 5,828 citations
- HuggingGPT: Solving AI Tasks with ChatGPT and its Friends in Hugging FaceYongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li et al.NeurIPS 2023 · 1,778 citations
Related papers
- Unveiling and Harnessing Hidden Attention Sinks: Enhancing Large Language Models without Training through Attention CalibrationZhongzhi Yu, Zheng Wang, Yonggan Fu, Huihong Shi et al.ICML 2024 · 63 citations
- LazyAttention: Efficient Retrieval-Augmented Generation with Deferred Positional EncodingHaocheng Xia, Mihir Pamnani, Hanxi Fang, Supawit Chockchowwat et al.ICML 2026
- Heads up! Large Language Models Can Perform Tasks Without Your Instruction via Selective Attention Head MaskingSenyu Han, Hongchuan Zeng, Kai Yu, Lu ChenICML 2025
- Attention Basin: Why Contextual Position Matters in Large Language ModelsZihao Yi, Zhenqing Ling, Delong Zeng, Haohao Luo et al.ACL 2026 · 2 citations
- Attention Entropy is a Key Factor: An Analysis of Parallel Context Encoding with Full-attention-based Pre-trained Language ModelsZhisong Zhang, Yan Wang, Xinting Huang, Tianqing Fang et al.ACL 2025 · 22 citations
