Re Reference AI

アーキテクチャ

MLAとは

Multi-head Latent Attention / マルチヘッド潜在注意

キーとバリューを低ランクの潜在ベクトルへ圧縮して保持し、KVキャッシュを削減する注意機構

Attention推論最適化

ひとことで言うと

AIが長い文章を扱うときに膨らむ「作業メモ(KVキャッシュ)」を、ぎゅっと圧縮して保存する工夫。DeepSeekのモデルで使われ有名になった。

概要

MLA(Multi-head Latent Attention)とは、注意機構のキーとバリューを低次元の潜在ベクトルへ圧縮して保持する方式を指す。 マルチヘッド注意ではヘッドごとにキーとバリューを保持するためKVキャッシュが肥大化し、長いコンテキストの推論でGPUメモリを圧迫する。 MLAはキーとバリューを共同で低ランクの潜在ベクトルへ圧縮して格納し、必要時に復元することで、精度の低下を抑えながらキャッシュ量を削減する。 DeepSeek-V2で導入されV3・R1にも採用された。GQAMQAに代わるKVキャッシュ削減手法として注目されている。

歴史

2024年5月にDeepSeekDeepSeek-V2の技術報告で発表した。

アーキテクチャ

通常のマルチヘッド注意は各ヘッドごとにキーとバリューを個別に保持するため、ヘッド数やコンテキスト長に比例してKVキャッシュが増大する。MLAはキーとバリューを共同で低ランクの潜在ベクトルへ圧縮する射影行列を学習し、推論時はこの圧縮済みの潜在ベクトルのみをキャッシュしておき、各ヘッドで利用する際に必要な次元へ復元する。位置エンコーディングRoPE)についても、圧縮と両立できるよう一部の次元を分離して扱う工夫が導入されている。

ワークフロー

入力トークンの表現からキー・バリューを共同で低ランクの潜在ベクトルへ圧縮する → この潜在ベクトルのみをKVキャッシュとして保持する → クエリとの注意計算時に、潜在ベクトルから各ヘッド用のキー・バリューを復元する → 通常のマルチヘッド注意と同様にアテンションスコアを計算する。

利点

  • 通常のマルチヘッド注意に比べ、同等以上の精度を保ちながらKVキャッシュのメモリ使用量を削減できる
  • キャッシュ削減により、同じGPUメモリでより長いコンテキストやより大きなバッチサイズを扱える
  • DeepSeek-V2・V3・R1で実運用され、大規模MoEモデルの推論コスト削減に寄与した実績がある

欠点

  • キーとバリューの圧縮・復元の計算が追加されるため、実装や学習の複雑さが通常のマルチヘッド注意より増す
  • 圧縮によって表現力が理論上は制約されるため、圧縮率の設定によっては精度への影響が生じうる
  • GQAMQAに比べて登場から日が浅く、他アーキテクチャでの採用実績はまだ限られる

比較

  • GQAGQAがキー・バリューのヘッド数を減らして共有することでキャッシュを削減するのに対し、MLAはキー・バリューを低ランクの潜在ベクトルへ圧縮して保持する

関連用語

KVキャッシュGQAMQASelf-AttentionDeepSeek

よくある質問

GQA/MQAとどちらが優れている?

DeepSeekの報告ではMLAはGQAより少ないキャッシュ量で同等以上の精度を達成したとされるが、実装の複雑さとのトレードオフがあり、採用実績の広がっているGQAも依然主流の選択肢になっている。

参考文献

関連Zenn記事