Latent Compaction with Test-Time Training

AI summary

Overview

Research area: Efficient long-context inference for large language models, specifically key–value (KV) cache compression, combined with test-time training (TTT).

Technical level: Advanced. The paper assumes familiarity with Transformer attention, grouped-query attention (GQA), KV caching, softmax normalization, knowledge distillation, and gradient-based optimization.

One-sentence scope: The paper proposes Latent Compaction, a method that learns a compact, document-specific KV cache by running gradient descent on cache entries (with model weights frozen) using generic prompts about the document, and evaluates it on RULER, LongBench, and Phonebook Lookup against budget-matched compression baselines.

What This Paper Is About

Long documents blow up the KV cache, which increases memory use and attention cost at every decoding step. Existing ways of shrinking that cache either drop entries (eviction) or replace context with a text summary, but both can lose details that only turn out to matter once a downstream question is asked. The paper's goal is to build a compact cache before any question is known, by treating the cache itself as a small trainable neural network and optimizing it per document so that it preserves what the full cache would have provided.

Key Contributions

1. A method for learning compact context representations trained with gradient descent and initialized through three complementary selection rules (preservers, anchors, and diversifiers).
2. A document-specific test-time training procedure that learns from the frozen model's own attention outputs and next-token predictions, using only generic prompts about the document and never downstream questions or answers.
3. Empirical results across three benchmarks (RULER, LongBench, Phonebook Lookup) and multiple model scales and families, including Qwen3-0.6B, 1.7B, 4B, and 8B, and MiniCPM5-1B and MiniCPM5-2B.
4. Ablations isolating the roles of initialization, both training stages, attention biases, and attention normalizer weighting.

Main Findings

• Large gains over compression baselines on RULER: On a RULER task, the method reaches 96.7% accuracy at 32× compression, versus 40.0% for KVZip+.
• Retrieval advantage on Phonebook Lookup: At 16× compression with Qwen3-1.7B, the default configuration scores 87.75 / 83.75 / 70.50 at 4K / 8K / 16K context (average 80.67), compared with an average of 20.50 when diversifiers are replaced with random document KVs and 19.75 with Gaussian random vectors.
• The advantage grows with compression, but with limits: Text summarization beats Latent Compaction on RULER at the most aggressive ratios, and all compressed methods approach zero Phonebook Lookup accuracy at 64× compression and beyond.
• LongBench results: Across Qwen3-0.6B, 1.7B, and 4B, Latent Compaction outperforms text summarization throughout and KVZip+ at 16× and 32×. At 4× its scores are close to KVZip+, and all compressed scores still decline as compression increases, leaving a gap to uncompressed models.
• Transfer across scales and families: Phonebook Lookup gains extend to Qwen3-0.6B, 4B, and 8B, and to MiniCPM5-1B and MiniCPM5-2B, with each model using its own teacher outputs.
• Stage 1 is essential: Removing stage 1 drops average Phonebook Lookup accuracy from 80.67 to 13.92. Letting preserver keys and values train in stage 1 drops it to 33.42, and removing the log-normalizer MSE term drops it to 36.50.
• Stage 2 adds measurable accuracy: Removing stage 2 reduces accuracy by 6.75, 15.50, and 7.75 percentage points at 4K, 8K, and 16K. Adding the language-modeling loss to the grouped KL improves those results by 4.75, 0.25, and 5.25 points.
• Grouped distillation is as good as full-vocabulary distillation: Replacing top-20-plus-tail targets with full-vocabulary targets leaves 4K accuracy unchanged, improves 8K by 0.50 points, and reduces 16K by 0.50 points, so it offers no consistent gain and requires more teacher-target storage.
• Initialization choices matter: The default 25% anchors / 75% diversifiers mixture beats every tested replacement initializer at each context length; substituting KVZip+ for the anchors reduces the 4K/8K/16K average from 80.67 to 57.17.
• Longer training does not help much: With stage 2 fixed at 1,000 updates, doubling stage 1 from 1,000 to 2,000 updates improves accuracy only from 61.00% to 61.25%; accuracy peaks at 1,000 stage-2 updates and then declines slightly, indicating overfitting.
• Attention normalizer weighting is critical at inference: Replacing the normalizer-weighted combination with an equal average of document and subsequent-token attention reduces accuracy to at most 0.25%, despite the stage-1 normalizer loss being retained. Removing learned attention biases reduces accuracy by 0.75, 3.00, and 0.25 percentage points at 4K, 8K, and 16K.
• Scope caveat: The results concern logical KV-pair budgets, not measured memory savings or inference speedups, and per-document fitting remains costly.

Methodology in Plain English

The aim is to replace a document's full set of key–value pairs with a much smaller set of "slots" that behave as if the document were still there. Model weights never change; only the per-document cache slots are trained.

1. Build the slot budget. Given an average budget of B KV pairs per layer and KV head, the method creates L×H×B slots total (counts can differ per layer–head pair). Each slot also carries a scalar attention bias that compensates for the attention mass lost through compression.

2. Initialize the slots in three ways. Preservers come from a global ranking of KVZip+ reconstruction scores pooled across all layers, KV heads, and document positions. Anchors are the first document position plus the original pairs with the highest normalized teacher attention (a quarter of the anchor–diversifier budget). Diversifiers are chosen by farthest-first traversal in key space, picking each new key that is farthest from its nearest already-selected key, with preservers excluded from the distance reference set.

3. Stage 1 — attention fitting. The frozen teacher model generates continuations for 480 generic prompts about the document (24 document requests × 20 response instructions, up to 24 tokens each). The queries captured during those rollouts serve as supervision. For each layer, the compacted cache is trained to match the full cache's attention outputs and log-normalizers, using an MSE loss on both, with statistics standardized per query head. Preservers stay frozen in this stage; anchors and diversifiers (their keys, values, and biases) are updated.

4. Stage 2 — end-to-end refinement. All cache entries are jointly optimized. For each teacher rollout, the compressed-cache model is teacher-forced on the teacher's own generated tokens and trained with a KL divergence against the teacher's distribution (top-20 tokens kept individually, remaining mass grouped into a single "tail" category) plus a standard next-token cross-entropy loss toward the teacher's greedy token. Both losses have equal weight.

5. Reuse. Once trained, the cache is fixed and reused across many downstream questions. Question and generated tokens are appended as ordinary, unbiased KV pairs outside the budget, and their attention is combined with the document attention using a shared softmax normalizer computed in log-space.

Training uses 1,000 stage-1 updates per layer and 1,000 stage-2 updates, with constant AdamW learning rates of 3×10⁻³ and 10⁻⁴, on Qwen3-1.7B by default, with BF16 model computation and an FP32 KV cache.

Why This Matters

Impact on research. The paper reframes KV cache compression as a per-document learning problem rather than a selection problem: instead of choosing which original entries to keep, it optimizes the entries themselves while keeping the model frozen. It also extends test-time training from adapting model parameters to adapting a memory artifact, and it shows that a grouped top-20-plus-tail distillation target retains accuracy comparable to full-vocabulary distillation with less storage. The reported ablations give a clear picture of which components carry the accuracy.

Real-world applications:

• Long-document question answering over a fixed corpus (legal filings, technical manuals, medical records), where one cache is prepared once and queried many times.
• Retrieval-augmented generation, where retrieved documents can be pre-compressed before the user's question is known.
• Multi-turn assistants with long histories, where the cache must serve unpredictable follow-up questions.
• On-device or memory-constrained deployment, where the KV cache, not the weights, is the binding constraint.

Industry relevance. KV cache size limits batch size and memory-bandwidth demand at every decoding step, so compression directly affects serving cost. The paper frames the target as a logical KV-pair budget and explicitly notes that it does not report measured memory savings or inference speedups, so the immediate industrial read is that the accuracy side is well supported while the systems side remains open. The method also works across Qwen3 and MiniCPM5 model scales, and the released code (https://www.github.com/THUNLP/latent-compaction) supports reproduction.

Future Directions

1. Reduce preparation cost. Per-document fitting is expensive: it takes 1,000 stage-1 updates per layer plus 1,000 stage-2 updates, using 480 teacher rollouts generated at up to 24 tokens each. The authors list reducing this cost as future work.
2. Recover accuracy under extreme compression. All compressed methods approach zero Phonebook Lookup accuracy at 64× compression and beyond, and Latent Compaction is beaten by text summarization on RULER at the most aggressive ratios.
3. Broader task coverage. The authors state that future work will study broader tasks beyond the three benchmarks evaluated here.
4. Measured systems benefits. The paper explicitly limits its claims to logical KV-pair budgets rather than measured memory savings or inference speedups, leaving the physical memory and latency side to be characterized.

Target Audience

Researchers and engineers working on efficient LLM inference, long-context modeling, and KV cache compression, who are comfortable with attention mechanics, GQA, and distillation. It is also relevant to practitioners deciding between eviction-based compression, text summarization, and learned cache representations, and to anyone studying test-time training as a technique for adapting memory rather than parameters.

Authors’ abstract

Long contexts enlarge key--value (KV) caches, increasing memory use and attention costs during generation. Compressing these caches must preserve information for unseen questions. We introduce **Latent Compaction**, which learns a compact document-specific cache through test-time training. We initialize cache entries using reconstruction importance, teacher attention, and key-space diversity. Training first fits attention outputs and log-normalizers to those of the full cache, with reconstruction-selected entries held fixed. It then jointly optimizes the entire cache through end-to-end distillation from the full-cache model. Model weights remain frozen throughout, and training uses only generic prompts about the document, without access to downstream questions or answers. Experiments on RULER, LongBench, and Phonebook Lookup demonstrate substantial gains over budget-matched compression baselines. On a RULER, for example, our method achieves 96.7% accuracy at 32$\times$ compression, compared with 40.0% for KVZip+. The method also transfers across model scales and families, with retrieval gains on both Qwen3 and MiniCPM5. These results show that treating the KV cache as a neural network and optimizing it with test-time training can preserve retrieval accuracy under aggressive compression. The code is released at https://www.github.com/THUNLP/latent-compaction