Latent Compaction with Test-Time Training
Long-context LLM inference is bottlenecked by expanding key-value caches. This paper proposes treating the KV cache as a learnable layer and optimizing it per document via test-time training while keeping model weights frozen. The two-stage procedure first matches teacher attention statistics and then distills teacher predictions, yielding large accuracy gains on synthetic and natural-document benchmarks under aggressive compression.
Latent Compaction presents a compelling conceptual advance: it formulates KV-cache compression as gradient-based optimization of document-specific slots with a frozen backbone, achieving 96.7% accuracy on RULER at $32\times$ compression versus 40.0% for KVZip+. However, the paper explicitly disclaims reporting physical memory or latency savings, and the expensive per-document optimization limits practical applicability. The method is strongest as a proof-of-concept that test-time training can preserve retrieval accuracy, but it is not yet a deployable compression system.
The two-stage training strategy and tri-partite initialization are rigorously validated. Ablations show that removing stage 1 collapses average Phonebook Lookup accuracy from 80.67% to 13.92% at $16\times$ compression, and the equal-average baseline in Equation (5) falls to near-zero accuracy, confirming that the learned normalizer weighting is essential. The design choices—preservers frozen during attention fitting, log-normalizer supervision, and per-entry attention biases—are all justified by controlled experiments.
The most significant limitation is the absence of measured memory savings or inference speedups; the paper studies only logical KV-pair budgets. Per-document preparation requires 480 teacher rollouts and 2,000 gradient updates, making the method costly for interactive use. Additionally, claims of transfer across model families actually describe reoptimizing the cache from scratch for each model rather than zero-shot cache transfer, and text summarization still surpasses the method on RULER at the most aggressive compression ratios. Extreme compression beyond $64\times$ collapses accuracy for all methods, including this one.
The evidence supports the central claim that learned latent compaction outperforms eviction-based baselines under comparable KV budgets. Comparisons to query-agnostic SnapKV and PyramidKV, as well as Attention Matching and KVZip+, are fair and well-specified. Nevertheless, the benchmark suite is narrow—centered on retrieval and synthetic long-context tasks—and natural-document gains on LongBench degrade substantially as compression increases, leaving a large gap to the uncompressed model.
Code, model checkpoints, and hyperparameters are publicly documented, which aids reproduction. Independent replication is still hindered by the need for a custom attention implementation that exposes log-normalizers and per-entry biases outside standard fused kernels like FlashAttention. Furthermore, the substantial compute cost for per-document optimization, the lack of reported random seeds or hardware details, and the reliance on 480 hand-designed generic prompts create practical barriers to exact reproduction.
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
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.