💬 NLP

End-to-End Context Compression at Scale

This paper introduces Latent Context Language Models (LCLMs), a family of encoder-decoder compressors trained on massive datasets that significantly improve the accuracy-efficiency frontier for long-context inference by reducing memory usage and compression time while maintaining high model quality.

Original authors: Ang Li, Sean McLeish, Haozhe Chen, Nimit Kalra, Zaiqian Chen, Artem Gazizov, Venkata Anoop Suhas Kumar Morisetty, Bhavya Kailkhura, Harshitha Menon, Zhuang Liu, Brian R. Bartoldson, Tom Goldstein, San

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ang Li, Sean McLeish, Haozhe Chen, Nimit Kalra, Zaiqian Chen, Artem Gazizov, Venkata Anoop Suhas Kumar Morisetty, Bhavya Kailkhura, Harshitha Menon, Zhuang Liu, Brian R. Bartoldson, Tom Goldstein, Sanae Lotfi, Micah Goldblum, Pavel Izmailov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, super-smart assistant (a Large Language Model) who can read millions of pages of text. The problem is, this assistant has a very small "working memory" (like a tiny notepad). Every time you give them a long document to read, they have to write down every single word on that notepad to remember it. If the document is too long, the notepad fills up, the assistant gets overwhelmed, and the process becomes incredibly slow and expensive.

This paper introduces a new tool called Latent Context Language Models (LCLMs) to solve this problem. Think of LCLMs as a super-efficient "summarizer" or "compression expert" that sits between you and the assistant.

Here is how it works, using simple analogies:

1. The Problem: The "Notepad" Bottleneck

Currently, if you ask an AI to read a 100-page book, it tries to keep every single word of those 100 pages in its active memory (the KV cache).

The Analogy: Imagine trying to carry a 100-page book in your pocket while running a marathon. It's heavy, it slows you down, and eventually, your pocket rips (the memory runs out).
Old Solutions: Previous methods tried to throw away pages they thought were unimportant (like deleting chapters) or tried to shrink the font size. But this often made the assistant forget crucial details or required so much time to "edit" the book that it wasn't worth it.

2. The Solution: The "Smart Compression" (LCLM)

The authors created a system that doesn't just delete words; it translates the book into a highly condensed, secret code.

The Encoder (The Translator): This is a smaller, specialized AI that reads chunks of your text (like a paragraph at a time) and turns them into a single, dense "thought token."
- Analogy: Instead of giving the assistant the whole 100-page book, the Encoder reads a page and writes down just one perfect sentence that captures the entire page's meaning. It does this for the whole book, turning 100 pages into just 10 sentences.
The Decoder (The Assistant): This is the main AI you want to use. It receives these 10 compressed sentences instead of the 100 pages. Because the "notepad" is now much smaller, the assistant can read it instantly, use less energy, and still understand the story.

3. How They Built It (The "Recipe")

The authors didn't just guess how to build this; they tested thousands of variations to find the perfect recipe.

The "Mean" vs. "Concat" Test: They tried different ways to combine the information.
- Analogy: Imagine you have 16 ingredients for a soup.
  - Option A (Concat): You keep all 16 ingredients separate in a big bowl.
  - Option B (Mean): You blend all 16 ingredients into one smooth, rich broth.
- They found that for very long texts, blending them (Mean Pooling) worked best, creating a smooth, information-dense "broth" that the assistant could easily digest.
The Training: They taught this system by feeding it massive amounts of data (350 billion words), alternating between reading normal text and reading compressed text. This taught the system how to keep the "flavor" of the original text even when it was shrunk down.

4. The Results: Faster, Lighter, and Smarter

The paper claims their new system creates a "new frontier" where you get the best of both worlds:

Speed: It is significantly faster to process long documents.
- Analogy: Instead of walking through a library to find a book, the LCLM gives the assistant a map that points directly to the right shelf.
Memory: It uses much less computer memory.
- Analogy: You can now carry the entire library in your backpack instead of a wheelbarrow.
Accuracy: Unlike older methods that made the assistant "dumb" or forgetful, LCLMs keep the assistant's intelligence high. They can still answer tricky questions and find specific details hidden in the text.

5. The "Agent" Feature: The "Expand" Button

The paper also shows how this works for AI agents (robots that do tasks).

The Scenario: Imagine an agent needs to find a specific bug in a massive codebase (a huge software project).
The Trick: The agent first reads the compressed version of the whole codebase to get a "bird's-eye view." It sees the general structure and knows roughly where the problem might be.
The "Expand" Tool: Once the agent spots a suspicious area in the compressed view, it can hit a button to "Expand" that specific section. The system then unpacks that tiny chunk back into full, detailed text so the agent can fix the bug precisely.
Analogy: It's like looking at a map of a city to find the right neighborhood, and then zooming in to see the street address. You don't need to look at every street in the city at once; you just zoom in when you need to.

Summary

The paper presents a new way to handle massive amounts of text for AI. Instead of forcing the AI to carry a heavy, unwieldy load of raw data, they use a smart "compression expert" to shrink the data into a lightweight, high-quality summary. This allows AI to read longer documents, faster, with less memory, without losing its ability to understand the details.

Technical Summary: End-to-End Context Compression at Scale

Problem Statement

Long-context inference in Large Language Models (LLMs) is increasingly constrained by memory and latency due to the linear growth of the Key-Value (KV) cache with context length. While recent KV cache compression techniques attempt to mitigate this by evicting cache entries or training compact caches, they face significant limitations:

Quality vs. Efficiency Trade-off: Many methods substantially degrade model performance or require domain-specific training that does not generalize.
Inference Engine Compatibility: Techniques that evict entries non-uniformly across attention heads or layers are difficult to integrate with modern production inference engines (e.g., vLLM, SGLang) which assume shared sequence lengths.
Prefill Overhead: Many approaches require the full context to be prefilled before compression, negating potential latency gains, or require substantial additional compute time.
Soft-Token Limitations: Existing encoder-decoder "soft-token" methods often rely on offline preprocessing, fail to preserve the base model's broad in-context behavior, or are evaluated only on narrow, domain-specific tasks rather than information-dense long-context benchmarks.

The central question addressed is whether a general, task-agnostic compressor can be trained to preserve a model's original capabilities while significantly reducing memory usage and inference time.

Methodology

The authors introduce Latent Context Language Models (LCLMs), a family of end-to-end trained encoder-decoder compressors. The core architecture maps a sequence of input tokens $x_{1:T}$ into a shorter sequence of latent embeddings $z_{1:M}$ via an encoder, which are then projected by an adapter and consumed by a decoder in place of the original context.

Architecture Design

Through a large-scale architecture search from scratch, the authors identified optimal design choices:

Pooling: Mean pooling was found to be superior to token-based (EOS/CLS) pooling and competitive with concatenation, particularly at high compression ratios (16×).
Encoding Granularity: An encoder window size ( $W$ ) of 1024 tokens was adopted, allowing the encoder to attend over a larger local context than the compression block size ( $N$ ) without prohibitive memory overhead.
Attention Masking: Causal masking within the encoder consistently outperformed bidirectional masking.
Adapter: A lightweight MLP adapter (without self-attention) proved more effective and efficient than attention-based adapters for projecting encoder latents to the decoder dimension.
Initialization: Initializing the encoder from a pre-trained embedding model (Qwen3-Embedding-0.6B) rather than a standard LLM provided better initial representations, though the gap narrowed with extensive training.

Training Recipe

The authors developed a multi-stage training recipe to prevent catastrophic forgetting and ensure stability:

Stage 0 (Adapter Warmup): Freeze encoder and decoder; train only the adapter.
Stage 1 (Encoder Training): Unfreeze encoder; keep decoder frozen.
Stage 2 (Continual Pre-training): Unfreeze decoder with a small learning rate; train end-to-end on a mixture of next-token prediction and reconstruction data.
Stage 3 (Supervised Fine-Tuning): Train on SFT data to enhance instruction following and reasoning with compressed inputs.

The training data includes a carefully curated mixture of code, text, reasoning, and long-context documents. Crucially, the pre-training data uses an interleaved format where compressed segments (wrapped in special tokens) alternate with uncompressed segments, forcing the model to condition on latent context at multiple positions rather than just the beginning.

Agent Integration

The authors propose an agentic system where the LCLM agent receives the entire compressed context. The agent can selectively invoke an EXPAND tool to retrieve the original uncompressed text for specific segments, enabling a "skim-and-zoom" strategy for long-horizon tasks.

Key Results

The LCLM family (specifically a 0.6B encoder and 4B decoder) was evaluated at compression ratios of 4×, 8×, and 16× on benchmarks including RULER, LongBench, LongHealth, and GSM8K.

Pareto Frontier: LCLMs establish a new Pareto frontier on the accuracy-efficiency trade-off. They achieve significantly faster Time-to-First-Token (TTFT) and lower peak GPU memory compared to KV cache compression baselines (e.g., SnapKV, KVzip, Attention Matching), especially at high compression ratios.
Accuracy Preservation: At 16× compression, LCLMs maintain high accuracy on long-context benchmarks (e.g., ~65% on RULER 16k vs. ~93% for full KV, significantly outperforming baselines which often drop below 20-40%).
Memory Efficiency: Unlike KV cache methods where memory scales with the full context length, LCLM memory usage plateaus at longer contexts because the encoder processes fixed-size windows. This allows LCLMs to handle contexts up to 1M tokens on a single H200 GPU, whereas baselines run out of memory at 512K or 1M tokens.
Fine-Grained Performance: On GSM8K (short, dense contexts), LCLMs achieve the highest accuracy across all compression ratios, demonstrating they are not limited to long-context QA but can compress information-dense inputs effectively.
Agent Performance: The agentic system with adaptive expansion significantly improves exact string-match accuracy on "needle-in-a-haystack" tasks, often matching the performance of the uncompressed model while retaining the efficiency of the compressed context.

Significance and Claims

The paper claims that LCLMs offer a practical, scalable solution for long-context inference that is compatible with standard production inference engines. By shifting the compression burden to a smaller, parallelizable encoder, LCLMs decouple the memory and latency costs of context length from the decoder's generation process.

The authors emphasize that their approach:

Preserves General Capabilities: Unlike prior soft-token methods that degrade performance or require task-specific tuning, LCLMs maintain strong performance across general tasks, reasoning, and instruction following.
Enables Adaptive Retrieval: The integration with agentic frameworks allows for dynamic expansion of relevant context, bridging the gap between global corpus visibility and fine-grained information retrieval.
Scalability: The method scales effectively to million-token contexts on commodity hardware (H200 GPUs), overcoming the memory bottlenecks that currently limit the deployment of long-horizon LLM systems.

The work suggests that learned, end-to-end compression is a viable alternative to heuristic KV cache eviction, providing a robust substrate for future long-horizon agents with persistent working memories.