通过低秩残差优化注意力缓存，显著降低内存占用且训练更快。
Low-Rank Key Value Attention
- 共享全秩投影+头特定低秩残差，平衡缓存与独立性。
- 仅用45-53%的缓存，测试损失低于标准MHA等方法。
- 训练速度提升18-25%，下游任务表现全面领先。
键值（KV）缓存是Transformer模型的主要内存瓶颈。我们提出低秩键值（LRKV）注意力机制，通过利用注意力头间的冗余性降低KV缓存内存占用，同时保持计算效率。每层采用共享全秩KV投影，并添加低秩、头特定的残差项，实现完全共享与完全独立之间的连续权衡。在128M至6.3B参数的预训练模型上，LRKV在标准MHA、MQA/GQA和MLA中均达到最低测试损失，仅需MHA 45-53%的KV缓存。在相同训练步数下，其性能达到基线水平的速度快18-25%。经监督微调后，LRKV在ARC-Easy、ARC-Challenge、MMLU、GSM8K和HumanEval等多个基准上表现最佳。
原文摘要 · Abstract (English)
The key-value (KV) cache is a primary memory bottleneck in Transformers. We propose Low-Rank Key-Value (LRKV) attention, which reduces KV cache memory by exploiting redundancy across attention heads, while being compute efficient. Each layer uses a shared full-rank KV projection augmented with low-rank, head-specific residuals, providing a continuous trade-off between complete sharing and full independence. After pretraining models of size 128M to 6.3B parameters, LRKV consistently achieves the lowest test loss among standard MHA, MQA/GQA, and MLA while using only 45-53\% of MHA's KV cache. LRKV reaches equivalent baseline quality 18-25\% faster (measured in training steps). After supervised midtraining, LRKV achieves the highest downstream task performance across ARC-Easy, ARC-Challenge, MMLU, GSM8K, and HumanEval benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。