トランスフォーマーの論文を読む⑥多頭
前回に引き続き、今回も論文を読みながらアテンションの仕組みを解説します。特に、Multi-Head Attention（マルチヘッド・アテンション、多頭アテンション）にフォーカスします。
多頭アテンションは、前回解説した「スケールされた内積アテンション」を複数使うことで、トークン間の様々な関係性に基づいたアテンションの処理を行います。
この辺りまで読み進めると、クエリ、キー、バリューに関して、その意味がより明確になってきます。
単純なアテンション
3.2.2 Multi-Head Attention は次のように始まります。
アテンションの計算を次元数$${d_\text{model}}$$のクエリ、キー、バリューに対して一度だけ実行するよりもクエリ、キー、バリューを$${h}$$個の異なる（学習可能な）線形変換によって次元数$${d_k}$$、$${d_k}$$、$${d_v}$$に投影することが有益であるとわかりました。
Instead of performing a single attention function with $${d_\text{model}}$$-dimensional keys, values and queries, we found it beneficial to linearly project the queries, keys and values $${h}$$ times with different, learned linear projections to $${d_k}$$, $${d_k}$$ and $${d_v}$$ dimensions, respectively.
ここで初めて、クエリ、キー、バリューが何を意味しているのかが明らかになりました。まず、最初に次元数$${d_\text{model}}$$に注目してください。この次元数は、セクション3.1 で述べられているようにトークンの埋め込みの次元数で論文では、$${d_\text{model}=512}$$が例として挙げられています。そして、全てのサブレイヤーや残差結合が出力する次元数となっています。
つまり、「アテンションの計算を次元数$${d_\text{model}}$$のクエリ、キー、バリューに対して一度だけ実行するよりも」という部分では、クエリ、キー、バリューは埋め込みベクトルそのものを意味しています。
単純化した例として「犬が歩く」という文章を考えます。これをトークン化したものが「犬」、「が」、「歩く」だとします。そして、この三つのトークンそれぞれが512次元のベクトルで表現されているとします。
我々が知りたいのは、各トークン間の関係の強さです。
まず、「犬」から文章内の各トークンとの関係の強さを計算すると、「犬」と「犬」、「犬」と「が」、「犬」と「歩く」の間の関係の強さを計算することになります。
「犬」がクエリの場合、「犬」と「が」と「歩く」がキーです。これらのキーに対するバリューも、それぞれ「犬」と「が」と「歩く」となります。後で見るように、実際にはもう１ステップあるのですが、ここでは埋め込みベクトルどうしで直接に内積を計算しているとしています。
以上から、一つのトークンに対して、512次元のクエリが一つ、512次元のキーが三つ、512次元のバリューが三つとなります。また、ここでのキーとバリューは同じベクトルです。この条件でアテンションを計算することになります。
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\dfrac{QK^\top}{\sqrt{d_k}}\right) V
$$
ここでのQ、K、Vは行列ですが、この例では、Qが「犬」の埋め込みを列ベクトル$${\vec{q}_1}$$として持っており、Kが「犬」、「が」、「歩く」の埋め込みを三列のベクトル$${\vec{k_1}, \vec{k_2}, \vec{k_3}}$$として持ち、Vも犬」、「が」、「歩く」の埋め込みを三列のベクトル$${\vec{v_1}, \vec{v_2}, \vec{v_3}}$$として持っていることになります。ただし、$${k_i = v_i \ \ (i = {1, 2, 3})}$$。
$$
Q = \begin{bmatrix}
- \vec{q}_1 -
\end{bmatrix}
$$
$$
K = \begin{bmatrix}
- \vec{k}_1 - \\
- \vec{k}_2 - \\
- \vec{k}_3 -
\end{bmatrix}
$$
$$
V = \begin{bmatrix}
- \vec{v}_1 - \\
- \vec{v}_2 - \\
- \vec{v}_3 -
\end{bmatrix}
$$
よって、「犬」というクエリだけを考えたアテンションの計算は次のようになります。
$$
\begin{aligned}
\text{Attention}(Q, K, V) &= \text{softmax}\left(\dfrac{QK^\top}{\sqrt{d_k}}\right) V \\[2ex]
&= \text{softmax}\Biggl(\dfrac{
\begin{bmatrix}
- \vec{q}_1 -
\end{bmatrix}\begin{bmatrix}
| & | & | \\
\vec{k}_1 & \vec{k}_2 & \vec{k}_3 \\
| & | & | \\
\end{bmatrix}
}{\sqrt{d_k}}\Biggr)\begin{bmatrix}
- \vec{v}_1 - \\
- \vec{v}_2 - \\
- \vec{v}_3 -
\end{bmatrix} \\[5ex]
&= \text{softmax}\left(\dfrac{
\begin{bmatrix}
\vec{q}_1 \cdot \vec{k}_1 & \vec{q}_2 \cdot \vec{k}_2 & \vec{q}_3 \cdot \vec{k}_3 \\
\end{bmatrix}
}{\sqrt{d_k}}\right)\begin{bmatrix}
- \vec{v}_1 - \\
- \vec{v}_2 - \\
- \vec{v}_3 -
\end{bmatrix} \\[5ex]
&= \begin{bmatrix}
w_1 & w_2 & w_3
\end{bmatrix} \begin{bmatrix}
- \vec{v}_1 - \\
- \vec{v}_2 - \\
- \vec{v}_3 -
\end{bmatrix} \\[5ex]
&= w_1 \vec{v}_1 + w_2 \vec{v}_2 + w_3 \vec{v}_3
\end{aligned}
$$
つまり、「犬」の埋め込みベクトルとシーケンス内の全ての埋め込みベクトル間の関係の強さによって、重み$${w_1, w_2, w_3}$$が計算され、この重みに従ってバリューの埋め込みベクトルを加重平均したものが、更新された「犬」の埋め込みベクトルとなります。
もちろん、「犬」と「犬」の関係が一番強いので、「犬」の情報の多くは「犬」から継続されます。そして、「が」と「歩く」からもアテンションが計算する重みによって情報が追加されます。
アテンションの計算は他のトークンからの観点からも行われます。よって、それぞれのトークンの埋め込みにシーケンス内のトークンからの関連情報が抽出されていきます。
また、アテンションの計算は、行列によってすべてのトークンに対して並列に行われます。
$$
\begin{aligned}
\text{Attention}(Q, K, V) &= \text{softmax}\left(\dfrac{QK^\top}{\sqrt{d_k}}\right) V \\[2ex]
&= \text{softmax}\Biggl(\dfrac{
\begin{bmatrix}
- \vec{q}_1 - \\
- \vec{q}_2 - \\
- \vec{q}_3 -
\end{bmatrix}\begin{bmatrix}
| & | & | \\
\vec{k}_1 & \vec{k}_2 & \vec{k}_3 \\
| & | & | \\
\end{bmatrix}
}{\sqrt{d_k}}\Biggr)\begin{bmatrix}
- \vec{v}_1 - \\
- \vec{v}_2 - \\
- \vec{v}_3 -
\end{bmatrix} \\[5ex]
&= \text{softmax}\Biggl(\dfrac{
\begin{bmatrix}
\vec{q}_1 \cdot \vec{k}_1 & \vec{q}_1 \cdot \vec{k}_2 & \vec{q}_1 \cdot \vec{k}_3 \\[2ex]
\vec{q}_2 \cdot \vec{k}_1 & \vec{q}_2 \cdot \vec{k}_2 & \vec{q}_2 \cdot \vec{k}_3 \\[2ex]
\vec{q}_3 \cdot \vec{k}_1 & \vec{q}_3 \cdot \vec{k}_2 & \vec{q}_3 \cdot \vec{k}_3 \\[2ex]
\end{bmatrix}}{\sqrt{d_k}}\Biggr)\begin{bmatrix}
- \vec{v}_1 - \\
- \vec{v}_2 - \\
- \vec{v}_3 -
\end{bmatrix} \\
&= \begin{bmatrix}
w_{11} & w_{12} & w_{13} \\
w_{21} & w_{22} & w_{23} \\
w_{31} & w_{32} & w_{33}
\end{bmatrix}\begin{bmatrix}
- \vec{v}_1 - \\
- \vec{v}_2 - \\
- \vec{v}_3 -
\end{bmatrix} \\[5ex]
&= \begin{bmatrix}
w_{11} \vec{v}_1 + w_{12} \vec{v}_2 + w_{13} \vec{v}_3 \\
w_{21} \vec{v}_1 + w_{22} \vec{v}_2 + w_{23} \vec{v}_3 \\
w_{31} \vec{v}_1 + w_{32} \vec{v}_2 + w_{33} \vec{v}_3
\end{bmatrix}
\end{aligned}
$$
このようにして、各トークンの埋め込みベクトルにはシーケンス内の関連情報が集約されていきます。つまり、各トークンの埋め込みベクトルの情報が強化されます。
また、エンコーダやデコーダではアテンションの計算がN=6回繰り返されます。よって、アテンションの計算が行われるたびに、強化された埋め込みベクトル同士の関係を使ってさらに各トークンの埋め込みベクトルの情報が強化されていきます。
しかし、論文では「アテンションの計算を次元数$${d_\text{model}}$$のクエリ、キー、バリューに対して一度だけ実行するよりも」ということで、これでは不十分だと主張しているわけです。
なぜならば、上述のような単純に埋め込みベクトルをそのまま使ったアテンションの計算では複雑なトークン間の関係を捉えきれないからです。そもそもトークンの埋め込みを直接使って内積を計算するのは、埋め込みベクトルが似ているかどうかを判断しているに過ぎません。
例えば、「犬」と「猫」の埋め込みベクトルはある意味似ている部分もあるでしょう。しかし、「犬」と「歩く」はどうでしょうか。そこには文法的に主語と述語の関係がありますが、単純な内積でそのような関係が浮かび上がってくるのでしょうか。
これに対して論文は多頭アテンションを提案してきます。
多頭アテンション
論文は「クエリ、キー、バリューを$${h}$$個の異なる（学習可能な）線形変換によって次元数$${d_k}$$、$${d_k}$$、$${d_v}$$に投影することが有益であるとわかりました」と説明を続けています。
つまり、シーケンス内の埋め込みベクトルをそのまま使った単純なアテンションの計算はあまり役に立たないと主張しています。代わりに、シーケンス内の埋め込みベクトルに対して線形変換をしたものに対してアテンションを計算するとより良い結果が生まれるのがわかったからです。
では、具体的に$${h}$$個や$${d_k}$$、$${d_v}$$の値は何なのでしょうか。セクション3.2.2の最後に次のように書いてあります。
この研究では、$${h = 8}$$の並列アテンション層（ヘッド）を使用しています。これらの各ヘッドに対して、$${d_k = d_v = d_{\text{model}} / h = 64}$$を使用しています。各ヘッドの次元が削減されているため、総計算コストは全次元を持つ単一ヘッドアテンションのそれと似ています。
In this work we employ $${h = 8}$$ parallel attention layers, or heads. For each of these we use $${d_k = d_v = d_\text{model} /h = 64}$$. Due to the reduced dimension of each head, the total computational cost is similar to that of single-head attention with full dimensionality.
つまり、512次元の（512個の要素を持つ）埋め込みベクトルを線形変換して、64次元の（64個の要素を持つ）埋め込みベクトルに変換します。
8個の別々な変換によって、一つの埋め込みベクトルに対して8個の異なる側面（クエリ）が抽出される事になります。
論文の続きを読みます。
これらの射影されたバージョンのクエリ、キー、バリューに対して、アテンション機能を並列に実行し、$${d_v}$$次元の出力値を生成します。これらは連結され、再び射影され、図2に示されているように、最終的な値を得ます。
On each of these projected versions of queries, keys and values we then perform the attention function in parallel, yielding $${d_v}$$-dimensional output values. These are concatenated and once again projected, resulting in the final values, as depicted in Figure 2.
（左）スケールされた内積アテンション
（右）多頭アテンション
上図の右の中央を見ると「Scaled Dot-Product Attention」（スケールされた内積アテンション）が$${h}$$個の並列処理として描かれているのがわかります。この一つ一つで上図左のスケールされた内積アテンションの処理が行われているわけです。
上図の右の下部を見ると「Linear」（線形変換）が並列処理として描かれています。これも$${h}$$個あり事になります。$${h=8}$$で考えると、線形変換を行うための行列が8個あるので、結果として一つの512次元の埋め込みベクトルから8個の64次元の埋め込みベクトルを生成します。この処理は、クエリ、キー、バリューに対してそれぞれ行われます。
つまり、一つの埋め込みベクトルから8個の異なるクエリ（ベクトル）、キー（ベクトル）、バリュー（ベクトル）が生成した上で、アテンションの計算は、これらの生成された埋め込みベクトル間で行われる事になり、様々な側面からの関係の強さが数値化される事になります。
最終的にスケールされた内積アテンションの出力が「Concat」（連結）されます。これは8個の64次元のベクトル出力をまとめて512次元に戻すことを指しています。連結された埋め込みベクトルにさらに線形変換を加えて最終的な出力とします。以上の処理を行うことでアテンション層からの出力は常に512次元となります。
ただし、$${64 \times 8 = 512}$$なので、実際には512次元から512次元への線形変換を行った上で、8個の部分に分けて扱うことで計算をさらに高速化することができます。
論文は、続いて多頭アテンションの利点を述べています。
マルチヘッドアテンション（多頭アテンション）により、モデルは異なる表現部分空間の異なる位置からの情報に同時に注目することができます。単一のアテンションヘッドでは、平均化によってこれが妨げられます。
Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this.
ここで表現部分空間とはデータの特徴や情報が埋め込まれている多次元空間のことを指します。つまり、元々の埋め込みベクトルの空間から複数の異なるベクトル空間が作り出されていることを意味します。これは、元々の512次元の空間から8個の64次元の異なる部分空間が生成されたことを簡潔に表現しています。
自然言語処理において、単語や文の意味は、しばしば高次元のベクトル空間（つまり、埋め込み空間）にエンコードされます。異なる「表現部分空間」とは、この高次元空間の中で異なる方向性や次元を指し、異なる種類の情報（例えば、文法的特徴、意味的特徴など）を捉えることができます。
多頭アテンションでは、複数の「ヘッド」を使用して、異なる「表現部分空間」での情報と、文中の異なる「位置」からの情報に同時に注目します。これにより、モデルは同じデータから複数の異なる特徴や文脈を捉えることができ、単一のアテンションを使用する場合に比べて、より豊富な情報を抽出し、より複雑な関係を理解することが可能になります。
このような処理を回帰でやろうとしたらかなり複雑な処理となることが想像できます。また、並列処理ができないので処理速度もかなり遅くなるでしょう。
論文は次の数式で多頭アテンションを表現しています。
$$
\begin{aligned}
\text{MultiHead}(Q, K, V) &= \text{Concat}(\text{head}_1, \dots, \ \text{head}_h)W^O \\
\text{where } \text{head}_i &= \text{Attention}(QW^Q_i, \ KW^K_i, \ VW^V_i)
\end{aligned}
$$
ここで、$${W^Q_i}$$、$${W^K_i}$$、$${W^V_i}$$は、それぞれクエリ、キー、バリューに対する線形変換を意味します。$${i}$$は、ヘッドの番号で１から８となります。
$${W^Q_i \in \mathbb{R}^{d_\text{model} \times d_k}}$$は、クエリ線形変換
$${W^K_i \in \mathbb{R}^{d_\text{model} \times d_k}}$$は、キー線形変換
$${W^V_i \in \mathbb{R}^{d_\text{model} \times d_v}}$$は、バリュー線形変換
$${W^O \in \mathbb{R}^{hd_v \times d_\text{model}}}$$は、最終出力の線形変換
前述したように、実際には$${d_k = d_v = 64}$$が使われました。
アテンションの適用
続いて論文のセクション3.2.3では、Transformerモデルにおける多頭アテンションの3つの異なる適用について説明しています。
エンコーダ・デコーダ・アテンション
論文はまず、エンコーダ・デコーダ・アテンションを解説しています。
「エンコーダ・デコーダ・アテンション」層では、クエリは前のデコーダー層から来ており、キーとバリューはエンコーダーの出力から来ます。これにより、デコーダー内のすべての位置が入力シーケンス内のすべての位置に注目することができます。これは、シーケンス・ツー・シーケンスモデル（sequence-to-sequence）（例えば、[38, 2, 9]など）における典型的なエンコーダ・デコーダ・アテンションメカニズムを模倣しています。
In "encoder-decoder attention" layers, the queries come from the previous decoder layer, and the memory keys and values come from the output of the encoder. This allows every position in the decoder to attend over all positions in the input sequence. This mimics the typical encoder-decoder attention mechanisms in sequence-to-sequence models such as [38, 2, 9].
「エンコーダ・デコーダ・アテンション」とは、エンコーダが抽出した文脈がデコーダへと連結した部分を指します。
エンコーダから来る二つの入力は、エンコーダの出力である文脈（埋め込みシーケンス）からのキーとバリューです。また、デコーダ内からのクエリがあるので、デコーダーが生成しているシーケンスのすべての位置がエンコーダからの入力シーケンス内のすべての位置に注目することができます。
つまりは、出力シーケンスが入力文章からの文脈を取り入れるためにエンコーダ・デコーダ・アテンションが使われています。また、このような仕組みは、以前に解説したRNNエンコーダ・デコーダでよく使われたものが由来となっています。
自己アテンション
次に論文は自己アテンションを解説しています。
エンコーダーには自己アテンション層が含まれています。自己アテンション層では、すべてのキー、バリュー、クエリが同じ場所から来ます。この場合、それはエンコーダーの前の層の出力です。エンコーダー内の各位置は、エンコーダーの前の層のすべての位置に注目することができます。
The encoder contains self-attention layers. In a self-attention layer all of the keys, values and queries come from the same place, in this case, the output of the previous layer in the encoder. Each position in the encoder can attend to all positions in the previous layer of the encoder.
エンコーダの自己アテンションは、入力文章のシーケンスからの文脈を抽出する中心的な役割を果たしています。ここでは、キー、バリュー、クエリが同じシーケンスから来ています。
よって、各位置（トークン）は入力シーケンスのすべての位置（トークン）に注目することができ、入力シーケンスの文脈（関連情報）をそれぞれのトークンの埋め込みベクトルへと取り込んでいきます。
マスクされた自己アテンション
最後に論文は、デコーダが行う自己アテンションについて解説しています。
同様に、デコーダー内の自己アテンション層は、デコーダー内の各位置がその位置までのデコーダー内のすべての位置に注目することを可能にします。自己回帰性質を保持するために、デコーダー内の左方向への情報流を防ぐ必要があります。これを実装するために、スケールされた内積アテンションで、ソフトマックスの入力において不正な接続に対応するすべての値をマスキングアウト（−∞に設定）します。図2を参照してください。
Similarly, self-attention layers in the decoder allow each position in the decoder to attend to all positions in the decoder up to and including that position. We need to prevent leftward information flow in the decoder to preserve the auto-regressive property. We implement this inside of scaled dot-product attention by masking out (setting to −∞) all values in the input of the softmax which correspond to illegal connections. See Figure 2.
マスクされた多頭アテンションについては以前に解説しましたが、このマスクは生成文章が将来のトークンからの情報がアテンションの計算に入ってくるのを防ぐ役割を果たしています。「自己回帰性質を保持するため」というのは、RNNのように自らが予測した情報に基づいて次のトークンを予測するという仕組みを指しています。
以前に解説しましたが、マスクされる位置に対してソフトマックスの入力を大きな負の値にすることでソフトマックスが計算する重みの値がほぼゼロになることで、最終的なバリューの加重平均から除外されるのがマスクをする仕組みとなっています。論文では、−∞に設定すると書いてありますが、実際には$${-10^9}$$などの十分に大きな負の値を使います。
マスクする以外は、自己アテンションとしての仕組みはエンコーダのものと同様で、生成シーケンス内の文脈を抽出するためのものです。よって、キー、バリュー、クエリはすべて同じシーケンスから来ています。
以上、アテンションの仕組みとどのように使われているかの解説になります。見事に回帰を必要とせずに文脈の処理をしていることが理解できました。改めて論文図１の全体像を見るとエンコーダとデコーダと多くの部分をカバーしてきました。
しかし、逆に位置ごとのフィードフォワード（Feed Forwardと書いてある青い箱）やエンコーダとデコーダの箱の外側はまだカバーしきれていません。
これらの詳細は、論文の続きを読むことである程度は明らかとなってきます。
次回予告
次回は、セクション３.３の位置ごとのフィードフォワードを読みながら解説を続けます。
お楽しみに！
