论文

Latent Compaction:用测试时训练压缩 KV cache,32 倍压缩下 RULER 准确率达 96.7%

Latent Compaction with Test-Time Training

精选理由

清华 THUNLP 开源了个 KV cache 压缩方法,压到 1/32 还能在 RULER 上保持 96.7% 准确率,比 KVZip+ 的 40% 高出一大截,长文档检索场景可以直接试。

长上下文推理中 KV cache 膨胀会推高显存和注意力开销。Latent Compaction 通过测试时训练学习针对单篇文档的紧凑 cache,初始化时综合重构重要性、teacher 注意力和 key 空间多样性三项信号,训练过程先拟合全量 cache 的注意力输出与 log-normalizer,再从全量模型端到端蒸馏优化整个 cache,模型权重全程冻结,且训练只用通用提示、不接触下游问题。在 RULER 基准上,该方法在 32 倍压缩下取得 96.7% 准确率,同期 KVZip+ 只有 40.0%。方法还在 Qwen3 和 MiniCPM5 两个模型家族上带来检索增益,代码已在 THUNLP 的 GitHub 开源。

原文 · arXiv cs.LG

Latent Compaction with Test-Time Training

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