---
author:
- 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
bibliography:
- refs.bib
title: End-to-End Context Compression at Scale
---
\newcommand{\TODO}[1]{\textcolor{red}{\textbf{TODO} #1}}
\newcommand{\inlineimage}[1]{\raisebox{-0.05\height}{\includegraphics[height=2.5em]{#1}}}
\newcommand{\todo}{{\color{red}{todo}}\xspace}
\DeclareRobustCommand\onedot{\futurelet\@let@token\@onedot}
\def\@onedot{\ifx\@let@token.\else.\null\fi\xspace}
\def\aka{\emph{a.k.a}\onedot}
\def\Eg{\emph{A.k.a}\onedot}
\def\eg{\emph{e.g}\onedot}
\def\Eg{\emph{E.g}\onedot}
\def\ie{\emph{i.e}\onedot}
\def\Ie{\emph{I.e}\onedot}
\def\cf{\emph{c.f}\onedot}
\def\Cf{\emph{C.f}\onedot}
\def\etc{\emph{etc}\onedot}
\def\vs{\emph{vs}\onedot}
\def\wrt{w.r.t\onedot}
\def\dof{d.o.f\onedot}
\def\etal{\emph{et al}\onedot}
\newcommand{\xxcomment}[4]{\textcolor{#1}{[$^{\textsc{#2}}_{\textsc{#3}}$ #4]}}
\newcommand{\micah}[1]{\xxcomment{red}{M}{G}{#1}}
\newcommand{\df}[1]{\xxcomment{teal}{D}{Fu}{#1}}
\newcommand{\bz}[1]{\xxcomment{violet}{B}{Z}{#1}}
\newcommand{\leon}[1]{{\color[rgb]{0.505, 0, 0.921}{Leon: \footnotesize\sf[#1]}}}
\renewcommand{\paragraph}[1]{\textbf{#1}}
\renewcommand{\paragraph}[1]{\sansbf{#1}}
\newcommand{\tokcount}{350B\xspace}
\newcommand{\enctokcount}{165B\xspace}
\newcommand{\llmtokcount}{185B\xspace}
\newcommand{\names}{LCLMs\xspace}
\newcommand{\name}{LCLM\xspace}
\newcommand{\fullnames}{Latent Context Language Models\xspace}
\newcommand{\fullname}{Latent Context Language Model\xspace}
\newcommand{\nyu}{\ensuremath{\textcolor[HTML]{57068C}{\spadesuit}}}
\newcommand{\umd}{\mbox{\textcolor[HTML]{E21833}{\scalebox{0.6}{\faHeart}}}}
\newcommand{\princeton}{\ensuremath{\textcolor[HTML]{E77500}{\clubsuit}}}
\newcommand{\columbia}{\ensuremath{\textcolor[HTML]{009EFF}{\blacklozenge}}}
\newcommand{\modal}{\ensuremath{\textcolor[HTML]{00A86B}{\scalebox{1.0}{$\blacktriangle$}}}}
\newcommand{\meta}{\ensuremath{\textcolor[HTML]{0064E0}{\scalebox{0.7}{$\blacksquare$}}}}
\newcommand{\llnl}{\mbox{\textcolor[HTML]{003C71}{\raisebox{-0.15ex}{\scalebox{1.25}{$\bullet$}}}}}
\newcommand{\harvard}{\ensuremath{\textcolor[HTML]{A51C30}{\bigstar}}}
\maketitle
\blankfootnote{$^{1}$Work completed during an internship at Modal.}
\blankfootnote{$^{2}$Served solely in an advisory role for this paper. All model training and data processing were conducted entirely outside of Meta.}
# Introduction {#sec:intro}
Reasoning over long contexts is a crucial capability for state-of-the-art Large Language Models (LLMs), as it enables them to parse long documents, engage in multi-turn conversations, and perform long-horizon agentic tasks. However, in production systems, the input context, working horizon, and memory can grow to millions of tokens, making inference increasingly constrained by memory and latency due to the growth of the KV cache [@hooper2024kvquant]. Even when inputs fit within the model's maximum context window, models often struggle to reliably use information distributed across long contexts [@liu2024lost; @an2025why]. As a result, context and memory management have emerged as critical systems and modeling challenges for deploying long-horizon LLM systems.
A natural response is to reduce the size of the KV cache directly. KV cache compression methods reduce the memory footprint of cached activations by evicting cache entries [@li2024snapkv; @devoto2025expected; @kim2025kvzip; @zweiger2026fast] or training a compact cache offline [@cartridges]. However, they have important limitations for general long-context inference. Many methods still require the full context to be prefilled before compression, and some require substantial additional time or memory beyond the original prefill. Query-dependent methods produce caches that are difficult to reuse across turns [@li2024scbench], while methods that evict caches non-uniformly across attention heads and layers [@kim2025kvzip] are difficult to integrate with modern inference engines that assume a shared sequence length across the cache.
Our achieve high quality compression while being fast and memory efficient. We show RULER (left) and LongBench (right) accuracy vs. TTFT (time to first token) per sample on a single H200. All methods are based on the same decoder Qwen3-4B-Instruct-2507 . We see that our models lie on a new Pareto frontier, compressing contexts much faster while maintaining high accuracy, especially for high compression ratios. KV cache compression baselines appear as vertical lines because the context is prefilled the same way regardless of compression ratio, and the eviction operation is much faster than the prefill time.
Encoder-decoder soft-token methods [e.g. @refrag; @in-context-autoencoder; @auto-compressor; @e2llm; @cepe] offer a promising way to address the limitations of KV cache compression. Rather than manipulating the KV cache directly, these models encode the raw input tokens into a much shorter sequence of continuous embeddings, which are then provided to the LLM in place of the original context. In principle, this approach is highly attractive: compression is parallelizable, supported by standard LLM inference engines, and works with standard LLM decoding while extending the decoder far beyond its native context length.
However, existing soft-token compression methods typically face a trade-off: they either substantially degrade the capabilities of the base language model or depend on domain-specific training to remain effective. This motivates a fundamental question: **can we train a general, task-agnostic compressor that preserves a model's original capabilities?** To this end, we introduce *`\fullnames`{=latex}* (`\names`{=latex}), a family of encoder-decoder compressors trained end-to-end at scale. We initialize both the encoder and decoder with pre-trained models and jointly train them on up to a total of `\tokcount `{=latex}tokens for the combined encoder-decoder. Crucially, we demonstrate that `\names `{=latex}can preserve decoder performance while reducing memory used for input context.
Our contributions are summarized as follows:
1. **Training recipe:** We develop a fine-grained training recipe for encoder-decoder context compression. Accordingly, we curate continual pre-training and finetuning data for training compressors and optimize training details across multiple training stages. See `\Cref{sec:training}`{=latex}.
2. **Comprehensive architectural search:** We conduct a large-scale architecture search to identify optimal design choices for encoder-decoder compression models in a controlled setting. See `\Cref{sec:arch}`{=latex} and `\Cref{fig:arch_sweep}`{=latex}.
3. **New memory/speed-performance frontier:** We train a family of models across different compression ratios at scale. Our experiments identify a Pareto frontier that effectively balances memory, latency, and downstream accuracy. We open-source this suite of models and all code. See `\Cref{sec:results}`{=latex} and `\Cref{fig:cover,fig:pareto_optimal,fig:more_benchmarks}`{=latex}.
4. **Agent with latent context:** We create an agentic system with a high compression ratio, where the agent can select which compressed chunk to expand. This agentic harness substantially enhances the model's performance on challenging needle-in-the-haystack tasks. See `\Cref{sec:agentic}`{=latex} and `\Cref{fig:agentic_bar_plot}`{=latex}.
# Related Work {#sec:related}
Context compression approaches generally fall into three categories: hard-token, soft-token, and KV cache compression. We defer the discussion of hard-token approaches to `\Cref{app-sec:extended-rel-work}`{=latex}, since they generally underperform compared to the latter two.
**KV cache compression.** KV cache compression methods evict entries from the KV cache. Most methods rely on hand-crafted or learned heuristics to select which entries to drop.
Prompt-agnostic methods [@kim2025kvzip; @devoto2025expected; @zweiger2026fast] prune the context without knowledge of the query, whereas prompt-dependent methods such as SnapKV [@li2024snapkv] require explicit context-prompt pairs and yield query-specific caches that limit multi-turn applicability [@li2024scbench; @kim2025kvzip]. An alternative is to anticipate future queries at compression time: Cartridges [@cartridges] distills a fixed-size KV cache per corpus. However, this comes at a substantial training cost per corpus, with an ICL-quality cache requiring approximately 30 minutes on an 8$\times$H100 node for an 8B model.
Although KV cache compression methods are well-studied, they are not widely adopted in inference engines such as vLLM [@kwon2023efficient] or SGLang [@zheng2024sglang]. Compression methods such as KVzip allocate eviction budgets non-uniformly across attention heads and layers; therefore, they cannot reduce the sequence-length dimension of the KV cache. Instead, KV cache compression methods mask evicted positions in the attention computation, forfeiting the memory and throughput benefits of these paged-attention engines. Some methods further require prefill passes that are two to three times longer than the original input [@kim2025kvzip]. We note that soft token compression methods are fully compatible with these open-source inference frameworks.
**Soft-token compression.** Despite promising results, 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 [@GMSA; @mean-pooling; @dai2025pretraining; @refrag; @xrag]. To the best of our knowledge, these methods are not evaluated on long-context benchmarks with information-dense tasks that require fine-grained details throughout the whole context, such as RULER [@hsieh2024ruler], to demonstrate complex long-context understanding.
In `\Cref{sec:arch}`{=latex}, we discuss the architectural decisions of prior work on soft-token compression as part of our architectural search and analysis. Our work shows that, with large-scale staged training, a single online soft-token compressor can robustly handle general heterogeneous inputs while closely matching uncompressed behavior.
# The `\fullname `{=latex}Architecture
An `\name `{=latex}consists of an encoder with a pooling method that maps token chunks to soft tokens, an adapter to project to the decoder's embedding dimension, and a decoder that consumes the soft tokens as latent context. Given a sequence of input tokens $x_{1:T}$ and a compression ratio $N$, the encoder maps each contiguous block of $N$ input tokens to a single latent token. Let $W$ denote the encoder window size, i.e., the number of input tokens processed by the encoder in one forward pass. We split the input sequence into $I=\lceil T/W\rceil$ encoder windows, each containing at most $W$ tokens: $$\begin{equation}
w_i
=
x_{
(i-1)W+1
\; : \;
\min(iW,T)
},
\qquad
i = 1,\dots,I.
\end{equation}$$ For each window $w_i$, the encoder computes the final hidden states $h^{(i)}_{1:|w_i|}$ for each input token in the window. A pooling operator then aggregates these hidden states into $M_i = \left\lceil \frac{|w_i|}{N} \right\rceil$ latent tokens $z^{(i)}_{1:M_i}$. When $W=N$, each encoder window contains a single compression block, yielding one latent token per forward pass. When $W=T$, the encoder processes the entire input sequence in one forward pass and produces $\lceil T/N\rceil$ latent tokens. We concatenate the latent tokens from all encoder windows to obtain the full compressed latent sequence $z_{1:M}$, where $M$ is the total number of latent tokens after compression.
Because the encoder and decoder may have different hidden dimensions, an adapter $a(\cdot)$ projects the compressed latent sequence $z_{1:M}$ from the encoder hidden dimension into the decoder hidden dimension, producing latent tokens $s_{1:M}$. The decoder then consumes these projected latent tokens in place of the corresponding input tokens.
# Training Recipe {#sec:training}
We now describe our multi-stage recipe and data for training. In contrast to prior context-compression work that trains specialized models on small-scale in-domain datasets [@GMSA; @mean-pooling; @arc-encoder; @xrag; @e2llm], our goal is to preserve the strong performance of a powerful LLM across downstream tasks. To this end, we curate three types of data: continual pre-training data, Supervised Fine-Tuning (SFT) data, and auxiliary reconstruction data.
Examples of the three data types used to train . We curate continual pre-training data with interleaved compressed/uncompressed blocks of tokens, supervised fine-tuning data with compressed prompts and long documents, and auxiliary reconstruction data that asks the decoder to reproduce the original context from the compressed latents.
## Training Data
#### Continual pre-training dataset.
We construct a carefully curated high-quality pre-training mixture covering Common Crawl text, code, science and reasoning, long-context data, and instruction-style data. We partition each sequence into multiple segments and alternate between compressed segments and standard token segments. We then compute the next-token prediction loss on the uncompressed tokens *only*. This interleaved format differs from the simple first-half compressed, second-half trainable split used in prior work [@refrag; @in-context-autoencoder; @auto-compressor; @cepe]. By distributing compressed spans throughout the sequence, the model learns to condition on latent context at multiple positions rather than only at the beginning of the context. We provide a visual example of this data format in `\Cref{fig:data_example}`{=latex}. The exact pre-training mixture is described in `\Cref{app-subsec:cont-pre-train-data}`{=latex}, with additional details on compression formatting in `\Cref{compression-data-details}`{=latex}.
#### SFT data.
To recover performance after continual pre-training [@gao2025train] and further enhance instruction following capability with latent context, we further post-train the model with supervised fine-tuning to improve reasoning, instruction following, and long-context understanding when conditioning on compressed inputs. We construct an SFT mixture spanning three components: (i) reasoning, (ii) long-context instruction following, and (iii) general instruction following and multi-turn conversation. Some completions in the original SFT data were generated by outdated models. To improve response quality, we relabel a subset of the SFT data using `Qwen3-30B-A3B-Instruct-2507` and `Qwen3-235B-A22B-Instruct-2507`, whose model checkpoints are released under the Apache-2.0 license. We report the exact SFT mixture and other details in `\Cref{app-subsec:posttrain-data}`{=latex}.
#### Auxiliary reconstruction data.
To accelerate training and encourage the latent representations to preserve fine-grained information, we introduce an auxiliary reconstruction task. Given a sampled document, we compress the text and ask the model to repeat the original content. To increase diversity, we include a variety of sources spanning code, text, long documents, math, and LaTeX, as described in `\Cref{app-subsec:aux-recon-data}`{=latex}. To reduce prompt overfitting, we use a bank of $100$ prompt templates per source. This auxiliary task encourages the compressor not only to support semantic understanding but also to retain fine-grained details needed for tasks such as exact retrieval. We include auxiliary reconstruction data in both continual pre-training and SFT.
## Training Stages
We adopt a multi-stage training recipe to improve stability, preserve the original model's capabilities, and mitigate catastrophic forgetting. The stages progressively increase the number of trainable components, starting with adapter and encoder warm up, followed by end-to-end continual pre-training and SFT.
- **Stage 0: Adapter warmup.** We freeze both the encoder and decoder and update only the adapter.
- **Stage 1: Encoder training.** We unfreeze the encoder while keeping the decoder frozen.
- **Stage 2: End-to-end continual pre-training.** We unfreeze the decoder with a small learning rate.
- **Stage 3: Supervised fine-tuning.** We train on the SFT mixture and increase the decoder learning rate.
This staged procedure is analogous to Vision-Language Model (VLM) training pipelines that first align pre-trained representations across modalities [@llava; @tong2024cambrian] and then fine-tune for downstream tasks. For stages 0, 1, and 2, we uniformly sample from the continual pre-training mixture; for stage 3, we train on the SFT mixture. In early experiments, we consider directly training the full model end-to-end from the beginning, but find that this approach underperforms the multi-stage recipe. At the beginning of training, the decoder is not accustomed to taking in the output of the embedding model, so both encoder and decoder parameter gradients can be large, leading to model degradation. The staged approach, where we first train the adapter before unfreezing the encoder and decoder, keeps gradients smooth and training more stable, preventing this degradation. We therefore adopt the multi-stage recipe for all scaled experiments. We also consider parameter-efficient variants that keep the decoder fully frozen [@in-context-autoencoder; @xrag; @cepe] or use LoRA [@e2llm; @GMSA; @mean-pooling], but we find them to substantially underperform compared to full-parameter training in early experiments; we therefore focus on end-to-end training. Details of the continual pre-training and SFT data mixtures are provided in `\Cref{app-tab:data-mixture-pre-training}`{=latex} and `\Cref{app-tab:SFT-mixture}`{=latex}, respectively. All other training details, including token budgets, learning rates, and sequence lengths, are provided in `\Cref{tab:training_recipe}`{=latex}.
\begin{table*}[t!]\small
\setlength{\tabcolsep}{8pt}
\caption{\textbf{Training recipe for compression model for four different stages.} We report, for each stage, which module(s) are optimized (adapter / encoder / decoder), peak learning rates, learning-rate schedule, optimizer, global batch size (in tokens), sequence length, training-token budget, and the data source used.}
\label{tab:training_recipe}
\resizebox{\linewidth}{!}{
\begin{tabular}{l|cccc}
\toprule
& \textbf{Stage 0} & \textbf{Stage 1} & \textbf{Stage 2} & \textbf{Stage 3} \\
& \textbf{Adapter Training} & \textbf{Encoder Training} & \textbf{LLM Training} & \textbf{SFT} \\
\hline
\\[-4pt]
\textbf{Training Config} \\
Adapter Peak LR & $1.0 \times 10^{-3}$ & $6.0 \times 10^{-5}$ & $6.0 \times 10^{-5}$ & $3.0 \times 10^{-5}$ \\
Encoder Peak LR & NA & $6.0 \times 10^{-5}$ & $6.0 \times 10^{-5}$ & $3.0 \times 10^{-5}$ \\
LLM Peak LR & NA & NA & $1.0 \times 10^{-5}$ & $3.0 \times 10^{-5}$ \\
LR scheduler & 5\% warmup steps with cosine decay to $1.0 \times 10^{-6}$ \\
Optimizer & AdamW ($\beta_1=0.9$, $\beta_2=0.95$) \\
\midrule
\textbf{Data} \\
LLM Batch size (tokens) & 4 Million \\
LLM Sequence Length & 16384 \\
Encoder Tokens & $17.54$ B & $35.07$ B & $82.43$ B & $29.73$ B \\
LLM Tokens (16x) & $21.29$ B & $42.58$ B & $100.08$ B & $21.32$ B \\
Total Tokens (16x) & $38.83$ B & $77.65$ B & $182.51$ B & $51.05$ B \\
Data Source & pre-training Mixture & pre-training Mixture & pre-training Mixture & SFT \\
\bottomrule
\end{tabular}
}
\end{table*}
A from-scratch pre-training sweep identifies the best encoder-decoder compressor architecture.Left: We explore how to best train an encoder LM to compress raw text into latent vectors and a decoder LM that uses the vectors as latent context. Right: We sweep many variants of encoder-decoder compressors end-to-end with 38B training tokens at a compression ratio of 16× and compare pre-training loss to find the best architecture. Since loss is computed only on the uncompressed tokens, the pre-training loss is substantially lower than a standard next-token prediction loss over all tokens. Extended experiments are in .
# Architectural Design Space of Latent Encoders {#sec:arch}
We conduct a controlled architecture search over encoder-decoder context compressors by pre-training `\name `{=latex}variants from scratch. Our cleanroom setting lets us isolate the effect of architectural choices without confounding from pre-trained initialization. This is particularly important for choices such as the encoder attention mask and pooling operators: for example, the pre-trained `Qwen3-Embedding-0.6B` encoder uses causal attention and EOS pooling, which could bias the architectural search. Specifically, we use `Qwen3-0.6B` as the base architecture for both the encoder and decoder, randomly initializing all weights with the same seed, and training all components end-to-end with no frozen parameters. Each variant is trained on the pre-training mixture described in `\Cref{sec:training}`{=latex} for 38B tokens at a compression ratio $N=16$. This comprehensive search covers nearly all major architectural choices explored in prior soft-token compression work, along with additional underexplored choices. We describe each design axis and its empirical effect below.
#### Pooling operator.
Prior work on context compression has primarily focused on two pooling strategies. Token-based pooling appends or prepends designated tokens, such as `EOS` or `CLS`, and uses their final hidden states as compressed representations [@refrag; @e2llm; @auto-compressor]. Mean pooling averages encoder hidden states over each block of $N$ input tokens [@in-context-autoencoder; @GMSA], a design that is common in vision and multimodal models [@gemma4_2025]. We additionally consider concatenation, a widely used alternative in vision-style architectures but less explored for context compression [@qwen3vl; @tong2024cambrian]. Concatenation preserves the $N$ sequential encoder hidden states by stacking them into a single wide vector, after which the adapter projects the resulting representation back to the LLM hidden dimension. The formal definitions of pooling operators can be found in `\Cref{app-subsec:encoder-defs}`{=latex}.
`\Cref{fig:arch_sweep}`{=latex} compares the three pooling operators in our pre-training sweep. Mean pooling consistently improves pre-training loss over token-based pooling, but is empirically indistinguishable from concatenation. This agrees with prior evidence that mean pooling can outperform token-based pooling for compression [@mean-pooling], and with findings that mean pooling can be surprisingly effective despite being a lossy operation in text embedding models [@hara2026mean]. It is also consistent with related observations in vision architectures [@chu2021conditional]. While concatenation can, in principle, preserve distinct representations that mean loses, and has been found to outperform mean pooling in other representation-learning settings [@chen2025comodo], we do not observe this advantage in our small-scale experiment. We compare these two operators at scale in `\Cref{app-fig:pool-abl-at-scale,app-subsubsec:pooling-operator}`{=latex}, where concatenation outperforms mean pooling at low compression ratios and mean pooling outperforms concatenation at high compression ratios, while the gap between the two narrows as context length increases.
#### Encoding granularity.
Ideally, one would set the encoder window size to the full context length, allowing the encoder to jointly contextualize all input tokens, analogous to ViT-style encoders [@dosovitskiy2020image] that process the entire image at once. For language inputs, however, the context length can be substantially larger, making full-context encoding prohibitive in both memory and time.
To study this trade-off, we vary the encoder window size over $W \in \{N, 256, 1024\}$ for mean pooling in `\Cref{fig:arch_sweep}`{=latex}. We find that increasing $W$ from $N\;(16)$ [@refrag; @e2llm] to $256$ leads to a large improvement, while increasing $W$ further from $256$ to $1024$ yields an additional smaller gain. Since $W=1024$ does not introduce significant memory or runtime overhead in our setting, we adopt $W=1024$ as the default. This setting enables the encoder to attend over a larger local context and produce richer compressed representations. One limitation of using $W < T$ is that the input must be chunked into multiple windows, so information crossing a boundary is split across two encoder forward passes. To address this, we also design a boundary-overlap variant that allows each window to see neighboring tokens during encoding. However, this approach does not improve pre-training loss while increasing encoder compute, so we use no overlap by default and defer details to `\Cref{app-subsec:arch-overlap}`{=latex}.
#### Encoder attention mask.
Unlike standard decoder language modeling, where causal masking is required to match autoregressive inference, context compression is applied to the prompt before decoding. The encoder can therefore use either causal or bidirectional attention, as in prefix-LMs [@raffel2020exploring]. In `\Cref{fig:arch_sweep,app-fig:adapter-ov-loss}`{=latex}, we compare these two masking choices and find that causal masking consistently achieves lower pre-training loss.
#### Adapter design.
The adapter projects encoder latents from dimension $d_{\mathrm{enc}}$ to the decoder embedding dimension $d_{\mathrm{dec}}$. Prior encoder-decoder compression work often assumes matching encoder and decoder hidden sizes, but a wide range of encoder-decoder models pair smaller encoders with larger decoders, making the adapter a non-trivial design choice. We compare a lightweight MLP projection against an attention-based adapter that adds a single self-attention layer over the latent sequence before the MLP projection. In contrast to prior findings [@GMSA], we observe in `\Cref{app-fig:adapter-ov-loss}`{=latex} that the MLP-only adapter achieves lower pre-training loss while using less compute, and we adopt it as the default. Formal definitions are provided in `\Cref{app-subsec:adapter}`{=latex}.
have lower TTFT and peak GPU memory as context length increases. We plot time to first token (left) and peak GPU memory (right) for a single sample. We note the baselines use the same time and memory regardless of compression ratio; hence, there is only one line per baseline method. We record all measurements on the same H200 GPU, and truncate lines when the method cannot be performed within the 141GB of memory. For peak GPU memory, our method plateaus at longer contexts for larger compression ratios: in this regime, the encoder’s batched processing dominates memory usage over the decoder. Memory begins to grow again once the decoder’s prefill memory surpasses the encoder’s activation per forward pass.
#### Optimal architecture at scale. {#par:optimal-arch-at-scale}
Overall, our from-scratch pre-training experiments in `\Cref{fig:arch_sweep}`{=latex} serve as a lightweight architecture sweep, allowing us to narrow the search space to a smaller set of promising design choices. We then scale these candidates using the full training pipeline, both to identify the strongest architecture at scale and to validate that the small-scale pre-training findings remain predictive under large-scale training. For our main experiments, we use `Qwen3-Embedding-0.6B` as the encoder with `Qwen3-4B-Instruct-2507` as the decoder. Across the continual pre-training loss curves in `\Cref{app-fig:pool-abl-at-scale,app-fig:adapter-ov-loss,app-fig:large-scale-sweep}`{=latex} and the downstream ablations in `\Cref{app-subsec:arch-sweep}`{=latex}, we find that the optimal architecture at scale uses an encoder window $W=1024$, causal masking, and an MLP adapter. The choice of pooling operator depends on the compression ratio: mean pooling has a small edge over concatenation at $16\times$ compression, while concatenation slightly outperforms mean pooling at $4\times$. Since we primarily focus on $16\times$ compression in our analysis, we adopt mean pooling as the default in all subsequent experiments. Full results and per-task breakdowns are reported in `\Cref{app-sec:training-recipe-sweeps}`{=latex}.
We also study alternative model choices, including whether to initialize the encoder from an embedding model or a language model, in `\Cref{app-subsec:encoder_lm_or_emb}`{=latex}. In `\Cref{app-sec:scaling model size}`{=latex}, we analyze the performance of scaling the sizes of the encoder and decoder.
# Results: Improving Latency- and Memory-Performance Tradeoffs {#sec:results}
We predominantly focus on long-context benchmarks: we evaluate on RULER [@hsieh2024ruler], LongBench [@bai2024longbench], and LongHealth [@adams2025longhealth]. For the long-context evaluation suite, we maintain instructions as uncompressed tokens, while we compress long-context segments. We report which parts of each benchmark are compressed and which are left as standard uncompressed tokens in `\Cref{app-tab:eval-benchmarks}`{=latex}.
We measure two efficiency axes: peak memory during compression, and compression time. Peak memory includes both the encoder and decoder for a single sample at each context length. We report the mean over 30 measurements on the same H200 GPU, using 20 warmup steps to reduce kernel-level noise. We define compression time as time-to-first-token (TTFT), i.e., the time required for a method to reach the point at which it can generate the first token. We adapt the KVPress [@devoto2025expected] timing code[^1] to support our method. All methods are benchmarked using Hugging Face Transformers [@hf_transformers] implementations.
Since `\names `{=latex}process input tokens in encoder-window chunks, we can parallelize compression by batching encoder forward passes before passing the resulting soft tokens to the decoder. For all figures in this paper, we use an encoder batch size of $128$; with encoder window size $W=1024$, this corresponds to processing $131{,}072$ input tokens per batched encoder pass. We find this batch size provides a good trade-off between compression speed and encoder memory usage, improving GPU utilization without exceeding memory limits for ultra-long contexts. Moreover, because our method preserves the standard KV cache structure, prefill and generation can be further accelerated with standard inference frameworks such as vLLM [@kwon2023efficient] and SGLang [@zheng2024sglang], as well as efficient attention implementations [@linearattention; @kimilinear; @deepseekv2; @deepseekv4]. Therefore, the throughput measurements reported here are conservative relative to what optimized serving implementations could achieve.
establish a new Pareto frontier on long-context benchmarks in terms of time to first token. We plot accuracy across RULER and LongHealth, and find that our compress samples faster with equivalent or higher accuracy than other compression methods. Full results can be viewed in .
## Baselines
We consider a range of KV cache compression baselines, including Expected Attention [@devoto2025expected], Attention Matching [@zweiger2026fast], KVzip [@kim2025kvzip], FastKVzip [@kim2026fast], and SnapKV [@li2024snapkv]. Each method has different native support for query-aware, query-agnostic, and self-study compression configurations.
Query-aware methods, such as SnapKV, are designed to compress a context given a known query; however, the resulting compressed cache may not generalize well to unseen queries or future turns. Query-agnostic methods, such as KVzip and FastKVzip, compress a cache without access to the downstream query, often using reconstruction-style objectives such as repeating the original context. Self-study, as introduced by @cartridges, performs compression using synthetic queries to expand coverage over potential inference-time queries. These synthetic queries are typically produced by prompting the target decoder to generate questions about the context. In Attention Matching [@zweiger2026fast], the authors experiment with a mixed query-agnostic and self-study configuration and find that self-study has limited impact on the repeat-prefill prompt. Since self-study introduces additional compression-time overhead, we do not apply it to RULER.
It is also important to distinguish algorithmic and theoretical cache reduction from systems-realizable cache reduction. Some KV cache methods produce an actual smaller cache that can, in principle, reduce decoding memory and latency. Others, such as @kim2025kvzip [@kim2026fast], are evaluated by masking or subsampling entries from a full cache, or by using non-uniform eviction patterns across heads and layers. These approaches are useful for measuring the quality of a compressed cache, but they do not necessarily translate into wall-clock speedups or memory savings in standard inference engines without specialized kernels that have not been released or even implemented.
We use the KVPress [@devoto2025expected] implementation for SnapKV, KVzip, FastKVzip, and Expected Attention. We use the official implementation for Attention Matching. We discuss all baselines further in `\Cref{sec:related}`{=latex} and `\Cref{app-sec:extended-rel-work}`{=latex}.
## A New Pareto Frontier
We train `\names `{=latex}with compression ratios of $16\times$, $8\times$, and $4\times$. In `\Cref{fig:cover,fig:more_benchmarks}`{=latex}, `\names `{=latex}establish a new Pareto frontier in compression time and accuracy over KV cache compression baselines on RULER, LongBench, and LongHealth. KV cache compression methods appear as nearly vertical lines in these plots because their compression time is largely independent of the target compression ratio. These methods first materialize the full KV cache and then perform eviction, masking, or compaction, whose cost is small relative to the full-context prefill. In contrast, `\names `{=latex}reduce the sequence length before decoder prefill, so higher compression ratios directly reduce the amount of decoder-side computation and memory.
**Time and memory scaling with context length.** In `\Cref{fig:pareto_optimal}`{=latex}, we visualize compression time and peak GPU memory for our method and the baselines across context lengths from 4K to 1M. `\names `{=latex}achieve the fastest compression time and, at longer context lengths, substantially lower peak GPU memory. Attention Matching runs out of memory at 1M tokens and fails at 512K tokens due to numerical instability in the linear solver. All other methods run out of memory at 512K and 1M tokens. This efficiency stems from two properties of `\names`{=latex}: compression is performed over fixed-size encoder windows rather than through a full-context decoder prefill, and the encoder is much smaller than the decoder. Since our encoder processes at most 128K input tokens per batched forward pass, peak memory is initially dominated by encoder activations. As a result, peak memory remains nearly flat for the $16\times$ model from 128K to 512K tokens, and for the $8\times$ model from 128K to 256K tokens. At longer contexts, the decoder prefill over the compressed latent sequence becomes the dominant memory cost, causing peak memory to increase.
**Fine-Grained Compression.** In `\Cref{app-fig:gsm}`{=latex}, we use GSM8K to analyze fine-grained compression quality, where we compress the entire prompt and context. We find that `\names `{=latex}achieve the highest accuracy across all compression ratios on GSM8K, with particularly strong gains over baselines at higher compression ratios.
# Agent Scaffolding With Latent Context {#sec:agentic}
Agentic systems are commonly used to handle large amounts of information [@zhang2025rlm]. However, agents usually rely on lexical or semantic search to locate the information needed. These information retrieval methods can fail when the search keyword is not obvious before reading through the information. For example, in a large codebase, a bug reported in the \`\`dashboard login flow" may actually originate in an indirectly called entitlement module that never mentions \`\`dashboard" or \`\`login". Ideally, an agent should read all information at once before acting, but context limits of existing models prevent this from happening.
can use tools to retrieve compressed context and improve exact string-match accuracy.Left: We equip with a tool that retrieves compressed context and inserts it into their working context uncompressed. Right: We test the agent on challenging NIAH tasks, and the agent substantially improves exact string-match accuracy on retrieval tasks.
LCLMs' more efficient long-context capacity makes it possible to place, for example, an entire code repository in an agent's context. We further enhance the agent's capability to locate precise information by allowing the `\name `{=latex}agent to expand compressed segments into raw text on demand. We segment the input into fixed-size chunks of 512 tokens, compress each chunk, and assign it an integer identifier. The model receives the entire compressed sequence in a single prompt, together with one `EXPAND` tool. At each turn, the model can decide to make a tool call of the form `EXPAND(i)`, and the tool returns the original text of the expanded segment.
We use the needle-in-a-haystack tasks from RULER as a testbed for evaluating this agentic retrieval mechanism with latent context. As shown in `\Cref{fig:agentic_bar_plot}`{=latex}, our agent substantially improves performance over the raw `\name `{=latex}with $16\times$ compressed context, and in some settings matches the performance of the original uncompressed context. These results suggest that compressed latent context can provide broad corpus-level visibility, while tool-based expansion supplies the exact fine-grained information needed. The marked success of this initial agentic attempt suggests the promise of adaptive expansion. LCLMs can skim a large quantity of context globally before deciding which relevant subset to zoom in on and read more carefully. Future work may explore new mechanisms for determining which context to expand or how to learn this behavior end-to-end.
# Conclusions {#sec:discuss}
Across long-context benchmarks, standard knowledge and instruction-following tasks, and agentic settings, our family of `\fullnames `{=latex}demonstrates that learned compression can preserve strong in-context capabilities while providing substantial efficiency gains, making it a practical building block for efficient long-context systems.
`\fullnames `{=latex}open a broad design space for future work and are naturally compatible with agentic frameworks such as Recursive Language Models [@zhang2025rlm]. More generally, `\names `{=latex}provide a promising substrate for long-horizon agents with large, persistent working memories by dramatically reducing the size of the inputs at scale. Future iterations of `\names `{=latex}could further improve the quality-efficiency trade-off by compressing inputs at multiple granularities and dynamically allocating capacity based on information density or input perplexity. Such adaptive compression could allow models to preserve fine-grained details where needed while maintaining a compact global context, reducing reliance on explicit expansion tools. Another promising direction is to extend compression beyond static input context to the model's generated state, including long chain-of-thought, tool observations, and an agent's accumulated working history, which can grow to dominate the context budget in long-horizon tasks.
# Acknowledgments
MG and ZC were supported by IBM, the Google Research Award, and the NVIDIA Academic Grant Program for Researchers. SM and TG are supported by DOE Office of Science's ASCR AI for Science initiative, the NSF TRAILS Institute (2229885), and Coefficient Giving. NK is supported by the NSF Graduate Research Fellowship.
Computing resources for this project were supported in part by the Swiss National Supercomputing Center (CSCS). Prepared in collaboration with LLNL under Contract DE-AC52-07NA27344 and supported by the LLNL-LDRD Program under Project No. 24-ERD-010 (LLNL-CONF-2019739). This manuscript has been authored by Lawrence Livermore National Security, LLC under Contract No. DE-AC52-07NA27344 with the U.S. Department of Energy. The United States Government retains, and the publisher, by accepting the article for publication, acknowledges that the United States Government retains a non-exclusive, paid-up, irrevocable, world-wide license to publish or reproduce the published form of this manuscript, or allow others to do so, for United States Government purposes. This material is based upon work supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, through solicitation DE-FOA-0003264, \`\`Advancements in Artificial Intelligence for Science," under Award Number DE-SC0025598.
Special thanks to Shengbang Tong for extensive discussions on architecture design and encoder-decoder model training. We would like to thank Preston Zhou and Eitan Borgnia for their early discussions and support for the project. We also thank Vadim Bereznyuk, John Kirchenbauer, and Neel Jain for providing feedback on the draft.
\bibliographystyle{plainnat}
\newpage
\appendix
\onecolumn
# Impact Statement {#app-sec:impact}
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be highlighted here.
# Extended Related Work {#app-sec:extended-rel-work}
#### Hard-token compression.
Hard-token methods achieve compression by deleting or rewriting input tokens. Token pruning approaches remove tokens via importance heuristics [@li2023compressing; @jiang2023llmlingua; @jiang2024longllmlingua], while summarization and rewriting methods compress by paraphrasing into fewer tokens [@chuang2024learning; @yoon2024compact]. While widely used in practice [e.g., context compaction in Claude Code and Codex; @claude_code; @codex] and very effective for reducing prompt length, these methods are intrinsically lossy: once information is dropped or paraphrased in the discrete token space, exact lexical and structural detail is unrecoverable. This limits their ability to support code fidelity, format-sensitive reasoning, and lookup-style in-context learning in long-horizon settings.
Relatedly, many agent systems externalize memory via summarization, retrieval, or structured stores. Examples include hierarchical memory management [@packer2023memgpt], long-term conversational memory [@zhong2023memorybank; @wang2023scm], and agentic memory frameworks [@xu2025amem]. While effective in extending interaction horizons, these approaches typically rely on lossy summaries or task-specific memory representations that are external to the base model.
#### KV cache compression.
KV cache compression reduces the size of the KV cache by evicting entries; most methods rely on a hand-crafted or learned heuristic to select which entries to drop. SnapKV [@li2024snapkv] compresses the cache by removing entries with low aggregated query attention on a per-head basis. SnapKV can be either prompt-agnostic, where pruning is performed without knowledge of the prompt, or prompt-dependent, where pruning is performed with knowledge of the prompt. Prompt-dependent methods have the downside of requiring labeling of the context and prompt for inputs, and pruning is highly specific to each individual input, limiting multi-turn chat applicability [@li2024scbench]. To overcome these limitations, self-study [@zweiger2026fast] can be used to generate synthetic query signals from the context itself, enabling prompt-agnostic pruning without requiring labeled prompt-context pairs.
KVzip [@kim2025kvzip] performs query-agnostic KV cache compression by using teacher-forced context reconstruction as a proxy objective, assigning each KV pair an importance score from the maximum attention it receives during reconstruction, and pruning low-scoring entries. Fast KV Compaction via Attention Matching [@zweiger2026fast] compresses the KV cache by fitting a smaller set of keys and values, together with per-entry bias terms, that match the original cache's attention outputs and attention mass on reference queries. Expected Attention [@devoto2025expected] predicts KV importance by approximating future queries with a Gaussian distribution, computing each key's expected attention in closed form, and pruning entries with the lowest expected contribution. There are many other KV cache compression techniques and we refer the reader to Awesome KV cache Compression Github [@October2001_AwesomeKVCacheCompression_2025] for a comprehensive review.
We note that KV cache compression methods also pool before eviction: both @li2024snapkv and @zweiger2026fast use max pooling, as does @devoto2025expected. @li2024snapkv further note that the choice of pooling operator does not significantly impact their performance.
It is also possible to do targeted KV cache compression offline for specific prompts. This approach applies to the setting where it is logical to expend more compute to compress the cache once, as this cost will be amortized over many user queries. @synk2025exploiting compute the full KV cache before applying top-k cache compression and reuse the smaller cache. @cartridges use self-study to learn small KV caches per corpus.
#### Soft-token compression.
A complementary direction encodes long contexts into fewer learned continuous tokens. The most successful examples of soft-token approaches to compression require training on the target context, such as prefix tuning [@li2021prefix; @gist; @lloco]. These methods can achieve comparable in-context performance to the base model with large compression ratios for the specific contexts they are trained on. Other soft-token approaches that do not require training on individual contexts are in-context autoencoders [@in-context-autoencoder] and encoder-decoder compression frameworks [@auto-compressor; @e2llm; @dai2025pretraining; @li2025500xcompressor].
#### Efficient long sequence modeling.
A complementary line of work targets the cost of long-context inference by modifying the underlying sequence model, the attention mechanism, or the positional encoding, rather than compressing the input. Early work scales Transformers to long inputs through segment-level recurrence and external memory: Transformer-XL [@transformerxl] caches hidden states from previous segments to extend the effective receptive field, and the Compressive Transformer [@rae2019compressive] further compresses older memory entries into a smaller compressed memory.
A second direction replaces softmax attention with formulations that scale linearly in sequence length. Linear-attention Transformers [@linearattention] cast attention as a kernel feature map to obtain $O(T)$ inference, state-space models such as S4 [@s4] and Mamba [@mamba] use recurrent or selective-scan formulations with constant-size hidden state, and more recent hybrids such as Kimi Linear [@kimilinear] interleave linear-attention blocks with sparse full-attention blocks to recover long-range recall while keeping most of the per-layer cost linear. These approaches lower the asymptotic cost of attention but typically trade off exact long-range recall on retrieval-style tasks.
Within the Transformer family, recent DeepSeek models compress the attention computation itself: Multi-head Latent Attention in DeepSeek-V2 [@deepseekv2] projects keys and values into a low-rank latent KV cache, and DeepSeek-V4 [@deepseekv4] introduces a hybrid attention architecture that interleaves Compressed Sparse Attention (CSA), which consolidates every $m$ tokens of the KV cache into a single entry and then applies sparse top-$k$ selection over those entries, with Heavily Compressed Attention (HCA), which uses a much larger compression rate $m' \gg m$ but retains dense attention over the resulting entries, jointly reducing both KV memory and per-token attention FLOPs at million-token context lengths.
A separate thread extends the usable context of pre-trained models through their positional encoding: rotary position embeddings [@rope] provide relative attention with smooth extrapolation, and YaRN [@yarn] rescales RoPE frequencies to extend the context window with minimal additional training.
Finally, parallel context encoding methods such as CEPE [@cepe] take a different route: a small encoder produces dense representations of long-context chunks in parallel, and a frozen decoder consumes them through inserted cross-attention layers, exposing long contexts to the decoder without full self-attention over the raw tokens.
These approaches reduce the per-token cost of attention, the cost of positional indexing, or the effective sequence length seen by self-attention, rather than the number of tokens reaching the decoder; they are largely orthogonal to LCLMs and can in principle be composed.
# Dataset Details {#app-sec:data}
## Continual Pre-training Dataset Curation {#app-subsec:cont-pre-train-data}
In early experiments, we only train on reconstruction data and hope the model can learn generalizable patterns. However, we find that training on such data alone allows the model to reconstruct the full sequence with negligible loss, but it generalizes poorly beyond reconstruction-style tasks. In fact, it cannot perform any other tasks, even when the LLM decoder is kept frozen. This might not seem surprising, as we are essentially performing a general version of prefix tuning. Thus, the learned representation, though general for different contexts, collapses to the task of reconstruction. On the other hand, training solely on next-token prediction enables the model to leverage compressed context for downstream generation and reasoning. However, the resulting representations often lack fine-grained fidelity (e.g., exact-string details), leading to weaker reconstruction and brittleness on tasks requiring precise recovery.
Motivated by this trade-off, we train on a mixture of next-token prediction and reconstruction data. Our working hypothesis is that next-token prediction provides task-aligned learning signals that teach the decoder to operate on compressed context, while reconstruction acts as an auxiliary objective that encourages information preservation and can accelerate early-stage representation learning. Theoretically, sufficiently large-scale next-token training alone may recover fine-grained fidelity, but we find that in practice, the reconstruction data serves as an auxiliary data source that accelerates the training of the general compressor.
\begin{table*}[ht]
\small
\setlength{\tabcolsep}{5pt}
\renewcommand{\arraystretch}{1.12}
\caption{
Dataset mixture for continual pre-training. Pre/Post-Comp are token counts before/after compression.
We refer to the unique Nvidia license as ``Nvidia,'' the license can be found at \url{https://huggingface.co/datasets/nvidia/Nemotron-CC-v2.1/blob/main/LICENSE.md}.
}
\label{app-tab:data-mixture-pre-training}
\resizebox{\linewidth}{!}{
\begin{tabular}{l l l r r r r r}
\toprule
\textbf{Data Name} &
\textbf{License} &
\textbf{Data Description} &
\textbf{Num Samples} &
\textbf{Tokens} &
\textbf{\%} \\
\cmidrule(lr){5-7}
& & & & \textbf{Pre-Comp} & \textbf{Post-Comp} & \textbf{Trainable} & \\
\midrule
Nemotron CC & Nvidia & Continual pre-training (CC Text) & 84.84M & 60.69B & 30.38B & 30.08B & 20.37 \\
Nemotron Code & Nvidia & Continual pre-training (Code) & 46.85M & 51.18B & 25.21B & 25.01B & 16.90 \\
Nemotron Reasoning & Nvidia & Continual pre-training (Reasoning) & 16.62M & 52.32B & 25.61B & 25.50B & 17.17 \\
Longmino & odc-by & Continual pre-training (Long-Context) & 0.548M & 26.62B & 4.24B & 4.23B & 2.84 \\
SFT & cc-by-4.0 & SFT with Compressed Prompt & 20.67M & 42.21B & 38.02B & 37.54B & 25.49 \\
Reconstruction & Nvidia & Reconstruction from Compression & 23.53M & 50.75B & 25.69B & 25.18B & 17.23 \\
\midrule
\textbf{Total} & -- & & \textbf{192.06M} & \textbf{283.78B} & \textbf{149.16B} & \textbf{147.54B} & \textbf{100.00} \\
\bottomrule
\end{tabular}
}
\end{table*}
#### Curation.
Our primary sources are NVIDIA's Nemotron pre-training datasets [@nvidia2025nvidianemotronnano2; @nvidia_nemotron_nano_v3_2025]. We sample 50M data points from Nemotron's pre-training code and cc code, 50M examples from high-quality and synthetic high-quality cc text, and 50M examples from Nemotron specialized SFT mixtures that cover diverse reasoning traces. A substantial portion of the synthetic data from Nemotron consists of rewrites or generations produced by strong teacher models (e.g., `Qwen3-30B-A3B`, `Qwen3-235B-A22B` [@qwen3], `GPT-oss-120B` [@agarwal2025gpt], and `DeepSeek-R1` [@guo2025deepseek]), which alleviate the problem of forgetting. Furthermore, to expose the compressor to longer documents, we sample approximately 500K examples from the OLMo-3 longmino collection [@olmo2025olmo3]. Finally, to better preserve instruction-following and chat-template behavior, we also include Nemotron pre-training SFT data and convert it into a standardized instruction tuning and multi-turn chat format. Exact dataset composition is reported in `\Cref{app-tab:data-mixture-pre-training}`{=latex}.
## Auxiliary Reconstruction Dataset Curation {#app-subsec:aux-recon-data}
We construct reconstruction data spanning three broad categories: text, code, and *LaTeX*. Text is sampled from Nemotron CC / math, FineWiki [@penedo2025finewiki], and RedPajama ArXiv [@together2023redpajama]; code is sampled from RedPajama GitHub, The Stack v2 (we concatenate files from the repo into one document) [@lozhkov2024starcoder], and CodeParrot; and LaTeX is sampled from TexTeller-OCR. Our goal is to maximize data diversity so that the compressor retains sufficient coverage for downstream tasks that may involve any of these modalities.
## SFT Dataset Curation {#app-subsec:posttrain-data}
For reasoning, we curate math, code, and science supervision from the Nemotron SFT datasets and OLMo-3 Dolci-Think [@NemotronPostTrainingDatasetV2; @NemotronPostTrainingDatasetV1; @nvidia_nemotron_nano_v3_2025; @olmo2025olmo3].
For long-context instruction tuning, we combine long-context and RAG-style QA datasets [@he2025scaling; @jin2019pubmedqa; @liu2024chatqa; @xu2024chatqa; @yang2018hotpotqa; @liu2024raginstructboostingllmsdiverse; @joshi2017triviaqa] with long-context code instruction datasets, including SWE-Fixer [@xie2025swe], NextCoder [@aggarwal2025nextcoder], PyResBugs [@11052783], CommitPack [@muennighoff2023octopack], and RepoBench [@liu2023repobench]. Since some public datasets contain outdated or low-quality responses, we regenerate targets by labeling them with `Qwen3-235B-A22B-Instruct-2507`. In addition, we generate three synthetic long-context corpora: (1) repository summarization, using summarization question templates as instructions and concatenated repository files from The Stack v2 [@lozhkov2024starcoder] as the document; (2) document summarization, using long documents from FineWiki [@penedo2025finewiki] and RedPajama-ArXiv [@together2023redpajama]; and (3) repository code generation, where we hold out one file and ask the model to reconstruct it given the remaining repository context. All synthetic completions are generated by `Qwen3-30B-A3B-Instruct-2507`.
For instruction following, we include multi-turn conversation data from WildChat and LMSYS-Chat-1M [@zhao2024wildchat; @zheng2023lmsys], along with additional alignment-style data obtained by re-labeling with `Qwen3-30B-A3B-Instruct-2507`. We also add OLMo-3 instruction tuning, Nemotron instruction tuning, and the Tulu-3 mixture [@lambert2024tulu3]; since these datasets are already generated by strong teacher models, we keep their original responses unchanged. Finally, we retain a small amount of reconstruction data to stabilize training and mitigate drift when optimizing on compressed inputs. Exact dataset composition is reported in `\Cref{app-tab:SFT-mixture}`{=latex}.
\begin{table*}[t]
\small
\setlength{\tabcolsep}{5pt}
\renewcommand{\arraystretch}{1.12}
\caption{SFT data mixture statistics. Pre/Post-Comp are token counts before/after compression; Weight (\%) is computed over total Post-Comp tokens. During packing, some ultra-long examples are truncated out, leaving us training with 30B tokens for the encoder and 20B for the decoder.
We refer to the unique Nvidia license as ``Nvidia,'' the license can be found at \url{https://huggingface.co/datasets/nvidia/Nemotron-CC-v2.1/blob/main/LICENSE.md}.
}
\label{app-tab:SFT-mixture}
\resizebox{\linewidth}{!}{
\begin{tabular}{l l l r r r r r}
\toprule
\textbf{Data Name} &
\textbf{License} &
\textbf{Description} &
\textbf{Num Samples} &
\textbf{Tokens} &
\textbf{(\%)} \\
\cmidrule(lr){5-7}
& & & & \textbf{Pre-Comp} & \textbf{Post-Comp} & \textbf{Trainable} & \\
\midrule
\textbf{Reasoning} & & & \textbf{2.67M} & \textbf{4.57B} & \textbf{3.87B} & \textbf{3.83B} & \textbf{17.88} \\
\quad Nemotron Reasoning & cc-by-4.0 & Nemotron reasoning data & 2.20M & 3.90B & 3.33B & 3.31B & 15.41 \\
\quad Dolci Think & odc-by & OLMo reasoning data & 0.36M & 0.50B & 0.40B & 0.40B & 1.87 \\
\quad Nemotron Science & cc-by-4.0 & Nemotron science & 0.11M & 0.16B & 0.13B & 0.13B & 0.60 \\
\midrule
\textbf{Long-Context Instruction Tuning} & & & \textbf{7.32M} & \textbf{47.25B} & \textbf{10.63B} & \textbf{9.90B} & \textbf{49.16} \\
\quad Repo Summarization & Apache & Synthetic repo-level summarization & 1.35M & 7.93B & 2.45B & 2.35B & 11.34 \\
\quad Document Summarization & Apache & Synthetic long-document summarization & 1.75M & 18.49B & 3.14B & 3.09B & 14.52 \\
\quad Repo Code Generation & Apache & Synthetic repo-level code generation & 2.02M & 7.83B & 3.54B & 3.38B & 16.36 \\
\quad Long-Context QA & Apache & Long-context / RAG QA & 0.84M & 6.21B & 0.30B & 0.18B & 1.37 \\
\quad Long-Context Code Instruct & Apache & Long-context code instruction tuning & 1.36M & 6.80B & 1.21B & 0.89B & 5.57 \\
\midrule
\textbf{General Instruction Following} & & & \textbf{5.26M} & \textbf{11.68B} & \textbf{4.52B} & \textbf{4.29B} & \textbf{20.91} \\
\quad Multi-turn Conversation & odc-by & WildChat + LMSYS Chat 1M & 2.48M & 7.68B & 1.95B & 1.80B & 9.00 \\
\quad Dolci Instruct & odc-by & OLMo instruction tuning & 1.46M & 2.06B & 1.76B & 1.73B & 8.15 \\
\quad Tulu3 SFT Mixture & odc-by & Tulu3 instruction tuning & 0.84M & 0.69B & 0.40B & 0.38B & 1.83 \\
\quad Nemotron Instruct & Nvidia & Nemotron instruction tuning & 0.35M & 1.07B & 0.32B & 0.28B & 1.46 \\
\quad HelpSteer & CC-4.0 & Alignment & 0.13M & 0.17B & 0.10B & 0.10B & 0.48 \\
\midrule
\textbf{Reconstruction Data} & & & \textbf{2.08M} & \textbf{5.13B} & \textbf{2.61B} & \textbf{2.55B} & \textbf{12.05} \\
\quad Reconstruction & Nvidia & Small reconstruction mixture & 2.08M & 5.13B & 2.61B & 2.55B & 12.05 \\
\midrule
\textbf{Total} & & & \textbf{17.32M} & \textbf{68.62B} & \textbf{21.63B} & \textbf{20.56B} & \textbf{100.00} \\
\bottomrule
\end{tabular}
}
\end{table*}
## Compression Data Format {#compression-data-details}
#### Continual pre-training.
For the continual pre-training data, given a sequence, we split it into multiple segments, wrap every other segment with the special tokens `<|memory_start|>` and `<|memory_end|>`, and treat the wrapped spans as compressed content while leaving the remaining segments uncompressed.
For raw pre-training documents, we construct a fixed alternating pattern of wrapped and unwrapped segments, always beginning with a wrapped segment and ending with an unwrapped segment. Documents longer than 16,384 tokens retain 8,192 trainable (unwrapped) tokens, while shorter documents retain half of their tokens as trainable. The number of segment pairs depends on document length: documents shorter than 2,048 tokens receive 1 pair; documents between 2,048 and 4,096 tokens receive 1-2 pairs; documents between 4,096 and 8,192 tokens receive 2-4 pairs; documents between 8,192 and 16,384 tokens receive 4-6 pairs; and documents longer than 16,384 tokens receive 4-8 pairs. Within each sequence, the first wrapped segment contains 30%-40% of the total wrapped tokens, the last unwrapped segment contains 40%-50% of the total trainable tokens, and the remaining segments are allocated with a $\pm 20\%$ jitter. No segment is allowed to be shorter than 128 tokens.
#### Instruction tuning.
For conversational data, we apply wrapping independently to each message. With 50% probability, the entire message content is enclosed within a single pair of tags. Otherwise, with the remaining 50% probability, we insert 1-3 wrapped regions inside the message and leave the remaining tokens outside the tags. The number of wrapped regions depends on message length: messages shorter than 32 tokens are always fully wrapped; messages between 32 and 200 tokens receive 1 region; messages between 200 and 600 tokens receive 1-2 regions; and messages of 600 tokens or more receive 1-3 regions. Wrapped regions cover 80%-100% of the message tokens in aggregate, and we enforce a minimum 5% unwrapped gap whenever feasible. If the realized wrapped coverage falls below 70%, we fall back to fully wrapping the message. Since this decision is made independently for each message, multi-turn conversations naturally contain a mixture of fully wrapped turns and turns with small unwrapped spans, with a bias toward high overall coverage.
# Extended Architecture Search {#app-sec:arch}
Comparison between mean and concat at scale for all ratios. From , we identify mean and concat as the most promising pooling operators, with mean slightly outperforming concat in training loss. We train at larger scale and find that mean remains marginally better than concat at higher compression ratios, but loses to concat at lower compression ratios.
## Formal Definitions of Pooling Operators {#app-subsec:encoder-defs}
We now give the formal definitions of pooling operators, and illustrate them visually in `\Cref{app-fig:arch_sweep}`{=latex}.
#### Token-based pooling.
The most common method is appending one or more learned special pooling tokens per latent output to each encoder window. Let $p_{i,1},\dots,p_{i,M_i}$ denote these pooling tokens. We encode $$\begin{equation}
(\tilde h^{(i)}_1,\dots,\tilde h^{(i)}_{|w_i|+M_i})
=
\mathrm{Enc}_\phi([w_i;p_{i,1};\dots;p_{i,M_i}]),
\end{equation}$$ and use the hidden states at the pooling-token positions: $$\begin{equation}
z^{(i)}_k
=
\tilde h^{(i)}_{|w_i|+k},
\qquad
k=1,\dots,M_i.
\end{equation}$$ This generalizes the common EOS-pooling strategy: when $W=N$, a single pooling token is appended for each compression block; when $W>N$, several pooling tokens summarize different regions of the same encoder window.
#### Mean pooling.
For each window $w_i$, we partition its hidden states into consecutive groups of size $N$. For latent index $k=1,\dots,M_i$, define $$\begin{equation}
G_{i,k}
=
\{(k-1)N+1,\dots,\min(kN,|w_i|)\}.
\end{equation}$$ The corresponding latent vector is the average hidden state over this group: $$\begin{equation}
z^{(i)}_k
=
\frac{1}{|G_{i,k}|}
\sum_{j\in G_{i,k}} h^{(i)}_j.
\end{equation}$$ Mean pooling directly aggregates the token-level representations assigned to each compression block. Across our architecture search, mean pooling consistently outperforms EOS-style token pooling, matching observations from prior text-compression and vision-encoder work.
#### Concat pooling.
Rather than averaging or selecting a designated token, concat pooling preserves the full per-token representation by concatenating consecutive encoder hidden states within each compression block. For window $w_i$ and latent index $k=1,\dots,M_i$, let $G_{i,k}=\{(k-1)N+1,\dots,\min(kN,|w_i|)\}$ as in mean pooling. The concatenated latent is $$\begin{equation}
z^{(i)}_k
=
\bigl[\,h^{(i)}_{j}\,\bigr]_{j\in G_{i,k}}\in\mathbb{R}^{N\,d_{\mathrm{enc}}},
\end{equation}$$ i.e., a vector of dimension $N\,d_{\mathrm{enc}}$ formed by stacking the $N$ encoder hidden states. The adapter is widened accordingly so that its first MLP layer takes inputs of dimension $N\,d_{\mathrm{enc}}$ and projects back to the decoder hidden dimension $d_{\mathrm{dec}}$, analogous to the patch-merger / spatial-concat scheme used in @qwen3vl.
MLP adapters, causal masking, and no boundary overlap give the lowest pre-training loss across the architecture sweep. We sweep combinations of adapter, boundary-overlap, and encoder mask, and find that (1) causal masking is strictly better than bidirectional masking, (2) boundary overlap does not improve over non-overlapping windows while substantially increasing compute, and (3) the MLP adapter outperforms attention-based adapters. The same trends hold at scale (plot d).
## Adapter Design {#app-subsec:adapter}
The adapter $a(\cdot)$ maps the compressed latent sequence $z_{1:M}\in\mathbb{R}^{M\times d_{\mathrm{enc}}}$ into a sequence $s_{1:M}\in\mathbb{R}^{M\times d_{\mathrm{dec}}}$ in the decoder embedding space. Prior encoder-decoder compression work often assumes $d_{\mathrm{enc}}=d_{\mathrm{dec}}$, but a wide range of encoder-decoder models pair smaller encoders with larger decoders, making the adapter an important design choice. We formally define the two adapter families in our sweep below.
#### MLP adapter.
The MLP adapter is a standard two-layer feed-forward network with a GELU non-linearity [@hendrycks2016gaussian] and an RMSNorm [@zhang2019root] pre-normalization, following the projection design used in LLaVA-style VLMs [@llava]: $$\begin{align}
\hat z_m \;&=\; \mathrm{RMSNorm}(z_m), \\
s_m \;&=\; W_2\,\mathrm{GELU}\!\bigl(W_1 \hat z_m + b_1\bigr) + b_2,
\qquad m = 1,\dots,M,
\end{align}$$ where $W_1\in\mathbb{R}^{d_h\times d_{\mathrm{enc}}}$ and $W_2\in\mathbb{R}^{d_{\mathrm{dec}}\times d_h}$ with hidden dimension $d_h$. Each latent token is projected independently and the adapter does not mix information across the latent sequence. When the encoder uses concat pooling (`\Cref{app-subsec:encoder-defs}`{=latex}), each $z_m$ has dimension $N\,d_{\mathrm{enc}}$ instead of $d_{\mathrm{enc}}$, so the input dimension of $W_1$ is widened to $N\,d_{\mathrm{enc}}$ and the first MLP layer is correspondingly $N\times$ larger.
#### Attention-based adapter.
The attention-based adapter adds a single multi-head self-attention layer over the latent sequence before the MLP projection, with an RMSNorm pre-normalization on the latents: $$\begin{align}
\hat z_m \;&=\; \mathrm{RMSNorm}(z_m), \\
\tilde z_{1:M} \;&=\; \mathrm{MHSA}\bigl(\hat z_{1:M}\bigr), \\
s_m \;&=\; W_2\,\mathrm{GELU}\!\bigl(W_1 \tilde z_m + b_1\bigr) + b_2,
\qquad m = 1,\dots,M.
\end{align}$$ The self-attention layer operates in the encoder hidden dimension $d_{\mathrm{enc}}$ and lets information mix across latent tokens before projection, at the cost of additional parameters and a compute term quadratic in $M$.
#### Comparison.
Across the adapter sweep in `\Cref{app-fig:adapter-ov-loss}`{=latex}, the MLP adapter achieves lower pre-training loss than the attention-based variant while adding less computation, in contrast to prior findings [@GMSA] that attention-based adapters perform better. We therefore use the MLP adapter as the default in all subsequent experiments.
## Boundary Context {#app-subsec:arch-overlap}
A potential limitation of fixed encoder windows is that information crossing a window boundary is split across two encoder forward passes. Increasing the window size $W$ reduces the frequency of such boundaries, but also increases compression-time cost. As an alternative, we implement boundary context overlapping, where each encoder window is extended with $O$ neighboring tokens while the number of latent tokens is kept unchanged.
For a target window $$w_i = x_{(i-1)W+1:\min(iW,T)},$$ we construct an overlapped encoder input $\tilde w_i$ by adding neighboring context tokens. For bidirectional attention, we include both left and right overlap when available: $$\begin{equation}
\tilde w_i
=
x_{\max(1,(i-1)W-O+1):\min(iW+O,T)}.
\end{equation}$$ We provide a visualization in the rightmost figure in `\Cref{app-fig:arch_sweep}`{=latex}. For causal attention, future context is unavailable, so we include only left overlap: $$\begin{equation}
\tilde w_i
=
x_{\max(1,(i-1)W-O+1):\min(iW,T)}.
\end{equation}$$
After encoding $\tilde w_i$, we apply pooling only to hidden states corresponding to the original non-overlapped region $w_i$, and discard pooled outputs that correspond to overlap tokens. Thus, overlap changes the encoder context available to each window but does not change the number of latent tokens passed to the decoder.
We sweep $O\in\{0,32,256\}$ for $W\in\{256,1024\}$ in `\Cref{app-fig:adapter-ov-loss}`{=latex}. Overlap does not noticeably improve pre-training loss in most settings, while increasing training cost because boundary tokens are encoded multiple times. For example, with $W=1024$ and $O=256$, each interior bidirectional window processes $1536$ tokens instead of $1024$. We therefore do not use overlap in the default architecture.
# Full-Scale Training Recipe Sweeps {#app-sec:training-recipe-sweeps}
## Pre-trained Encoder Representations {#app-subsec:encoder_lm_or_emb}
For the encoder initialization, to determine whether an embedding model's representation [@refrag; @e2llm; @cepe] is better than that of a standard language model [@GMSA; @auto-compressor; @mean-pooling], we conduct an experiment with the EOS pooling token. We can see from `\Cref{app-fig:large-scale-sweep}`{=latex} (d) that using the embedding model as the encoder outperforms using an LLM as the encoder. Note the embedding model was initialized from the LLM [@qwen3-embedding]. This shows that embedding model training provides a better representation for the encoder, though the gap gradually narrows with longer training. We also show downstream evaluation performance in `\Cref{tab:abl-encoder}`{=latex}.
Architecture choices selected by the small-scale sweep continue to hold at scale. Continual pre-training loss for the full-pipeline runs at compression ratio 16×. Across (a) encoder window size, (b) encoder attention mask, (c) encoder initialization (LM vs. embedding), and (d) adapter family, the configuration chosen from the from-scratch sweep (W = 1024, causal masking, embedding-initialized encoder, MLP adapter) yields the lowest pre-training loss.
## LR Sweep {#app-sec:continunal pre-training lr sweep}
Language models are prone to catastrophic forgetting during continual pre-training and SFT [@luo2025empirical; @li2024revisiting]. To actively avoid forgetting, we sweep the peak learning rate for stages 2 and 3 of training. To test for catastrophic forgetting, we then benchmark the language model without compression after SFT and check that its performance is comparable to the benchmark performance of the language model before training. For stage 2, we sweep from $1\times10^{-6}$ to $1\times10^{-5}$ and find that a larger learning rate improves downstream benchmarks without causing forgetting. We show these results in `\Cref{app-fig:large-scale-sweep}`{=latex} (c). For stage 3, we sweep from $2\times10^{-5}$ to $5\times10^{-5}$. We do not observe significant differences in those SFT LRs across different models; we report the detailed numbers on downstream benchmarks in `\Cref{app-subsec:lr-sweep}`{=latex}.
Scaling the decoder helps more than scaling the encoder, and lower compression ratios yield lower pre-training loss.Left: pre-training loss at the end of stage 2 with compression ratio 16× across encoder/decoder size combinations; increasing the decoder size produces a much larger drop in loss than increasing the encoder size. Right: pre-training loss at the end of stage 2 for our 0.6B-4B models across compression ratios; lower compression ratios yield lower pre-training loss. Note that the 4B-encoder and 8B-decoder runs use 256 ranks rather than 128 ranks. Because our stateful Parquet dataloader assigns work by rank, these runs complete slightly fewer optimization steps. For a clean comparison, we therefore truncate all loss curves at the same optimization step in the plot above.
## Scaling Up Model Size {#app-sec:scaling model size}
To study scaling behavior, we ask two questions: (i) does a larger encoder improve compression quality, and (ii) does a larger decoder, which has a wider hidden dimension that could compress more information per token, improve performance under compression? Accordingly, we add two model configurations: (1) using the larger `Qwen3-Embedding-4B` as the encoder with the same size `Qwen3-4B-Instruct-2507` decoder; and (2) using the same size `Qwen3-Embedding-0.6B` encoder and the larger `Qwen3-8B` as the decoder. In `\Cref{fig:stage_2_loss}`{=latex} (a), we see that increasing the size of the decoder is much more beneficial in terms of pre-training loss than increasing the size of the encoder.
However, the scaling results in `\Cref{app-tab:scaling}`{=latex} are mixed. The $0.6$B encoder performs best across the RULER tasks, while the $4$B encoder performs best on the remaining evaluations. In contrast, the $8$B decoder achieves substantially lower pre-training loss but does not yield the downstream gains we expected. We suspect this is partly due to a mismatch between the training mixture and the decoder initialization: our data curation and recipe were tuned around the $4$B instruct decoder, whereas the $8$B decoder is a hybrid-thinking model and may require a different data distribution. Moreover, the $4$B instruct model already outperforms the $8$B model in the full-cache setting on several evaluations, which may limit the apparent benefit of scaling the decoder in our current setup.
\small
\setlength{\tabcolsep}{6pt}
Enc / Dec RULER 4K RULER 8K RULER 16K LongBench LongHealth5 GSM8K
----------- --------------------- --------------------- --------------------- --------------------- --------------------- ---------------------
0.6B / 4B **75.06** **70.23** **65.91** 37.33 [67.50]{.underline} [81.05]{.underline}
0.6B / 8B 71.20 [67.40]{.underline} [62.22]{.underline} [37.80]{.underline} 64.80 77.30
4B / 4B [71.30]{.underline} 66.90 62.00 **39.00** **69.80** **83.20**
: Encoder/decoder scaling at 16$\times$ compression {#app-tab:scaling}
## Packed Sequences
Since our corpus contains many long-context examples with highly variable lengths, conventional batching would waste substantial computation on padding. To improve training throughput, we pack multiple examples into a single sequence. Because we continually pre-train an instruction-tuned model, we prevent cross-example information leakage by resetting the attention mask at example boundaries, yielding a block-diagonal attention pattern. We implement this efficiently using variable-length attention kernels [@dao2023flashattention; @dong2024flex]. `\clearpage`{=latex}
# Compute {#app-subsec:compute}
All pre-training-from-scratch experiments are run on 16 nodes of MI300A GPUs, each containing 4 GPUs. Full-scale continual pre-training experiments for 0.6B-4B models are trained on $32$ nodes of MI300A GPUs, each containing $4$ GPUs.
All evaluations are conducted on H200 GPUs.
All scaling experiments (`Qwen3-Embedding-4B` encoder, `Qwen3-8B` decoder) are trained on H200 clusters with 32 nodes to alleviate memory issues. `\clearpage`{=latex}
\FloatBarrier
\FloatBarrier
\clearpage
Benchmarks and Results
GSM8K Performance vs Compression Ratio
The long-context benchmarks show that can compress extended contexts efficiently. We next evaluate on GSM8K to test whether can also handle short, information-dense inputs, where nearly every token may be relevant to the answer. In , achieve the highest accuracy across all compression ratios on GSM8K, with particularly strong gains over baselines at higher compression ratios. These are aggressive compression settings: compression ratios of 16× and 8× remove 93.75% and 87.5% of the input tokens, respectively. Strong performance in this setting demonstrates that are not specialized only for long-context QA; they can also compress short, dense contexts, making them practical for general-purpose deployment.
