通过复用首层值头,让模型更聪明且节省一半缓存。
Improving Model Representation and Reducing KV Cache via Skip Connections with First Value Heads
- 从第二层起复用首层一半的值头,减少计算与缓存
- 缓存降低约25%,困惑度优于标准Transformer
- 可低开销升级现有模型,兼容先进优化方法
Transformer模型在语言任务中表现卓越,但扩展时需大量内存和计算资源,尤其是自回归解码中的键值(KV)缓存。跳连机制有望提升表示能力而不增加资源消耗,但现有方法或只增强表达力而未减缓存,或降缓存却牺牲性能。本文提出SkipV1Former,通过将第一层的值头信息直接传递至后续层,强化模型表示并减少缓存。具体地,从第二块开始,每层仅重新计算一半值头,另一半复用首层结果,使值投影和缓存近乎减少50%。理论分析表明,该设计恢复了压缩丢失的信息,并加速Transformer在自回归任务中的隐式元优化。实验显示,不同规模模型上,该方法稳定降低约25%的KV缓存,同时提升困惑度,优于标准多头注意力(MHA)及部分先进变体。此外,我们提出一种仅需10-15%额外计算量的微调方案,可将现有MHA模型高效升级为SkipV1Former。结合组查询注意力(Group-Query Attention)和多潜空间注意力(Multi-Latent Attention)等技术,可进一步节省缓存,与YOCO结合时实现近50%的缓存缩减并保持性能提升。
原文摘要 · Abstract (English)
Transformer models have driven breakthroughs across various language tasks by their strong capability to learn rich contextual representations. Scaling them to improve representation, however, often demands substantial memory and compute costs, such as the Key-Value (KV) cache used during auto-regressive decoding. Skip connections offer a promising way to improve representation without bloating resource usage, yet most prior works either improve expressivity while leaving KV costs unchanged, or reduce memory at the cost of weaker representation. In this work, we propose SkipV1Former, a Transformer variant that uses skip connections from the first layer's Value heads to strengthen model representation and reduce KV cache. Specifically, from the second block onward, each layer reuses half of its Value heads from the very first layer, while computing the other half as usual-cutting Value projections and V cache by nearly 50 \%. Theoretically, we show that routing uncompressed first-layer Values into deeper layers restores information lost to compression and accelerates the model's implicit mesa-optimization-a key pattern of Transformer in auto-regressive tasks. Empirically, across different model scales, SkipV1Former delivers consistent reductions of approximately 25 \% in KV cache while improving perplexity relative to standard Multi-Head Attention (MHA) Transformers and some advanced variants. Moreover, we propose a recipe for uptraining existing MHA Transformer checkpoints to SkipV1Former with only 10-15\% additional compute. Finally, SkipV1Former can seamlessly combine advanced methods like Group-Query Attention and Multi-Latent Attention to achieve further KV cache savings and performance improvement. When combined with YOCO, it cuts KV cache size by nearly 50 \% while still improving performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。