通过跨层融合提升长序列注意力效率,内存减半且性能更优。
Reconstructing KV Caches with Cross-layer Fusion For Enhanced Transformers
- Top层缓存融合底层与中层的键值信息,保留位置关系。
- 332M至4B参数模型下,缓存内存减少50%,困惑度更低。
- 适合长序列生成任务,尤其对显存受限场景友好。
Transformer解码器在各类任务中表现优异,但长序列下的键值(KV)缓存占用内存过大。尽管跨层共享(如YOCO、CLA)可缓解此瓶颈,但通常性能低于同层方法(如GQA)。我们分析顶层键值的信息来源发现:值主要来自底层,而键则同时依赖底层和中层。基于此,提出FusedKV,其顶层缓存为底层与中层最相关信息的可学习融合,直接作用于旋转编码后的键,保持相对位置信息且无需重新计算旋转嵌入。进一步提出FusedKV-Lite,顶层缓存直接由底层值与中层键构成,降低输入输出开销,仅轻微增加困惑度。在332M至4B参数的LLM上,该方法实现50%缓存内存减少,且验证困惑度低于标准Transformer,成为高效高能的架构替代方案。
原文摘要 · Abstract (English)
Transformer decoders have achieved strong results across tasks, but the memory required for the KV cache becomes prohibitive at long sequence lengths. Although Cross-layer KV Cache sharing (e.g., YOCO, CLA) offers a path to mitigate KV Cache bottleneck, it typically underperforms within-layer methods like GQA. To understand the root cause, we investigate the information flow of keys and values of the top-layers. Our preliminary reveals a clear distribution: values are predominantly derived from the bottom layer, while keys draw more information from both bottom and middle layers. Building upon this, we propose FusedKV, whose top-layer KV caches are a learnable fusion of the most informative ones from the bottom and middle layers. This fusion operates directly on post-RoPE keys, preserving relative positional information without the computational cost of re-applying rotary embeddings. To further improve efficiency, we propose FusedKV-Lite, an cross-layer sharing approach, where top-layer KV caches are directly derived from the bottom-layer values and the middle-layer keys. Compared to FusedKV, FusedKV-Lite reduces I/O overhead at the cost of a slight increase in perplexity. In experiments on LLMs ranging from 332M to 4B parameters, our proposed method reduce 50\% cache memory while achieving lower validation perplexity than the standard Transformer decoder, establishing it as a memory-efficient, high-performance architectural alternative.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。