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 THUNLP/latent-compaction
Comments: 19 pages, 21 figures
Mentions
THUNLP/latent-compaction is trending on GitHub. THUNLP/latent-compaction is trending on GitHub. Learn how Latent Compaction achieves 96.7% accuracy at 32x KV cache compression using test-time training. Repo: THUNLP/latent-compaction — Follow NexBit for the dev tools worth knowing. #machinelearning #artificialintelligence #deeplearning #llm #largelanguagemodels #airesearch #python #developer #techcommunity #softwareengineering — 👉 Follow Nexbit 🌐 Website: nexbit.art Discord community: discord.gg Telegram channel: t.me 𝕏 X: x.com 📸 Instagram: instagram.com 🧵 Threads: threads.com 👍 Facebook: facebook.com #Shorts
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 THUNLP/latent-compaction Comments: 19 pages, 21 figures