はじめに
個人的な学習として、KV Cache に関連する技術を調べています。
KV Cache は LLM 推論においてメモリボトルネックの主要因となり得ます。前回の記事では、MQA や GQA によって KV Cache の容量を削減できることを見てきました。
本稿では、さらに大きく KV Cache を削減する手法として、DeepSeek-V2 で採用された MLA (Multi-Head Latent Attention) を見ていきます。
MHA における KV Cache の問題
Decoding Multi-Head Latent Attention (Part 1): The KV Cache Memory Bottleneck, Solved. という記事を参考にしています。1
LLM の主要なアーキテクチャである Transformer は Multi-Head Attention というメカニズムでコンテキストを理解しています。Query (Q) と Key (K) の類似度を計算し、その重みに応じて Value (V) を加重平均します。
テキスト生成のタスクでは、前の単語全てに基づいて、トークン・単語ごとにテキストを生成します。これは膨大な計算負荷となります。例えば、モデルは現在のトークンだけでなく、シーケンス内の先行するトークン全てに対して Attention を計算する必要があるからです。
この処理を軽減するために KV Cache という手法があります。これは、各ステップで以前のすべての Key, Value を再計算するのではなく、一度計算した値をキャッシュしておくというものです。モデルが新しいトークンを生成するたびに、そのトークンに対応する Key と Value だけを新たに計算し、KV Cache に追加します。そして、次のトークンでは、キャッシュに既に存在する Key と Value を再利用し、新しい Query のみを計算します。これは冗長な計算を回避し、推論速度を大きく向上させます。
しかしながら、KV Cache にはメモリに大きなコストがかかるという問題点があります。以前の記事で示したように、レイヤーごとの KV Cache の容量はおおよそ以下のようになります。
Size = 2 \times L \times HiddenSize
- シーケンス長(L): 会話履歴または入力ドキュメントの全長。つまり、これまでに生成されたトークンの合計
- HiddenSize = n_h × d_h
- アテンションヘッド数(n_h): アテンションメカニズムが使用する異なる「視点」の数
- d_h : 1ヘッドの次元
Key, Value なので 2 が係数としてかかっています。また、レイヤー数を NLayers とすると、$ NLayers \times 2 \times L \times d \times nh $ となります。
DeepSeek-V2 の場合は、12万8000トークンというコンテキスト長となり、その容量は膨大です。
しかしながら、DeepSeek-V2 では、その前のモデルと比較して KV Cache の容量を大幅に削減することに成功しています。これには、MLA (Multi-Head Latent Attention) というメカニズムが関係しています。
DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model 2
Multi-Head Latent Attention
Multi-Head Latent Attention (MLA) は DeepSeek-AI が開発した DeepSeek-V2 で導入された技術です。DeepSeek-V2 は2024年5月に発表されたモデルですが、推論時のメモリ消費を大幅に削減しつつ、高い処理速度を実現しました。論文2の評価実験では、従来モデルと比較して学習コストを42.5%削減し、生成スループットを5.76倍に向上させたと報告されています。
MLA の特徴は以下です。
- 低ランク行列による圧縮（Low-Rank Compression）: Key と Value の次元数を、各ヘッドごとの Key と Value を保持する代わりに、低次元の潜在表現（Latent Vector）に次元削減します
- キャッシュの削減: 各ヘッドの Key と Value ではなく、この共通化・圧縮された潜在空間の行列 $(C_{KV})$ のみをキャッシュします。これにより、キャッシュサイズを小さく抑えます
- Absorption Trick（吸収トリック）: Query を計算する重み行列と、Key を復元するための重み行列を事前に掛け合わせておく Absorption Trick を採用しています。これにより、Query 側とキャッシュ側で独立して計算が可能になり、アテンションスコアが効率的に計算できます
- また、RoPE は低ランク圧縮と相性が悪いため、MLA では RoPE を適用するための専用のクエリと共有キーを別途用意する「分離型 RoPE」を採用しています
Multi-Head Latent Attention (MLA) 2
Low-Rank Key-Value Joint Compression
MLAの核心は、Key と Value を単一の潜在ベクトルへ圧縮することにあります。
圧縮計算
入力ベクトル $h_t$ は、低ランクの潜在ベクトル $ C^{KV}_{t} $ に射影 (projection) されます。
C^{KV}_{t} = W^{DKV} h_t
ここで、$ W^{DKV} $ はダウンプロジェクション行列と呼ばれています。
復元
この潜在ベクトルから、Attention 計算に必要な Key と Value が生成されます。
\begin{align}
k^{c}_{t} &= W^{UK} c^{KV}_t \\
v^{c}_{t} &= W^{UV} c^{KV}_t
\end{align}
$ W^{UK} $ と $ W^{UV} $ はアッププロジェクション行列です。
推論時には、この $c^{KV}_t$ のみをキャッシュすればよいため、KVキャッシュが劇的に削減されます。
Query の圧縮
学習時のアクティベーションメモリを節約するため、Query に対しても同様の圧縮が行われます。
\begin{align}
c^{Q}_{t} &= W^{DQ} h_t \\
q^{Q}_{t} &= W^{UQ} c^{Q}_{t}
\end{align}
RoPE の扱い
RoPE は位置情報を回転行列として埋め込むため、単純な低ランク圧縮と組み合わせると、後述する Attention 計算時の最適化 (Absorption Trick) が適用できなくなります。MLA では RoPE を適用する成分を別途用意します。
\begin{align}
q_{t,i}^{R} &= RoPE(W^{QR}c_{t}^{Q}) \\
k_{t}^{R} &= RoPE(W^{KR}h_{t})
\end{align}
ここで、$k_{t}^{R}$ はすべてのヘッドで共有される Key です。
最終的な Key と Query の結合
コンテンツ成分（圧縮されたもの）と位置成分（RoPE を適用したもの）を結合します。
\begin{align}
q_{t,i} &= \big[ q_{t,i}^{C};q_{t,i}^{R} \big] \\
k_{t,i} &= \big[ k_{t,i}^{C};k_{t}^{R} \big]
\end{align}
Attention 計算
MLA の最大の工夫は、推論時に Key と Value を明示的に復元せず、そのまま Attention を計算できる点です。
通常は潜在ベクトルから Key と Value を復元して Attention を計算します。しかし MLA では、行列積の結合法則を利用することで、この復元処理そのものを省略できます。
連結された Query と Key を用いて、最終的なスコアと出力を計算します。
ここで、$d_h$ は各ヘッドの次元、$d^R_h$ は RoPE 用ベクトルの次元です。
\begin{align}
o_{t,i} &= \sum^{t}_{j=1}softmax_{j} \big( \frac{q^{T}_{t,i} k_{j,i}}{\sqrt{d_{h}+d^{R}_{h}}} \big) v^{C}_{j,i} \\
u_{t} &= W^{O}\big[ o_{t,1}; o_{t,2}; ...; o_{t,n_h}\big]
\end{align}
Absorption Trick
推論時には、アッププロジェクション行列 $W_{UK}$ を Query 側の射影行列 $ W_{UQ} $ へ、また $W_{UV}$ を出力射影 $W_{O}$ にあらかじめ組み込む（吸収する）ことができます。
これにより、潜在ベクトル $ c^{KV}_{t} $ から Key と Value を明示的に復元することなく、直接 Attention 計算を行うことが可能になり、計算コストを抑えることができます。
以下では、Key の計算を例に、Absorption Trick の仕組みを見ていきます。
通常のステップ
まず、潜在ベクトルから Key を復元します。
k^{C}_{t} = W_{UK} c^{KV}_{t}
続いて、この Key と Query の内積を計算して Attention スコアを求めます。
q^{T}_{t}k^{C}_{t} = q^{T}_{t} \big( W_{UK}c^{KV}_{t} \big)
つまり、通常は 一度高次元の Key を生成してから Query との内積を計算しています。
Absorption Trick による最適化
ここで、行列積の結合法則
(A B) C = A (B C)
を利用すると、
q^{T}_{t} \big( W_{UK}c^{KV}_{t} \big) = \big(q^{T}_{t} W_{UK} \big) c^{KV}_{t}
と書き換えることができます。
この変形により、$ W_{UK} $ をあらかじめ Query 側の射影行列へ吸収しておけば、高次元の $ k^{C}_{t} $ を明示的に生成する必要がなくなります。
その結果、潜在ベクトル $c^{KV}_{t} $ に対して直接 Attention スコアを計算できるため、計算量とメモリアクセスの療法を削減できます。
まとめ
本稿では、DeepSeek-V2 で採用された Multi-Head Latent Attention (MLA) について調べました。
MLA の最大の特徴は、KV Cache のメモリ使用量を削減しながら、Attention 計算を効率化している点にあります。具体的には、次の 2 つの工夫が導入されています。
低ランク圧縮：各ヘッドの Key・Value を保持する代わりに、低次元の潜在ベクトルのみを KV Cache に保持することで、メモリ使用量を大幅に削減する。
Absorption Trick により、潜在ベクトルから Key や Value を明示的に復元することなく、低次元のまま Attention を計算できるため、計算量とメモリアクセスの両方を削減する。
MQA や GQA は 「KV Cache を共有することで容量を削減する」 アプローチでした。一方、MLA は 「KV Cache 自体を低ランクの潜在表現へ圧縮する」 という異なるアプローチを採用しており、長いコンテキスト長を扱う大規模言語モデルにおいて、高いメモリ効率と推論性能を実現しています。