X-EcoMLA：将预训练注意力升级为 MLA 以实现高效且极致的 KV 压缩
计算与语言
2025-09-09 v4
摘要
多头潜在注意力（MLA）旨在通过低秩键值联合压缩来优化 KV 缓存内存。MLA 不再单独缓存键和值，而是存储它们的压缩潜在表示，从而在保持性能的同时降低内存开销。虽然 MLA 在不损害语言模型准确性的情况下提高了内存效率，但其主要局限性在于需要在预训练阶段进行集成，要求模型从头开始训练。这引出了一个关键问题：我们能否在已经使用不同注意力机制预训练的模型中完全或部分地利用 MLA 的优势？在本文中，我们提出 X-EcoMLA，采用训练后蒸馏，通过轻量级的训练后适配，将基于 Transformer 的注意力升级为高效的混合 MLA 变体，从而无需进行大量的预训练。我们证明，利用训练良好的模型的暗知识可以提高训练精度，并在不损害模型性能的情况下实现 MLA 中极致的 KV 缓存压缩。实验结果表明，我们提出的方法可以在保持基准性能的同时有效压缩 KV 缓存；具体而言，对于 Llama3.2-1B-Instruct 基线，在 AMD MI300 上仅使用 3.6B 训练 token 和 70 GPU 小时即可实现 6.4 倍压缩并保持相同的平均分数，而使用 7B 训练 token 和 140 GPU 小时实现 10.6 倍压缩时，平均分数下降不到 0.1%。该工作的代码可在 https://github.com/AMD-AGI/AMD-Hybrid-Models 获取。
引用
@article{arxiv.2503.11132,
title = {X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression},
author = {Guihong Li and Mehdi Rezagholizadeh and Mingyu Yang and Vikram Appia and Emad Barsoum},
journal= {arXiv preprint arXiv:2503.11132},
year = {2025}
}