[论文] Latent Compaction with Test-Time Training

论文概要

研究领域: NLP
作者: Yingfa Chen, Xingyu Shen, Chaojun Xiao, Xu Han
发布时间: 2026-10-14
arXiv: 2610.20582

中文摘要

长上下文会增大键值（KV）缓存，增加生成过程中的内存使用和注意力开销。压缩这些缓存必须保留足够信息以应对未见问题。我们提出潜在压缩（Latent Compaction），通过测试时训练学习一个紧凑的文档特定缓存。我们基于重建重要性、教师注意力和键空间多样性初始化缓存条目。训练首先将注意力输出和对数归一化项拟合到完整缓存的对应值，同时冻结重建选中的条目；随后通过端到端蒸馏联合优化整个缓存。模型权重始终保持冻结，训练仅使用关于文档的通用提示，不访问下游问答。在 RULER、LongBench 和 Phonebook Lookup 上的实验表明，该方法显著优于同等预算的压缩基线。例如在 RULER 上，本文方法在 32 倍压缩下达到 96.7% 的准确率，而 KVZip+ 仅为 40.0%。该方法还能跨模型规模和家族迁移，在 Qwen3 和 MiniCPM5 上都实现了检索增益。这些结果说明，将 KV 缓存视为神经网络并用测试时训练进行优化，可以在激进压缩下保持检索准确率。代码已开源：https://www.github.com/THUNLP/latent-compaction

原文摘要

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 RU...

自动采集于 2026-10-16

#论文 #arXiv #NLP #小凯