[论文解读] TransMLA: Multi-Head Latent Attention Is All You Need

Fanxu Meng, Tang, Pingzhi+4 more|ArXiv.org|Feb 11, 2025

Topic Modeling被引用 5

一句话总结

论文证明在相同 KV 缓存下 MLA 相较于 GQA 具有更强的表达能力，并引入 TransMLA，一种后训练方法将基于 GQA 的预训练模型转换为基于 MLA 的模型，随后进行微调以在不增加 KV 缓存大小的情况下提升下游性能。

ABSTRACT

In this paper, we present TransMLA, a framework that seamlessly converts any GQA-based pre-trained model into an MLA-based model. Our approach enables direct compatibility with DeepSeek's codebase, allowing these models to fully leverage DeepSeek-specific optimizations such as vLLM and SGlang. By compressing 93% of the KV cache in LLaMA-2-7B, TransMLA achieves a 10.6x inference speedup at an 8K context length while preserving meaningful output quality. Additionally, the model requires only 6 billion tokens for fine-tuning to regain performance on par with the original across multiple benchmarks. TransMLA offers a practical solution for migrating GQA-based models to the MLA structure. When combined with DeepSeek's advanced features, such as FP8 quantization and Multi-Token Prediction, even greater inference acceleration can be realized.

研究动机与目标

• 动机：降低大语言模型自注意力中的 KV 缓存瓶颈。
• 在相同 KV 缓存成本下，确立 MLA 相较于 GQA 的理论表达能力优势。
• 提供一个实用的转换工作流（TransMLA），将流行的基于 GQA 的模型转换为基于 MLA 的模型。
• 证明微调后的 TransMLA 模型在下游任务上优于其 GQA 对应模型，且 KV 缓存开销最小。

提出的方法

• 定义并在 KV 缓存与表示方面比较 MHA、GQA、MQA 和 MLA。
• 证明通过复制与低秩分解等论证，GQA 可以在相同 KV 缓存下转换为 MLA。
• 揭示基于 SVD 的低秩分解，将 GQA 表示为具有潜在 KV 表示的 MLA。
• 引入 TransMLA 作为一种后训练转换，在扩展表达能力的同时通过 W_K^a、W_K^b、W_V^a、W_V^b 参数化保持 KV 缓存固定。
• 描述一种吸收（absorb）操作，将某些矩阵合并以在推理时维持固定的潜在 KV 维度。
• 给出将 Qwen2.5 及类似模型从 GQA 转换到 MLA 并在 SmolTalk 上对数学/代码任务进行微调的实验设置。

(a) Group Query Attention (GQA)

实验结果

研究问题

• RQ1 在相同 KV 缓存大小下，MLA 是否能够匹配或超越 GQA 的表现？
• RQ2 是否有可能在不增加 KV 缓存开销的前提下，将基于 GQA 的预训练模型转换为基于 MLA 的模型？
• RQ3 转换后的 TransMLA 进行微调，是否在数学与编码等下游任务上获得可测量的提升？
• RQ4 正交分解在 TransMLA 的性能提升中起到的作用是什么？

主要发现

• 当 KV 缓存大小固定时，MLA 比 GQA 更具表达能力，且任意 GQA 配置都可以转换为 MLA 表示。
• 一种实用、低开销的转换（TransMLA）可以将流行的基于 GQA 的模型（如 LLaMA、Qwen、Mixtral）在不增加 KV 缓存大小的情况下转换为 MLA。
• 微调后的 TransMLA 模型在下游任务上相较于其 GQA 对应模型表现更佳，尤其在数学和编码任务中。
• 性能提升归因于潜在 KV 表示带来的表达能力扩展以及因子分解中的正交分解，而不仅仅是参数数量的增加。
• 在没有正交分解的同等扩展（类似恒等映射的维度扩展）的情况下，收益仅为边际，凸显潜在分解方法的重要性。

(b) Multi-Head Attention (MHA)