通过随机跨层注意力,实现模型深度维度的动态缓存共享,大幅降低推理内存开销。
Stochastic KV Routing: Enabling Adaptive Depth-Wise Cache Sharing
- 训练时随机选择跨层注意力,让模型适应不同深度缓存策略。
- 在不损失性能前提下,显著减少键值缓存内存占用。
- 特别适合资源受限的大模型部署,兼具稳定性与可扩展性。
高吞吐量自回归生成需缓存键值(KV)以避免重复计算,但其内存开销巨大,影响服务成本。现有方法多从时间轴压缩或淘汰缓存入手,本文提出从深度维度优化的新思路。尽管已有研究指出每层全量缓存冗余,但跨层共享仍面临吞吐下降或首词延迟增加的问题。本文证明舍弃某层缓存亦可实现无信息损失的优化。提出一种简单训练方法:训练中各层随机选择使用自身或前层的KV状态进行注意力计算。该随机过程使模型对多种深度缓存共享策略具备鲁棒性,适应部署时未知硬件约束。实验表明,此方案在预训练或微调阶段应用,可支持多种模型族实现深度缓存共享;尤其在数据有限的大模型场景中,具类似正则化效果,常保持甚至提升性能,同时显著降低缓存内存占用。
原文摘要 · Abstract (English)
Serving transformer language models with high throughput requires caching Key-Values (KVs) to avoid redundant computation during autoregressive generation. The memory footprint of KV caching is significant and heavily impacts serving costs. This work proposes to lessen these memory requirements. While recent work has largely addressed KV cache reduction via compression and eviction along the temporal axis, we argue that the \emph{depth} dimension offers an orthogonal and robust avenue for optimization. Although prior research suggests that a full cache for every layer is redundant, implementing cross-layer cache sharing remains a practical challenge; existing methods typically suffer from reduced throughput or increased time-to-first-token. In this paper, we demonstrate that dropping a layer's cache offers efficient optimization without information loss. We propose a simple training approach: random cross-layer attention. During training, layers randomly choose to attend either to their own KV states or those of a preceding layer. This stochastic process adapts the model to be robust to various depth-wise cache sharing strategies, ensuring flexibility for unknown hardware constraints at deployment time. Our evaluations show that applying this scheme during pre-training or fine-tuning enables depth-wise cache sharing for various model families. Furthermore, for larger models in data-constrained settings, this approach is suggestive of a regularization-like effect, frequently preserving or improving performance while significantly reducing the cache's memory footprint.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。