PaperlayerView exact revision on arXiv

Source-linked AI summary

End-to-End Context Compression at Scale

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

arXiv:2606.09659v1cs.CLcs.AIcs.LG

TL;DR

Long-context inference is constrained by KV-cache memory and latency, while existing soft-token compressors do not reliably preserve broad model capabilities. This paper introduces end-to-end trained Latent Context Language Models, which preserve strong in-context capabilities while establishing improved memory-, speed-, and performance tradeoffs.

  • Problem

    Long-context inference needs efficient context management because growing KV caches constrain memory and latency, while existing soft-token compressors may degrade capabilities or require domain-specific training.

  • Method

    The paper trains encoder-decoder Latent Context Language Models end-to-end at scale to compress input tokens into shorter latent sequences consumed by a decoder.

  • Results

    LCLMs establish a new Pareto frontier across compression time, peak memory, and accuracy while preserving strong in-context capabilities across long-context and agentic settings.

  • Takeaways & Limitations

    LCLMs provide a practical building block for efficient long-context systems, including agents that selectively expand relevant compressed context.

  • Takeaways & Limitations

    Prior soft-token methods have mostly been evaluated with domain- or task-specific finetuning and lack convincing evidence on information-dense long-context benchmarks.

Abstract

from arXiv · show

Long-context language model inference is bottlenecked by memory, as the KV cache grows with context length. Recent techniques to compress the KV cache fall short: they either degrade model quality substantially or require considerable time and compute to compress a single long prompt. Furthermore, many methods require the input to fit within the target model's context window, and are generally incompatible with modern production inference engines. Encoder-decoder compressors, which map a long token sequence to a shorter sequence of latent embeddings consumed by a decoder, are an appealing alternative in principle. However, existing approaches are not competitive with KV cache compression on the accuracy-efficiency frontier. In this work, we revisit encoder-decoder compression and close this gap. We first perform an architecture search, pre-training many variants from scratch to determine how best to design and train encoder-decoder compressors. Guided by our findings, we continually pre-train a family of 0.6B-encoder, 4B-decoder models on over 350B tokens each, at compression ratios of 1:4, 1:8, and 1:16. We introduce Latent Context Language Models (LCLMs), a family of compressors that improve the Pareto frontier across general-task performance, compression speed, and peak memory usage. We demonstrate that LCLMs serve as efficient backbones for long-horizon agents, letting the agent skim through a compressed long context and adaptively expand relevant segments on demand.

1 Introduction

Long-context inference is increasingly constrained by KV-cache memory and latency, while existing compression methods impose important quality, compute, and compatibility limitations. This work introduces Latent Context Language Models (LCLMs), developed through large-scale training and architecture search to improve the memory/speed-performance frontier and support adaptive context expansion for agents.

  • Production contexts can reach millions of tokens, making inference increasingly constrained by memory and latency as the KV cache grows.
  • KV-cache compression reduces memory but may require full-context prefill or substantial time and compute, limiting general long-context inference.
  • Encoder-decoder soft-token methods encode raw tokens into shorter continuous embeddings, offering a parallelizable alternative to direct KV-cache manipulation.
  • LCLMs are end-to-end, task-agnostic encoder-decoder compressors introduced to preserve base-model capabilities without relying on domain-specific training.
  • The work contributes a training recipe, comprehensive architecture search, a memory/speed-performance Pareto frontier across compression ratios, and an agent that selectively expands compressed chunks.

2 Related Work

Context compression methods comprise hard-token, soft-token, and KV cache approaches, but prior KV and soft-token methods face deployment, efficiency, or generalization limitations. This work addresses these gaps with a single online soft-token compressor trained at scale to handle heterogeneous inputs while closely matching uncompressed behavior.

  • Compression categories: Context compression approaches fall into hard-token, soft-token, and KV cache compression categories.Hard-token approaches are deferred because they generally underperform the latter two categories.
  • KV cache compression: KV cache compression methods evict entries using mostly hand-crafted or learned heuristics to select which entries to drop.Prompt-agnostic methods prune without query knowledge, while prompt-dependent methods require explicit context-prompt pairs and produce query-specific caches.
  • KV cache compression: KV cache compression is not widely adopted in inference engines such as vLLM or SGLang because non-uniform eviction budgets cannot reduce the cache’s sequence-length dimension.These methods mask evicted positions during attention, forfeiting memory and throughput benefits.
  • Soft-token compression: Most prior soft-token methods rely on offline preprocessing and do not convincingly preserve the base model’s broad in-context behavior.Prior work is mostly evaluated on domain- or task-specific finetuning that is not transferable across tasks, with limited evaluation on long-context benchmarks with informativeness.
  • Soft-token compression: Large-scale staged training enables a single online soft-token compressor to robustly handle general heterogeneous inputs while closely matching uncompressed behavior.The paper discusses these architectural decisions and findings through its architectural search and analysis.

3 The Latent Context Language Model Architecture

LCLMs encode contiguous token chunks into latent tokens, concatenate them across encoder windows, project them to the decoder’s dimension, and use them as compressed context. The architecture supports both block-wise and whole-sequence encoding, with training data combining compressed and uncompressed contexts, long-document supervision, and reconstruction objectives.

  • Architecture: LCLMs use an encoder, pooling operator, adapter, and decoder to convert token chunks into soft latent tokens consumed as context.The encoder maps each contiguous block of N input tokens to one latent token, while the adapter aligns encoder and decoder hidden dimensions.
  • Windowed Encoding: The input is divided into I = ⌈T/W⌉ encoder windows, each containing at most W tokens, and latent tokens from all windows are concatenated into z1:M.Each window is encoded separately before its pooled representations are combined into the full compressed sequence.
  • Windowed Encoding: When W = N, each forward pass yields one latent token, whereas W = T processes the entire input in one pass and produces ⌈T/N⌉ latent tokens.These settings respectively represent block-wise and whole-sequence encoder processing.
  • Training Data: LCLM training uses interleaved compressed/uncompressed blocks, compressed prompts with long documents, and reconstruction targets for reproducing original context.The three data types support continual pre-training, supervised fine-tuning, and auxiliary reconstruction objectives.

4 Training Recipe

The LCLM training recipe uses staged optimization and three complementary data sources to preserve capabilities while improving reasoning, instruction following, long-context understanding, and fine-grained information retention. Continual pre-training interleaves compressed and standard segments, while supervised fine-tuning and auxiliary reconstruction target downstream performance with compressed inputs.

  • Training objectives: The recipe is designed to preserve a powerful LLM’s downstream performance rather than specialize on small in-domain datasets.The stated goals are to improve stability, preserve original capabilities, and mitigate catastrophic forgetting.
  • Continual pre-training data: The continual pre-training mixture covers web text, code, science and reasoning, long-context data, and instruction-style data.Sequences are partitioned into segments that alternate between compressed and standard tokens, with next-token loss computed only on uncompressed tokens.
  • SFT data: SFT post-training targets reasoning, instruction following, and long-context understanding when the model conditions on compressed inputs.Its mixture includes reasoning, long-context instruction following, and general instruction following with multi-turn content.
  • Auxiliary reconstruction data: Auxiliary reconstruction trains the model to reproduce compressed documents, encouraging latent representations to retain fine-grained details for exact retrieval.The data spans code, text, long documents, math, and LATEX, uses 100 prompt templates per source, and appears in both continual pre-training and SFT.
  • Multi-stage training: Training progresses through four stages, from adapter-only warmup to encoder training, end-to-end continual pre-training, and supervised fine-tuning.Stages progressively unfreeze components: the decoder remains frozen initially, then is trained with a small learning rate before receiving a higher rate during SFT.

5 Architectural Design Space of Latent Encoders

The section presents a from-scratch architecture search for encoder-decoder compressors, isolating choices including pooling, encoding granularity, attention masking, and adapter design. The experiments identify mean pooling, a larger encoder window, and causal masking as important design findings, while scaling candidates validates small-scale results.

  • Architecture search: From-scratch pre-training isolates architectural effects without confounding from pre-trained initialization.The search specifically examines encoder attention masks and pooling operators.
  • Pooling operator: Mean pooling consistently improves pre-training loss over token-based pooling and is empirically indistinguishable from concatenation.The comparison covers token-based pooling, mean pooling, and concatenation.
  • Encoding granularity: Increasing the encoder window from W = N (16) to 256 produces a large improvement, while increasing it to 1024 yields a smaller additional gain.Because W = 1024 adds no significant memory or runtime overhead in this setting, it is adopted as the default.
  • Encoder attention mask: Causal masking consistently achieves lower pre-training loss than bidirectional masking in the reported comparisons.Unlike decoder language modeling, prompt compression does not require causal masking to match autoregressive inference.
  • Optimal architecture at scale: The architecture sweep narrows the search to promising choices, which are then scaled with the full training pipeline to identify the strongest architecture and test predictive validity.The large-scale experiments assess whether small-scale pre-training findings remain predictive.

6 Results: Improving Latency- and Memory-Performance Tradeoffs

LCLMs improve the latency–accuracy and memory tradeoffs for long-context compression, establishing a new Pareto frontier across multiple benchmarks. Their efficiency comes from fixed-size encoder-window processing and a smaller encoder, while maintaining strong accuracy under high compression ratios.

  • Pareto frontier: 16×, 8×, and 4× compression ratios establish a new Pareto frontier in compression time and accuracy over KV cache baselines.The result holds across RULER, LongBench, and LongHealth; KV cache methods have compression time largely independent of target compression ratio.
  • Scaling with context length: LCLMs achieve the fastest compression time and substantially lower peak GPU memory at longer context lengths from 4K to 1M.Attention Matching runs out of memory at 1M tokens and fails at 512K tokens due to numerical instability; other methods also run out of memory at 512K and 1M.
  • Scaling with context length: Fixed-size encoder-window processing and a smaller encoder reduce memory demands compared with full-context decoder prefill.The encoder processes at most 128K input tokens per batched forward pass, and peak memory remains nearly flat for the 16× model from 128K to 512K tokens.
  • Fine-Grained Compression: LCLMs achieve the highest GSM8K accuracy across all compression ratios, with particularly strong gains over baselines at higher ratios.The analysis compresses the entire prompt and context.

7 Agent Scaffolding With Latent Context

LCLM agents can provide corpus-level visibility by compressing large inputs and selectively expanding segments into raw text when precise details are needed. On needle-in-a-haystack tasks, this adaptive expansion substantially improves over compressed-context LCLMs and can sometimes match uncompressed-context performance.

  • Motivation: Agents may miss relevant information when lexical or semantic search terms are not obvious from the underlying content.A bug described as part of a dashboard login flow may originate in an indirectly called entitlement module that mentions neither term.
  • Method: LCLM agents compress fixed-size 512-token chunks, assign integer identifiers, and can expand selected segments into raw text on demand.The agent receives the entire compressed sequence in one prompt along with an EXPAND mechanism.
  • Evaluation: Needle-in-a-haystack evaluations from RULER test the agentic retrieval mechanism with latent context.The testbed evaluates whether compressed latent context can support retrieval across large contexts.
  • Results: Adaptive expansion substantially improves performance over raw LCLM with 16× compressed context and sometimes matches the original uncompressed context.The approach combines broad global skimming with fine-grained reading of selected relevant subsets.

8 Conclusions · B Extended Related Work · C Dataset Details

LCLMs preserve strong in-context capabilities while improving efficiency across long-context, general-task, and agentic settings. The related work spans hard- and soft-token compression, KV-cache methods, efficient sequence modeling, positional encoding, and parallel context encoding, with these approaches largely orthogonal to LCLMs.

  • 8 Conclusions: LCLMs preserve strong in-context capabilities while providing substantial efficiency gains across long-context benchmarks, standard tasks, and agentic settings.The paper presents learned compression as a practical building block for efficient long-context systems.
  • 8 Conclusions: LCLMs reduce input size at scale and offer a promising substrate for long-horizon agents with large, persistent working memories.They are naturally compatible with agentic frameworks such as Recursive Language Models, while future work may compress inputs at multiple granularities and dynamically.
  • B Extended Related Work: Hard-token compression deletes or rewrites input tokens through importance-based pruning, summarization, or paraphrasing, effectively reducing prompt length.These methods include token pruning and context compaction used in systems such as Claude Code and Codex.
  • B Extended Related Work: Agent memory systems extend interaction horizons through summarization, retrieval, or structured stores, but typically use lossy summaries or task-specific representations external to the base model.Examples include hierarchical memory management, long-term conversational memory, and agentic memory frameworks.
  • B Extended Related Work: KV-cache compression evicts or reconstructs entries using learned or hand-crafted importance heuristics, including attention-based pruning, cache reconstruction, pooling, and offline prompt-specific compression.Methods include SnapKV, KVzip, Fast KV Compaction via Attention Matching, Synk et al., and Eyuboglu et al.
  • B Extended Related Work: Soft-token compression encodes long contexts into fewer learned continuous tokens, with the strongest approaches often training on the target context and achieving high compression ratios for those contexts.Prefix-tuning methods can achieve comparable in-context performance to the base model on contexts used for training.
  • B Extended Related Work: Efficient long-sequence modeling modifies sequence models, attention, or positional encodings, while parallel context encoding exposes compressed chunk representations to a decoder.These directions include recurrence, linear attention, state-space models, latent or sparse attention, RoPE extensions, and CEPE; they reduce attention cost, positional-indexing cost, or effective sequence length rather than decoder input tokens, and can in principle compose with LCLMs.

C.1 Continual Pre-training Dataset Curation … C.4 Compression Data Format

The compressor is trained on diverse next-token, reconstruction, and instruction-tuning data because reconstruction alone generalizes poorly beyond reconstruction tasks. Compression data uses alternating wrapped and unwrapped segments, with distinct wrapping schemes for documents and conversations.

  • C.1 Continual Pre-training Dataset Curation: Reconstruction-only training preserves sequences but generalizes poorly, failing to perform other tasks even when the LLM decoder remains frozen.This resembles a general form of prefix tuning.
  • C.1 Continual Pre-training Dataset Curation: A mixture of next-token prediction and reconstruction trains task-aligned decoder behavior while encouraging information preservation and accelerating early representation learning.The authors note that sufficiently large-scale next-token training alone may recover fine-grained fidelity.
  • C.1 Continual Pre-training Dataset Curation: Continual-pretraining sources include Nemotron code, text, and specialized SFT mixtures, plus approximately 500K OLMo-3 longmino examples to expose the compressor to longer documents.Nemotron synthetic data includes teacher-model rewrites or generations intended to alleviate forgetting.
  • C.2 Auxiliary Reconstruction Dataset Curation: Reconstruction data spans text, code, and LATEX from diverse corpora to preserve coverage for downstream tasks.Sources include Nemotron CC/math, FineWiki, RedPajama ArXiv, RedPajama GitHub, The Stack v2, CodeParrot, and TexTeller-OCR.
  • C.3 SFT Dataset Curation: Reasoning supervision is curated from math, code, and science data in Nemotron SFT datasets and OLMo-3 Dolci-Think.These sources provide supervision for reasoning-oriented instruction tuning.
  • C.3 SFT Dataset Curation: Long-context instruction tuning combines long-context and RAG-style QA with code-instruction datasets such as SWE-Fixer, NextCoder, PyResBugs, CommitPack, and RepoBench.The cited code datasets target software engineering and repository-level tasks.
  • C.3 SFT Dataset Curation: Instruction-following data combines multi-turn conversations, relabeled alignment data, OLMo-3 and Nemotron instruction tuning, and the Tulu-3 mixture.Responses from datasets already generated by strong teacher models are kept unchanged.
  • C.4 Compression Data Format: Continual-pretraining sequences alternate compressed and uncompressed segments by wrapping every other segment with <|memory_start|> and <|memory_end|>.Raw documents begin with a wrapped segment and end with an unwrapped segment; documents longer than 16,384 tokens retain 8,192 trainable tokens, while shorter documents retain half.

D Extended Architecture Search … D.3 Boundary Context

The extended architecture search compares pooling, adapter, masking, and boundary-context choices for encoder-decoder compression. It favors mean pooling, MLP adapters, causal masking, and non-overlapping windows based on loss and compute tradeoffs.

  • D.1 Formal Definitions of Pooling Operators: At scale, mean pooling remains marginally better than concat at higher compression ratios, whereas concat wins at lower compression ratios.The comparison identifies mean and concat as the most promising pooling operators, with mean slightly better in training loss initially.
  • D Extended Architecture Search: The architecture sweep finds causal masking better than bidirectional masking, non-overlapping windows preferable to boundary overlap, and MLP adapters better than attention-based adapters.These trends also hold at scale, while attention-based adapters add quadratic-in-M latent mixing computation.
  • D.1 Formal Definitions of Pooling Operators: Mean pooling averages hidden states within each compression block and consistently outperforms EOS-style token pooling across the architecture search.For latent index k, the pooling group contains consecutive positions from (k − 1)N + 1 through min(kN, |wi|).
  • D.1 Formal Definitions of Pooling Operators: Concat pooling preserves each token representation by concatenating consecutive hidden states, producing a latent vector of dimension N denc.Its adapter input dimension is widened accordingly before projection to the decoder hidden dimension ddec.
  • D.2 Adapter Design: The MLP adapter achieves lower pre-training loss than the attention-based variant with less computation and becomes the default for subsequent experiments.The attention-based adapter adds self-attention over latent tokens before MLP projection, whereas the MLP processes each latent independently.
  • D.3 Boundary Context: Boundary overlap extends each encoder window with O neighboring tokens while keeping the number of latent tokens unchanged, reducing information splitting across window boundaries.Pooling is applied only to the original non-overlapped region, and overlap outputs are discarded.
  • D.3 Boundary Context: With W = 1024 and O = 256, each interior bidirectional window processes 1536 tokens instead of 1024, while overlap usually fails to improve pre-training loss.The sweep uses O ∈ {0, 32, 256} for W ∈ {256, 1024}; overlap is therefore excluded from the default architecture.

E Full-Scale Training Recipe Sweeps … E.4 Packed Sequences

The full-scale recipe sweeps identify effective encoder initialization, learning-rate ranges, model-scaling behavior, and packed-sequence training procedures. At scale, embedding-initialized encoders and the selected architecture configuration minimize pre-training loss, while scaling outcomes vary by evaluation and packed attention prevents cross-example leakage.

  • E.1 Pre-trained Encoder Representations: Embedding-model initialization outperforms language-model initialization for the encoder when using the EOS pooling token.The embedding model was initialized from the language model.
  • E.1 Pre-trained Encoder Representations: At 16× compression, W = 1024, causal masking, an embedding-initialized encoder, and an MLP adapter yield the lowest continual pre-training loss.Figure 10 reports that the configuration selected by the from-scratch sweep continues to perform best in full-pipeline runs.
  • E.3 Scaling Up Model Size: Scaling experiments compare a Qwen3-Embedding-4B encoder with a Qwen3-4B-Instruct-2507 decoder against a Qwen3-Embedding-0.6B encoder with a Qwen3-8B decoder.The experiments test whether larger encoders or decoders improve compression quality and performance.
  • E.3 Scaling Up Model Size: The 0.6B encoder performs best across RULER tasks, whereas the 4B encoder performs best on the remaining evaluations.These mixed scaling results are reported for 16× compression.
  • E.3 Scaling Up Model Size: The 8B decoder substantially lowers pre-training loss but does not deliver the expected downstream gains.The paper attributes this partly to a mismatch between the training mixture and decoder initialization, which were tuned around the 4B instruct decoder.
  • E.4 Packed Sequences: Packing multiple variable-length examples into one sequence improves training throughput by reducing computation wasted on padding.Attention masks reset at example boundaries, producing block-diagonal attention and preventing cross-example information leakage.

F Compute … G.3.1 Continual Pre-training LR sweep

The paper reports compute configurations, benchmark coverage, GSM8K performance under aggressive compression, aggregate and per-task results, and continual-pretraining learning-rate sweeps. LCLMs achieve the highest GSM8K accuracy across all tested compression ratios while maintaining strong performance at larger ratios.

  • F Compute: From-scratch pre-training uses 16 MI300A nodes, while full-scale continual pre-training of 0.6B–4B models uses 32 nodes.Each MI300A node contains 4 GPUs.
  • F Compute: Evaluations run on H200 GPUs, while scaling experiments use 32-node H200 clusters to alleviate memory issues.The scaling setup uses a Qwen3-Embedding-4B encoder and Qwen3-8B decoder.
  • G Benchmarks and Results: The evaluation suite covers long-context, general reasoning, knowledge, instruction-following, and agentic coding benchmarks, with specified compressible input components.Instructions and other remaining components stay as hard tokens when applicable.
  • G.1 GSM8K Performance vs Compression Ratio: LCLMs achieve the highest GSM8K accuracy across all compression ratios, with particularly strong gains over baselines at higher compression ratios.These settings include aggressive compression ratios of 16× and beyond.
  • G.1 GSM8K Performance vs Compression Ratio: Figure 12 shows that LCLMs maintain much higher GSM8K accuracy at larger compression ratios.Compression time is not reported because GSM8K contexts are too small for reliable measurements.
  • G.3.1 Continual Pre-training LR sweep: The continual-pretraining learning-rate sweep is summarized for 16 × compression using a bidirectional mask, MLP adapter, and O = 0.Additional tables provide RULER, LongBench English, and LongBench Chinese per-task or per-subtask results.

G.3.2 Post-training LR sweep … G.5 NIAH Agent results

The appendix evaluates post-training learning rates and multiple architectural choices for 16× compression, including pooling, adapters, overlap, masks, windows, and encoder representations. It also reports NIAH retrieval results with and without an agentic expand-or-answer loop across context lengths.

  • G.3.2 Post-training LR sweep: The post-training learning-rate sweep fixes 16× compression, causal masking, an MLP adapter, and continual pre-training learning rate 1.0 × 10−5.The sweep is summarized overall and by RULER, LongBench English, and LongBench Chinese subtasks.
  • G.4 Architecture comparison: The architecture comparison examines pooling operators, adapter and overlap choices, mask types, window sizes, and encoder representations under 16× compression.The ablations use the configurations specified in their respective tables, including mean or EOS pooling and varying attention settings.
  • G.4.1 Pooling Operator: Pooling is compared between Mean and Concat with W = 1024, causal attention, and continual pre-training learning rate 1.0 × 10−5.The comparison includes aggregate results and per-task or per-subtask evaluations for RULER and both LongBench language settings.
  • G.4.2 Adapter & Overlap: Adapter and overlap configurations are compared at 16× compression with mean pooling, W = 1024, and continual pre-training learning rate 1.0 × 10−5.The evaluation reports aggregate, RULER per-task, and LongBench English and Chinese per-subtask results.
  • G.4.3 Mask-type Comparison: Mask type is ablated at 16× compression with mean pooling, W=1024, and continual pre-training learning rate 1.0 × 10−5.This isolates the effect of masking under a fixed compression and pooling configuration.
  • G.4.4 Window-size Comparison: Window size is ablated at 16× compression with mean pooling and continual pre-training learning rate 1.0 × 10−6.W=16 and W=256 use causal attention, whereas W=1024 uses bidirectional attention; W=1024 causal attention is not run at this learning rate.
  • G.4.5 Encoder-representation Comparison: Encoder representations are compared at 16× compression using EOS pooling, causal attention, and continual pre-training learning rate 1.0 × 10−6.This ablation isolates the representation supplied by the encoder under a fixed pooling and attention setup.
  • G.5 NIAH Agent results: NIAH retrieval is evaluated for the LCLM 0.6B-4B model at 16× compression with and without an agentic expand-or-answer loop across context lengths.∆agent denotes the cell-wise gain over the non-agentic baseline.
Loading 2606.09659v1…