arXiv:2503.11132cs.CL2025-03被引 4

将预训练模型的注意力机制改造为高效压缩的MLA,无需重训练。

X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression

  • 通过轻量级微调将已有模型注意力转为混合式MLA结构
  • 实现6.4倍压缩时仅用36亿参数训练,10.6倍压缩性能损失小于0.1%
  • 适合需要极致内存优化的推理场景,尤其适配大模型部署

多头潜在注意力(MLA)通过联合低秩压缩键值对来优化KV缓存内存。与分别缓存键值不同,MLA存储其压缩后的潜在表示,降低内存开销同时保持语言模型性能。然而,现有MLA需在预训练阶段集成,导致必须从头训练。本文提出X-EcoMLA,通过后训练蒸馏实现将基于Transformer的注意力向量“回收升级”为高效的混合式MLA变体,无需大规模预训练。实验表明,利用已训练模型的隐含知识可提升训练精度,并在不损失模型性能的前提下实现极端KV缓存压缩。以Llama3.2-1B-Instruct为基准,6.4倍压缩仅需3.6B训练参数和70 GPU小时(AMD MI300),10.6倍压缩的平均得分下降不足0.1%,耗时140 GPU小时,使用7B训练参数。代码已开源:https://github.com/AMD-AGI/AMD-Hybrid-Models。

原文摘要 · Abstract (English)

Multi-head latent attention (MLA) is designed to optimize KV cache memory through low-rank key-value joint compression. Rather than caching keys and values separately, MLA stores their compressed latent representations, reducing memory overhead while maintaining the performance. While MLA improves memory efficiency without compromising language model accuracy, its major limitation lies in its integration during the pre-training phase, requiring models to be trained from scratch. This raises a key question: can we use MLA's benefits fully or partially in models that have already been pre-trained with different attention mechanisms? In this paper, we propose X-EcoMLA to deploy post training distillation to enable the upcycling of Transformer-based attention into an efficient hybrid MLA variant through lightweight post-training adaptation, bypassing the need for extensive pre-training. We demonstrate that leveraging the dark knowledge of a well-trained model can enhance training accuracy and enable extreme KV cache compression in MLA without compromising model performance. The experimental results show that our proposed method can effectively compress the KV cache while preserving the performance on the benchmarks; specifically, for Llama3.2-1B-Instruct baseline, a 6.4x compression achieves the same average score by using only 3.6B training tokens and 70 GPU hours on AMD MI300, whereas a 10.6x compression have less than 0.1% average score drop with 7B training tokens and 140 GPU hours. The code for this work is available at https://github.com/AMD-AGI/AMD-Hybrid-Models.

KV压缩轻量化后训练MLA

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