Infinite Retrieval: Attention Enhanced LLMs in Long-Context Processing

TL;DR

The paper tackles the challenge of processing inputs beyond typical LLM context windows by proposing InfiniRetri, a training-free method that leverages the model's own attention to enable retrieval over infinitely long inputs. It introduces a segment-and-merge workflow, a retrieval-in-attention mechanism, and a cache of sentence-level token IDs to dynamically assemble context across segments, avoiding external embeddings or heavy retraining. Empirical results show state-of-the-art performance on the Needle-In-a-Haystack task with a 0.5B model on 1M tokens and significant improvements on LongBench, while achieving substantial reductions in latency and compute overhead compared to traditional KV-cache methods. While highly effective for retrieval and QA tasks, the approach is less transformative for long-document summarization, indicating a direction for further refinement and task-specific adaptations. Overall, InfiniRetri demonstrates a practical, scalable paradigm for long-context processing by exploiting latent retrieval signals in LLM attention and can be applied to any Transformer-based model without additional training.

Abstract

Limited by the context window size of Large Language Models(LLMs), handling various tasks with input tokens exceeding the upper limit has been challenging, whether it is a simple direct retrieval task or a complex multi-hop reasoning task. Although various methods have been proposed to enhance the long-context processing capabilities of LLMs, they either incur substantial post-training costs, or require additional tool modules(e.g.,RAG), or have not shown significant improvement in realistic tasks. Our work observes the correlation between the attention distribution and generated answers across each layer, and establishes the attention allocation aligns with retrieval-augmented capabilities through experiments. Drawing on the above insights, we propose a novel method InfiniRetri that leverages the LLMs's own attention information to enable accurate retrieval across inputs of infinitely length. Our evaluations indicate that InfiniRetri achieves 100% accuracy in the Needle-In-a-Haystack(NIH) test over 1M tokens using a 0.5B parameter model, surpassing other method or larger models and setting a new state-of-the-art(SOTA). Moreover, our method achieves significant performance improvements on real-world benchmarks, with a maximum 288% improvement. In addition, InfiniRetri can be applied to any Transformer-based LLMs without additional training and substantially reduces inference latency and compute overhead in long texts. In summary, our comprehensive studies show InfiniRetri's potential for practical applications and creates a paradigm for retrievaling information using LLMs own capabilities under infinite-length tokens. Code will be released in link.

Figures (10)

• Figure 2: Visual 0, 1, 2, 26, 27 layers of Attention Scores Heatmap from using Qwen2-7B-Instruct inference in a QA Sample
• Figure 3: The retrieval accuracy on LLMs each layers
• Figure 4: Entire Workflow of Our Method InfiniRetri for Enhancing Long-Context Processing in LLMs
• Figure 5: Humans are limited by their field of vision, but can read the entire book page by page.
• Figure 6: Performance Comparison on the Needle in a Haystack Task Using Llama3-8B-Instruct.