マルチヘッドアテンション

1回のアテンションを複数に分ける発想

アテンション関数は、クエリ・キー・値を受け取って出力を計算する仕組みです。素朴なやり方では、モデルの内部表現と同じd_model次元のクエリ・キー・値を使い、アテンション関数を1回だけ実行します。

論文 Attention Is All You Need は、この1回の計算の代わりに、別のやり方の方が有益だと報告しました。クエリ・キー・値を、学習によって得られる互いに異なる線形写像でh回写し、それぞれd_k次元、d_k次元、d_v次元の表現にします。そして写した各組に対して、アテンション関数を並列に実行します。この1組ごとの計算をヘッドと呼び、複数のヘッドを並べた構成をマルチヘッドアテンションと呼びます。

なぜ複数のヘッドが役立つのか

単一のアテンションヘッドでは、注意の結果が1つの平均にまとめられます。論文は、この平均化が、性質の異なる情報を同時に拾うことを妨げると述べています。

マルチヘッドアテンションでは、ヘッドごとに異なる線形写像を通すため、各ヘッドは異なる表現部分空間を見ることになります。その結果、モデルは異なる位置にある、異なる表現部分空間の情報に同時に注意を向けられます。

計算の手順と数式

各ヘッドの計算は次の式で表されます。

head_i = Attention(Q W_i^Q, K W_i^K, V W_i^V)

W_i^Q、W_i^K、W_i^Vはヘッドiに固有の学習可能なパラメータ行列で、大きさはそれぞれd_model×d_k、d_model×d_k、d_model×d_vです。各ヘッドはd_v次元の出力を返します。

次に、h個のヘッドの出力を連結し、もう一度線形写像を掛けて最終的な値を得ます。

MultiHead(Q, K, V) = Concat(head_1からhead_hまでのすべてのヘッドの出力) W^O

W^Oの大きさはh・d_v×d_modelで、連結したh・d_v次元の表現をd_model次元へ写す役割を持ちます。手順をまとめると、写像、並列のアテンション、連結、再写像の4段階です。

論文での設定と計算コスト

論文では、h=8個の並列なアテンション層、つまり8個のヘッドを使いました。各ヘッドの次元はd_k=d_v=d_model/h=64です。

ヘッドを増やしても計算量が大きく膨らまないのは、各ヘッドの次元をd_model/hに縮めているためです。このため全体の計算コストは、フル次元で単一ヘッドのアテンションを行う場合と同程度に収まります。複数の視点を持たせながら、計算の負担はほぼ変えない点がこの設計の要点です。

学習のポイント

マルチヘッドアテンションはシラバス2024のキーワードに含まれる語です。式の形だけでなく、なぜ分けるのか(平均化の回避)と、分けても計算コストが同程度で済む理由(各ヘッドの次元の縮小)を、組にして押さえておくと理解が安定します。