第10章 マルチヘッドアテンションと出力射影:複数の視点で同時に見る

こんにちは。ゆうせいです。

第9章で、自己注意という仕組みが一通り完成しました。しかし、実際のTransformerでは、この自己注意を一つだけ動かしているわけではありません。同じ処理を複数、並列に動かしています。これがマルチヘッドアテンションです。Transformer Explainerを開くと、Head 1、Head 2といったタブや切り替えが用意されており、ヘッドを切り替えるたびに、注目のパターンがまったく異なることに気づくはずです。なぜ一つでは足りないのでしょうか。複数動かした結果は、どのように一つにまとめられるのでしょうか。第10章では、この二つを解説します。ここまでで自己注意の中身は理解できていますので、この章は、その使い方の話になります。

なぜ一つの注意機構では足りないのか

一つの重み行列では、一つの観点しか学習できない

第9章までで見てきた自己注意では、各トークンについて、注意の重みが一組だけ計算されました。3番目のトークンは、1番目に20パーセント、2番目に50パーセント、自分自身に30パーセント注目する、という具合です。

ここで問題になるのは、文中のトークンどうしの関係は、一種類ではないという点です。

「The tired cat quickly ate the fresh fish」という文を考えます。

「ate」というトークンは、複数の異なる観点から、他のトークンと関係を持っています。

文法的な主語として、「cat」と関係しています。

文法的な目的語として、「fish」と関係しています。

動作の様態として、「quickly」と関係しています。

これらは、性質のまったく異なる関係です。一組の注意の重みでは、これらすべてを同時に表現できません。

主語に強く注目すれば目的語への注目が薄れ、目的語に強く注目すれば主語への注目が薄れます。合計が1になるという制約があるため、複数の関係を同時に強く保持することが、構造的に難しいのです。

複数の観点を並列に持たせる

そこで採られる解決策が、自己注意の仕組みを複数用意することです。

一つの自己注意の単位を、ヘッド(Head)と呼びます。

ヘッド1は主語と動詞の関係を捉える、ヘッド2は動詞と目的語の関係を捉える、ヘッド3は修飾関係を捉える、というように、それぞれが異なる観点を担当します。

ただし、「どのヘッドが何を担当するか」は、人間が指定するものではありません。それぞれのヘッドが独立した重み行列を持ち、学習の過程で、自然と異なる観点を担当するようになります。

なぜ自然と分かれるのかというと、すべてのヘッドが同じ観点を学習しても、モデル全体の性能向上に寄与しないためです。異なる情報を捉えるヘッドがあるほうが、次のトークンの予測精度が上がります。この圧力によって、ヘッドの役割が自然に分化していきます。

ヘッド数と次元数の関係

ここが、初学者が最もつまずく箇所です。丁寧に確認します。

次元を分割する方式

マルチヘッドにする際、素朴に考えると、768次元の自己注意を12個用意する、という設計が思い浮かびます。しかし、実際はそうではありません。

768次元を12個のヘッドに分割し、各ヘッドは64次元だけを担当します。

\frac{768}{12} = 64

第5章で、この64という数字が出てきたことを思い出してください。第7章のスケーリングで \sqrt{64} = 8 を使ったのも、この分割によるものです。

なぜ分割するのか

分割せずに768次元のヘッドを12個作れば、表現力は上がりそうに思えます。しかし、そうしない理由があります。

第一の理由は、計算量です。768次元のヘッドを12個作ると、パラメータ数も計算量も、およそ12倍になります。分割方式であれば、ヘッドを増やしても、全体の計算量はほぼ変わりません。

第二の理由は、実験的に、分割方式でも十分な性能が得られると確認されているためです。1ヘッドあたりの表現力は下がりますが、複数の観点を持てる利点が、それを上回ります。

つまり、マルチヘッドは、計算資源を増やす仕組みではなく、同じ計算資源をより効率的に使う仕組みなのです。

実装上の処理の流れ

実際の処理は、次の順序で行われます。

第一に、入力ベクトル(768次元)から、Query、Key、Valueを生成します。この時点では、いずれも768次元です。

第二に、それぞれを12個の塊に分割します。各塊は64次元です。

第三に、12組のQuery、Key、Valueについて、第6章から第9章までの処理を、それぞれ独立に実行します。

第四に、12個の出力(各64次元)を、横につなげて768次元に戻します。

第五に、出力射影と呼ばれる変換を適用します。

第一から第四までは、これまでの内容の繰り返しです。新しいのは、第五の出力射影だけです。

分割は並列に計算される

12個のヘッドの計算は、互いに独立しています。ヘッド1の計算結果は、ヘッド2の計算に影響しません。

したがって、12個の計算を同時に実行できます。GPUのような並列計算に強いハードウェアでは、この性質が大きな利点になります。

実装上は、ループで12回計算するのではなく、テンソルの形を工夫することで、一度の行列演算としてまとめて処理されます。

出力射影とは何か

12個のヘッドの出力を横につなげると、768次元のベクトルに戻ります。この段階で処理を終えてもよさそうですが、もう一つ、変換が加えられます。

これが、出力射影(Output Projection)です。

出力 = Concat(head_1, head_2, \ldots, head_{12}) \times W_O

W_O は、768行768列の重み行列です。これも学習によって獲得されます。

なぜ出力射影が必要なのか

単に横につなげただけの状態には、問題があります。

つなげただけのベクトルは、0番目から63番目の要素がヘッド1の出力、64番目から127番目がヘッド2の出力、という具合に、区画が分かれた状態です。

各ヘッドの出力は、それぞれ独立した観点の情報です。しかし、これらは、統合されて初めて意味を持ちます。「主語との関係」と「目的語との関係」を、組み合わせて解釈する必要があるのです。

出力射影は、この区画を混ぜ合わせ、ヘッドをまたいだ情報の統合を行う役割を持ちます。

行列の掛け算では、出力の各要素が、入力のすべての要素の重み付き和になります。したがって、出力射影を通した後のベクトルでは、各要素が、12個すべてのヘッドの情報を含んだものになります。

第5章のアダマール積との対比

第5章までの内容を踏まえると、ここで通常の行列積が使われる理由が理解できます。

もしアダマール積を使えば、要素ごとの独立した調整になり、ヘッドをまたいだ混合が起きません。区画が分かれたままです。

新しい特徴を作り出したい、情報を混ぜ合わせたい場面では、通常の行列積が使われます。出力射影は、まさにその場面にあたります。

パラメータ数を確認する

1つのTransformerブロックにおける、注意機構全体のパラメータ数を計算します。

Query生成用の行列:768 \times 768 = 589824

Key生成用の行列:768 \times 768 = 589824

Value生成用の行列:768 \times 768 = 589824

出力射影用の行列:768 \times 768 = 589824

合計:589824 \times 4 = 2359296

約236万個です。第5章で計算した約177万個に、出力射影の約59万個が加わった形です。

ブロックが12個ありますから、モデル全体では、次のようになります。

2359296 \times 12 = 28311552

約2831万個のパラメータが、注意機構に使われていることになります。

なお、ここではバイアス項を省略しています。実際のGPT-2では、各線形変換にバイアスが加わるため、実際の値はこれよりわずかに多くなります。

Transformer Explainerで観察できること

Transformer Explainerでは、ヘッドを切り替えて、それぞれの注意パターンを比較できます。

次の点を確認してみてください。

一つ目は、ヘッドごとに、注目のパターンがまったく異なるという点です。あるヘッドでは直前のトークンに強く注目し、別のヘッドでは離れた位置のトークンに注目する、といった違いが観察できます。

二つ目は、対角線に沿った注目パターンを持つヘッドが存在するという点です。これは、直前のトークンに注目するヘッドであり、多くのモデルで観察される典型的なパターンです。

三つ目は、最初のトークンに強く注目するヘッドが存在するという点です。文頭のトークンは、特に参照する先がないため、注目の受け皿として使われることがあります。この現象は、アテンションシンクと呼ばれ、近年、研究の対象になっています。

四つ目は、明確なパターンを持たないヘッドも存在するという点です。すべてのヘッドが解釈しやすい役割を持つわけではありません。

初学者がつまずきやすい点

つまずき1:ヘッド数を増やせば性能が上がると考える

ヘッド数を増やすと、1ヘッドあたりの次元数が減ります。

GPT-2(small)で、ヘッド数を24にすると、1ヘッドあたりは32次元になります。

\frac{768}{24} = 32

次元数が減ると、1つのヘッドが表現できる情報が少なくなります。観点は増えますが、それぞれの観点の解像度が下がる、というトレードオフです。

適切なヘッド数は、モデルの規模や課題によって異なります。研究では、学習後に一部のヘッドを削除しても性能がほとんど低下しない、という報告もあります。すべてのヘッドが等しく重要なわけではない、という点は押さえておいてください。

つまずき2:ヘッドの役割を決めつける

「このヘッドは主語を担当している」といった解釈は、可視化の結果から後付けで行われるものです。

学習の前に役割が決まっているわけではありませんし、すべてのヘッドが人間に理解できる役割を持つわけでもありません。

分かりやすいパターンを持つヘッドは、可視化の説明でよく取り上げられますが、それは全体の一部にすぎない、と理解しておいてください。

つまずき3:QKVの生成が3つの別々の行列だと思い込む

概念としては、Query、Key、Valueそれぞれに専用の行列があります。

しかし、実際のGPT-2の実装では、これら三つの行列が、一つの大きな行列としてまとめられています。768行2304列(768の3倍)の行列を一度掛けてから、結果を三つに分割する方式です。

768 \times 3 = 2304

数学的には同じ計算ですが、行列積を一度で済ませられるため、計算効率が上がります。

第5章と第9章の演習コードで、c_attn という一つの層からQKVを取り出して分割していたのは、この実装によるものです。

演習

演習1:次元数を計算する

次の条件のモデルについて、1ヘッドあたりの次元数と、スケーリングに使う値を計算してください。

埋め込み次元が1024、ヘッド数が16のモデル。

埋め込み次元が2048、ヘッド数が16のモデル。

演習2:マルチヘッドアテンションを実装する

第9章の演習3を拡張し、マルチヘッド化してみてください。

import torch
import torch.nn.functional as F
import math

torch.manual_seed(0)

# 設定
seq_len = 5
d_model = 768
n_heads = 12
d_k = d_model // n_heads

print(f"埋め込み次元: {d_model}")
print(f"ヘッド数: {n_heads}")
print(f"1ヘッドあたりの次元: {d_k}")
print(f"スケーリング係数: {math.sqrt(d_k):.2f}")

# 入力
x = torch.randn(seq_len, d_model)

# QKVを生成するための行列
W_q = torch.randn(d_model, d_model) * 0.02
W_k = torch.randn(d_model, d_model) * 0.02
W_v = torch.randn(d_model, d_model) * 0.02
W_o = torch.randn(d_model, d_model) * 0.02

# QKVを生成
Q = x @ W_q
K = x @ W_k
V = x @ W_v
print(f"\nQKV生成後の形: {Q.shape}")

# ヘッドに分割
Q = Q.view(seq_len, n_heads, d_k).transpose(0, 1)
K = K.view(seq_len, n_heads, d_k).transpose(0, 1)
V = V.view(seq_len, n_heads, d_k).transpose(0, 1)
print(f"ヘッド分割後の形: {Q.shape}")

# 各ヘッドで自己注意を計算
scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k)
mask = torch.tril(torch.ones(seq_len, seq_len))
scores = scores.masked_fill(mask == 0, float('-inf'))
weights = F.softmax(scores, dim=-1)
head_outputs = weights @ V
print(f"各ヘッドの出力の形: {head_outputs.shape}")

# ヘッドを結合
concatenated = head_outputs.transpose(0, 1).contiguous().view(seq_len, d_model)
print(f"結合後の形: {concatenated.shape}")

# 出力射影
output = concatenated @ W_o
print(f"出力射影後の形: {output.shape}")

# ヘッドごとの注意パターンの違いを確認
print("\n=== ヘッドごとの注意パターン(最終トークンの行)===")
for h in [0, 1, 2]:
    print(f"ヘッド{h}: {weights[h, -1].numpy().round(3)}")

演習3:実際のGPT-2で注意パターンを観察する

from transformers import GPT2Tokenizer, GPT2Model
import torch

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2Model.from_pretrained("gpt2", output_attentions=True)

text = "The cat sat on the mat"
inputs = tokenizer(text, return_tensors="pt")
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

outputs = model(**inputs)
attentions = outputs.attentions

print(f"ブロック数: {len(attentions)}")
print(f"各ブロックの注意行列の形: {attentions[0].shape}")
print(f"トークン: {tokens}\n")

# 最初のブロックの、いくつかのヘッドを比較
layer = 0
for head in [0, 3, 7, 11]:
    attn = attentions[layer][0, head]
    last_row = attn[-1]
    top_idx = last_row.argmax().item()
    print(f"層{layer} ヘッド{head}: 最終トークンが最も注目するのは「{tokens[top_idx]}」")
    print(f"  重み分布: {last_row.detach().numpy().round(3)}")

演習の解答例と解説

演習1の解答は、次の通りです。

埋め込み次元1024、ヘッド数16の場合。

1ヘッドあたりの次元数:\frac{1024}{16} = 64

スケーリング係数:\sqrt{64} = 8

埋め込み次元2048、ヘッド数16の場合。

1ヘッドあたりの次元数:\frac{2048}{16} = 128

スケーリング係数:\sqrt{128} \approx 11.31

一つ目の例では、埋め込み次元がGPT-2(small)より大きいにもかかわらず、1ヘッドあたりの次元数は同じ64になります。ヘッド数も同時に増やしているためです。

多くのモデルで、1ヘッドあたりの次元数を64または128に保つ設計が採られています。この値が、経験的に扱いやすい範囲だとされています。

演習2では、テンソルの形が段階的に変化していく様子が確認できます。特に、view と transpose によってヘッドに分割し、また元に戻す部分が、実装の要点です。

また、ヘッドごとに注意の重みが異なることが、数値として確認できます。同じ入力から出発しているにもかかわらず、重み行列が異なるため、まったく違う注目パターンになります。

演習3では、実際に学習されたGPT-2のヘッドが、どのような注目パターンを持つかを観察できます。層やヘッドを変えて、どのような違いが現れるかを試してみてください。

講師向けの補足

この章で最も重要なのは、「なぜ複数のヘッドが必要か」という動機の部分です。

「主語との関係」と「目的語との関係」を同時に強く保持できない、という制約を、実際の文章を使って示してください。合計が1になるという制約があるため、複数の関係を同時に扱えない、という説明が最も分かりやすいものです。

次元の分割については、混乱が生じやすい箇所です。「768次元のヘッドが12個」ではなく「768次元を12分割」であることを、図で示すことを推奨します。ホワイトボードに768の長方形を描き、それを12個に区切る図が有効です。

演習3で、実際のGPT-2の注意パターンを観察する時間は、必ず確保してください。ヘッドによって注目先が異なることを、自分の目で確認する体験が、マルチヘッドの意義を実感させます。

理解度を確認する問い

第10章の内容を、次の問いで確認してみてください。

一つの注意機構では不十分である理由を、具体的な文章の例を使って説明してみてください。

768次元のヘッドを12個作るのではなく、768次元を12分割する方式が採られている理由は何でしょうか。

出力射影が必要な理由を、ヘッドの出力を横につなげただけの状態が持つ問題から説明してみてください。

ヘッド数を増やすと、何が増えて、何が減るでしょうか。

まとめ

マルチヘッドアテンションとは、自己注意の仕組みを複数、並列に動かす構造です。1つの単位をヘッドと呼びます。

複数のヘッドが必要な理由は、トークンどうしの関係が一種類ではないためです。注意の重みは合計が1になるという制約を持つため、一組の重みでは、主語との関係、目的語との関係、修飾関係といった複数の観点を同時に強く保持できません。

GPT-2(small)では、768次元を12個のヘッドに分割し、各ヘッドが64次元を担当します。ヘッドごとに768次元を持たせるのではなく分割する方式を採ることで、計算量を増やさずに複数の観点を獲得できます。

各ヘッドの出力は、横につなげて768次元に戻された後、出力射影と呼ばれる行列変換を通ります。この変換により、ヘッドごとに区画が分かれていた情報が混ぜ合わされ、統合された表現になります。

注意機構全体のパラメータ数は、1ブロックあたり約236万個、モデル全体では約2831万個です。

次の第11章では、残差接続を扱います。ここまでで完成した注意機構の出力が、そのまま次に渡されるのではなく、入力と足し合わされてから渡される仕組みと、その理由を解説します。

第11章に進む前に、Transformer Explainerでヘッドを切り替えながら、注意パターンの違いを観察しておいてください。ヘッドによってまったく異なるパターンが現れることを実感しておくと、なぜ出力射影で統合が必要なのかが、理解しやすくなります。

セイ・コンサルティング・グループでは新人エンジニア研修のアシスタント講師を募集しています。

投稿者プロフィール

山崎講師
山崎講師代表取締役
セイ・コンサルティング・グループ株式会社代表取締役。
岐阜県出身。
海外放浪の末、2000年創業、2004年会社設立。
IT企業向け人材育成研修歴業界歴20年以上。
すべての無駄を省いた費用対効果の高い「筋肉質」な研修を提供します!
この記事に間違い等ありましたらぜひお知らせください。

学生時代は趣味と実益を兼ねてリゾートバイトにいそしむ。長野県白馬村に始まり、志賀高原でのスキーインストラクター、沖縄石垣島、北海道トマム。高じてオーストラリアのゴールドコーストでツアーガイドなど。現在は野菜作りにはまっている。