Multi-Head Attentionとは？同じ文脈を複数の重みで読む
Self-Attentionの計算を一度だけ行うと、各トークンから見た「どこをどれだけ参照するか」は一組の重みになる。Multi-Head Attentionは、Query・Key・Valueを別々の学習済み行列で複数組へ変換し、Attentionを並行して計算する。最後に各組の結果を連結し、もう一度線形変換して次の層へ渡す仕組みだ。
「ヘッドが8個なら、同じ計算を8回繰り返しているだけでは？」という疑問が出る。違いは各ヘッドに使う射影行列だ。同じ入力でも別の行列を通るため、参照先の重みや集める情報が一致する必要はない。ただし「このヘッドは必ず文法、このヘッドは必ず意味を担当する」と事前に役割が固定されるわけでもない。
1つの入力から、ヘッドごとに異なるQ・K・Vを作る
原論文『Attention Is All You Need』では、入力のQuery・Key・Valueに対し、ヘッドごとに異なる学習済み線形射影を適用する。ヘッド の計算は次の形だ。
各ヘッドが同じトークン列を受け取っても、、、 が違えば、照合に使う表現も結果も変わる。各ヘッドの出力を連結し、出力用の学習済み行列 を掛ける。
連結したまま終わりではない。最後の射影は、複数ヘッドから集めた情報をモデルが使う次元へ戻し、混ぜ合わせる役割を持つ。
2ヘッドの重みが違うと、集める値も変わる
次はAttentionの中身だけを1次元・2位置へ縮めた説明用の例だ。各ヘッドへ入るQueryを 1 とし、Keyから得たスコアが第1ヘッドで [2, 0]、第2ヘッドで [0, 2] になるように置く。両ヘッドのValueは説明のため [10, 0] とする。実際の射影行列を学習するコードではなく、射影後にスコアが異なると何が起きるかを見る計算だ。
import math
for scores in ([2.0, 0.0], [0.0, 2.0]): values = [10.0, 0.0] unnormalized = [math.exp(score) for score in scores] total = sum(unnormalized) weights = [value / total for value in unnormalized] output = sum(weight * value for weight, value in zip(weights, values)) print([round(weight, 3) for weight in weights], round(output, 3))実行結果:
[0.881, 0.119] 8.808[0.119, 0.881] 1.192第1ヘッドは最初の位置へ約88%の重みを置き、第2ヘッドは2番目の位置へ約88%を置く。同じValueの並びでも出力は約8.808と約1.192に分かれる。完全なMulti-Head Attentionなら、この2つを連結したあと、さらに出力射影を適用する。例の数字は仕組みを示す手作りの値で、実在するLLMのAttention重みではない。
ヘッドを増やせば単純に性能が上がるわけではない
原論文の基本モデルでは、モデル全体の幅512を8ヘッドに分け、各ヘッドのKey・Valueを64次元にした。単一ヘッドを8個、幅512のまま複製したわけではない。幅を分けることで、原論文は全体の計算量を幅512の単一ヘッドと同程度に保ったと説明している。
ヘッド数は設計上の選択であり、多ければ必ず良いとは言えない。各ヘッドの次元、モデル全体の幅、学習済み射影、利用するAttentionの方式も結果に関わる。PyTorchの MultiheadAttention は embed_dim と num_heads を指定し、埋め込み次元をヘッド数で分ける実装だ。別方式のモデルまで、すべてこの分け方だと考えないほうがよい。
Self-Attentionや因果マスクとの関係
Self-Attentionは「同じ系列からQuery・Key・Valueを作り、系列内を参照する」という入力間の関係を表す。Multi-Headは、そのAttentionを複数の射影で並行に行う構成を表す。両者は対立する方式ではない。Transformerのデコーダーでは、複数ヘッドのSelf-Attentionに因果マスクをかけ、未来の位置を参照させない。
KV Cacheで保存するKeyとValueも、Attentionの層とヘッドに対応する。ここで保存されるのはトークンの文字列ではなく、次の位置の計算で再利用する数値表現だ。
Multi-Head Attentionは、同じ文脈を単一の重み付けだけで集約しないための構成だ。複数の学習済み射影による出力を作り、連結と出力射影で組み合わせる。各ヘッドの役割を人が最初から指定する仕組みではないことも、併せて押さえておきたい。