Lighter And Better: Towards Flexible Context Adaptation For Retrieval Augmented Generation

TL;DR

FlexRAG tackles the high cost and limited effectiveness of retrieval-augmented generation by compressing retrieved contexts into learnable embeddings and down-sampling them with a flexible budget. The approach combines a compressive encoder, a selective compression mechanism, and a two-stage training workflow (unlabeled pre-training followed by task-specific fine-tuning) to optimize RAG performance without modifying the base LLM. Empirical results across multiple QA datasets show superior cost efficiency and robustness, with strong gains in open-domain and long-sequence multi-document QA, especially under tight context budgets. The framework offers adaptable compression ratios and methods, enabling practical deployment with varied hardware and task requirements while preserving general task capabilities.

Abstract

The existing Retrieval-Augmented Generation (RAG) systems face significant challenges in terms of cost and effectiveness. On one hand, they need to encode the lengthy retrieved contexts before responding to the input tasks, which imposes substantial computational overhead. On the other hand, directly using generic Large Language Models (LLMs) often leads to sub-optimal answers, while task-specific fine-tuning may compromise the LLMs' general capabilities. To address these challenges, we introduce a novel approach called FlexRAG (Flexible Context Adaptation for RAG). In this approach, the retrieved contexts are compressed into compact embeddings before being encoded by the LLMs. Simultaneously, these compressed embeddings are optimized to enhance downstream RAG performance. A key feature of FlexRAG is its flexibility, which enables effective support for diverse compression ratios and selective preservation of important contexts. Thanks to these technical designs, FlexRAG achieves superior generation quality while significantly reducing running costs. Comprehensive experiments on various question-answering datasets validate our approach as a cost-effective and flexible solution for RAG systems.

Figures (3)

• Figure 1: Comparison of related techniques. 1) Context compression: token embeddings are compressed into compact summary vectors. 2) Context filtering: important token embeddings are filtered from the input prompt. 3) Prompt tuning: soft-prompt is learned to improve the downstream task. 4) FlexRAG: unifying all functions in one framework, with compressive embeddings (summary vectors) down-sampled by importance (filtering) and learned to optimize the RAG performance (prompt tuning).
• Figure 2: Architecture of FlexRAG. It transforms the retrieved contexts $\mathbf{X}_{retr}$ into compressive embeddings $\mathbf{E}_{retr}$ using the compressive encoder. With estimated importance, it down-samples $\mathbf{E}_{retr}$ into $\mathbf{E}'_{retr}$ as the compressed context for RAG.
• Figure 3: Performance on NQ and TriviaQA with varied #Doc.