arXiv:2410.14731cs.LGcs.AI2024-10ICLR被引 26

通过可训练正交投影实现动态KV缓存压缩,显著降低内存占用且保持高精度。

MatryoshkaKV: Adaptive KV Compression via Trainable Orthogonal Projection

  • 用可训练的正交投影矩阵压缩KV缓存特征维度,替代传统PCA方法。
  • 在LLaMA2-7B和Mistral-7B上实现平均60%压缩率,性能损失低于10%。
  • 支持不同层和注意力头自适应压缩,适合资源受限场景部署。

KV缓存已成为大语言模型推理中的标准技术,其张量形状为(层数,头数,序列长度,特征维度),用于缓存自注意力的历史信息。随着模型和数据规模增长,KV缓存成为存储与内存传输的瓶颈。以往研究主要关注缓存张量前三个轴的压缩,本文则聚焦特征维度轴,利用低秩投影矩阵将缓存特征映射到低维空间。我们首先考察基于主成分分析(PCA)的正交投影方法,发现其在低压缩率下性能明显下降。为此,提出采用蒸馏目标与精心设计的套娃式训练策略直接优化正交投影矩阵。训练后,根据不同的压缩预算,自适应地为各层和头搜索最优压缩率。相比已有方法,本方案可轻松适配预训练大模型,并在性能与压缩率间实现平滑权衡。实验表明训练过程数据效率高,可在主流模型如LLaMA2-7B-base和Mistral-7B-v0.3-base上实现平均60%的KV缓存压缩率(极端情况下达75%),同时维持超过90%的原始性能。

原文摘要 · Abstract (English)

KV cache has become a de facto technique for the inference of large language models (LLMs), where tensors of shape (layer number, head number, sequence length, feature dimension) are introduced to cache historical information for self-attention. As the size of the model and data grows, the KV cache can quickly become a bottleneck within the system in both storage and memory transfer. To address this, prior studies usually focus on the first three axes of the cache tensors for compression. This paper supplements them, focusing on the feature dimension axis, by utilizing low-rank projection matrices to transform the cache features into spaces with reduced dimensions. We begin by investigating the canonical orthogonal projection method for data compression through principal component analysis (PCA). We observe the issue with PCA projection where significant performance degradation is observed at low compression rates. To bridge the gap, we propose to directly tune the orthogonal projection matrices with a distillation objective using an elaborate Matryoshka training strategy. After training, we adaptively search for the optimal compression rates for various layers and heads given varying compression budgets. Compared to previous works, our method can easily embrace pre-trained LLMs and hold a smooth tradeoff between performance and compression rate. We empirically witness the high data efficiency of our training procedure and find that our method can sustain over 90% performance with an average KV cache compression rate of 60% (and up to 75% in certain extreme scenarios) for popular LLMs like LLaMA2-7B-base and Mistral-7B-v0.3-base.

KV缓存模型压缩正交投影大模型推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。