PaperlayerView exact revision on arXiv

Source-linked AI summary

Context Embeddings for Efficient Answer Generation in RAG

David Rau, Shuai Wang, Hervé Déjean, Stéphane Clinchant

arXiv:2407.09252v3cs.CLcs.IR

TL;DR

Long RAG inputs slow answer generation, while existing compression approaches have limitations in effectiveness, compression-rate flexibility, and multi-document support. COCOM compresses contexts into embeddings and reports faster inference with higher performance than existing context-compression methods. It also supports varying compression rates and multiple contexts.

  • Problem

    Long retrieved contexts slow RAG generation, and current embedding-compression approaches are described as limited in effectiveness, compression-rate flexibility, and multi-document support.

    Source: §1 · pp. 1–2

  • Method

    COCOM compresses contexts into context embeddings that can be used by the LLM for answer generation.

    Source: §3.2 · p. 3, §7 · pp. 9–10

  • Results

    Up to 5.69× lower inference time and 22× fewer GFLOPs are reported while maintaining high performance and outperforming existing context-compression methods.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    Varying compression rates and handling multiple contexts allow COCOM to trade answer quality against generation speed within the reported RAG setting.

    Source: §7 · p. 10

    Experiments were limited to a 7-billion-parameter model and five documents, leaving performance with larger models and larger collections unexplored.

    Source: §8 · p. 10

Abstract

from arXiv · show

Retrieval-Augmented Generation (RAG) allows overcoming the limited knowledge of LLMs by extending the input with external information. As a consequence, the contextual inputs to the model become much longer which slows down decoding time directly translating to the time a user has to wait for an answer. We address this challenge by presenting COCOM, an effective context compression method, reducing long contexts to only a handful of Context Embeddings speeding up the generation time by a large margin. Our method allows for different compression rates trading off decoding time for answer quality. Compared to earlier methods, COCOM allows for handling multiple contexts more effectively, significantly reducing decoding time for long inputs. Our method demonstrates a speed-up of up to 5.69 $\times$ while achieving higher performance compared to existing efficient context compression methods.

1 INTRODUCTION

RAG can improve knowledge-intensive question answering with retrieved context, but long inputs slow generation and current compression methods leave gaps in effectiveness, compression-rate flexibility, and multi-document support. COCOM compresses contexts into embeddings and reports faster inference alongside high performance.

1 INTRODUCTION

  • RAG supplies external context for knowledge-intensive tasks, but adding more context considerably slows generation during inference.The paper connects the slowdown to transformer self-attention’s space and memory requirements growing exponentially with input length.Source: §1 · p. 1
  • Existing embedding-compression approaches are described as limited by large compressors, low effectiveness, fixed compression rates, and single-document support.The paper notes that these methods may leave the decoder untuned and hypothesizes that freezing it hinders use of compressed contexts.Source: §1 · p. 2
  • COCOM compresses contexts into a small number of context embeddings that are input to the LLM for answer generation.The paper describes it as a multi-context compression method using a single model for compression and answer generation.Source: §1 · p. 2
  • The paper presents COCOM as more effective than current context-compression approaches and offers different compression rates to trade inference time against answer quality.It also says COCOM supports multiple contexts, unlike prior approaches described as limited to one document.Source: §1 · p. 2
  • Up to 5.69× lower inference time and 22× fewer GFLOPs are reported while maintaining high performance.The introduction also reports an ablation of pretraining collection, pretraining, fine-tuning, decoder freezing, target dataset, and decoder training.Source: §1 · p. 2

2 RELATED WORK

Related work frames context compression as a way to reduce RAG input size and generation time, with lexical- and embedding-based approaches as two main lines. The discussion compares their mechanisms and limits, especially whether they fit multi-document RAG.

2 RELATED WORK

  • Embeddings are motivated as a way to provide passage content to an LLM without supplying the original context in token form.The paper situates this approach among concurrent embedding-based and lexical-based compression research aimed at reducing RAG input size and computational cost.Source: §2 · p. 2

2.1 Lexical-based Compression.

  • Lexical methods select or summarize context tokens, while document-level selection includes or excludes retrieved documents according to the query.Because these approaches rely on query inputs, some compression must occur online, which the paper says slows generation.Source: §2.1 · p. 2

2.2 Embedding-based Compression.

  • Embedding-based compression represents context with summary embeddings, but prior methods differ in training, scope, and suitability for RAG.AutoCompressor uses iterative chunk aggregation and at least 50 embeddings; ICAE uses the decoder model but freezes its decoder, while GridLM is described as single-context and not speeding decoding.Source: §2.2 · pp. 2–3
  • xRAG reuses frozen ranking representations, but the paper says these representations may lack information needed to recover context and that lightweight compressors can have low effectiveness.The cited discussion contrasts retrieval representations with representations learned for compression, which it says can reconstruct original context.Source: §2.2 · p. 3

2.3 Overview

  • The paper notes that many prior embedding-compression works use one relevant document, so their reported effectiveness may not transfer directly to RAG with potentially irrelevant, multiple retrieved documents.It argues that fine-tuning the decoder is a simple yet necessary solution for handling multiple context representations.Source: §2.3 · p. 3

3 METHODOLOGY

The RAG task uses a ranking system to retrieve context segments for a user input, then generates a response from that context and input. Because context can be much longer than the query, it substantially expands the LLM input.

3.1 Task Definition: RAG

  • A ranking system builds an index from a collection and searches it at request time to return context segments relevant to the user input.The task definition allows the ranking system to be multistaged.Source: §3.1 · p. 3
  • The LLM generates a response based on both retrieved context and the user input.Source: §3.1 · p. 3
  • RAG context dramatically increases the LLM input because the context length is much greater than the user-input length.Source: §3.1 · p. 3

3.2 COCOM: Effective Context Compression

COCOM replaces long context token sequences with smaller sets of context embeddings that an LLM uses to generate responses. Its variable compression rate controls the embedding count, and it can compress multiple contexts independently for use by the LLM.

3.2 COCOM: Effective Context Compression

  • COCOM compresses context tokens into a smaller set of embeddings, which the LLM combines with the user input to generate a response.The compressor is trained to represent input content in compressed form, and the LLM is trained jointly to decode those embeddings for answering queries.Source: §3.2 · pp. 3–4
  • Independent context compression allows embeddings to be precomputed and stored offline, reducing inference computation and answer-generation input size.The paper says this smaller input leads to a large generation speed-up.Source: §3.2 · p. 4
  • COCOM may also use an embedding model as compressor, including lightweight encoder-only models such as BERT.Source: §3.2 · p. 4

3.2.1 Adaptable Compression Rate.

  • The number of context embeddings can vary with compression rate and tokenized input length; for a 128-token context at rate 64, COCOM uses two embeddings.Source: §3.2.1 · p. 4

3.2.2 Multiple Contexts.

  • For multiple retrieved passages, COCOM compresses each context independently and inserts [SEP] tokens between their embeddings before passing them to the LLM.Source: §3.2.2 · p. 4

3.3 Pre-training Context Embeddings

COCOM pre-training uses two autoregressive variations of next-token prediction to compress context into context embeddings and train the LLM to use them. The tasks pair reconstruction of compressed inputs with generation conditioned on compressed context.

3.3.1 Auto-encoding with Context Embeddings.

  • Auto-encoding trains the compressor and LLM jointly to recover the original input tokens from compressed context embeddings.This is a preliminary step toward answering questions from context embeddings: first the model learns to compress and decompress the same input effectively.Source: §3.3.1 · pp. 4–5

3.3 Pre-training Context Embeddings

  • The pre-training objective combines auto-encoding and language modeling so the model learns both to reconstruct compressed input and to use context embeddings.The language-modeling task complements auto-encoding, which alone could bias the LLM toward recovering the input rather than leveraging its content.Source: §3.3.1 · p. 5, §3.3.2 · p. 5

3.3.2 Language Modeling from Context Embeddings.

  • Language modeling trains the LLM to generate a continuation conditioned on compressed representations of the preceding input.The input is split into a compressed first part and a continuation; this task is intended to teach the model to leverage the embeddings rather than only reconstruct the input.Source: §3.3.2 · p. 5

3.4 Fine-tuning

For downstream RAG, the model is fine-tuned to generate target answers from a question and retrieved context compressed into context embeddings.

3.4 Fine-tuning

  • Instruction fine-tuning trains the LLM to generate target responses from questions and retrieved contexts compressed into context embeddings.The training data combines publicly available QA datasets, and updates are based only on the target responses.Source: §3.4 · p. 5

4 EXPERIMENTAL SETUP

The experimental setup specifies the models, training data and procedures, evaluation tasks and metrics, and comparison baselines for COCOM. It evaluates question answering with retrieved contexts and context-compression alternatives.

4.1 Implementation Details

  • COCOM uses Mistral-7B-Instruct-v0.24 for answer generation and compression, while COCOM-light uses bert-base-uncased; experiments use compression rates ξ=1, 16, and 128 and top-5 retrieved documents.Retrieval uses SPLADE-v3 [18] with top-50 reranking by DeBERTa-v3 [9].Source: §4.1 · p. 5

4.2 Training

  • Both pre-training and fine-tuning use parameter-efficient LoRA tuning..Source: §4.2 · p. 5

4.2.1 Pre-training.

  • Pre-training uses equal-probability samples from auto-encoding and language-modeling tasks, with 10m Wikipedia-KILT [27] samples formed from 128-token chunks.The authors report trying other task ratios but finding equal sampling performed best.Source: §4.2.1 · p. 5

4.2.2 Fine-tuning.

  • Fine-tuning concurrently combines training samples from ten question-answering datasets, then filters out queries with more than 128 tokens and labels with more than 64 tokens.The datasets are Natural Questions, MS MARCO, adversarial QA, HotpotQA [35], WikiQA, SCIQ [15], ASQA [30], TriviaQA [16], Freebase QA, and squad [28].Source: §4.2.2 · p. 5

4.3 Evaluation

  • Evaluation covers Natural Questions, TriviaQA [16], HotpotQA [35], ASQA [30], and PopQA [22], using Exact Match as the main metric and Match for comparison with prior work.Match checks whether the label appears as an exact match within the generated answer.Source: §4.3 · p. 5, §4.3.1 · p. 5

4.4 Baselines without Context Compression

  • The uncompressed RAG baseline receives five retrieved contexts and serves as an upper bound, while the closed-book baseline answers without context and serves as a lower bound.The base model for these baselines is Mistral-7B-Instruct-v0.2.Source: §4.4 · pp. 5–6

4.5 Baselines with Context Compression

  • Prior context-compression baselines leave decoder LLMs untuned, which the authors argue limits their ability to leverage context embeddings for question answering.They note that answering from context embeddings differs fundamentally from standard language modeling.Source: §4.5 · p. 6
  • Baseline constraints include single-context compression for AutoCompressor [4] and xRAG, while ICAE [8] receives concatenated top-five contexts truncated to 512 tokens.ICAE's maximum compressed output is 128 tokens, approximately a compression rate of 4; AutoCompressor and xRAG use the top retrieved document.Source: §4.5 · p. 6

5 RESULTS

COCOM and COCOM-light are evaluated for answer effectiveness and computational efficiency, with ablations examining context count, pretraining, decoder tuning, and fine-tuning data. The results show trade-offs between compression, answer quality, and efficiency.

5.1 Main Results

  • COCOM significantly outperforms existing context-compression methods across tested datasets and rates, including xRAG Mixtral-8x7B, while compressing contexts by up to 128 times with relatively high average performance versus uncompressed RAG.Average performance decreases 4 points at ξ=4 and 10 points at ξ=128; Match results also retain COCOM’s significant advantage over other methods despite higher scores for methods with untuned decoders.Source: §5.1 · p. 6, Table 2 · p. 7

5.2 COCOM-light

  • COCOM-light is an effective alternative for resource-constrained settings at lower compression rates, but its performance drops considerably at ξ=128.It uses BERT as a computationally lighter context compressor.Source: §5.2 · pp. 6–7

5.3 Computational Efficiency

  • COCOM reduces answer-generation inference time by up to 5.69× and operations by up to 22× versus RAG without compression; COCOM-light compresses documents up to 89× faster than COCOM.Higher compression rates reduce index storage requirements, but the associated trade-off lowers answer-generation quality.Source: §5.3 · p. 7, Table 3 · p. 7, Table 4 · p. 7

5.4.1 Handling multiple contexts.

  • Using the default k=5 instead of k=1 yields substantial gains on NQ and ASQA at every compression rate; k=1 still beats single-document baselines ICAE and xRAG.The comparison measures Exact Match.Source: §5.4.1 · p. 8, Table 5 · p. 8

5.4.2 Pre-training Context Compression.

  • Dedicated auto-encoding and language-modeling-from-context-embedding pretraining improves downstream QA performance, possibly because compression is too complex to learn alongside the downstream task or needs larger fine-tuning datasets.Source: §5.4.2 · p. 8, Table 7 · p. 9

5.4.3 Pre-training Corpus.

  • Pretraining on Fine-Web’s sample-10BT subset causes a slight decrease in downstream performance, while results indicate robustness to variation in the pretraining corpus.The tested subset contains 10 million tokens and differs from the target corpus.Source: §5.4.3 · p. 8

5.4.4 Decoder LLM Tuning.

  • Tuning the decoder is critical for high effectiveness when using context embeddings, supporting the authors’ hypothesis that those embeddings differ from the decoder’s trained input-token embeddings.The ablation compares tuning all components with freezing the decoder and tuning only the compressor.Source: NQ ASQA · p. 8, §5.4.4 · p. 8, Table 6 · p. 8

5.4.5 Fine-tuning Data.

  • Fine-tuning on multiple QA datasets improves transferability and average performance across datasets, although training only on NQ performs slightly better on NQ itself.The comparison evaluates zero-shot transfer to other datasets after fine-tuning on NQ alone or multiple datasets.Source: §5.4.5 · p. 8, Figure 3 · p. 8

6 ANALYSIS

The analysis examines why performance falls at high compression rates and compares answers from different methods on one NQ question. Both COCOM variants reconstruct inputs effectively at lower compression rates, but struggle at 128; in the case study, both answer correctly at rate 4 but have difficulty at 128.

6.1 Context compression

  • At compression ratio ξ=128, both full and lightweight models have significant difficulty reconstructing inputs, with the problem notably more pronounced for COCOM than COCOM-light.At lower compression factors (ξ=4, 16), both effectively master auto-encoding. On the second pre-training task, COCOM consistently outperforms COCOM-light, a finding that correlates with final question-answering effectiveness.Source: §6.1 · p. 9

6.2 Case Study Answer Quality

  • At rate 4, COCOM and COCOM-light answer the selected NQ question correctly; at rate 128, both encounter difficulties.The outputs at rate 4 are Rosalind Bailey for both models; at rate 128, COCOM returns Alison Steadman and COCOM-light returns Rosalind Elliott.Source: §6.2 · p. 9
  • Without RAG, the LLM gives an irrelevant name; xRAG returns an incorrect entity, and ICAE gives an unreasonable answer.The paper attributes xRAG’s error likely to limitations in reading compressed embeddings accurately; it notes that the xRAG response was intentionally capped at 30 tokens.Source: §6.2 · p. 9

7 CONCLUSION

The conclusion presents COCOM as a method for compressing multiple contexts into embeddings that augment answer generation. It reports faster generation and a flexible efficiency–effectiveness trade-off, particularly relevant for long inputs.

7 CONCLUSION

  • COCOM compresses multiple contexts into context embeddings that can be precomputed to augment answer generation.The conclusion also says the method tunes all model components and outperforms existing context-compression methods in RAG.Source: §7 · pp. 9–10
  • The number of compression tokens can be varied to trade answer quality against generation speed, with the distinction from token-form context most apparent for very long inputs.The conclusion says multiple contexts enhance generation quality while making use of reduced decoding time.Source: §7 · p. 10
  • The authors hope the work will inspire further research and support efficient and effective RAG deployment.Source: §7 · p. 10

8 LIMITATIONS

The authors identify resource constraints that limit the model scale and document count in their experiments, as well as the tasks and languages evaluated.

8 LIMITATIONS

  • Experiments were limited to a 7-billion-parameter model, leaving larger models such as LLaMA70B and Mixtral7x8B unexplored.The authors note that those models might offer enhanced performance but require substantial training and inference resources.Source: §8 · p. 10
  • Experiments used only five documents, so they may not reflect effectiveness with larger document collections.The authors state that efficiency gains could be more pronounced at larger scale, but this remains untested in their experiments.Source: §8 · p. 10
  • Evaluation covered only QA tasks and English corpora, leaving performance on diverse tasks and multilingual datasets unassessed.Source: §8 · p. 10

A APPENDIX

The appendix includes tables on comparison results, pretraining and fine-tuning hyperparameters, and the datasets used for fine-tuning COCOM variants.

A APPENDIX

  • Table 9 compares COCOM variants with other context-compression methods using Match, with all methods using five context passages unless noted otherwise.The caption marks single-context methods, upper and lower baselines, and statistical non-significance relative to COCOM at ξ=4.Source: Table 9 · p. 12
  • Tables 10 and 11 list pretraining and fine-tuning hyperparameters, while Table 12 describes the fine-tuning dataset collection and its query and label-length filters.Table 12 excludes queries longer than 128 tokens and labels longer than 64 tokens.Source: Table 10 · p. 12, Table 11 · p. 13, Table 12 · p. 13

Related papers

Loading 2407.09252v3…