Latent Compaction with Test-Time Training
Yingfa Chen, Xingyu Shen, Chaojun Xiao, Xu Han
arXiv:2610.20582v1Long-context language models are powerful, but they come with a very practical problem: the longer the input, the larger the key-value cache becomes during generation. That cache is what lets the model attend back to earlier tokens efficiently, but it also becomes a major memory and bandwidth bottleneck. In real use, that means lower batch sizes, slower decoding, and a harder time keeping long-document systems responsive. This paper is about a very specific version of that problem: before a user asks any downstream question, can we compress a document’s cache in a way that still preserves the information needed for questions we have not seen yet? That question matters because a lot of existing context-compaction methods are either too lossy or too tied to the immediate prompt. Some systems summarize old context in text, which is useful for conversation but can drop details that only become important later. Others evict key-value pairs from the cache using heuristics like attention scores or reconstruction scores. Those methods can work well, but they are still basically selecting from the original cache. This paper asks a more ambitious question: instead of only choosing which original entries to keep, can we treat the cache itself as something we can train, like a tiny document-specific neural memory, while keeping the base model frozen? The method is called Latent Compaction. The core idea is to build a compact cache for one document, optimize it at test time, and then reuse that trained cache across many downstream questions. Importantly, the model’s weights never change. Only the document cache entries are learned. That makes the method closer in spirit to test-time training than to ordinary fine-tuning: the system adapts to the document, not to the downstream task labels, and it does so without seeing the actual questions or answers that will come later. To understand the method, it helps to recall what the KV cache is doing. In a Transformer, each earlier token contributes a key and a value. When a new token arrives, the query attends over those keys, produces attention weights, and uses them to mix the values into an output. If you compress the cache too aggressively, you risk changing both the attention distribution and the resulting representation. So the paper’s design tries to preserve not just the most “important” tokens, but the behavior of the attention mechanism itself. Latent Compaction does that with a compact set of learned slots per layer and per KV head. Each slot has a key, a value, and also an attention bias. Those biases are useful because compression changes the total attention mass assigned to the document, not just the relative ranking of entries. By learning biases, the method can compensate for that shift instead of only trying to approximate the original keys and values. Initialization is a big part of the method. The authors use three complementary kinds of slots. First are preservers: these are original KV pairs chosen globally by a reconstruction-based importance score, specifically KVZip+ scores. The highest-scoring pairs across the whole document and all layers and heads are kept as preservers. Second are anchors: these are selected using teacher attention, including the first document token as a possible attention sink. Third are diversifiers: these are chosen to spread out in key space using a farthest-first rule, so the compact cache covers a broader range of document content rather than collapsing onto a few similar entries. The intuition is simple: preserve the strongest original evidence, seed the cache with attention-relevant points, and make sure the learned memory has geometric diversity. Training then happens in two stages. In stage one, the method fits the compact cache to match the full cache’s attention outputs and log-normalizers. This is done using teacher queries captured from generic prompts about the document, not downstream questions. The full document cache serves as the teacher, and the compact cache is trained to reproduce what the full cache would have produced. Preservers stay frozen in this stage, while anchors and diversifiers are updated. Matching the log-normalizer is especially important because it helps calibrate how much attention the document should receive relative to later question tokens or generated tokens. Stage two is a joint refinement stage. Now all cache entries are optimized together. Instead of only matching attention, the method distills the full model’s token predictions on a teacher rollout. The loss compares the probability mass on the teacher’s top tokens and on the tail mass, plus a standard language-model loss on the teacher’s greedy token. This is still done without any access to downstream questions or answers. The idea is that the compact cache should not only imitate internal attention patterns, but also preserve the model’s predictive behavior on generic document-conditioned continuations. That is the paper’s main conceptual move: treat the KV cache as a trainable latent representation and optimize it directly. The experimental evaluation spans three benchmarks: RULER, LongBench, and Phonebook Lookup. RULER is used for synthetic retrieval-style long-context tasks, LongBench covers broader long-context understanding, and Phonebook Lookup is a clean retrieval setting that also lets the authors test transfer across model families and sizes. The paper compares Latent Compaction against budget-matched compression baselines, including eviction-style methods and query-agnostic adaptations. The comparison is intentionally strict: results are reported under the same logical KV-pair budget, so the question is how well each method preserves information under the same compression pressure. The headline result is on RULER. At 32-times compression, Latent Compaction reaches 96.7% accuracy, while KVZip+ gets 40.0%. That is a very large gap, and it suggests that learned cache optimization can do much more than selecting a subset of original entries. Across the broader benchmark suite, the method also improves over the baselines at different context lengths and compression ratios. On LongBench, the gains indicate that the approach is not limited to synthetic retrieval alone; it also helps on more natural long-context tasks. On Phonebook Lookup, the method transfers across model scales and families, with retrieval gains reported for multiple Qwen3 sizes and MiniCPM5 models. So the benefit is not just tied to one architecture. The ablations are also informative. They show that initialization matters, that both training stages matter, and that attention normalization matters. In other words, the method is not simply “do gradient descent on a smaller cache and hope for the best.” It depends on a carefully constructed starting point and on training objectives that preserve both attention structure and predictive behavior. The paper also notes that these results concern logical KV-pair budgets, not measured memory savings or inference speedups, so the reported gains are about retention quality under compression rather than end-to-end systems throughput. What should we take away from this? First, the paper makes a strong case that a compressed cache can be more than a pruned subset of the original cache. By training document-specific slots at test time, the model can recover much of the information needed for later questions even under aggressive compression. Second, the combination of teacher attention fitting and end-to-end distillation seems to be doing complementary work: one stage aligns internal attention behavior, and the other aligns the model’s outputs. Third, the approach appears to transfer across tasks and model families, which is encouraging if you think of cache compression as a reusable inference-time tool rather than a benchmark-specific trick. There are also clear limitations to keep in mind. The method requires test-time optimization before reuse, so there is extra upfront computation. The paper emphasizes that the cache is trained on generic document prompts, but it still needs a teacher rollout and gradient updates, which may be expensive for some deployments. The results are framed around logical cache budgets, so you still need to ask how the method translates into actual memory, latency, and engineering constraints. And while the method is designed to avoid downstream questions and answers, it still depends on the full-cache model as a teacher, so its quality is bounded by the teacher’s own behavior. Overall, Latent Compaction is an interesting step beyond pure eviction. It treats the KV cache as a trainable memory object, uses test-time training to shape that memory before the questions arrive, and shows that this can preserve retrieval accuracy surprisingly well under severe compression. The main message is that for long-context inference, the best compact representation may not be the one you merely keep, but the one you actively learn.
Also spotted that day
Previous daily papers