PaperlayerView exact revision on arXiv

Source-linked AI summary

REFRAG: Rethinking RAG based Decoding

Xiaoqiang Lin, Aritra Ghosh, Bryan Kian Hsiang Low, Anshumali Shrivastava, Vijai Mohan

arXiv:2509.01092v2cs.CLcs.AIcs.LG

TL;DR

Long-context inputs raise inference latency and KV-cache memory, and RAG contexts contain many passages with sparse cross-chunk attention. REFRAG compresses context representations and selectively expands chunks; it reports up to 30.85× TTFT acceleration over LLaMA and 3.75× over CEPE without loss in perplexity or downstream accuracy.

  • Problem

    Long prompts increase inference latency and KV-cache memory, while RAG contexts contain many passages with sparse cross-chunk attention.

    Source: §1 · pp. 1–2

  • Method

    REFRAG uses compressed chunk embeddings as decoder inputs and selectively expands chunks while preserving autoregressive decoding.

    Source: §2 · p. 3

  • Results

    Up to 30.85× TTFT acceleration over LLaMA and 3.75× over CEPE is reported without loss in perplexity or downstream accuracy.

    Source: §4 · p. 5, §7 · p. 11

  • Takeaways & Limitations

    REFRAG extends context handling beyond the original LLaMA-2-7B 4k context window while maintaining performance at longer contexts.

    Source: §4 · p. 5

    Ablation findings suggest a practical limit to the compression rate beyond which model capability is significantly reduced.

    Source: §4.1 · p. 7

Abstract

from arXiv · show

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.

1 Introduction

Long-context prompts increase inference latency and KV-cache memory, while RAG contexts contain many retrieved passages that may contribute little to a query. REFRAG targets this structure by compressing context chunks and selectively expanding them during decoding.

1 Introduction

  • Longer prompts increase KV-cache memory linearly and TTFT quadratically, while TTIT grows linearly with prompt length.These costs reduce throughput in applications that require high throughput and low latency.Source: §1 · p. 1
  • RAG contexts often include uninformative or reusable passages, discard retrieval-time encoding information, and have sparse cross-attention between chunks.Diversity and deduplication can make many context chunks unrelated.Source: §1 · p. 2

1.1 Our Contributions

  • REFRAG reduces decoding latency, TTFT, and memory usage without modifying the LLM architecture or adding decoder parameters.The framework is designed for RAG applications.Source: §1.1 · p. 2
  • REFRAG replaces retrieved passage tokens with precomputed compressed chunk embeddings and uses a lightweight RL policy to select chunks for full-token expansion.This supports compression at arbitrary positions while preserving autoregressive decoding.Source: §1.1 · p. 2
  • 30.75× TTFT acceleration without loss in perplexity is reported for REFRAG, 3.75× the acceleration of the previous method.The evaluation also covers long-context applications including RAG, multi-turn conversation with RAG, and long-document summarization.Source: §1.1 · p. 2

2 Model Architecture

REFRAG encodes context chunks as shorter decoder inputs and selectively expands chunks. With k = 16, it accelerates TTFT over CEPE at a context length of 16384.

2 Model Architecture

  • REFRAG chunks the context, encodes each chunk into a projected embedding, and feeds those embeddings with question-token embeddings to the decoder.For RAG, where context dominates the input, this shortens the decoder input by approximately a factor of k and reduces latency and memory use.Source: §2 · pp. 2–3, Figure 1 · p. 3
  • At context length 16384, REFRAG with k = 16 achieves 16.53× TTFT acceleration with cache and 8.59× without cache, surpassing CEPE’s 2.01× and 1.04×, respectively.Source: §2 · p. 3
  • At k = 32, TTFT acceleration reaches 32.99× versus LLaMA and 3.75× versus CEPE, while performance remains similar to CEPE.REFRAG also achieves up to 6.78× throughput acceleration compared with LLaMA.Source: §2 · p. 3

3 Methodology

REFRAG aligns its encoder and decoder through continual pre-training, then fine-tunes for downstream tasks. Its training recipe combines reconstruction, curriculum learning, and selective compression.

3 Methodology

  • Continual pre-training uses next-paragraph prediction to align encoder–decoder generations with those from a decoder given the full context.The workflow then applies supervised fine-tuning for downstream tasks such as RAG and multi-turn conversation.Source: §3 · p. 4

3.1 Continual Pre-training Recipe

  • The reconstruction task trains the encoder and projection layer, with the decoder frozen, to compress chunks and map embeddings into the decoder’s token space.After this alignment, the decoder is unfrozen for continual pre-training.Source: §3.1 · p. 4
  • Curriculum learning addresses the growing difficulty of reconstructing longer chunks and multiple chunk embeddings by gradually increasing task difficulty.Training shifts from easier single-chunk examples toward examples with more chunk embeddings.Source: §3.1 · p. 4
  • A lightweight RL policy uses next-paragraph prediction perplexity as a negative reward to select chunks for uncompressed expansion.The encoder and decoder are fine-tuned to handle mixed compressed and uncompressed inputs while preserving autoregressive decoding.Source: §3.1 · p. 4

4 Experimental Results

REFRAG improves log-perplexity and time-to-first-token (TTFT) across evaluated settings, while chunk embeddings support longer contexts. Ablations show that reconstruction training and RL-based selective compression contribute to performance, with compression rate 64 diminishing it.

4 Experimental Results

  • Across output lengths of 512, 1024, and 2048 tokens, REFRAG8 and REFRAG16 outperform other baselines in nearly all settings and have lower latency than CEPE.The evaluation uses held-out Book and ArXiv data plus PG19 and Proof-pile; baselines are primarily based on LLaMA-2-7B.Source: §4 · p. 5, Table 1 · p. 6
  • At context lengths of 4096, 8192, and 16384 tokens, REFRAG8 and REFRAG16 maintain superior performance despite training at a 6144-token combined context and output length.Chunk embeddings enable extrapolation beyond LLaMA-2-7B’s 4k context window.Source: §4 · p. 5, Table 2 · p. 6
  • At compression rate 32, REFRAG matches CEPE in log-perplexity while reaching 30.85× TTFT acceleration over LLaMA and 3.75× over CEPE.At rate 16, it improves average log-perplexity over CEPE by 9.3% across four datasets and is 16.53× faster than LLaMA in TTFT.Source: §4 · p. 5

4.1 Ablation Study

  • RL-based selective compression consistently outperforms heuristic and random chunk selection across compression rates, and REFRAG16 at rate 8 outperforms natively trained REFRAG8 across datasets and context lengths.The REFRAG16 configuration achieves rate 8 by selectively compressing chunks without recomputing chunk embeddings.Source: §4 · p. 5, §4.1 · pp. 6–7, Figure 3 · p. 6
  • Curriculum learning is essential to the reconstruction task, and reconstruction-task initialization is important for successful continual pre-training.The reconstruction task involves recovering the input from chunk embeddings.Source: §4.1 · pp. 5–6
  • Performance regresses as compression increases; rate 32 remains competitive, while rate 64 is overly aggressive and diminishes performance.These results indicate a practical limit beyond which model capability is significantly reduced.Source: §4.1 · p. 7
  • Increasing decoder size substantially reduces loss, whereas a larger encoder yields only modest improvement and may not help with limited continual-pretraining data.Across additional decoder models, the training recipe achieves performance comparable to full context, and longer context lowers perplexity.Source: §4.1 · p. 7

5 Contextual Learning Applications

REFRAG is evaluated on retrieval-augmented generation and multi-turn RAG, with reported performance gains over LLaMA-based fine-tuning across passage settings. The section also describes the datasets and retrieval setup used for these evaluations.

5 Contextual Learning Applications

  • The contextual-learning applications use instruction-tuning datasets curated for RAG, long-document summarization, and multi-turn conversation with RAG.For RAG fine-tuning, the authors use 1.1 million data points across five domains and evaluate on held-out data and commonly used RAG datasets, under strong- and weak-retriever settings.Source: §5 · p. 7, §5.1 · pp. 7–8

5.1 Retrieval Augmented Generation

  • For RAG, REFRAG16 and REFRAG32 outperform LLaMAFT despite using 2× and 4× fewer decoder tokens, respectively.The analysis reports the same pattern in short- and long-context scenarios.Source: §5.1 · p. 8, Table 3 · p. 9
  • At equal latency, REFRAG improves average performance over LLaMA by 1.22% across 16 RAG tasks with a strong retriever and 1.93% with a weak retriever.With a strong retriever and at most 10 passages, it matches LLaMA while achieving 5.26× TTFT speedup; with a weak retriever and 10 passages, it improves performance by 0.71% and accelerates TTFT by 5.26×.Source: §5.1 · p. 8

5.2 Multi-Turn Conversation

  • The multi-turn evaluation uses TopiOCQA, ORConvQA, and QReCC, retrieving K passages for each turn with the same retriever and corpus as Section 5.1.Source: §5.2 · p. 8
  • In multi-turn RAG, REFRAG outperforms LLaMAFT on two of three datasets with five passages and all three with ten passages.The authors attribute the gains to LLaMAFT’s 4k-token context window, which can require truncating conversation history; REFRAG maintains performance with many passages through compression.Source: §5.2 · p. 8, Table 4 · p. 9

6 Related Works

Related work covers retrieval-augmented language modeling, efficient long-context methods, compression, and prompt compression. The authors distinguish REFRAG by its support for precomputed chunk embeddings at arbitrary prompt positions and adaptive compression.

Retrieval-Augmented Language Modeling

  • Retrieval-augmented language modeling includes retrieved passages prepended to prompts and fusion-in-decoder, which removes cross-document attention but does not compress in the decoder.Source: §6 · pp. 9–10

Efficient Long-Context LLMs

  • Prior long-context approaches target different bottlenecks: compressed attention reduces attention complexity, StreamingLLM reduces KV-cache memory, and CEPE and PCC are limited to prefix contexts.The paper describes these methods as complementary or unsuitable for tasks such as multi-turn RAG and summarization; Kuratov et al. (2025) study compression into a single embedding.Source: §6 · p. 10

Compressive transformer

  • Earlier compressive-transformer methods reduce KV-cache memory or recursively compress context, but the cited approaches do not reduce decoding latency or support precomputed chunks at arbitrary prompt positions.The authors characterize REFRAG as enabling precomputation, arbitrary-position use, and adaptive compression rates without recomputing chunk embeddings.Source: §6 · p. 10

Prompt compression

  • Prompt-compression methods such as LLMLingua and LongLLMLingua reduce input length, while ranking and pruning approaches remove tokens or sentences estimated to be less informative.The paper describes these approaches as complementary and potentially integrable with REFRAG.Source: §6 · p. 10

7 Conclusion

The conclusion presents REFRAG as a RAG-focused decoding framework that reduces latency and memory by compressing, sensing, and expanding context representations. Across RAG, multi-turn conversations, and long-document summarization, it reports faster TTFT without loss in perplexity or downstream accuracy.

7 Conclusion

  • REFRAG uses RAG context sparsity and block-diagonal attention patterns to reduce memory and inference latency, achieving up to 30.85× TTFT acceleration without loss in perplexity or downstream accuracy.The conclusion reports 3.75× acceleration over previous state-of-the-art methods across RAG, multi-turn conversations, and long-document summarization.Source: §7 · p. 11

A Additional Discussion

REFRAG’s latency and throughput gains depend on context length, while selective compression uses a learned policy to preserve important chunks and compress others. The policy is trained to choose chunks based on their contribution to predicting output tokens.

A Additional Discussion

  • With short contexts, REFRAG achieves k× acceleration in TTFT and up to k× in throughput; with longer contexts, both reach up to k2×.The analysis assumes precomputed, cached chunk embeddings and focuses on LLaMA-2-7B; the authors state the results should generalize to other models.Source: A Additional Discussion · p. 16, Table 6 · p. 16
  • At context length 16384, REFRAG achieves 16.53× TTFT acceleration with cache and 8.59× without cache, versus CEPE’s 2.01× and 1.04×.With longer context, TTFT acceleration reaches 32.99×; long-context throughput acceleration reaches 6.78× and 6.06×.Source: A Additional Discussion · p. 16

A.1 Modeling REFRAG Selective Compression

  • For context of s tokens split into L chunks, pretraining randomly retains T′ = pL chunks as original tokens and compresses the rest, enabling mixed inputs at arbitrary positions.The decoder predicts the output from the mixed sequence while preserving its autoregressive property.Source: §A.1 · pp. 16–17
  • The policy selects T′ chunks from L to maximize a reward, framing selective compression as a discrete optimization problem with NP-hard exact solutions.The authors adopt sequential decision-making and reinforcement learning to train the selection policy.Source: §A.1 · p. 17
  • A two-layer transformer scores chunk embeddings, and the policy samples T′ chunk indices sequentially; logits are not recomputed after each selection to speed training.The PPO-style objective uses an advantage based on negative log-perplexity over the output tokens.Source: §A.1 · p. 17

B Additional Details on Experimental Settings

The experimental settings define REFRAG’s evaluation inputs and compare it with LLaMA-based and retrieval-augmented baselines. The setup includes both compressed-chunk variants and models that use different amounts or forms of context.

B.1 Additional Details on Baselines

  • Unless otherwise specified, baselines use LLaMA-2-7B to match prior methods, and REFRAG uses RoBERTa-large to encode context into compressed chunk embeddings.REFRAGk denotes a compression rate k; REFRAGRL uses the RL policy for selective compression.Source: §B.1 · pp. 17–18
  • Perplexity evaluations use 4096-token examples with 2048 context tokens followed by 2048 output tokens.For REFRAG, the context is encoded and output-token perplexity is evaluated; the cited setup also includes other encoder-decoder combinations in ablations.Source: §B.1 · p. 18
  • LLaMA-No Context omits context, LLaMA-Full Context uses the whole sequence, and LLaMAK adds the last K context tokens for comparison at matched context length.Full-context perplexity serves as a reference upper bound, while LLaMAK is described as a strong baseline when context-token counts are matched.Source: §B.1 · p. 18
  • CEPE encodes the context with a T5-like architecture; CEPED is its instruction-fine-tuned variant, while LLaMA-32K extends LLaMA-2-7B context from 4K to 32K.Source: §B.1 · p. 18
  • REPLUG ensembles generation across retrieved contexts, with REPLUGChat and REPLUGFT denoting its LLaMA chat and downstream-task-fine-tuned variants.Source: §B.1 · p. 18

B.2 Additional Details on Hyperparameters and Experimental Settings for CPT

The appendix specifies training settings and a curriculum that shifts toward longer examples, then details how latency and throughput are calculated and measured. Its calculations distinguish computation-bound prefill from memory-bound generation.

B.2 Additional Details on Hyperparameters and Experimental Settings for CPT

  • Training continues from a continual-pretraining checkpoint, uses p = 0.1 to retain 10% of chunks as original tokens, and runs four epochs with a 4% warm-up and batch size 256.The setup uses AdamW, a cosine learning-rate schedule, and a two-layer projection MLP; training uses bfloat16 and FSDP on 8 nodes with 8 H100 cards each.Source: §B.2 · p. 19

B.3 Curriculum learning data mixture

  • The nine-stage geometric curriculum gradually increases the share of longer-sequence data, allocating more samples to longer-context examples.Table 8 describes the stage mixtures and total sample counts for each data type.Source: §B.3 · p. 19, Table 8 · p. 20, Figure 6 · p. 20

B.4 Detailed Calculation of Acceleration in Latency and Throughput of Our Model

  • For LLaMA-2-7B loaded in bfloat16, TTFT is primarily computation-limited by attention and feedforward operations; at sequence length s, l layers, and batch size b, prefill computation is (24d2+4ds)lbs.Source: §B.4 · pp. 20–21
  • Generation latency is modeled as memory-bound, accounting for KV-cache transfers of 4dlb(s + o) bytes and model-parameter transfers of 2n bytes.Throughput is defined as bo divided by TTFT plus data latency; Table 9 compares KV-cache memory, TTFT, generation latency, and throughput.Source: §B.4 · p. 21
  • Empirical resource measurements use batch size 1 on one A100, with a shared modeling file and hyperparameters across baselines to isolate differences between model designs.The controlled setup aims to reduce environmental factors that could favor a method.Source: §B.4 · p. 21

C Additional Experimental Results

Additional experiments examine attention across retrieved passages, compression rates, model combinations, and performance in RAG and summarization. The reported RAG comparisons find similar performance among baselines at matched context counts.

C Additional Experimental Results

C.1 Additional Contextual Application - Summarization Task

The summarization evaluation tests REFRAG on long scientific articles from Arxiv and Pubmed. Under equal decoder-token counts, REFRAG achieves the best reported performance, with REFRAG16 outperforming REFRAG8 at 128 tokens.

C.1 Additional Contextual Application - Summarization Task

  • The task generates an abstract from an entire long scientific article, using Arxiv and Pubmed articles from the Cohan et al. (2018) dataset.REFRAG and LLaMA are fine-tuned on the datasets and evaluated on their validation sets.Source: §C.1 · p. 24
  • The evaluation probes whether summarization benefits more from additional document information in compressed form or less information without compression.The authors motivate this comparison by noting that correct summarization requires information from the whole document.Source: §C.1 · p. 24
  • Under the same decoder-token count (and latency), REFRAG achieves the best performance; at 128 tokens, REFRAG16 outperforms REFRAG8.The authors attribute the latter comparison to REFRAG16 incorporating more document information at a higher compression rate.Source: §C.1 · p. 24, Table 21 · p. 30

Related papers

Loading 2509.01092v2…