Attention Reveals More Than Tokens: Training-Free Long-Context Reasoning with Attention-guided Retrieval
This addresses a key limitation in LLMs for users needing long-context reasoning, though it is incremental as it builds on existing CoT methods.
The paper tackled the problem of LLMs having shorter effective context lengths than claimed for complex reasoning tasks, and proposed Attrieval, a training-free algorithm that uses attention weights to retrieve implicit facts, enhancing performance on synthetic and real-world QA datasets.
Large Language Models (LLMs) often exhibit substantially shorter effective context lengths than their claimed capacities, especially when handling complex reasoning tasks that require integrating information from multiple parts of a long context and performing multi-step reasoning. Although Chain-of-Thought (CoT) prompting has shown promise in reducing task complexity, our empirical analysis reveals that it does not fully resolve this limitation. Through controlled experiments, we identify poor recall of implicit facts as the primary cause of failure, which significantly hampers reasoning performance. Interestingly, we observe that the internal attention weights from the generated CoT tokens can effectively ground implicit facts, even when these facts are not explicitly recalled. Building on this insight, we propose a novel training-free algorithm, Attrieval, which leverages attention weights to retrieve relevant facts from the long context and incorporates them into the reasoning process. Additionally, we find that selecting context tokens from CoT tokens further improves performance. Our results demonstrate that Attrieval enhances long-context reasoning capability notably on both synthetic and real-world QA datasets with various models.