ARC-Encoder: learning compressed text representations for large language models

AI summary

Overview

Research area: Natural Language Processing / large language model inference efficiency, specifically context (prompt) compression using soft continuous representations.

Technical level: Advanced. The paper assumes familiarity with transformer self-attention, embedding spaces, encoder-decoder training, and benchmark conventions such as exact match and ROUGE.

Scope: This paper introduces ARC-Encoder, a trainable text encoder plus a small MLP projector that compresses a document into fewer continuous vectors which replace token embeddings in an unmodified, frozen decoder LLM.

What This Paper Is About

LLMs increasingly need to read long inputs, such as retrieved documents in retrieval-augmented generation (RAG) or long user prompts, which is expensive because transformer attention scales quadratically and can also exceed the model's context window. Existing compression methods either compress weakly, or achieve strong compression by fine-tuning and even modifying the target decoder, which risks degrading the model's general abilities and requires a separate dedicated encoder per decoder. ARC-Encoder instead learns an external compressor whose output can be injected into an untouched decoder exactly where token embeddings normally go, so the decoder keeps its original capabilities and few-shot behavior.

Key Contributions

1. An external compressor needing no decoder modification. ARC-Encoder produces "pooled tokens" that replace raw text input in a frozen decoder. It reduces input sequence length while preserving strong performance across benchmarks, including in-context learning (ICL).

2. One encoder, multiple decoders. A single shared encoder can serve more than one decoder by specializing only a projector layer and special tokens, which account for less than 1% of the encoder's weights (15M parameters per decoder in the reported experiments).

3. Context window extension via chunked compression. ARC-Encoder can be trained to compress chunks of a long document in parallel, showing competitive results on long-context benchmarks and letting a decoder process up to 8 times more input than its original context window.

4. Evidence that pretraining and fine-tuning are both essential. The paper ablates training objectives, architecture, and pooling design, and reports that compressed representations of Wikipedia require memory on the same order as storing the raw text, which permits precomputing representations.

Main Findings

• ARC-Encoder beats the closed-book baseline and approaches open-book. With a Mistral 7B decoder at a 4x pooling factor, ARC4-Encoder^M scores 39.0 NQ, 68.9 TQA, 45.1 HQA, 71.1 SQuAD, 31.0 FLORES, 23.8 CNN (average 46.5), versus open-book 39.9 / 70.5 / 48.3 / 77.7 / 31.3 / 27.2 (average 49.2). With a Llama3.1 8B decoder, ARC4-Encoder^L scores 39.7 / 70.1 / 46.9 / 74.0 / 33.7 / 23.7 (average 48.0) versus open-book 38.6 / 67.1 / 47.1 / 72.2 / 32.8 / 26.5 (average 47.4).

• Larger pooling factors cost accuracy. ARC8-Encoder^M averages 43.4 on Mistral 7B and ARC8-Encoder^L averages 45.0 on Llama3.1 8B, both below their 4x counterparts.

• Compression gains differ by task. ARC-Encoder outperforms baselines more clearly on tasks where the decoder cannot rely on parametric memory, such as reading comprehension (SQuAD) and summarization (CNN). It is slightly behind PISCO-like methods on FLORES because those sequences are very short (under 30 tokens), where fixed 32-memory-token methods end up with a compression ratio below 1.0, effectively expanding the input.

• Constant pooling factor versus fixed token counts. ARC-Encoder maintains a fixed pooling factor regardless of sequence length, while fixed-count methods provide more capacity on short inputs and effectively expand them.

• Cross-model sharing is nearly lossless. The shared encoder (ARC-Encoder with the shared-token marker) loses less than 1.0 point on average compared to decoder-specialized versions.

• Adaptation to a brand-new decoder is cheap but weaker. Feeding OLMo-7B with a frozen encoder plus only 15M newly trained parameters gives ARC4-Encoder^M 31.5 NQ, 62.5 TQA, 26.5 HQA, 46.4 SQuAD, 17.2 FLORES, 19.1 CNN (average 33.9) versus OLMo-7B open-book average 39.6 and closed-book 19.5/48.3/17.8/11.2. On HotpotQA it beats open-book.

• Long-context extension works on some tasks. ARC8-Encoder with Llama2 7B Chat (context truncated to 32k tokens, 32 chunks of 1024 tokens, decoder side 4k, noted as 32k//8) scores 27.5 NQA, 28.3 QASPER, 14.1 GovReport, 19.1 QM-Sum, versus Llama2 Chat at 16.1 / 17.2 / 15.7 / 19.8 and Llama2-32k Instruct at 14.2 / 16.4 / 17.8 / 17.6. The GovReport and QM-Sum results suggest a distribution mismatch with the synthetic fine-tuning data.

• The optimal pooling factor changes between stages. Pretraining with a higher pooling factor than fine-tuning gives improved models (Appendix A.2); the final ARC4-Encoder shares the pretrained encoder with ARC8-Encoder and is only fine-tuned at a pooling factor of 4.

• Pretraining length matters greatly. After 20k pretraining steps there is roughly a +16 point improvement in average score, and after 80k steps roughly +19 points, compared with fine-tuning directly without pretraining.

• Both pretraining tasks matter. With 0% reconstruction batches the average is 39.8; 20% gives 41.6; 50% gives 41.5; 100% gives 37.5. Omitting reconstruction or using too little continuation hurts substantially.

• Encoder truncation trades size for accuracy. Default 3.0B averages 41.7 ± 0.4; truncating 0 layers (3.2B) gives 42.4 ± 0.3; 4 layers (2.8B) gives 41.6 ± 0.3; 21 layers (1.1B) gives 38.6 ± 0.1.

• Average pooling over consecutive queries is the most robust pooling choice. At 8x, average queries give 41.7, last queries 42.1, k-means merged queries 38.6, and memory tokens (about 16) 38.4. At 4x, average queries rise to 45.5 and last queries to 44.5. Averaging was chosen because it was more robust in the multi-decoder setting.

• Pooling position matters. Default placement averages 41.7; pooling between self-attention and MLP averages 40.8; before queries averages 41.4; two layers earlier averages 41.7.

• LoRA and causality variants underperform the default. LoRA with rank 128 averages 41.1 ± 0.4 and keeping causality averages 41.3 ± 0.2, against 41.7 ± 0.4 for the default.

• Prefilling efficiency. ARC4-Encoder nearly matches the open-book baseline while achieving x1.8 gains in prefilling FLOPs, as profiled in Section 4.5 of the paper.

• Reconstruction is nearly perfect when trained for it. The auto-encoder architecture achieves near-perfect reconstruction at a 16x pooling factor on relatively short sequences (up to around 128 tokens), though such representations are not well exploited on downstream tasks.

Methodology in Plain English

The system has three parts. First, a text encoder built from the Llama3.2 3B backbone with its output head removed, its causal mask removed, and its last two layers removed, which converts a document into hidden-state vectors. Second, a pooling mechanism that shrinks the sequence: in the last self-attention layer, consecutive query vectors are averaged in groups, while keys and values stay unchanged, so the pooled queries attend the full uncompressed context. Third, a two-layer MLP projector with a dimensional bottleneck (3072 to 2048 to 4096) that maps encoder outputs into the decoder's embedding dimension. Only the encoder and MLP are trained; the decoder is frozen and simply receives the compressed vectors where token embeddings would normally sit.

Training proceeds in two phases. Pretraining alternates between two teacher-forced tasks: reconstruction, where the decoder must regenerate the original full text from compressed tokens, and continuation, where subsequences are replaced by their compressed form and the decoder must predict what follows. Special learned tokens <Cont> and <Rec> are appended after each compressed sequence to signal which task is active. Fine-tuning then specializes the compressor to downstream behaviors such as in-context learning or long-context understanding, sometimes by interleaving compressed documents with full-text queries and answers in a few-shot template and masking all loss tokens except the final answer. The pooling factor can differ between pretraining and fine-tuning.

For multi-decoder support, the encoder is shared while each decoder gets its own MLP projector and special tokens (under 1% of encoder weights, 15M parameters per decoder). At each training step one decoder is sampled uniformly, and only its projector plus the shared encoder are updated. For long context, ARC8-Encoder is paired with an instruct decoder, long documents are split into fixed-size chunks encoded in parallel, and the compressed chunk tokens are concatenated and fed to the decoder; special tokens are dropped in this setting because instruction prompts play a similar role.

Evaluation uses 5-shot setups across question answering (Natural Questions, TriviaQA, HotpotQA in the distractor setting, SQuAD, with top-5 passages retrieved by NV-Embed v2 from Atlas Wikipedia chunks), translation (FLORES, BLEU averaged over English to Danish, French, German and Spanish), and summarization (CNN-DailyMail, ROUGE-L), plus long-context tasks from ZeroSCROLLS (NarrativeQA and QASPER F1, GovReport and QM-Sum Rouge-L). Fine-tuning mixes exclude the benchmark training sets to test generalization.

Why This Matters

Impact on research. The paper argues that strong compression does not require touching the decoder, contrasting with gist tokens, memory tokens, PISCO-style methods, and approaches that add internal modules or fine-tune the whole model. It also challenges the claim that pretraining is unnecessary, showing it is essential for aligning compressed representations with a decoder's hidden-state space, and demonstrates that compressed-token usage need not interfere with standard-token usage. It reports the ablation details, pooling placements, and pooling schedules that make soft compression work, and notes its use of exact match with base decoders and exclusion of benchmark training sets, which differs from earlier reporting conventions.

Real-world applications (implied by the paper's scenarios):

• RAG pipelines where retrieved documents must be processed on the fly and inference cost is dominated by long contexts.
• Question answering and reading comprehension with retrieved passages, especially tasks where the decoder cannot rely on parametric memory.
• Summarization and translation of documents, where the reported gains over compression baselines were largest or most consistent.
• Extending the usable input length of a model with a small context window, such as Llama2 7B Chat's 4096-token window, without altering that model.
• Precomputing compressed representations of a corpus such as Wikipedia, which the paper notes requires memory on the same order as raw text.

Industry relevance. A portable compressor that supports multiple LLMs with only small MLPs means one encoder can serve a fleet of models instead of one dedicated compressor per model. Because the decoder is untouched, teams avoid the forgetting risk and retraining cost of modifying a production model, and they retain the option of mixing compressed and standard tokens. The reported x1.8 gain in prefilling FLOPs at a 4x pooling factor is a direct serving-cost argument.

Future Directions

• Scaling to much longer contexts. The paper explicitly leaves scaling to contexts such as 128k tokens for future work, noting that it would require specialized datasets and substantial computational budgets.
• Extending beyond two decoders. The multi-decoder experiments used two decoders for engineering simplicity, and the authors state their adaptation results suggest the method can extend to more than two.
• Better merging logic. Dynamic merging strategies aimed at a specific average pooling factor gave variable results and need further refinement.
• Improving training-data alignment for specific tasks. Underperformance on GovReport and QM-Sum is attributed to distribution mismatch with the synthetic fine-tuning data, implying better aligned data would help.
• Exploring alternative pooling schedules. Pooling every two tokens in the last layers performed better than the default at the same pooling factor, suggesting pooling strategies could be adapted to the target pooling factor.

Target Audience

Researchers and engineers working on context compression, retrieval-augmented generation, inference efficiency, and long-context language modeling. It is most useful to readers already comfortable with transformer internals and LLM training pipelines, and to practitioners who need LLM inference cost reduction without modifying or retraining their decoder. Readers looking for production latency or throughput measurements should note that the paper reports a prefilling FLOPs gain (x1.

Authors’ abstract

Recent techniques such as retrieval-augmented generation or chain-of-thought reasoning have led to longer contexts and increased inference costs. Context compression techniques can reduce these costs, but the most effective approaches require fine-tuning the target model or even modifying its architecture. This can degrade its general abilities when not used for this specific purpose. Here we explore an alternative approach: an encoder that compresses the context into continuous representations which replace token embeddings in decoder LLMs. First, we perform a systematic study of training strategies and architecture choices for the encoder. Our findings led to the design of an Adaptable text Representations Compressor, named ARC-Encoder, which outputs $x$-times fewer continuous representations (typically $x\!\in\!\{4,8\}$) than text tokens. We evaluate ARC-Encoder across a variety of LLM usage scenarios, ranging from in-context learning to context window extension, on both instruct and base decoders. Results show that ARC-Encoder achieves state-of-the-art performance on several benchmarks while improving computational efficiency at inference. Finally, we demonstrate that our models can be adapted to multiple decoders simultaneously, allowing a single encoder to generalize across different decoder LLMs. This makes ARC-Encoder a flexible and efficient solution for portable encoders that work seamlessly with multiple LLMs. We release a training code at https://github.com/kyutai-labs/ARC-Encoder , fine-tuning dataset and pretrained models are available at https://huggingface.co/collections/kyutai/arc-encoders-68ee18787301407d60a57047 .