Rethinking RAG-based Decoding
Free while signed in. Answers cite the passages they came from.
REFRAG replaces most retrieved tokens with precomputed chunk embeddings at decode time, then selectively expands only the few chunks that matter. This exploits block-diagonal attention in RAG prompts to cut latency and memory while preserving accuracy across RAG, multi-turn dialog, and long-doc summarization.
Core idea: Chunk the retrieved context, encode each chunk with a lightweight encoder, project to the decoder’s embedding size, and feed embeddings directly alongside the user query; an RL policy decides which chunks to keep uncompressed (“compress anywhere,” not only in the prefix).
Big speedups without accuracy loss: Up to 30.85× time-to-first-token acceleration vs LLaMA (and 3.75× over CEPE) at high compression rates, with comparable perplexity; throughput gains up to 6.78×.
Longer effective context: Compression lets the model handle much larger contexts (reported 16× extension) while maintaining or improving perplexity as sequence length grows.
RAG wins under fixed latency: With the same latency budget, REFRAG uses more passages and outperforms a LLaMA baseline on 16 RAG tasks. Aggregated plots and detailed results show gains for both strong and weak retrievers.
Generalization across applications: On multi-turn conversational QA, REFRAG preserves longer history and improves scores as passages and turns increase. On long-document summarization, it achieves the best ROUGE at matched decoder tokens.
Get next week’s papers.
The same picks and the same summaries, in your inbox. Free, and 176 issues deep.
Subscribe on Substack