Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AttentionRAG: Attention-Guided Context Pruning in Retrieval-Augmented Generation

Yixiong Fang, Tianran Sun, Yuling Shi, Xiaodong Gu

TL;DR

AttentionRAG addresses the problem of excessive and redundant retrieved contexts in Retrieval-Augmented Generation by introducing an attention-guided context-pruning method that reformulates queries into a next-token-prediction template to focus attention on a single anchor token. The approach processes long contexts in chunks, computes chunk-level attention features, and selects sentences containing the most-attended tokens to form a compressed context, enabling high efficiency without requiring extra training. Empirical results on LongBench and Babilong show up to 6.3x compression with equal or improved task performance compared to uncompressed contexts, and competitive results versus state-of-the-art prompt-compression methods. The method demonstrates strong transferability across models and tasks, with practical speedups via batching and quantization, and offers a flexible framework for scalable, long-context RAG in real-world applications.

Abstract

While RAG demonstrates remarkable capabilities in LLM applications, its effectiveness is hindered by the ever-increasing length of retrieved contexts, which introduces information redundancy and substantial computational overhead. Existing context pruning methods, such as LLMLingua, lack contextual awareness and offer limited flexibility in controlling compression rates, often resulting in either insufficient pruning or excessive information loss. In this paper, we propose AttentionRAG, an attention-guided context pruning method for RAG systems. The core idea of AttentionRAG lies in its attention focus mechanism, which reformulates RAG queries into a next-token prediction paradigm. This mechanism isolates the query's semantic focus to a single token, enabling precise and efficient attention calculation between queries and retrieved contexts. Extensive experiments on LongBench and Babilong benchmarks show that AttentionRAG achieves up to 6.3$\times$ context compression while outperforming LLMLingua methods by around 10\% in key metrics.

AttentionRAG: Attention-Guided Context Pruning in Retrieval-Augmented Generation

TL;DR

AttentionRAG addresses the problem of excessive and redundant retrieved contexts in Retrieval-Augmented Generation by introducing an attention-guided context-pruning method that reformulates queries into a next-token-prediction template to focus attention on a single anchor token. The approach processes long contexts in chunks, computes chunk-level attention features, and selects sentences containing the most-attended tokens to form a compressed context, enabling high efficiency without requiring extra training. Empirical results on LongBench and Babilong show up to 6.3x compression with equal or improved task performance compared to uncompressed contexts, and competitive results versus state-of-the-art prompt-compression methods. The method demonstrates strong transferability across models and tasks, with practical speedups via batching and quantization, and offers a flexible framework for scalable, long-context RAG in real-world applications.

Abstract

While RAG demonstrates remarkable capabilities in LLM applications, its effectiveness is hindered by the ever-increasing length of retrieved contexts, which introduces information redundancy and substantial computational overhead. Existing context pruning methods, such as LLMLingua, lack contextual awareness and offer limited flexibility in controlling compression rates, often resulting in either insufficient pruning or excessive information loss. In this paper, we propose AttentionRAG, an attention-guided context pruning method for RAG systems. The core idea of AttentionRAG lies in its attention focus mechanism, which reformulates RAG queries into a next-token prediction paradigm. This mechanism isolates the query's semantic focus to a single token, enabling precise and efficient attention calculation between queries and retrieved contexts. Extensive experiments on LongBench and Babilong benchmarks show that AttentionRAG achieves up to 6.3 context compression while outperforming LLMLingua methods by around 10\% in key metrics.

Paper Structure

This paper contains 34 sections, 3 equations, 2 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Retrieve-Augmented Generation
  4. Attention Mechanism in LLMs
  5. Method
  6. Problem Formulation
  7. Construct Answer Hint Prefix
  8. Compute Attention Features from LLM
  9. Compress with Attention
  10. Time Efficiency and Batch Generation
  11. Experimental Setup
  12. Datasets
  13. Long Context Reasoning
  14. Babilong
  15. Long Context Summarization
  16. ...and 19 more sections

Figures (2)

  • Figure 1: Illustration of RAG context compression
  • Figure 2: Illustration of AttentionRAG: 1) We first generate a hint prefix using the original query; if generation fails, we fallback to a fixed hint prefix. 2) The long context is then divided into smaller chunks for localized processing. 3) For each chunk, we batch the input with the hint prefix (1 token only) and use an LLM to compute attention features over the full context, focusing on the generated anchor token. 4) Using these attention features, we identify and extract the most relevant content from each chunk to construct a compressed context.