💻 computer science

InfiniRetri: Attention-Guided Retrieval for Enhancing LLMs' Long-Context Processing

InfiniRetri is a training-free method that leverages a pre-trained LLM's intrinsic attention mechanisms to iteratively distill relevant context from arbitrarily long inputs, enabling efficient and highly accurate long-context retrieval and reasoning without the need for external embedding models or costly retraining.

Original authors: XiaoJu YE, Fang Wang, FangFang Wang, YuanRui Ren

Published 2026-10-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: XiaoJu YE, Fang Wang, FangFang Wang, YuanRui Ren

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Modern artificial intelligence systems, known as large language models, have become remarkably capable at understanding and generating human language. They can write stories, solve math problems, and answer complex questions. However, these systems face a significant physical limitation: they can only hold a certain amount of information in their "working memory" at one time. When a user asks the model to read a very long document, such as a full novel or a collection of legal records, the system often struggles. It tends to forget details from the beginning of the text or lose track of the main point as the input grows longer. This happens because the mathematical process the model uses to connect words to one another becomes incredibly expensive and slow as the text gets bigger. To fix this, researchers have tried two main approaches. One is to force the model to learn how to handle longer texts through expensive and time-consuming retraining. The other is to use a helper system that searches for relevant parts of the text before the model reads it, but this helper often relies on simple word-matching that misses deeper meaning.

A team of researchers from WenHua College in China has proposed a different solution that requires no retraining and no external search tools. They discovered that the model already possesses the ability to find the right information within a long text; it just needs to be guided to use it. By observing how the model's internal attention mechanism works, they found that the model naturally focuses its attention on the specific sentences that contain the answers to a question, even when those sentences are buried deep within a massive document. This pattern is most visible in the final layer of the model's processing network. The researchers built a method called InfiniRetri that exploits this natural behavior. Instead of trying to compress the text into smaller, unreadable fragments or relying on a separate search engine, their method iteratively reads chunks of the long document, asks the model to identify the most relevant sentences based on its own internal focus, and saves only those complete sentences into a small, efficient memory cache.

The process works like a careful reader going through a thick book page by page. The researchers split the long text into manageable sections. For each section, they feed the text along with the user's question into the model. The model then looks at how much attention its internal mechanisms pay to different parts of that section. The researchers extract a map of this attention, specifically looking at how the question words connect to the words in the text. They found that the model's attention scores act as a precise signal, highlighting exactly which sentences hold the answer. The system then selects the top-scoring sentences and saves them into a lightweight cache. This cycle repeats for every chunk of the document. By the time the entire text has been processed, the cache contains only the most critical sentences needed to answer the question, effectively filtering out the rest of the noise. Finally, the model reads this small, curated collection of sentences and generates the final answer.

The results of this approach are striking. In tests where the model had to find a hidden fact within a massive document, a very small version of the model using this method achieved perfect accuracy even when the text contained over one million words. This is a significant improvement over previous methods, which often failed when the text exceeded 32,000 words. The method also proved highly effective at answering questions that required finding information across multiple documents, outperforming other advanced techniques that rely on complex training or external databases. Remarkably, this high level of performance was achieved using only two gigabytes of computer memory, a size small enough to run on standard consumer hardware. The researchers noted that the method works best when there is a clear question to ask, as the system relies on the question to guide its attention. While it showed less improvement on tasks like summarizing a whole book, where the goal is to condense everything rather than find a specific fact, the ability to retrieve precise information from virtually unlimited text lengths without expensive training represents a major step forward. The study suggests that by trusting the model's own internal signals, we can unlock its ability to read and reason over texts of any size.

Technical Summary: InfiniRetri

Problem Statement

Large Language Models (LLMs) suffer from consistent performance degradation as input lengths increase, affecting both direct retrieval and complex reasoning over long documents. Existing solutions face significant limitations:

Context Window Extension: Extending context windows via continued training is prohibitively expensive due to the quadratic complexity of attention mechanisms.
Prompt Compression: Current methods (e.g., SelectiveContext, LLMLingua) often operate at the token level, discarding semantic units and producing unreadable text, which compromises performance without training.
KV Cache Compression: Methods like H2O or SnapKV retain only partial Key-Value states, failing to compress all states in a training-free manner and often overlooking fine-grained attention patterns.
Retrieval-Augmented Generation (RAG): Standard RAG relies on external embedding models that capture only surface-level lexical relevance, struggling to bridge semantic gaps or establish associations between multiple retrieved passages.

Methodology: InfiniRetri

The authors propose InfiniRetri, a training-free method that leverages the intrinsic attention mechanisms of pre-trained LLMs to retrieve and cache relevant information from arbitrarily long inputs. The approach is built upon the discovery of the Attention-Aligned Retrieval (AAR) pattern.

Core Insight: Attention-Aligned Retrieval (AAR)

Through layer-wise visualization and analysis, the authors observed that:

Retrieval Capability: The question-to-context attention sub-matrix of a pre-trained LLM (specifically in the last transformer layer) provides a precise, interpretable signal for identifying content relevant to a query. The model inherently "knows" which context spans are relevant.
Layer Depth: This retrieval pattern strengthens with depth, peaking in the final transformer layer, where attention between relevant and irrelevant context is most discriminative.

Workflow

InfiniRetri processes long contexts through an iterative loop consisting of five steps:

Chunking: The long context is split into segments at sentence boundaries to preserve semantic integrity.
Merging: A lightweight external cache of previously retrieved sentences is prepended to the current chunk.
LLM Inference: The merged context and the user question are fed into the LLM.
Retrieval in Attention:
- The last-layer attention matrix is extracted.
- The question-to-context sub-matrix is sliced out (excluding attention sinks).
- Attention scores are aggregated at the phrase level using 1-D sum-pooling.
- Top-K most relevant positions are identified based on the aggregated scores.
Cache Update: Instead of storing individual tokens, the complete sentences containing the top-K tokens are retrieved from the original context and added to the cache.

After processing all chunks, the final answer is generated using the accumulated cache and the original question.

Key Contributions

Identification of AAR Pattern: The paper identifies that the question-to-context attention sub-matrix in the last layer of pre-trained LLMs serves as an accurate, reliable, and interpretable retrieval signal, eliminating the need for external embedding models.
InfiniRetri Framework: A training-free method that iteratively distills complete, relevant sentences into a lightweight cache. It aligns with prompt compression goals but operates at the sentence level for better interpretability and semantic coherence.
Comprehensive Evaluation: The method is evaluated against three families of baselines: prompt compression, KV cache compression, and RAG, demonstrating superior performance across multiple benchmarks.

Experimental Results

InfiniRetri was evaluated on LongBench, LongBenchV2, $\infty$ Bench, and the Needle-In-a-Haystack (NIH) test.

LongBench Performance:
- Prompt Compression: Outperformed the previous training-free SOTA (DAC) by 11.4% in average score (42.05 vs. 37.76) with a significantly lower compression rate (0.26 vs. 0.5).
- KV Cache Compression: Surpassed the strongest prior method, CompilerKV (which requires offline calibration), by 15.9% overall on Qwen2-7B. Notably, it achieved a 238.7% improvement in multi-document QA tasks (13.40 $\to$ 45.38).
- RAG: Outperformed leading embedding-based retrievers (BGE-M3, Stella-v5) and advanced RAG systems (GraphRAG, RQ-RAG). It trailed only MemoRAG, which benefits from full-stage training.
Ultra-Long Context Capabilities:
- NIH Test: Achieved 100% accuracy on the NIH test with a 0.5B-parameter model over 1 million tokens, requiring only 2GB of GPU memory.
- $\infty$ Bench: A 1.5B model achieved 100% accuracy on ultra-long retrieval tasks (214K tokens), surpassing larger models equipped with KV retrieval methods.
- LongBenchV2: A 7B model using InfiniRetri reached performance parity with its 72B counterpart, particularly on long-text tasks.
Complexity: The method reduces inference complexity from quadratic $O(N^2)$ to approximately linear $O(mN)$ , where $m$ is the chunk size, by pruning the context to a small cache of relevant sentences.

Significance and Claims

The paper claims that InfiniRetri establishes a new paradigm for information retrieval built upon the intrinsic capabilities of LLMs. By utilizing the model's own attention modules, it bridges the semantic gap between queries and documents without external retrieval components or training.

The authors emphasize that:

Interpretability: Unlike token-level compression, InfiniRetri retains complete sentences, making the retrieval process transparent and human-readable.
Efficiency: It drastically reduces memory overhead (e.g., 2GB for 1M tokens) and computational cost while maintaining or improving accuracy.
Generalizability: The method is applicable to any attention-based LLM and theoretically supports infinite context lengths.

Limitations Acknowledged:
The paper notes that InfiniRetri requires a predefined question (limiting open-ended scenarios without question generation) and shows limited gains on summarization tasks, which may require aggregating information across multiple windows beyond the current single-pass retrieval mechanism. Additionally, performance is sensitive to hyperparameter tuning, specifically the phrase size relative to the target answer length.