[Paper] 通过跨层融合重建 KV 缓存以增强 Transformers

Source: arXiv - 2512.03870v1

概览

Transformer 解码器驱动了当今的大型语言模型，但它们的键值（KV）缓存会随序列长度线性增长，导致在长输入上很快耗尽 GPU 显存。论文《通过跨层融合重建 KV 缓存以增强 Transformer》提出了一种巧妙的方法，将 KV 缓存减半，同时实际上 提升了困惑度。通过学习如何融合来自低层的最有信息量的键和值，作者构建了一种既节省显存又高性能的新解码器架构。

主要贡献

• 跨层 KV 融合 (FusedKV)： 引入可学习的融合模块，将底层的值与中层的键组合，形成顶层的 KV 缓存。
• FusedKV‑Lite： 一个精简变体，直接复用底层值和中层键，消除额外的 I/O 并进一步降低显存开销。
• 经验洞察： 发现在深层解码器中，值 主要来源于底层，而 键 则从底层和中层都获取有用信号。
• 显存削减： 在 332 M 到 4 B 参数的模型上实现约 50 % 的 KV 缓存显存节省。
• 性能提升： 尽管缓存减小，验证困惑度仍低于原始 Transformer 解码器。

方法论

1. KV 流向诊断 – 作者对标准解码器进行仪器化，追踪每个顶层键/值的来源。热力图显示出明显的分割：值高度偏向底层，键则是底层和中层的混合。
2. FusedKV 设计 –
   • 融合模块： 一个小型线性层（或 MLP）学习权重，以混合底层值 (V_bottom) 和中层键 (K_mid)。
   • 后 RoPE 融合： 融合在旋转位置嵌入（RoPE）之后进行，此时相对位置信息已嵌入向量，无需重新计算。
3. FusedKV‑Lite 变体 – 跳过可学习的融合，直接将 V_bottom 与 K_mid 复制到顶层缓存中。这样去除了额外的读写步骤，以略微增加的困惑度换取更低的延迟。
4. 训练与集成 – 融合参数与语言模型一起端到端训练，使用标准的下一个 token 预测。注意力计算本身无需修改；解码器只读取更小的、已融合的缓存。

结果与发现

• 显存： 在所有规模上，KV 缓存大约减半，直接转化为能够在相同硬件上翻倍上下文长度或使用更大批次。
• 质量： 验证困惑度始终下降（数值越低越好），说明融合后的信息比单纯的缓存共享更丰富。
• FusedKV‑Lite： 额外节省约 5 % 的 I/O 带宽；相对于完整的 FusedKV，困惑度仅上升约 0.2 %，仍优于基线。

实际意义

• 长上下文推理： 开发者现在可以在单个 GPU 上运行 8 k 或 16 k token 的提示，而之前只能到约 4 k token，开启文档级摘要、代码库分析等场景。
• 成本降低： KV 显存减半意味着相同上下文长度下的显存需求减半，可使用更便宜的 GPU 实例（例如 A100‑40 GB 替代 A100‑80 GB）来提供相同工作负载。
• 部署简便： 由于融合在模型图内部完成，无需额外的缓存管理代码——只需替换解码器类即可。这使得在现有推理框架（如 Hugging Face Transformers、vLLM）中的集成非常直接。
• 微调潜力： 轻量的融合参数可以在特定领域数据上微调，为在不增加显存的情况下适配大模型提供了一条低成本路径。

局限性与未来工作

• 架构范围： 本研究聚焦于仅解码器的 Transformer；编码器‑解码器或纯编码器模型的 KV 动态可能不同。
• 训练开销： 引入融合模块会增加少量参数并在前向传播时多一次小的计算，可能会在超大模型上略微延长训练时间。
• 对极端规模的泛化： 实验止步于 4 B 参数；尚不清楚该方法在 30 B+ 模型上是否仍然有效，KV 模式可能会发生变化。
• 未来方向： 作者建议探索自适应融合（对每个 token 使用不同权重）以及将该思路扩展到多查询注意力或稀疏缓存等场景。

结论： 通过认识到键和值在 Transformer 解码器中分布于不同层，作者提出了一种简单却强大的跨层融合技术，将 KV 缓存显存削减一半，同时实际 提升 了模型质量。对于任何需要长上下文的 LLM 服务构建或部署者，FusedKV（及其 lite 变体）都是值得尝试的实用升级。

作者

• Hongzhan Lin
• Zhiqi Bai
• Xinmiao Zhang
• Sen Yang
• Xiang Li
• Siran Yang
• Yunlong Xu
• Jiaheng Liu
• Yongchi Zhao
• Jiamang Wang
• Yuchi Xu
• Wenbo Su
• Bo Zheng

论文信息

• arXiv ID: 2512.03870v1
• 分类: cs.CL
• 发布日期: 2025 年 12 月 3 日
• PDF: Download PDF