Preprint
Machine Learning

Rethinking RAG-based Decoding

Xiaoqiang Lin, Aritra Ghosh, K. H. Low, Anshumali Shrivastava, Vijai Mohan
September 1, 2025arXiv.org7 citations
arXivPDFDOI

7

Citations

3

Influential Citations

arXiv.org

Venue

2025

Year

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities in leveraging extensive external knowledge to enhance responses in multi-turn and agentic applications, such as retrieval-augmented generation (RAG). However, processing long-context inputs introduces significant system latency and demands substantial memory for the key-value cache, resulting in reduced throughput and a fundamental trade-off between knowledge enrichment and system efficiency. While minimizing latency for long-context inputs is a primary objective for LLMs, we contend that RAG require specialized consideration. In RAG, much of the LLM context consists of concatenated passages from retrieval, with only a small subset directly relevant to the query. These passages often exhibit low semantic similarity due to diversity or deduplication during re-ranking, leading to block-diagonal attention patterns that differ from those in standard LLM generation tasks. Based on this observation, we argue that most computations over the RAG context during decoding are unnecessary and can be eliminated with minimal impact on performance. To this end, we propose REFRAG, an efficient decoding framework that compresses, senses, and expands to improve latency in RAG applications. By exploiting the sparsity structure, we demonstrate a 30.85 the time-to-first-token acceleration (3.75 improvement to previous work) without loss in perplexity. In addition, our optimization framework for large context enables REFRAG to extend the context size of LLMs by 16. We provide rigorous validation of REFRAG across diverse long-context tasks, including RAG, multi-turn conversations, and long document summarization, spanning a wide range of datasets. Experimental results confirm that REFRAG delivers substantial speedup with no loss in accuracy compared to LLaMA models and other state-of-the-art baselines across various context sizes.

Analysis

Why This Paper Matters

Retrieval-augmented generation (RAG) has become a cornerstone of modern LLM applications, allowing models to incorporate vast external knowledge without retraining. However, the standard practice of concatenating multiple retrieved passages into the context introduces severe latency and memory overhead, particularly for the key-value cache during decoding. This paper identifies a crucial insight: in RAG, attention patterns are block-diagonal, meaning only a small fraction of the context is actually relevant to the query at each decoding step. This observation challenges the conventional wisdom that all context tokens must be processed equally, and opens the door to significant efficiency gains.

The practical importance of this work is hard to overstate. As LLMs are deployed in agentic and multi-turn settings, the ability to handle long contexts efficiently directly impacts user experience and system cost. REFRAG's 30.85x speedup in time-to-first-token and 16x context extension are not incremental improvements—they represent a paradigm shift in how RAG decoding can be optimized. For practitioners building real-time RAG systems, this could mean the difference between a responsive assistant and a frustratingly slow one.

Technical Contributions

  • Block-diagonal attention sparsity identification: The authors empirically show that RAG contexts exhibit a unique sparsity pattern where attention is concentrated within small diagonal blocks corresponding to relevant passages, unlike standard long-context tasks.
  • REFRAG framework: A three-stage pipeline—compress (prune irrelevant KV cache entries), sense (determine which tokens are needed for the current step), and expand (reconstruct only necessary context)—that exploits this sparsity.
  • Optimization for context extension: By reducing memory footprint, REFRAG allows LLMs to handle contexts up to 16x longer than their original capacity without additional hardware.
  • Generalizability: The method is validated across multiple tasks (RAG, multi-turn, summarization) and models (LLaMA), showing it is not limited to a specific architecture or dataset.

Results

REFRAG achieves a 30.85x acceleration in time-to-first-token (TTFT) compared to standard RAG decoding, and a 3.75x improvement over the previous state-of-the-art. These gains come with no degradation in perplexity or accuracy across all evaluated tasks. The context extension capability is demonstrated by increasing effective context size by 16x, enabling LLMs to process much longer documents than their original training length. The paper reports consistent results across multiple datasets and LLaMA model sizes, with thorough ablation studies confirming the importance of each component in the REFRAG pipeline.

Significance

This work has immediate implications for the deployment of RAG systems in production. By dramatically reducing latency and memory requirements, REFRAG makes it feasible to use larger context windows and more retrieved passages without sacrificing responsiveness. The insight about block-diagonal attention sparsity may also inspire further research into context-aware decoding strategies for other tasks. For the AI community, this paper demonstrates that careful analysis of model behavior (attention patterns) can lead to practical optimizations that bridge the gap between research and real-world application.

retrieval-augmented-generationllm-decodingattention-sparsitylatency-optimizationlong-contextefficient-inferencekey-value-cache
Large Language ModelsAttention MechanismsRetrieval Augmented Generation

Related Papers