TransMLA: マルチヘッド潜在的注意はすべて必要です

TransMLA: Multi-head Latent Attention Is All You Need

February 11, 2025

著者: Fanxu Meng, Zengwei Yao, Muhan Zhang

cs.AI

要旨

現代の大規模言語モデル（LLM）は、純粋な計算上の制約よりも、現在のハードウェア上で通信のボトルネックにしばしば遭遇します。マルチヘッド潜在注意（MLA）は、キー値（KV）レイヤーで低ランク行列を使用することで、圧縮された潜在的なKV状態をキャッシュできるようにし、この課題に取り組んでいます。このアプローチにより、従来のマルチヘッド注意と比較して、KVキャッシュサイズが大幅に削減され、推論が高速化されます。さらに、MLAは、表現力を向上させるためにアッププロジェクション行列を使用し、追加の計算を交換条件として通信オーバーヘッドを削減します。MLAはDeepseek V2/V3/R1で効率と効果を実証していますが、多くの主要なモデルプロバイダーは依然としてGroup Query Attention（GQA）に依存しており、MLAを採用する計画を発表していません。本論文では、GQAは常に同じKVキャッシュオーバーヘッドを維持しながらMLAで表現できることを示し、逆は成り立たないことを示します。MLAのより広範な使用を促進するために、**TransMLA**という、広く使用されているGQAベースの事前学習モデル（LLaMA、Qwen、Mixtralなど）をMLAベースのモデルに変換する事後トレーニング手法を紹介します。変換後、モデルはKVキャッシュサイズを増やさずに表現力を向上させるために追加のトレーニングを受けることができます。さらに、Deepseek R1の効率的な蒸留を可能にするために、MLA固有の推論加速技術を開発する予定です。

English

Modern large language models (LLMs) often encounter communication bottlenecks on current hardware, rather than purely computational constraints. Multi-head Latent Attention (MLA) tackles this challenge by using low-rank matrices in the key-value (KV) layers, thereby allowing compressed latent KV states to be cached. This approach significantly reduces the KV cache size relative to traditional multi-head attention, leading to faster inference. Moreover, MLA employs an up-projection matrix to increase expressiveness, trading additional computation for reduced communication overhead. Although MLA has demonstrated efficiency and effectiveness in Deepseek V2/V3/R1, many major model providers still rely on Group Query Attention (GQA) and have not announced any plans to adopt MLA. In this paper, we show that GQA can always be represented by MLA while maintaining the same KV cache overhead, but the converse does not hold. To encourage broader use of MLA, we introduce **TransMLA**, a post-training method that converts widely used GQA-based pre-trained models (e.g., LLaMA, Qwen, Mixtral) into MLA-based models. After conversion, the model can undergo additional training to boost expressiveness without increasing the KV cache size. Furthermore, we plan to develop MLA-specific inference acceleration techniques to preserve low latency in transformed models, thus enabling more efficient distillation of Deepseek R1.