摘要
arXiv:2502.07864v2 通信类型:跨域
摘要:现代大规模语言模型(LLMs)通常在当前硬件上遇到通信瓶颈,而不仅仅是计算约束。多头潜在注意(MLA)通过在键值(KV)层中使用低秩矩阵来应对这一挑战,从而允许压缩的潜在KV状态被缓存。这种方法相对于传统的多头注意力显著减少了KV缓存的大小,从而加快了推理速度。此外,MLA 使用上投影矩阵以增加表达能力,通过增加额外的计算来减少通信开销。尽管 MLA 在 Deepseek V2/V3/R1 中已经证明了其效率和有效性,但许多主要的模型提供商仍然依赖于组查询注意(GQA),并且没有宣布任何计划采用 MLA。在本文中,我们展示了在保持相同的 KV 缓存开销的同时,GQA 总是可以被表示为 MLA,但反之则不成立。为了促进 MLA 的更广泛应用,我们引入了 TransMLA,这是一种后训练方法,可以将广泛使用的基于 GQA 的预训练模型(如 LLaMA、Qwen、Mixtral)转换为基于 MLA 的模型。转换后,模型可以进行额外的训练以增强表达能力,而不增加 KV 缓存的大小。此外,我们计划开发专门针对 MLA 的推理加速技术,以在转换后的模型中保留低延迟,从而能够更有效地蒸馏 Deepseek R1。