训练Transformer让键值缓存更易压缩,提升长文本处理效率
Training Transformers for KV Cache Compressibility

- 训练时随机屏蔽部分键值缓存,引导模型学习可压缩表征
- 在长文本问答等任务中,压缩后性能提升显著,节省内存与计算
- 适合需要高效处理长上下文的模型部署场景
长上下文语言建模受限于键值(KV)缓存,其内存和解码访问成本随前缀长度线性增长。现有方法多为对固定预训练模型的后处理压缩,效果受限于模型内部表示的可压缩性。本文提出将KV可压缩性视为模型表征的属性,并证明几乎所有序列到向量函数都存在可压缩与不可压缩的Transformer实现,因此需在训练中引导模型学习可压缩表示。为此,我们提出KV压缩感知训练(KV-CAT),通过训练时对KV槽位进行掩码,迫使模型减少依赖,学习更利于事后压缩的表征。实验表明,KV-CAT在检索、长文本问答及压缩前缀续写困惑度评估中,均提升了压缩质量与预算之间的权衡表现。
原文摘要 · Abstract (English)
Long-context language modeling is increasingly constrained by the Key-Value (KV) cache, whose memory and decode-time access costs scale linearly with the prefix length. This bottleneck has motivated a range of context-compression methods, from token-level summarization to recent optimization-based KV compression methods. These post-hoc methods operate on the KV cache of a fixed pretrained model, so their effectiveness is fundamentally limited by how well the model's internal representations can be compressed. In this work, we formalize the notion of KV compressibility and show that it is a property of the learned representations, rather than of the context alone. We prove that almost any sequence-to-vector function admits both highly compressible and inherently non-compressible transformer implementations, highlighting the need to guide transformers toward compressible representations during training. Motivated by this, we propose KV-Compression Aware Training (KV-CAT), a continued pretraining procedure that incentivizes the emergence of compressible representations. We introduce a train-time KV sparsification policy that masks KV slots during training. This forces the model to use fewer KV slots and encourages it to learn representations amenable to post-hoc compression. Empirically, we show that KV-CAT improves the quality-budget tradeoff of downstream compression methods across retrieval, long-context question answering, and perplexity-based evaluation of compressed-prefix continuation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。