Lexical-based compression focuses on either selecting tokens from the context (Li, 2023) or summarizing contexts (Xu et al., 2023), both aiming to retain essential information while reducing overall context size. LLMLingua comprises a query-independent token filtering module that uses a LLM to first select important tokens in the context. Then, a query-dependent token classifier is used to select tokens to form the compressed context.

On the other hand, Zhu et al. (2024) do not consider compression at the term level, but at the document level. Retrieved documents are either included or excluded with respect to the query. Only the included documents form the context for answer generation. It is worth noting that current lexical-based compression approaches all rely on specific query inputs. Therefore, compression needs to be (partially) processed online not allowing to compress documents offline, slowing down generation time.

Embedding-based compression approaches focus on compressing the context into one or multiple summary embeddings that can be directly interpreted by the decoder model. This first work of this line is called AutoCompressor (Chevalier et al., 2023). This approach attempts to compress contextual information by segmenting it into randomly segmented chunks, subsequently aggregating these into summary embeddings through an iterative process until target embedding size is met. However, the training of the summary embeddings relies exclusively on next token prediction tasks, raising concerns about their ability to effectively encapsulate relevant contextual data. Furthermore, AutoCompressor is designed primarily for long contexts, generating a minimum of 50 summary embeddings. Such a configuration is not suitable for common RAG pipelines where short passages are retrieved, such as KILT.

Building up on AutoCompressor, ICAE by Ge et al. (2024) explores training a context compressor using the same LLM as the decoder model, and compress only once to get the summary embeddings. However, their approach limits the model’s capacity by using a frozen decoder module, preventing the accumulation of gradients from the decoder part during training. In this paper, we argue that decoder training is an important factor that strongly impacts the performance of the model. We illustrate this argument in Section 4.2.1.

Furthermore, GridLM Muennighoff et al. (2024) addresses the issue of double decoding the same context first for retrieval and then again as the provided context to the LLM. They use the same LLM for ranking and generation which allows them to cache all representations during encoding the contexts and to reuse them during generation. This approach compared to ours is limited to only a single context, does not speed up decoding time, and results in gigantic storage requirements.

Cheng et al. (2024) propose xRAG concurrently to our method. They directly reuse frozen ranking representations based on embedding models while freezing the decoder. Although this approach successfully resolves the double decoding problem, it suffers from low effectiveness because the representation is not trained prior to its application to compression tasks. This issue becomes particularly challenging when light-weight encoder models, such as DPR with 109 million parameters, are used as compressors. In such cases, the model achieves similar effectiveness to the Mistral-7b model when retrieval is not applied ¹¹1By default, xRAG uses a 7B SFR LLM-based ranking model as compressor. On the other hand, using retrieval representations from lightweight models for compression is counter-intuitive. Representations gathered from retrieval tasks may lack sufficient information to fully recover the context. Conversely, representation learned for compression demonstrate its capacity to reconstruct the original context (Ge et al., 2024). This suggest that, upon further adjustment, it may show a higher potential to serve as an effective retriever.

In Table 1 we contrast our method with the described related works on embedding-based compression. It is important to note that most previous works mentioned so far have only considered cases that may not directly apply to RAG settings but rather to long-context question answering. In their setting, only one relevant document is used for each query to fulfill the user request.

Therefore, such models are not naturally able to deal with effectively multiple documents. Furthermore, their reported effectiveness may not directly indicate the final performance in RAG systems, where the document may be potentially irrelevant, and often multiple top-retrieved documents are used. As a decoder model, by design, should be able to handle multiple context representations, we argue that fine-tuning the decoder is a simple yet necessary solution compared to existing works

RAG employs a ranking system $\mathcal{R}$ and a parametric generative language model $\theta_{LLM}$, where the ranking system can be multi-staged. First, the ranking system builds a search index $\mathcal{I}$ based on a collection. Then, at request time, the index $\mathcal{I}$ is searched yielding context segments²²2The segments can be at different granularities for instance sentences, passages, or entire documents. In this work, we focus on passages. $\mathcal{C}$ that are relevant to the user input $x$: $f_{\mathcal{I},\mathcal{R}}:\{x\}\rightarrow\mathcal{C}$.

Next, the LLM generates a response $r$ based on the context $\mathcal{C}$ and user input $x$:

Note how in RAG the context is added to the input of the LLM dramatically increasing the input to the LLM, as $|\mathcal{C}|\gg|x|$.

The main idea of COCOM is to enhance efficiency by compressing the context, which is typically given in surface form as input tokens into a smaller set of context embeddings which then serve as the input to the LLM. An overview of our entire pipeline is given in Figure 2. More formally, our approach can be described as follows:

Given a context $\mathcal{C}$ tokenized into a a sequence of tokens $\{t_{1},t_{2},$ $\dots,t_{n}\}$, a compressor model $\phi_{comp}$, we compress $\mathcal{C}$ into context embeddings $\mathcal{E}$, a smaller set of embeddings $\{e_{1},e_{2},\ldots,e_{k}\}$, where $k\ll n$. Each embedding $e_{i}\in\mathbb{R}^{d}$, with $d$ being the LLM’s hidden dimension.

Next, based on the compressed context embeddings $\mathcal{E}$ and the user input $x$ the LLM $\phi_{LLM}$ generates a response $r$:

The $\phi_{comp}$ model is trained to generate context embeddings that capture the content of the input tokens in a compressed form. As both models are trained jointly, $\theta_{LLM}$ learns to decode these context embeddings, extracting the relevant information required to answer user queries.

COCOM compresses the context-embeddings question independently. This means not only do individual contexts have to be contextualized by an LLM only once, but they can also be pre-computed offline and stored, drastically reducing computational costs of the LLM at inference time. Further, by only feeding a small number of context embeddings instead of the long context, the input size is reduced to a fraction leading to a massive speed-up for answer generation.

For COCOM, we utilize the same model for compression and answer generation $\phi_{comp}=\theta_{LLM}$. Therefore, we effectively train a single model on the two tasks. For the compression task, we prepend a special token <AE> to the input and depending on $\xi$ append a different number of context embedding tokens <CTX> at the end of the sequence. We directly use the representations of the last hidden layer as our context embeddings as input - to the same model - for the answer generation.

As demonstrated later in the experiments, our method also allows us to potentially employ any embedding model as a compressor; including more lightweight encoder-only models such as BERT ³³3See Section 5.2.

We propose two auto-regressive variations of the next-token prediction task to learn to compress context into context embeddings and to use these context embeddings as input to the LLM.

Following our earlier notation, the objective function for the standard next token prediction for input $\mathcal{X}=\{x_{1},x_{2},\dots,x_{T}\}$ can be written as:

For the downstream RAG application, we fine-tune the model on a question $q$, relevant context(s) retrieved by a retrieval system and compressed into context embeddings $\mathcal{E}$, which are combined into an instruction $I_{q,\mathcal{E}}$. We train the LLM to generate the target response $R=(r_{1},r_{2},\dots,t_{T})$. We fine-tune our models on a combined set of publicly available QA datasets. We employ instruction fine-tuning only updating the models based on the target responses.

We use Mistral-7B-Instruct-v0.2⁴⁴4https://https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2 as our backbone LLM for answer generation. For context compression in COCOM, we utilize the same model. For our more light-weight context compression, in COCOM-light, we employ bert-base-uncased⁵⁵5https://https://huggingface.co/google/bert-base-uncased. We apply three different compression rates: $\xi={1,16,128}$. We employ SPLADE-v3 (Lassance et al., 2024) with reranking top-50 using DeBERTa-v3 (He et al., 2021) as our retrieval system. For all our experiments we use top-5 documents as context.

For both pre-training and fine-tuning, we apply parameter-efficient LoRA tuning.

We evaluate our model on several widely used QA datasets. Natural Questions (Kwiatkowski et al., 2019), TriviaQA (Joshi et al., 2017), HotpotQA (Yang et al., 2018), ASQA (Stelmakh et al., 2022), and PopQA (Mallen et al., 2023).

We fine-tune the base model (Mistral-7B-Instruct-v0.2):

• •

  RAG - upper bound. The model receives the top-5 retrieved contexts, alongside the query and answers the question. This model serves as an upper bound in our experiment not applying context compression.

• •

  Closed Book - lower bound. (w/o RAG). The LLM generates an answer based on the query without any provided context. This serves as a lower-bound baseline.

We compare our models to the context compression methods mentioned below. As mentioned earlier these models tune only parts of their model components on the downstream data but leave their decoder LLM untuned applying it zero-shot. We argue this to be a major limitation, as answering questions from context embeddings differs fundamentally from the standard language modeling hindering the model to effectively leverage the context embeddings.

To ensure comparability among approaches we use the same retrieval system as mentioned earlier in Section 4.1.

• •

  Autocompressor  (Chevalier et al., 2023): We use the princeton-nlp/AutoCompressor-Llama-2-7b-6k checkpoint producing 50 summary vectors. As their model is limited to compressing one single context, we just use the top retrieved document as context.

• •

  ICAE.  (Ge et al., 2024): We use the Mistral-7B-Instruct-v0.2 LoRa-checkpoint⁹⁹9https://huggingface.co/sggetao/icae which uses the same base LLM as ours and is therefore directly comparable. ICAE is fine-tuned to compress a single long context, however, in our work we use multiple contexts. To alleviate this we concatenate the top five retrieved documents together as the context input for the model and truncate as the maximum length of 512 tokens. Note the model has a maximum output length of 128 compressed tokens, which approximately indicates a compression rate of 4 from its original concatenated context input.

• •

  xRAG. We utilize the xRAG-7b¹⁰¹⁰10https://huggingface.co/Hannibal046/xrag-7b, and 8x7B mixture-of-experts model ¹¹¹¹11https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1 alongside their strongest SFR compressor. The base model is again the same as ours, to ensure comparability. As their model is limited to compressing only a single context into a single compressed representation, we use the top retrieved document for the xRAG setting.¹²¹²12We also tested compressing five documents together, which yielded lower effectiveness. We apply their predefined stopping criteria for answer generation, which aims at cutting the verbose nature of a untuned decoder LLM.

The main results for COCOM are presented in Table 2. We measure performance following the standard practice for fine-tuned models using the Exact Match (EM) metric. Compared to existing context compression methods¹³¹³13Side note: As previously mentioned earlier in Section 2, existing context compression methods do not tune the decoder LLM and therefore compare their performance and make effectiveness claims against zero-shot baselines. However, we argue that tuning compression models while freezing the decoder LLM could not be considered zero-shot, as it involves tuning some parts of the model on the task data. This setting is akin to soft-prompt tuning (Cuconasu et al., 2024; Li and Liang, 2021), where the compressor model effectively parameterizes the soft prompt. Consequently, the performance of these methods should be regarded as intermediate between zero-shot and full decoder tuning and should be compared against similar tuning settings, such as soft prompt tuning. , our approach demonstrates a significantly (Tested with paired t-test (p¡0.05)). higher effectiveness across different compression rates for all datasets tested. COCOM even outperforms the much stronger xRAG Mixtral-8x7B model by a large margin having 8 times more parameters than COCOM. The highest performance is observed at a low compression rate ($\xi$=4). Increasing the compression rate results in a slight performance decline, which we will analyze further in Section 6.1.

Compared to our upper bound baseline RAG without compression, we reduce the context by up to 128 times while still maintaining relatively high performance on average over datasets.

Performance decreases on average 4 points for our strongest model (COCOM $\xi$ = 4) and 10 points for the highest compression rate (COCOM $\xi$ = 128). Compared to the lower bound baseline LLM without provided context we gain up to 17 points, adding only a small number of additional context embeddings to the input.

Note, while EM is a standard metric for evaluating tuned models, it might underestimate zero-shot decoder methods that do not adapt the decoder to generate answers. To address this, we also provide results using the Match metric in the appendix in Table 9. Although models that do not tune their decoder achieve relatively higher performance when measured in Match, our method’s effectiveness compared to other methods still remains consistently significantly higher.

Overall, considering the effectiveness and the efficiency gains from context compression (discussed further in Section 5.3), COCOM shows a very favorable trade-off.

Even though context compression has to be done only once offline, using a very large LLM can be costly, especially in resource-constraint settings. To this end, we propose COCOM-light, a computationally efficient context compression model based on BERT as a context compressor.

To alleviate the dimensional mismatch between the bert-based compressor and the - typically larger - LLM, we learn a linear projection layer $\bm{W}^{\gamma\times d}$, where $\gamma$ is the compression rate and $d$ the hidden dimension of the LLM. To obtain a set of Context Embeddings we leverage the last hidden representation of each input token. We simply split the hidden representations into blocks of length $\gamma$ and project each block into a single Context Embedding. This way, we learn a block-wise aggregation of the input representations that depending on the input length, and the compression rate $\gamma$ yields a different number of Context Embeddings per input. Note that a similar approach is applied in xRAG, where a projection layer is used on the embedding vector to resolve the dimensional mismatch. However, we argue that compressing using a single vector embedding could significantly restrict the compression quality, especially when using lightweight encoder models such as BERT. This restriction can result in much lower effectiveness compared to using a larger embedding model (Cheng et al., 2024).

We present the results in Table 2 measured in EM. Results for Match can be again found in the appendix in Table 9. We find that while being highly effective for small compression rates to drop considerably for the highest compression rate of $\xi$=128. COCOM-light, compared to other methods poses an effective alternative to it’s bigger counterpart COCOM, in resource-constrained settings.

We measure efficiency in answer generation time (ms), maximum GPU memory (GB), and number of operations per generated token (Giga FLOPs) using the torch profiler. We run the experiments on a single A100 40GB with a fixed batch size of 16¹⁴¹⁴14Maximum batch size that fits on GPU across models.. We load the model in half-precision and use the PyTorch inference mode. We discard the first warm-up batch from the measurement and measure the bare forward pass of the model. Note decoding results are independent of the compressor, therefore COCOM  and COCOM-light share efficiency results.

We show our efficiency results for answer generation in Table 3 for different compression rates $\xi$ and compare them to RAG without context compression. Context compression with COCOM reduces answer generation time , GPU memory, and the number of operations drastically with up to 5.69 $\times$ less inference time cost, 1.27 $\times$ GPU memory, and 22 $\times$ GFLOPs compared to no compression.

In addition, Table 4 presents the compression costs for all documents in the kilt-100w ( 24m contexts) collection using COCOM-light models at various compression rates. COCOM-light models demonstrate significantly faster compression speeds compared to the COCOM model by employing a much computationally lighter compressing module (up to 89 $\times$). Index size varies inversely with compression rate: higher compression rates result in smaller index storage requirements. However, this trade-off leads to lower quality in answer generation, as shown in Table 2.

In the following section, we run additional ablation experiments for COCOM and COCOM-light. Most results can be found in Table 6. We report performance in Exact Match on two datasets (NQ and ASQA).

In our earlier results in Section 5.1, we observe a decline in performance with higher compression rates, particularly for the lightweight compressor in COCOM-light. To investigate potential reasons for this drop, we assess the model’s ability to perform the two pre-training tasks: (i) compressing and decompressing input (auto-encoding) and (ii) language modeling from compressed representations after pre-training.

Our results in Table 7 indicate that both models effectively learn the auto-encoding task at lower compression rates ($\xi$=4, 16), but struggle to recover the input when the context is compressed into fewer embeddings ($\xi$=128), with this issue being more pronounced for the lightweight compression module.

We identify two possible explanations: First, compressing longer contexts is inherently more challenging and might require additional objectives. Second, decompressing information from a smaller set of Context Embeddings may be more difficult due to the sequential decoding nature of decoder-only models. Introducing additional pause tokens (Goyal et al., 2024) could help alleviate this issue, providing the model with a means to hierarchically decompress information, drawing on ideas from Chain-of-Thought prompting (Wei et al., 2022). We also experimented with pre-training on more samples but found that this did not improve downstream performance. Regarding the second pre-training task, it is noteworthy that COCOM-light outperforms its larger counterpart in language modeling from Context Embeddings. This analysis shows compressing and re-constructing longer texts is challenging and needs further investigation.

We investigate the answers generated with different models. For this, we randomly select a query from the NQ dataset and compare the responses generated by each method. Table 8 presents the responses to the selected question.

From the responses, we observe that without RAG, the LLM tends to hallucinate and provide an irrelevant name as an answer. On the other hand, XRAG understands the question but returns an incorrect named entity, likely due to limitations in reading compressed embeddings accurately. ICAE struggles to comprehend the question, resulting in an unreasonable answer. Both COCOM and COCOM-light successfully answer the question correctly at a compression rate of 4. However, they encounter difficulties when the compression rate is increased to 128.

It is also worth noting that the XRAG response was intentionally truncated to a maximum of 30 tokens in its original publication, with the stopping criteria involving halting at punctuation mark such as periods, commas, and colons.