Multi-head Latent Attention (MLA)の詳細【be absorbedとは？】

目的

DeepSeek V2論文のMulti-head Latent Attention (MLA)における、下記表現を理解するために「数式で」整理をすることが目的の記事です。
主に、「どの重みを」「どのように吸収するのか」を「数式で」整理します。
（そのうち詳細な解説付きの記事にするかも）

In addition, during inference, since $WUKW^{UK}$ can be absorbed into $WQW^Q$, and $WUVW^{UV}$ can be absorbed into $WOW^{O}$, we even do not need to compute keys and values out for attention.

本記事は、DeepSeek V2の論文中に記載されているMulti-head Latent Attention (MLA)について理解されていることが前提になります。
https://arxiv.org/abs/2405.04434

なお、RoPE導入後も成立するという主張に関しては、「今回は」取り扱いません。
ただ、基本的な主張は変わらないです。そのうち記事にするかもです。

私はまだまだ勉強中の身です。
誤った理論展開をしておりましたら、ぜひご指摘ください。

なお、私が理解できるように、どこが間違っているのか、なぜ間違っているのかをわかりやすくお伝えいただけますと幸いです。

本記事は、DeepSeek V2の論文を理解している方にご教授をいただくための記事でもあります。
記事の一番下に私が理解できなくて困っている疑問点があるので、ぜひ詳しい方はそちらご教授いただけますと嬉しいです！

とはいえ、論文中では行間がすごく広くなっているMLAに関して、詳細に記述しているので、誰かの理解の助けになれば嬉しいです。
あと、今のAIでは書けない記事を書きたかった。
（o1 proとかに解説を依頼したとしても、ここまで詳細には解説してくれないはずです）

なお、最近流行りのDeepSeek-R1に関しても論文解説記事を書いています。
本記事とは異なり、理論に詳しくない方が読んで面白い！と思っていただけるように記事を書いたつもりなので、ぜひ読んでいただけると嬉しいです。
https://zenn.dev/asap/articles/34237ad87f8511

数式でMLAの整理

Scaled Dot-Product Attentionは下記で表現されます。

$$Attention(i)(Q(i),K(i),V(i))=softmax(Q(i)K(i)⊤dh)V(i) \text{Attention}_{(i)}(Q_{(i)}, K_{(i)}, V_{(i)}) = \text{softmax} \left( \frac{Q_{(i)}K_{(i)}^\top}{\sqrt{d_h}} \right) V_{(i)}$$

上記を前提に、MLAの式展開を考察していきます。

記号の導入

MLAは主にKVキャッシュとクエリの圧縮を目的にしています。
したがって、圧縮されたKV表現と圧縮されたクエリ表現を取り扱うために、記号を導入していきます。

圧縮KV表現とKV

ここでは、圧縮KV表現と、元のKVについて記載していきます。

DeepSeek V2の論文中では、特定のtokenに限定した式を記載していますが、同じではつまらないと思ったので、ここでは、全tokenをまとめた行列で書き下してみます。

Attention Is All You Needの論文でも、上記の式でAttentionは記載していたので、それをリスペクトします。

MLAにおいてKVの圧縮表現$CKV∈RT×dc\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$を導入します。
$TT$は全体のtoken数、$dcd_c$は圧縮KV表現の次元数（論文では512次元）です。

この圧縮表現は下記のより作成されます。
（各tokenごとに下記の式変形が適用されます。つまりtoken数分だけ並列処理されます）

$$ctKV=WDKVht \begin{equation} \mathbf{c}_t^{KV} = W^{DKV} \mathbf{h_t} \end{equation}$$

$$CKV=[(c1KV)T(c2KV)T⋮(cTKV)T] \begin{equation} \mathbf{C}^{KV} = \begin{bmatrix} (\mathbf{c}_1^{KV})^T \\ (\mathbf{c}_2^{KV})^T \\ \vdots \\ (\mathbf{c}_T^{KV})^T \end{bmatrix} \end{equation}$$

ただし$tt$は、ある$tt$ token目の特徴量であることを示します。
また、$ht∈Rd\mathbf{h_t} \in \mathbb{R}^{d}$は、前の層の、ある$tt$ token目での出力で、$dd$は前の層の出力次元数です（論文中では5120次元）

また、$WDKV∈Rdc×dW^{DKV} \in \mathbb{R}^{d_c \times d}$はdown-projection行列です。
$ctKV∈Rdc\mathbf{c}_t^{KV} \in \mathbb{R}^{d_c}$は、ある$tt$ token目での、KVの圧縮表現です。
そして、$CKV∈RT×dc\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$は$ctKV\mathbf{c}_t^{KV}$を全てのtokenでまとめた行列になります。

KVの圧縮表現$CKV\mathbf{C}^{KV}$を利用すると、キーとバリューを復元できます。

$$K=(WUK(CKV)T)T \begin{equation} \mathbf{K} = \left(W^{UK} (\mathbf{C}^{KV})^T\right)^T \end{equation}$$

$$V=(WUV(CKV)T)T \begin{equation} \mathbf{V} = \left(W^{UV} (\mathbf{C}^{KV})^T\right)^T \end{equation}$$

$WUK,WUV∈Rdhnh×dcW^{UK},W^{UV} \in \mathbb{R}^{d_hn_h \times d_c}$はup-projection行列です。
圧縮されたKVから、マルチヘッドの全てのキー、バリュー表現を復元します。
$dhd_h$は1ヘッドあたりの次元数（論文では128次元）、$nhn_h$はマルチヘッドのヘッド数（論文では128個）
$K,V∈RT×dhnh\mathbf{K}, \mathbf{V} \in \mathbb{R}^{T \times d_hn_h}$は、それぞれ圧縮されていないキー、バリュー表現です。

圧縮クエリ表現とQ

続いて、クエリに関しても学習時のパラメータ削除のために圧縮します。

$$ctQ=WDQht \begin{equation} \mathbf{c}_t^{Q} = W^{DQ} \mathbf{h_t} \end{equation}$$

$$CQ=[(c1Q)T(c2Q)T⋮(cTQ)T] \begin{equation} \mathbf{C}^{Q} = \begin{bmatrix} (\mathbf{c}_1^{Q})^T \\ (\mathbf{c}_2^{Q})^T \\ \vdots \\ (\mathbf{c}_T^{Q})^T \end{bmatrix} \end{equation}$$

また、$WDQ∈Rdc′×dW^{DQ} \in \mathbb{R}^{d_c' \times d}$はdown-projection行列です。
$dc′d_c'$は圧縮クエリ表現の次元数（論文では1536次元）
$ctQ∈Rdc′\mathbf{c}_t^{Q} \in \mathbb{R}^{d_c'}$は、ある$tt$ token目での、クエリの圧縮表現です。
$CQ∈RT×dc′\mathbf{C}^{Q} \in \mathbb{R}^{T \times d_c'}$は、token全体でのクエリの圧縮表現になります。

このクエリの圧縮表現$CQ\mathbf{C}^{Q}$を利用すると、元のクエリを復元できます。

$$Q=(WUQ(CQ)T)T \begin{equation} \mathbf{Q} = \left(W^{UQ} (\mathbf{C}^{Q})^T\right)^T \end{equation}$$

$WUQ∈Rdhnh×dc′W^{UQ} \in \mathbb{R}^{d_hn_h \times d_c'}$はup-projection行列です。
クエリの圧縮表現から、マルチヘッドの全てのクエリ表現を復元します。
$Q∈RT×dhnh\mathbf{Q} \in \mathbb{R}^{T \times d_hn_h}$は、それぞれ圧縮されていないクエリ表現です。

通常のMHA（Multi-Head Attention）のQKVと、 MLAの違い

通常のMHAにおいて、QKVはそれぞれ下記で表現されます。

MHAでのキー表現

$$kt=WKht \begin{equation} \mathbf{k}_t = W^{K} \mathbf{h_t} \end{equation}$$

$$K=[(k1)T(k2)T⋮(kT)T] \begin{equation} \mathbf{K} = \begin{bmatrix} (\mathbf{k}_1)^T \\ (\mathbf{k}_2)^T \\ \vdots \\ (\mathbf{k}_T)^T \end{bmatrix} \end{equation}$$

したがって、

$$K=[(WKh1)T(WKh2)T⋮(WKhT)T] \begin{equation} \mathbf{K} = \begin{bmatrix} (W^{K} \mathbf{h_1})^T \\ (W^{K} \mathbf{h_2})^T \\ \vdots \\ (W^{K} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

なお、$kt∈Rdhnh\mathbf{k}_t \in \mathbb{R}^{d_hn_h}$、$WK∈Rdhnh×dW^{K} \in \mathbb{R}^{d_hn_h \times d}$、$K∈RT×dhnh\mathbf{K} \in \mathbb{R}^{T \times d_hn_h}$

以下はほぼ同じなので、隠しておきます。興味あればご覧ください。

MHAでのクエリ表現

MHAでのクエリ表現

$$qt=WQht \begin{equation} \mathbf{q}_t = W^{Q} \mathbf{h_t} \end{equation}$$

$$Q=[(q1)T(q2)T⋮(qT)T] \begin{equation} \mathbf{Q} = \begin{bmatrix} (\mathbf{q}_1)^T \\ (\mathbf{q}_2)^T \\ \vdots \\ (\mathbf{q}_T)^T \end{bmatrix} \end{equation}$$

したがって、

$$Q=[(WQh1)T(WQh2)T⋮(WQhT)T] \begin{equation} \mathbf{Q} = \begin{bmatrix} (W^{Q} \mathbf{h_1})^T \\ (W^{Q} \mathbf{h_2})^T \\ \vdots \\ (W^{Q} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

なお、$qt∈Rdhnh\mathbf{q}_t \in \mathbb{R}^{d_hn_h}$、$WQ∈Rdhnh×dW^{Q} \in \mathbb{R}^{d_hn_h \times d}$、$Q∈RT×dhnh\mathbf{Q} \in \mathbb{R}^{T \times d_hn_h}$

MHAでのバリュー表現

MHAでのバリュー表現

$$vt=WVht \begin{equation} \mathbf{v}_t = W^{V} \mathbf{h_t} \end{equation}$$

$$V=[(v1)T(v2)T⋮(vT)T] \begin{equation} \mathbf{V} = \begin{bmatrix} (\mathbf{v}_1)^T \\ (\mathbf{v}_2)^T \\ \vdots \\ (\mathbf{v}_T)^T \end{bmatrix} \end{equation}$$

したがって、

$$V=[(WVh1)T(WVh2)T⋮(WVhT)T] \begin{equation} \mathbf{V} = \begin{bmatrix} (W^{V} \mathbf{h_1})^T \\ (W^{V} \mathbf{h_2})^T \\ \vdots \\ (W^{V} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

なお、$vt∈Rdhnh\mathbf{v}_t \in \mathbb{R}^{d_hn_h}$、$WV∈Rdhnh×dW^{V} \in \mathbb{R}^{d_hn_h \times d}$、$V∈RT×dhnh\mathbf{V} \in \mathbb{R}^{T \times d_hn_h}$

見るとわかるように、それぞれほぼ同じように構成されます。
なので、キー表現に着目して、MHAとMLAの違いを見ていきます。

MLAにおけるキー表現

重要なのは、「圧縮KV表現とKV」の章で記載した下記の3式です。

$$ctKV=WDKVht \mathbf{c}_t^{KV} = W^{DKV} \mathbf{h_t}$$

$$CKV=[(c1KV)T(c2KV)T⋮(cTKV)T] \mathbf{C}^{KV} = \begin{bmatrix} (\mathbf{c}_1^{KV})^T \\ (\mathbf{c}_2^{KV})^T \\ \vdots \\ (\mathbf{c}_T^{KV})^T \end{bmatrix}$$

$$K=(WUK(CKV)T)T \mathbf{K} = \left(W^{UK} (\mathbf{C}^{KV})^T\right)^T$$

それぞれ式番号(1)(2)(3)に対応しています。
これをまとめていきます。

式(2)(3)から

$$K=(WUK[c1KV,c2KV,⋯ ,cTKV])T \begin{equation} \mathbf{K} = \left(W^{UK} [\mathbf{c}_1^{KV}, \mathbf{c}_2^{KV}, \cdots ,\mathbf{c}_T^{KV}]\right)^T \end{equation}$$

式(1)(11)から

$$K=(WUK[WDKVh1,WDKVh2,⋯ ,WDKVhT])T \mathbf{K} = \left(W^{UK} [W^{DKV} \mathbf{h_1} ,W^{DKV} \mathbf{h_2} ,\cdots ,W^{DKV} \mathbf{h_T}]\right)^T$$

$$=[WUKWDKVh1,WUKWDKVh2,⋯ ,WUKWDKVhT]T = [W^{UK}W^{DKV} \mathbf{h_1} ,W^{UK}W^{DKV} \mathbf{h_2} ,\cdots ,W^{UK}W^{DKV} \mathbf{h_T}]^T$$

となるため、これまでの書き方と表記を揃えると

$$K=[(WUKWDKVh1)T(WUKWDKVh2)T⋮(WUKWDKVhT)T] \begin{equation} \mathbf{K} = \begin{bmatrix} (W^{UK}W^{DKV} \mathbf{h_1})^T \\ (W^{UK}W^{DKV} \mathbf{h_2})^T \\ \vdots \\ (W^{UK}W^{DKV} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

なお、念のため記載すると、$K∈RT×dhnh\mathbf{K} \in \mathbb{R}^{T \times d_hn_h}$、$WUKWDKVht∈RdhnhW^{UK}W^{DKV} \mathbf{h_t} \in \mathbb{R}^{d_hn_h}$です。

ちなみに、MHAのキー表現は下記でした。

$$K=[(WKh1)T(WKh2)T⋮(WKhT)T] \begin{equation} \mathbf{K} = \begin{bmatrix} (W^{K} \mathbf{h_1})^T \\ (W^{K} \mathbf{h_2})^T \\ \vdots \\ (W^{K} \mathbf{h_T})^T \end{bmatrix} \end{equation}$$

したがって、適切に学習がなされた場合、式(12)(13)から、$WK=WUKWDKVW^{K} = W^{UK}W^{DKV}$となります。
このとき、$WKW^{K}$の行列のランクや$WUKW^{UK}$や$WDKVW^{DKV}$の行列の自由度$dcd_c$によって、どの程度劣化するかが決まります。
万が一、$WKW^{K}$の行列のランクが、$WUKW^{UK}$や$WDKVW^{DKV}$の行列の自由度$dcd_c$よりも小さい場合は、適切に学習がなされることにより、劣化が生じません。

ついでに、$WK∈Rdhnh×dW^{K} \in \mathbb{R}^{d_hn_h \times d}$は、パラメータ数が$dhnh×dd_hn_h \times d$なのに対して、$WUK∈Rdhnh×dcW^{UK} \in \mathbb{R}^{d_hn_h \times d_c}$と$WDKV∈Rdc×dW^{DKV} \in \mathbb{R}^{d_c \times d}$に分割した場合、パラメータ数が$dhnh×dc+dc×dd_hn_h \times d_c + d_c \times d$となります。
ここで、$dd$は$dcd_c$と比較して非常に大きいため、前者のパラメータよりも、少ないパラメータにて近似が可能です。

深層学習の世界において、しばしば学習可能なパラメータ数を少なくして近似することで、精度の向上が見られたことがあり（画像処理でいうCNNの世界）、今回MHAよりもMLAが性能が良い結果になったのは、その観点も効いているのかもしれません。

be absorbedってどういうこと？？

さて、実際にここから、論文中の下記の内容が成立することを確かめます。

In addition, during inference, since $WUKW^{UK}$ can be absorbed into $WQW^Q$, and $WUVW^{UV}$ can be absorbed into $WOW^{O}$, we even do not need to compute keys and values out for attention.

W^{UK}の吸収

単一ヘッドを前提とした場合、あるヘッド$ii$に着目すると、Attentionは下記のように表すことができます。

$$SDPA(i)(Q(i),K(i),V(i))=softmax(Q(i)K(i)⊤dh)V(i) \begin{equation} \text{SDPA}_{(i)}(Q_{(i)}, K_{(i)}, V_{(i)}) = \text{softmax} \left( \frac{Q_{(i)}K_{(i)}^\top}{\sqrt{d_h}} \right) V_{(i)} \end{equation}$$

なお、$SDPA(i)∈RT×dh\text{SDPA}_{(i)} \in \mathbb{R}^{T \times d_h}$です。
SDPAはScaled Dot-Product Attentionのことです。

今回の

$WUKW^{UK}$ can be absorbed into $WQW^Q$

の部分は、$Q(i)K(i)⊤Q_{(i)}K_{(i)}^\top$に関連する部分なので、その部分に着目します。

ただし、ヘッドごとに見ていくため、これまでの記号は分割されて下記のようになります。
$Q(i)∈RT×dhQ_{(i)} \in \mathbb{R}^{T \times d_h}$、$K(i)∈RT×dhK_{(i)} \in \mathbb{R}^{T \times d_h}$、$W(i)UQ∈Rdh×dc′W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$、$W(i)UK∈Rdh×dcW^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$
下記は、ヘッドごとに共通で利用するため変わりません。
$CQ∈RT×dc′\mathbf{C}^{Q} \in \mathbb{R}^{T \times d_c'}$、$CKV∈RT×dc\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$

さて、MLAにおいて$Q(i)K(i)⊤Q_{(i)}K_{(i)}^\top$は下記のように書き下せます。

$$Q(i)K(i)⊤=(W(i)UQ(CQ)T)TW(i)UK(CKV)T Q_{(i)}K_{(i)}^\top = \left(W^{UQ}_{(i)} (\mathbf{C}^{Q})^T\right)^T W^{UK}_{(i)} (\mathbf{C}^{KV})^T$$

$$=CQ(W(i)UQ)TW(i)UK(CKV)T = \mathbf{C}^{Q} (W^{UQ}_{(i)})^T W^{UK}_{(i)} (\mathbf{C}^{KV})^T$$

$$=((W(i)UK)TW(i)UQ(CQ)T)T(CKV)T \begin{equation} = \left((W^{UK}_{(i)})^T W^{UQ}_{(i)} (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T \end{equation}$$

上記において、$W(i)UK∈Rdh×dcW^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$、$W(i)UQ∈Rdh×dc′W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$は学習時は分割して学習された重み行列になります。

そこで、推論時には、これらの重み行列から事前に、$(W(i)UK)TW(i)UQ∈Rdc×dc′(W^{UK}_{(i)})^T W^{UQ}_{(i)} \in \mathbb{R}^{d_c \times d_c'}$を計算しておき、新しい一つの重み行列にすることもできます。

そうすると、圧縮KV表現$CKV\mathbf{C}^{KV}$と圧縮クエリ表現$CQ\mathbf{C}^{Q}$から明示的にキーとクエリを復元することなく、一撃で$Q(i)K(i)⊤Q_{(i)}K_{(i)}^\top$を計算できます。

この主張が私的イチオシポイントです！とても面白いと思いました！

これが推論時に、「$WUKW^{UK}$は$WUQW^{UQ}$($WQW^{Q}$)に吸収される」という言葉の意味だと思われます。

W^{UV}の吸収

今回の、

$WUVW^{UV}$ can be absorbed into $WOW^{O}$

の部分は、Attentionマップを作成し、バリューとの行列積を計算した後、各Headごとの出力を$WOW^{O}$で一つにまとめる部分です。この部分は少し入り組んでいるので、少し整理していきます。

まず、Attentionマップの部分を$Atn(i)Atn_{(i)}$とまとめます。
具体的には下記です。

$$Atn(i)(Q(i),K(i))=softmax(Q(i)K(i)⊤dh) \begin{equation} \text{Atn}_{(i)}(Q_{(i)}, K_{(i)}) = \text{softmax} \left( \frac{Q_{(i)}K_{(i)}^\top}{\sqrt{d_h}} \right) \end{equation}$$

ここで、$Atn(i)∈RT×T\text{Atn}_{(i)} \in \mathbb{R}^{T \times T}$です。

すると、Scaled Dot-Product Attentionは下記のように表せます。

$$SDPA(i)=Atn(i)V(i) \begin{equation} \text{SDPA}_{(i)} = \text{Atn}_{(i)} V_{(i)} \end{equation}$$

その場合、MHA（Multi-Head Attention）は、下記のように表現されます。

$$MHA(Q,K,V)=WO[SDPA(1)TSDPA(2)T⋮SDPA(nh)T] \begin{equation} \text{MHA}(Q, K, V) = W^{O} \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \end{equation}$$

ここで、$WO∈Rd×dhnhW^{O} \in \mathbb{R}^{d \times d_hn_h}$、$SDPA(i)T∈Rdh×T{SDPA}_{(i)}^T \in \mathbb{R}^{d_h \times T}$であり、

$$[SDPA(1)TSDPA(2)T⋮SDPA(nh)T]∈Rdhnh×T \begin{equation} \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \in \mathbb{R}^{d_hn_h \times T} \end{equation}$$

となります。

ここで、ちょっとトリッキーですが、$WOW^{O}$を下記のように分解していきます。

$$MHA(Q,K,V)=[W(1)O,W(2)O,⋯ ,W(nh)O][SDPA(1)TSDPA(2)T⋮SDPA(nh)T] \begin{equation} \text{MHA}(Q, K, V) = [W^{O}_{(1)}, W^{O}_{(2)}, \cdots, W^{O}_{(n_h)}] \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \end{equation}$$

このとき、$W(i)O∈Rd×dhW^{O}_{(i)} \in \mathbb{R}^{d \times d_h}$となります。

このとき、行列積の性質として下記が成立します。

$$MHA(Q,K,V)=[W(1)O,W(2)O,⋯ ,W(nh)O][SDPA(1)TSDPA(2)T⋮SDPA(nh)T]=W(1)OSDPA(1)T+W(2)OSDPA(2)T+⋯+W(nh)OSDPA(nh)T \begin{equation} \begin{aligned} \text{MHA}(Q, K, V) &=[W^{O}_{(1)}, W^{O}_{(2)}, \cdots, W^{O}_{(n_h)}] \begin{bmatrix} \text{SDPA}_{(1)}^T \\ \text{SDPA}_{(2)}^T \\ \vdots \\ \text{SDPA}_{(n_h)}^T \end{bmatrix} \\ &= W^{O}_{(1)} \text{SDPA}_{(1)}^T \\ &\quad + W^{O}_{(2)} \text{SDPA}_{(2)}^T \\ &\quad + \cdots + W^{O}_{(n_h)} \text{SDPA}_{(n_h)}^T \end{aligned} \end{equation}$$

式(17)を展開して、下記のように書きます。

$$MHA(Q,K,V)=W(1)O(Atn(1)V(1))T+W(2)O(Atn(2)V(2))T+⋯+W(nh)O(Atn(nh)V(nh))T \begin{equation} \begin{aligned} \text{MHA}(Q, K, V) &= W^{O}_{(1)} \left(\text{Atn}_{(1)}V_{(1)}\right)^T \\ &\quad + W^{O}_{(2)} \left(\text{Atn}_{(2)}V_{(2)}\right)^T \\ &\quad + \cdots + W^{O}_{(n_h)} \left(\text{Atn}_{(n_h)} V_{(n_h)}\right)^T \end{aligned} \end{equation}$$

このとき、式(4)から下記にように式変形できます。

式(4)は下記です。

$$V=(WUV(CKV)T)T \mathbf{V} = \left(W^{UV} (\mathbf{C}^{KV})^T\right)^T$$

なお、$CKV∈RT×dc\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$、$WUV∈Rdhnh×dcW^{UV} \in \mathbb{R}^{d_hn_h \times d_c}$、$V∈RT×dhnh\mathbf{V} \in \mathbb{R}^{T \times d_hn_h}$

すると、各Headごとに分解すると下記にようになります。

$$V(i)=(W(i)UV(CKV)T)T \mathbf{V}_{(i)} = \left(W^{UV}_{(i)} (\mathbf{C}^{KV})^T\right)^T$$

$$=CKV(W(i)UV)T = \mathbf{C}^{KV} (W^{UV}_{(i)})^T$$

なお、$CKV∈RT×dc\mathbf{C}^{KV} \in \mathbb{R}^{T \times d_c}$、$W(i)UV∈Rdh×dcW^{UV}_{(i)} \in \mathbb{R}^{d_h \times d_c}$、$V(i)∈RT×dh\mathbf{V}_{(i)} \in \mathbb{R}^{T \times d_h}$

$$MHA(Q,K,V)=W(1)O(Atn(1)CKV(W(1)UV)T)T+W(2)O(Atn(2)CKV(W(2)UV)T)T+⋯+W(nh)O(Atn(nh)CKV(W(nh)UV)T)T \begin{aligned} \text{MHA}(Q, K, V) &= W^{O}_{(1)} \left(\text{Atn}_{(1)}\mathbf{C}^{KV} (W^{UV}_{(1)})^T \right)^T \\ &\quad + W^{O}_{(2)} \left(\text{Atn}_{(2)}\mathbf{C}^{KV} (W^{UV}_{(2)})^T \right)^T \\ &\quad + \cdots + W^{O}_{(n_h)} \left(\text{Atn}_{(n_h)} \mathbf{C}^{KV} (W^{UV}_{(n_h)})^T \right)^T \end{aligned}$$

$$MHA(Q,K,V)=W(1)OW(1)UV(CKV)TAtn(1)T+W(2)OW(2)UV(CKV)TAtn(2)T+⋯+W(nh)OW(nh)UV(CKV)TAtn(nh)T \begin{equation} \begin{aligned} \text{MHA}(Q, K, V) &= W^{O}_{(1)} W^{UV}_{(1)} (\mathbf{C}^{KV})^T \text{Atn}_{(1)}^T \\ &\quad + W^{O}_{(2)} W^{UV}_{(2)} (\mathbf{C}^{KV})^T \text{Atn}_{(2)}^T \\ &\quad + \cdots + W^{O}_{(n_h)} W^{UV}_{(n_h)} (\mathbf{C}^{KV})^T \text{Atn}_{(n_h)}^T \\ \end{aligned} \end{equation}$$

上記において、$W(i)O∈Rd×dhW^{O}_{(i)} \in \mathbb{R}^{d \times d_h}$、$W(i)UV∈Rdh×dcW^{UV}_{(i)} \in \mathbb{R}^{d_h \times d_c}$は学習時は分割して学習された重み行列になります。

そこで、推論時には、これらの重み行列から事前に、$W(i)OW(i)UV∈Rd×dcW^{O}_{(i)} W^{UV}_{(i)} \in \mathbb{R}^{d \times d_c}$を計算しておき、新しい一つの重み行列にすることもできます。

そうすると、圧縮KV表現$CKV\mathbf{C}^{KV}$から明示的にバリューを復元することなく、一撃で$MHA(Q,K,V)\text{MHA}(Q, K, V)$を計算できます。

これが推論時に、「$WUVW^{UV}$は$WOW^{O}$に吸収される」という言葉の意味だと思われます。

まとめて

「W^{UK}の吸収」のまとめ

「$WUKW^{UK}$の吸収」の章にて、式(15)から下記の式が成立しています。

$$Q(i)K(i)⊤=((W(i)Q)′(CQ)T)T(CKV)T Q_{(i)}K_{(i)}^\top = \left((W^{Q}_{(i)})' (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T$$

ただし、まとめた式（吸収した式）に、$(W(i)Q)′(W^{Q}_{(i)})'$という名前をつけています。
$(W(i)Q)′(W^{Q}_{(i)})'$は下記で定義されています。

$$(W(i)Q)′=(W(i)UK)TW(i)UQ (W^{Q}_{(i)})' = (W^{UK}_{(i)})^T W^{UQ}_{(i)}$$

ただし、$(W(i)Q)′∈Rdc×dc′(W^{Q}_{(i)})' \in \mathbb{R}^{d_c \times d_c'}$。
また、$W(i)UK∈Rdh×dcW^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$、$W(i)UQ∈Rdh×dc′W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$

その上で、Scaled Dot-Product Attentionにおける、Attentionマップの式は下記のようになります。

$$Atn(i)(Q(i),K(i))=softmax(((W(i)Q)′(CQ)T)T(CKV)Tdh) \text{Atn}_{(i)}(Q_{(i)}, K_{(i)}) = \text{softmax} \left( \frac{\left((W^{Q}_{(i)})' (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T}{\sqrt{d_h}} \right)$$

「W^{UV}の吸収」のまとめ

「$WUVW^{UV}$の吸収」の章にて、式(23)から下記が成立しています。

$$MHA(Q,K,V)=(W(1)O)′(CKV)TAtn(1)T+(W(2)O)′(CKV)TAtn(2)T+⋯+(W(nh)O)′(CKV)TAtn(nh)T \begin{aligned} \text{MHA}(Q, K, V) &= (W^{O}_{(1)})' (\mathbf{C}^{KV})^T \text{Atn}_{(1)}^T \\ &\quad + (W^{O}_{(2)})' (\mathbf{C}^{KV})^T \text{Atn}_{(2)}^T \\ &\quad + \cdots + (W^{O}_{(n_h)})' (\mathbf{C}^{KV})^T \text{Atn}_{(n_h)}^T \\ \end{aligned}$$

ただし、まとめた式（吸収した式）に、$(W(i)O)′(W^{O}_{(i)})'$という名前をつけています。
$(W(i)O)′(W^{O}_{(i)})'$は下記で定義されています。

$$(W(i)O)′=W(i)OW(i)UV (W^{O}_{(i)})' = W^{O}_{(i)} W^{UV}_{(i)}$$

ただし、$(W(i)O)′∈Rd×dc(W^{O}_{(i)})' \in \mathbb{R}^{d \times d_c}$。
また、$W(i)O∈Rd×dhW^{O}_{(i)} \in \mathbb{R}^{d \times d_h}$、$W(i)UV∈Rdh×dcW^{UV}_{(i)} \in \mathbb{R}^{d_h \times d_c}$。

全体のまとめ

Multi-head Latent Attention (MLA)を計算するにあたり、必要なのは下記に全て示ました。

$$Q(i)K(i)⊤=((W(i)Q)′(CQ)T)T(CKV)T Q_{(i)}K_{(i)}^\top = \left((W^{Q}_{(i)})' (\mathbf{C}^{Q})^T \right)^T (\mathbf{C}^{KV})^T$$

$$MHA(Q,K,V)=(W(1)O)′(CKV)TAtn(1)T+(W(2)O)′(CKV)TAtn(2)T+⋯+(W(nh)O)′(CKV)TAtn(nh)T \begin{aligned} \text{MHA}(Q, K, V) &= (W^{O}_{(1)})' (\mathbf{C}^{KV})^T \text{Atn}_{(1)}^T \\ &\quad + (W^{O}_{(2)})' (\mathbf{C}^{KV})^T \text{Atn}_{(2)}^T \\ &\quad + \cdots + (W^{O}_{(n_h)})' (\mathbf{C}^{KV})^T \text{Atn}_{(n_h)}^T \\ \end{aligned}$$

ここからわかることとして、推論時には$(W(i)Q)′(W^{Q}_{(i)})'$と$(W(i)O)′(W^{O}_{(i)})'$という新しい重み行列を事前に計算してモデルにおいておくことで、キーバリュークエリを明示的に計算する必要がなく、少ない計算回数にて MLAを計算することが可能だとわかりました。

ここ！めちゃくちゃ面白いですよね！

以上が、論文中の

In addition, during inference, since $WUKW^{UK}$ can be absorbed into $WQW^Q$, and $WUVW^{UV}$ can be absorbed into $WOW^{O}$, we even do not need to compute keys and values out for attention.

だと考えています。

私の疑問点

MLAは主にKVキャッシュの圧縮や、クエリの圧縮による、メモリ効率を向上させることに非常に重きを置いた手法に見えています。

一方で、計算回数を減らすために、行列を吸収させる場合、吸収後の行列のパラメータ数は非常に多くなるように思います。

例えば下記の式で考えます。

$$(W(i)Q)′=(W(i)UK)TW(i)UQ (W^{Q}_{(i)})' = (W^{UK}_{(i)})^T W^{UQ}_{(i)}$$

ただし、$(W(i)Q)′∈Rdc×dc′(W^{Q}_{(i)})' \in \mathbb{R}^{d_c \times d_c'}$であり、また、$W(i)UK∈Rdh×dcW^{UK}_{(i)} \in \mathbb{R}^{d_h \times d_c}$、$W(i)UQ∈Rdh×dc′W^{UQ}_{(i)} \in \mathbb{R}^{d_h \times d_c'}$です。

このとき、論文中の数字（$dc=512,dc′=1536,dh=128d_c = 512, d_c' = 1536, d_h = 128$）を代入すると、$(W(i)Q)′(W^{Q}_{(i)})'$のパラメータ数は$786,432786,432$、$W(i)UKW^{UK}_{(i)}$のパラメータ数は$65,53665,536$、$W(i)UQW^{UQ}_{(i)}$のパラメータ数は$196,608196,608$となります。

すなわち、吸収前のパラメータ数は$65,536+196,608=262,14465,536 + 196,608 = 262,144$
吸収後のパラメータ数は$786,432786,432$となってしまいます。

これでは、推論時の必要パラメータが多くなってしまうため、 MLAの利点であるメモリ効率向上と相反する形に見えます。
という意味で、この吸収は実装されているのか？されていないのか？
されている場合は、メモリ効率の低減は、無視ができる程度なのか？そうでないのか？

この辺り、まだ私の中でもあまり固まっていない疑問ですが、詳しい方に教えていただけると嬉しいです。

なお、「吸収して大きな行列を作っているのではなく、2つの行列のまま2回行列を掛け算している」という考えは、私はないと思っております。なぜからその方法であればRoPEを導入した際に問題にならないからです。

その場合、普通に2つの行列演算の間に回転行列をかければいいだけのはずで、わざわざRoPEを適用するQuery・Key、適用しないQuery・Keyを分ける必要はないですよね？

ぜひ詳しい方！教えていただけると嬉しいです！

まとめ

読んでくださってありがとうございました！
MLAの「be absorbed」の部分の検証をしたいと思っていた方の参考になれば幸いです。

また、ぜひ私の疑問点にも、ご回答をもらえると嬉しいです。

Attentionについてのおすすめ書籍

機械学習エンジニアのためのTransformers ―最先端の自然言語処理ライブラリによるモデル開発
Hugging Faceの開発者が書いた本ですので、信頼できます。
当然のことながら、Attentionレイヤーの詳細な解説に関しても記載してあり、理論・基礎を学ぶ上でも非常におすすめの書籍です。

大規模言語モデル入門
大規模言語モデル入門Ⅱ〜生成型LLMの実装と評価
よく紹介させていただいておりますが、こちらの書籍は、LLMのファインチューニングから、RLHF、RAG、分散学習にかけて、本当に幅広く解説されており、いつも参考にさせていただいております。
特に1冊目は、TransformerのAttentionレイヤーの数式に関してもわかりやすく説明されておりおすすめです。