1 Introduction↩︎

Reasoning over long documents is common in scenarios of retrieval-augmented generation (RAG). This is a computationally costly process, both as far as time and memory. Time is impacted by processing the document itself and self-attending over its computed representations in later parts of the generation process. Memory costs spike due to the key-value (KV) cache of the processed document. The common way to reduce these costs is soft context compression, where a sequence of continuous representations is pre-computed. This representation is compressed in the sense that its length is significantly lower than the document length, thereby reducing both time and memory costs of reasoning over the document. This compressed representation is computed once, and then retrieved as needed. This important problem is receiving significant attention [1]–[3].

We study both compressor model design and the compression training process, with simplicity in mind. The compression encoder we design is composed of an encoding LLM, and straightforward mean-pooling operations to collapse together representations to achieve a target compression ratio. This approach adds no parameters beyond the encoder LLM, and the computation beyond the encoding of the document is minimal. It also naturally allows for compression ratio flexibility, raising the question of the benefits or downsides of training the same compressor to support multiple compression ratios. We design a simple training objective and process to achieve this. This is motivated foremost by the benefit of training a single model to serve different compute budgets, rather than maintaining and training a separate model for each compression ratio. It also allows to examine if and when training to compress at multiple ratios can perform better than training for a single ratio.

We construct a rigorous evaluation suite using multiple question-answering (QA) datasets. We distinguish between datasets that are part of our training set, and these that are completely held-out, allowing us to better gauge generalization. We conduct a battery of experiments, across three model families, model scales, and variations of both compressor architecture and multi-ratio training.

We find that our approach consistently outperforms the conventional compression-tokens approach, while being more efficient. In addition, we also show that by simply altering the attention pattern in the conventional compression tokens method, one can mitigate the gap between the approaches significantly, albeit not entirely. Our multi-ratio training experiments reveal that it is possible to train and deploy only a single model for a wide range of compression ratios with only minor performance drops. Interestingly, our proposed enhancement for the compression-tokens approach even benefits from multi-ratio training. A comparison of compression performances for scales between 0.6B to 8B shows that compression quality increases with scale, amplifying the benefits of applying such compression methods to larger models. Our code is available at https://github.com/lil-lab/simple-context-compression.

2 Task Definition↩︎

Soft context compression is an approach where a document of length \(L\) is mapped to a sequence of vectors of length \(C\), where \(L \gg C\). While the original document can be described as a sequence of tokens, the compression is made of dense continuous vectors, hence soft. This process allows an LLM that uses the compressed version of the document to invest significantly less computation, both in time and key-value (KV) cache space, both reduced from dependence on \(L\) to dependence on \(C\). This benefit increases with repeated use of the document, as is likely in RAG scenarios.

Formally, we define soft context compression to support multiple compression ratios. Let \(\mathcal{M}\) be a language model and \(\mathcal{R} \subseteq \mathbb{N_+}\) the admissible set of compression ratios. Let \(\mathcal{V}\) denote the vocabulary and \(d\) the embedding dimension of \(\mathcal{M}\). The goal of learning is to construct a compression function \[f_c : \mathcal{V}^{L} \times \mathcal{R} \to \mathbb{R}^{C \times d}\;\;,\] which maps a token sequence \(T = (t_{1}, \ldots, t_{L})\), \(t_i \in \mathcal{V}\) of length \(L\) and a ratio \(r \in \mathcal{R}\) to a compressed representation of \(C\) vectors of dimension \(d\). The length \(C\) is determined by the specified ratio \(C = \left\lceil L / r \right\rceil\).

An ideal compressor \(f_c\) preserves the conditional distribution of the model using the compressed version for any prompt \(P\): \[p_{\mathcal{M}}(\cdot \mid T, P) \;\approx\; p_{\tilde{\mathcal{M}}}(\cdot \mid f_c(T; r), P)\;\;,\] where \(\tilde{\mathcal{M}}\) is an adapted version of \(\mathcal{M}\), for example augmented with lightweight parameters such as LoRA [4] modules that can be fused into the base model without altering its capacity.

3 Background and Related Work↩︎

3.0.0.1 Soft Context Compression

A dominant line of research on context compression adopts the use of artificial compression tokens. As shown in [fig:arch:tokens], a sequence of length \(L\) with a target compression ratio \(r\) is augmented with \(C = \lceil L/r \rceil\) additional identical tokens. ¹ The embedding of the compression token is learned. The final hidden state at the time steps of the compression tokens is taken as the compressed representation. A decoder, conditioned on this representation and a downstream prompt (e.g., a question), produces the output. Training typically combines a language modeling objective on the decoder with a distillation loss that encourages the compressor–decoder to approximate the behavior of a target LLM with access to the full uncompressed context ([fig:arch:regular95llm]). The decoder parameters are either tuned during learning or are frozen.

This paradigm has been explored extensively in recent work. AutoCompressors [5] introduce recursive compression by appending a fixed set of compression tokens and extracting their hidden states. They tie the encoder and decoder weights. The ICAE framework [1] adopts an encoder–decoder setup where the decoder is frozen and only the encoder is trained, with a two-stage process of autoencoding pretraining and task-specific finetuning. COCOM [6] extends this approach to retrieval-augmented QA, experimenting with lighter encoders and with training decoders to jointly process multiple compressed contexts. Other work has sought more aggressive reduction: xRAG [2] maps document retrieval embeddings directly into the decoder’s input space, achieving single-token compression but with severe constraints on sequence length and generality. PISCO [7] demonstrates that training compressors on LLM-generated answers improves downstream RAG performance, while PCC [3] decouples the compressor from the target LLM by learning a converter to project compressed representations into another model’s hidden space. Most recently, GMSA [8] proposed grouping hidden representations and learning a layer semantic alignment module that bridges the gap between the encoder’s final hidden states and the decoder’s first attention layers. Their approach is related to our study of pooling, but differs in the complexity of multi-stage reconstruction training and the use of compression-decoder adapters.²

In this work, we revisit some of these design choices and provide a systematic comparison with token-based architectures under both single- and multi-ratio training regimes.

3.0.0.2 KV Cache Compression

In contrast to representing contexts as input embeddings, another line of work compresses the entire set of key–value (KV) states. Some approaches remove or compress less informative entries in the KV cache without additional training [9]–[11], while others train the model to perform the compression explicitly [12], [13]. A different variant introduces compression tokens, but instead of retaining only the final hidden representation, all KV states are propagated to the decoder [14], [15]. Although these methods provide higher-capacity compressed representations that are well suited for efficient long-context comprehension, their increased size makes it impractical to store them for reuse in retrieval-augmented generation frameworks, where caching compressed representations could otherwise avoid recomputation.

3.0.0.3 Hard Prompt Compression

An alternative approach is to compress contexts directly in the token space. This has been done by removing unimportant tokens or lexical units [16]–[18] or generating concise summaries that preserve salient details [19]. While these methods can be more interpretable and storage-efficient, they are inherently constrained by their reliance on explicit tokens.

4 Methodology↩︎

We propose a compression model design where representations of adjacent tokens are pooled together to reduce the document representation length. We train via knowledge distillation to replicate the functionality of a teacher receiving the original (uncompressed) input.

4.1 Compression via Mean Pooling↩︎

We propose a simple compression architecture that relies only on mean pooling of encoded representations. [fig:arch:pool] illustrates our approach. Given a document, we compute its representation with an encoder, and apply a non-overlapping mean pooling operator with window size \(r\), the same size as the compression ratio, and stride \(r\) to generate continuous vectors as the output compression.

Formally, let \(h= (h_1,\ldots,h_L) \in \mathbb{R}^{L\times d}\) denote the sequence of hidden states produced by the encoder. For a compression ratio \(r\in \mathcal{R}\), we partition the sequence into \(k\) consecutive, non-overlapping blocks: \[S_k = \{(k-1)r+1,\ldots,\min(kr, L)\}\;\;, \qquad k=1,\ldots,\lceil L/r\rceil\;\;.\] The compressed representation of length \(\lceil L/r\rceil\) is obtained by averaging within each block: \[f_c(T,r) = (z_1,\ldots,z_{\lceil L/r\rceil}) \in \mathbb{R}^{\lceil L/r\rceil \times d} \text{s.t.}z_k = \frac{1}{|S_k|} \sum_{i \in S_k} h_i\;\;.\]

The encoder we use is Transformer-based. Critically, we use a full self-attention mask during encoding, allowing each encoded vector to include information from the entire context, and thereby each compressed segment to aggregate information across the entire context. In practice, we initialize with an autoregressive LLM, and remove the self-attention mask before learning.

Our pooling design introduces no additional parameters beyond those of the encoder backbone and the decoder (i.e., target LLM) using the compressed representation, and has low computational overhead. The compression-tokens method is slightly more expensive. It requires an encoder input of size \(L+ L/r\), while our approach only processes the original \(L\) tokens, with negligible overhead for the pooling operator.

4.2 Training Objective↩︎

Our objective is to approximate the behavior of the target LLM before any fine-tuning when provided with the full uncompressed context ([fig:arch:regular95llm]). We consider this LLM as the teacher model within a knowledge distillation framework. Each training instance consists of a context \(T\), a prompt \(P\), and a ground-truth answer \(A= (a_1,\ldots,a_m)\). For each answer position \(i = 1,\dots,m\), the teacher defines the distribution \(q(\cdot \mid T, P, A_{<i})\), where \(A_{<i} = (a_1,\ldots,a_{i-1})\) denotes the gold prefix of the answer. The compressor produces a compressed representation \(f_c(T,r)\), which is passed to the fine-tuned decoder together with \((P, A_{<i})\), yielding the student distribution \(p_\theta(\cdot \mid f_c(T,r), P, A_{<i})\). This combined encoder-decoder model is the student.

For a single compression ratio \(r\), the distillation loss is the token-level KL divergence between the teacher and student distributions: \[\mathcal{L}_{\text{KD}}(T,P,A; r) \;=\; \sum_{i=1}^m \mathrm{KL}\!\left( q(\cdot \mid T,P,A_{<i}) \;\|\; p_\theta(\cdot \mid f_c(T,r),P,A_{<i}) \right)\;\;.\]

We propose a unified training strategy in which a single compressor is trained to handle multiple ratios simultaneously. This is in contrast to most previous work, where a separate model is trained for each compression ratio. We generate compressed representations for all ratios \(r \in \mathcal{R}\) for each training instance. Each compressed representation is passed independently to the decoder, and the corresponding losses are computed. The final objective for one training instance is obtained by summing across the ratios: \[\mathcal{L}_{\text{multi}}(T,P,A) \;=\; \sum_{r\in \mathcal{R}} \mathcal{L}_{\text{KD}}(T,P,A; r)\;\;.\]

The iteration over ratios is performed within each batch, and a single parameter update is applied after aggregating the losses. Since the encoder computation is shared across all ratios, this procedure is substantially more efficient than training separate models. By relying exclusively on knowledge distillation, rather than combining it with a language modeling objective, we enable a direct and fair comparison between the original model and the compressor.

5 Experiments and Results↩︎

Our central objective is to isolate the contribution of context compression itself, without entanglement with retrieval noise or incomplete supervision. While compression methods are often demonstrated in retrieval-augmented generation (RAG) settings, these introduce extraneous challenges, such as when retrieved passages may not contain sufficient evidence, making performance conflate retrieval quality with compression quality. To avoid this confounder, we focus exclusively on reading comprehension, where given contexts are guaranteed to contain the necessary evidence to answer the question. This setup allows for a controlled, head-to-head comparison of different compression strategies, across a range of datasets that stress both single-hop and multi-hop reasoning.

5.1 Experimental Setup↩︎

5.1.0.1 Data

We curated our training set by mixing multiple context-based datasets, in tasks spanning reading comprehension (RC) and summarization. A detailed list of the datasets we incorporated in our training mixture can be found in 2. We evaluate with six reading comprehension benchmarks: SQuAD [20], NarrativeQA [21], HotpotQA [22], AdversarialQA [23], TriviaQA [24], and ParaphraseRC [25]. This selection covers a broad spectrum of reasoning styles, from factual extraction to adversarial paraphrasing, thereby testing the generality of compression. For TriviaQA, we restrict the evaluation to the verified subset, ensuring that every question has sufficient supporting evidence. The training mixture includes the train splits of SQuAD, NarrativeQA, and HotpotQA, which thus serve as in-domain testbeds. AdversarialQA, TriviaQA, and ParaphraseRC are excluded from the training mixture and instead used purely for out-of-domain evaluation.

5.1.0.2 Model Training

For each target language model, we first finetune it on the training mixture using LoRA [4]. This finetuned model is then fixed and used as the teacher in the distillation process. This ensures that performance differences stem solely from the compressor rather than mismatched finetuning (e.g., to the question domain). Both the compressor’s encoder and decoder are initialized from the same target LLM but are trained with separate LoRA weights. We always use the instruction-tuned model weights as our backbone. For multi-ratio training, we always train on the ratios \(\{4\times, 8\times, 16\times, 32\times, 64\times, 128\times\}\), unless stated otherwise. In addition, we found that applying a single linear layer slightly improves performance for both our method and the compression-tokens method, so for all experiments in this paper a learned matrix \(W\in\mathbb{R}^{d\times d}\) is applied to \(f_c(T,r)\) before the compressed representation is given as input to the decoder LLM, unless stated otherwise. Detailed training configuration and hyperparameter choices are provided in 7.

5.1.0.3 Implementation of the Compression-Tokens Approach

The main approach we compare ours against is using compression tokens. A central design consideration in our experiments is the attention pattern applied to compression tokens. Under the compression-tokens paradigm, the causal attention mask typically employed by transformer-based LLMs imposes a strong Matryoshka-style [26] constraint: compressed representations at smaller lengths must correspond to strict prefixes of those at larger lengths. We relax this restriction by allowing compression tokens to attend bidirectionally among themselves, while retaining causal attention over the original context. This simple, albeit not explored in prior work, modification makes the model aware of how many compression tokens are available (i.e., its compression budget), and therefore to allocate information differently for each ratio while still benefiting from shared computation and KV caching. We experiment with both the conventional causal attention mask and our bidirectional modification. Empirically, we observe that this modification significantly improves the approach’s performance. In addition, in our implementation of the compression-tokens models, we utilize only a single compression token that is appended \(\lceil L/r\rceil\) times to the context, rather than having several compression tokens. This enables us to compress any context to any arbitrary ratio we choose, while retaining an equal number of parameters in the model.

5.1.0.4 Metrics

We evaluate our models using the standard exact match (EM) and \(F_1\) metrics.³ Several recent works reported QA performance using a substring accuracy metric, which assigns a score of 1 if the exact match is a substring of the output and 0 otherwise. We opt against the adoption of this metric as it is easily exploitable.⁴ This forces us to exclude some baselines from our primary results.

In addition, for each metric, we also define its teacher-normalized version. Given a metric \(M\), a target language model \(\mathcal{M}\) and a compressor \(f_c\), we calculate the following: \(M_T\), the teacher’s score, by passing the original uncompressed contexts to the decoder model; \(M^{\varnothing}_T\), the no-context score, by passing only the question without any context; and \(M_{f_c}\), by using the compressed context. Then, the teacher-normalized score is given by \(\frac{M_{f_c}-M^\varnothing_T}{M_T -M^\varnothing_T}\). This score does two things. First, it scales the compressor’s score with respect to the teacher’s, allowing for a direct assessment of the amount of retained performance under compression. Second, it accounts for how easy it is for the model to answer the input question without any context. This latter consideration comes to account for cases where a question does not really require the context, potentially inflating the score of the compression model, while actually simply ignoring the compressed input.

5.1.0.5 Systems

Our main comparison point is our implementations of the compression tokens method, with causal or bidirectional attention. We compare against these two systems throughout our experiments. We also include comparisons to other soft context compression methods: ICAE [1] and PCC [3]. We also evaluate LLMLingua2 [18], a hard prompting compression approach, by passing LLMLingua’s compressed prompts to our finetuned Qwen3-8B teacher model.

Comparison between compression methods is challenging in general, due to inconsistencies in the training procedures. Our main goal is to build a systematic understanding of the architecture landscape, and we compare ourselves to other methods mainly to showcase the validity of our experiments. Showing the strength of performance compared to the prior state of the art is secondary, and not even necessarily feasible because many approaches do not release code or models, or adopt training methods that complicate the evaluation. This challenge is demonstrated well by the PISCO [7] method. Their training method is not well suited for the traditional EM/\(F_1\) metrics and performs poorly on these metrics, and indeed the authors evaluated their method only using a more forgiving substring accuracy score. We therefore omit this method from our primary comparisons, but include it and report substring accuracy in Appendix 8.1 ([tab:primary95results95accuracy]).

All experiments are conducted with context lengths up to 1,024 tokens, a practical limitation imposed by our computational budget.

5.2 Results↩︎

[tab:primary95results95hybrid95transpose] shows our main results. We demonstrate the robustness and generality of our findings by comparing six different models, spanning three model families and four model scales: Llama3.2-1B [27], Gemma2-2b [28], and Qwen3-0.6/1.7/4/8B [29]. We show the average \(F_1\) scores over all six evaluation datasets. Below the results table, we provide bar charts that summarize the results along specific dimensions that emphasize trends. The bar charts show the teacher-normalized \(F_1\) metric averaged across all models and datasets to get a single aggregated result per method. We provide the teacher model’s performance with the original context ("Original") and without context at all ("No Ctx"). In between we report the results for different compression ratios and for both the single- and multi-ratio training schemes.

Our mean-pooling method is consistently better than both the standard (causal) compression tokens architecture and the bidirectional variation. Simply adding bidirectional attention between the compression tokens dramatically improves performance. Comparing the single-ratio models to the multi-ratio models, we see that the bidirectional compression tokens approach significantly benefits from multi-ratio training, while for the mean-pooling approach a trade-off clearly exists. This is potentially because the bidirectional approach has a clear signal about the target compression ratio – the forward attention allows each time step to be aware of the total length of the compressed output. Without this, the model must produce a one-size-fits-all representation in each time step. Lastly, while we do see a performance drop in the multi-ratio setting at \(128\times\) for all methods, it must be taken in the context of the already relatively low performance retention at that ratio.

5.2.0.1 Compression Scaling

It is well known that LLM performance increases with scale, if scaled appropriately [30]. But does compression quality scale as well? In [tab:primary95results95hybrid95transpose] we can see that the teacher’s performance improves as the model size grows. Having a compressor increase in performance at the same rate as the teacher would actually tell us that the compressor does not scale well, since that would mean that the teacher-normalized scores stay constant throughout the scales. 2 shows the compression performances of the four Qwen3 model scales we use. We train the models under the multi-ratio setting to efficiently evaluate multiple ratios. We evaluate using the teacher-normalized \(F_1\) score, and present the average scores across all datasets. The results are excitingthe compressors clearly have desirable scaling properties. Critically, the efficiency gains of context compression are much larger as the model size increases. Therefore, observing that larger models retain higher levels of performance (higher relative F1) justifies this line of research. Similar trends are observed for the other methods ([tab:primary95results95hybrid95transpose]), demonstrating the generality of this finding.

5.2.0.2 In-Domain vs. Out-of-Domain

We construct our evaluation suite with both in-domain QA datasets and out-of-domain QA datasets (5.1). The training splits of the in-domain datasets are included in the training data mixture, while the out-of-domains datasets are not. It is expected that downstream performance will drop for out-of-domain datasets. Critical for our study, though, is the performance drop of the compressor itself. 3 (a) plots the in-domain and out-of-domain performance using the teacher-normalized \(F_1\) score for the Qwen3-8B model, averaged over the datasets in each category. We first observe that while the mean-pooling approach is superior for ratios up to \(16\times\) in all settings, its performance deteriorates as the compression ratio increases. To better understand the performance change due to the domain gap, we plot the differences between the in-domain and out-of-domain performance in 3 (b). The performance gap is higher for low ratios, and lower for higher ratios. One possible explanation is that at low compression ratios the compressed representations still retain much of the original contextual signal, so the model is more sensitive to domain-specific distributional shifts; differences between in-domain and out-of-domain language patterns thus manifest as a larger performance gap. By contrast, at higher compression ratios much of the fine-grained contextual detail is already lost to compression noise, which dominates over the domain gap. In this regime, both in-domain and out-of-domain datasets suffer similarly from the limited representational capacity, resulting in a smaller relative difference.

5.2.0.3 Model Ablations

We run all our ablations using the Gemma2-2B model, since it demonstrates strong performance while being relatively compute-efficient. 1 presents the ablation results. We conduct several ablations: (1) Fixed Decoder: the decoder is kept frozen and only the encoder is trained; (2) Fixed Encoder: the encoder is kept frozen and only the decoder is trained; (3) No Encoder: we remove the encoder entirely, and obtain the initial context representation using only the token embeddings of the decoder model; (4) w/o Linear Layer: we remove the linear layer that is applied after the pooling operation; (5) Ratio Sampling: instead of training on all ratios for each sample, a single ratio is randomly chosen for each sample during training.

Freezing the decoder results in considerable performance reduction, although not catastrophic. This is in line with findings of previous works ([7]). Freezing or removing the encoder is more detrimental, lowering performance by more than 12%. The effect of the linear layer is not very significant, its removal results in a reduction of only 0.7%. Finally, while randomly sampling a single ratio per sample would speed up training significantly, it does so at the cost of a small drop in performance (1.4%), likely because each ratio gets less updates during training.

6 Discussion↩︎

We provide a systematic study of soft context compression, showing that a simple mean-pooling approach consistently outperforms compression-tokens architectures while requiring no additional parameters. We further demonstrate that multi-ratio training is both feasible and effective, enabling a single compressor to support a wide range of compression budgets with only minor performance degradation. Interestingly, we observe that the bidirectional compression-tokens method consistently benefits from multi-ratio training. A plausible explanation is that, unlike mean-pooling or causal compression-tokens, this method has explicit access to the number of compression tokens available, allowing it to adapt to the target budget. Exploring how to incorporate similar signals into other architectures is an important direction for future work.

Finally, our study highlights a broader open problem: the evaluation of compression methods remains hindered by inconsistent setups, metrics, and benchmarks. By isolating compression quality from retrieval confounders and applying a uniform evaluation across models, scales, and domains, we aim to provide a rigorous basis for future research on soft context compression. We hope this work helps establish more standardized practices and clarifies the core principles that should guide the development of next-generation compressors.

Acknowledgements↩︎

This research was supported by a gift to the LinkedIn–Cornell Bowers Strategic Partnership, NSF under grant No. OAC-2311521, NASA under award No. 20-OSTFL20-0053, and a gift from Open Philanthropy. We thank Google for enabling experiments with Gemini through a gift. This research was supported with Cloud TPUs from Google’s TPU Research Cloud (TRC). We thank Nathan Godey for his comments.

7 Additional Experimental Setup Details↩︎

7.1 Data↩︎

2 and 3 provide detailed lists of our training data mixture and evaluation suite. We use the summaries as contexts for NarrativeQA, instead of the full stories. For HotpotQA, we only use the two gold paragraphs as contexts, and remove the distractors.

We increase the training data diversity by randomly sample a prompt template that fits the task, when training samples are composed of a context \(C\), question \(Q\) and answer \(A\). For example, for the extractive QA task, an example of a prompt template is: “<C>\n Extract the answer from the text above. \n Question: <Q>\n Answer: <A>”. Similar templates are defined for other tasks as well. We created approximately 100 prompt templates for each task. The full list of templates for each task is available along with our code.

7.2 Training Hyperparameters↩︎

8 Additional Results↩︎

We provide additional results from the same experiments conducted in the main body of the paper. Appendix 8.1 provides the primary results of the paper with the EM and substring accuracy metrics. Appendix 8.2 shows the \(F_1\) performance on each individual dataset from the evaluation suite.

8.1 Primary Results Additional Metrics↩︎

8.2 Results Per Dataset↩︎

8.2.1 In Domain Datasets Results↩︎

8.2.2 Out-of-Domain Datasets Results↩︎

9 LLM Usage↩︎

LLMs (specifically, ChatGPT) were used in the process of writing this paper for creating tables and figures, as well as proof-reading.

References↩︎