1

前言

Multi-Head Latent Attention（MLA）隨著DeepSeek的火爆，成為大家關注的熱點。然而DeepSeek V2原文中只通過消融實驗驗證MLA的訓練效果好于MHA，并沒有為此提供理論保障。研究人員紛紛下場研究不同設計的能力對比，例如蘇劍林提出三個猜想，并通過實驗驗證注意力頭的維度是關鍵因素[1]。

2

GQA v.s. MLA

本文直接通過理論證明了，給定相同的KV Cache下，MLA的表達能力總是超過目前廣泛使用的Group-Query Attention（GQA）。證明思路如圖1所示：GQA能夠等價轉換為MLA，反之無法用GQA表示所有的MLA。

圖1給定相同的KV Cache預算，GQA的表達能力小于MLA，小于MQA。

我們將Qwen2.5中的GQA等價轉換為MLA，在SmolTalk數據集上進行訓練。轉換前后的模型，在訓練過程中的Loss以及訓練后模型的效果如圖2所示。

圖2將Qwen2.5中的GQA等價轉換為MLA，在下游任務中訓練的效果對比。

從圖中可以看出，經過轉換的MLA模型在訓練過程中表現(xiàn)出更低的Loss值，表明其對訓練數據的擬合能力更強。最終基于MLA的模型在數學和代碼任務上的準確率顯著高于原始的基于GQA的模型。

3

TransMLA

GQA等價變?yōu)镸LA能提升模型的表達能力，但并未減少模型的KV Cache。因此并不能提升推理速度，阻礙了TransMLA的實際應用價值。因為通常我們更希望獲得一個預訓練時強大，推理時高效的模型。因此我們通過如圖3所示的RoRoPE、FreqFold和BKV-PCA三項創(chuàng)新，實現(xiàn)了：

1）壓縮LLaMA-2-7B 93%的KV Cache；

2）性能損失很小，通過少量訓練即可恢復；

3）無需優(yōu)化，直接使用DeepSeek模型加載，在多個硬件上實現(xiàn)5-10倍加速。

圖3TransMLA先將位置信息集中到第一個head，再對KV進行低秩壓縮。

整個轉換過程為：

1.分組合并。將分組的KV拼接為一個Latent表示，并將分組的RoPE拼接為RoPE'。

2. Decouple RoPE。將RoPE'分為RoPE和NoPE，將位置信息集中到RoPE head，去除NoPE部分的位置編碼。

3.壓縮KV。對K_rope與 V的大小進行平衡后，進行聯(lián)合低秩壓縮。其中RoRoPE表示我們可以在RoPE'兩端對QK進行旋轉，只要滿足：1）旋轉只發(fā)生在不同K head的相同維度，2）RoPE'的實部和虛部對應的維度需要使用相同的旋轉方式。RoRoPE通過這一特殊旋轉方式將K的主成分集中到一個頭，去除其他頭的位置編碼，重新使用一個標準的RoPE表示位置信息。

圖4RoRoPE能夠將多個K head的信息集中到第一個注意力頭中，從而以很小的誤差將其分割為K_rope和K_nope

如圖4（a）中黃線所示，使用RoRoPE后輸出分布集中在原來的第一個頭。綠線使用一種頻率近似的方法，使分布更集中在原來的第一個頭中。如圖4（b）使用RoRoPE裁剪至一個頭（128維）的效果顯著好于沒有進行主成分提取的MHA2MLA[2]。

4

實驗

TransMLA減少轉換過程的性能損失，能夠輕易通過訓練恢復效果

我們在SmolLM-1.7B和LLaMA-2-7B上驗證TransMLA的效果，使用同期工作MHA2MLA作為對比。由于TransMLA使用的RoRoPE、FreqFold和BKV-PCA顯著減少轉換時的誤差，如表1所示，裁剪LLaMA-2-7B 68.75%的KV Cache，無需訓練，在6個benchmark上只損失1.65%的效果，而MHA2MLA則損失約21.85%的效果。

表1直接將模型轉換為MLA，使用TransMLA和MHA2MLA的效果對比。

如表2所示，由于TransMLA轉換后對模型的破壞更小，因此只使用500M Tokens的訓練即可超過使用6B Tokens訓練的MHA2MLA。

表2通過少量訓練后，使用TransMLA和MHA2MLA的效果對比。

DeepSeek模型直接加載TransMLA的Checkpoint，輕易的使用vllm加速

不同于其他KV Cache壓縮方法需要專門定制推理框架，TransMLA將所有的模型都統(tǒng)一轉換為DeepSeek模型。利用其豐富的生態(tài)，只要能支持DeepSeek的硬件和環(huán)境，就能支持TransMLA的推理加速。目前我們實現(xiàn)了Transformers和vllm版本的代碼，未來將會在SGLang等其他框架上進行測試。

圖5TransMLA裁剪LLaMA-2-7B 92.97%的KV Cache，在不同的硬件環(huán)境上的加速比。

如圖5所示，僅僅將LLaMA-2-7B轉為MLA，就帶來了最多10.6x推理加速。未來我們將會結合DeepSeek的混合精度量化，MTP等技術進一步優(yōu)化模型推理速度。

TransMLA已經支持了主流的模型

TransMLA已經支持Llama、Qwen、Gemma、Mistral/Mixtral等主流模型，轉換代碼已經開源，近期將通過訓練恢復模型效果，發(fā)布MLA加速版本的基座模型。

TransMLA支持Grouped Latent Attention(GLA)

近期flash attention和Mamba作者發(fā)布的GLA[3]充分發(fā)揮了tensor parallel的優(yōu)勢，推理速度能比MLA快2倍。然而他們將這一方法定位為一個從頭預訓練的架構，從頭訓練GLA需要巨大的成本。我們指出，MLA模型，包括DeepSeek以及使用TransMLA轉化的模型，通通可以轉化為GLA模型。

DeepSeek-V2-Lite原始模型在wikitext2上的ppl為6.3102，直接使用原始GLA的實現(xiàn)加載的ppl為21.0546，這可能是其未提供直接加載DeepSeek模型效果的原因。我們通過解決tensor parallel時RMSNorm和Softmax分割的問題，將ppl降低到了7.2416，通過少量訓練即可恢復模型效果。接下來我們將會在DeepSeek V3/R1上進行實驗，盡量維持滿血DeepSeek能力，同時提升推理速度。

5

總結與展望

本文理論證明了MLA的表達能力大于GQA，呼吁基座模型全面轉為MLA架構。同時提供一種將存量GQA模型轉換為MLA模型的方法，減少遷移所需的成本。此外將會完善GQA/MHA/MLA轉GLA的方法，嘗試突破DeepSeek的能力邊界。

6

引用

[1] Transformer升級之路：20、MLA究竟好在哪里？https://kexue.fm/archives/10907

[2] Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs

[3] Hardware-Efficient Attention for Fast Decoding