[論文レビュー] 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 より表現力を持つことを証明し、GQA ベースの事前学習モデルを MLA ベースのモデルに変換する後処理学習法である TransMLA を導入し、その後の微調整により 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 より理論的に表現力の利点を持つことを確立する。
• 人気のある GQA ベースのモデルを MLA ベースのモデルへ変換する実践的な変換ワークフロー（TransMLA）を提供する。
• 微調整済みの TransMLA モデルが、KV キャッシュのオーバーヘッドを最小限に抑えつつ、下流タスクで GQA の対向モデルを上回ることを示す。

提案手法

• KV キャッシュと表現力の観点から MHA、GQA、MQA、MLA を定義・比較する。
• 複製と低ランク分解による議論を通じて、同じ KV キャッシュで GQA を MLA に変換できることを示す。
• 潜在 KV 表現を持つ MLA として GQA を表現するための SVD ベースの低ランク分解を明らかにする。
• W_K^a、W_K^b、W_V^a、W_V^b のパラメータ化によって表現力を拡張しつつ KV キャッシュを固定した後処理変換として TransMLA を導入する。
• 推論時に固定された潜在 KV 次元を維持するために、特定の行列を統合する absorb 演算を記述する。
• 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)