Self-Attention完全解剖(後編)――因果マスクとMulti-Headの必然。なぜ未来を目隠しし、32個の複眼で世界を見るのか?
1. 素朴な疑問:なぜ超能力を「目隠し」し、視線を「32分割」するのか?
前話(第4話)で、私たちはSelf-Attentionの美しい幾何学を解き明かしました。 $Q$ と $K$ の内積を取り、$\sqrt{d_k}$ でスケーリングし、Softmaxを通して $V$ を加重平均する。この一撃によって、静止していた単語は周囲の文脈を吸い込んだ「生きた動的ベクトル」へと進化します。
しかし、前話で実装した素朴なSelf-AttentionをそのままGPTのような生成AIに組み込むと、2つの致命的な問題によってモデルは即座に機能不全に陥ります。
- なぜ「未来」を見えなくするのか? 推論時には1文字ずつしか出力できないのに、訓練時にはなぜ文章全体を一瞬で並列学習できるのか? 未来の単語をどうやって物理的に遮断しているのか?
- なぜ「ヘッド(Head)」を細かく切り分けるのか? 4096次元の巨大な注意機構を1つ作るのと、128次元の小さな注意機構を32個並べるのとでは、一体何が違うのか?
今回は、現代のあらゆる言語モデル(GPT-4、Claude、Llama 3など)が例外なく採用している**「因果マスク(Causal Masking)」と「Multi-Head Attention(MHA)」**の必然性を、テンソルの変形軌跡とともに徹底的に暴きます。
2. 因果マスク(Causal Mask):訓練を100倍高速化する「時間の防壁」
LLMを訓練する目的は、第1話で学んだ通り「過去の文脈から、次に来る未知の1単語を予測すること」です。
もしモデルが「The cat sat on the mat」という文章を学習するとき、2単語目の「cat」を処理している最中に、3単語目の「sat」や6単語目の「mat」を視界に入れてしまったら何が起きるでしょうか?
【カンニングの発生】
入力: "The cat sat on the mat"
「cat」の次の単語を予測したい!
↓
素朴なSelf-Attention:「あ、後ろの配列の3番目に 'sat' って書いてあるのが見えます!」
↓
モデル:「じゃあ、後ろの文字をそのままオウム返しに出力すれば100点満点ですね」
モデルは文章の論理や意味を何一つ理解することなく、**「ただ後ろのインデックスの文字をそのまま書き写すだけのチート(カンニング)」**を覚えて学習を終えてしまいます。これでは、文章生成の知能など絶対に宿りません。
なぜ1文字ずつ隠して入力しては駄目なのか?
「カンニングを防ぎたいなら、訓練のときも人間と同じように1文字ずつ入力すればいいのではないか?」 そう思うかもしれません。
ステップ1: "The" → 次を予測
ステップ2: "The cat" → 次を予測
ステップ3: "The cat sat" → 次を予測
しかし、この逐次入力方式(RNNやLSTMがやっていた方式)を行うと、GPUの超並列アーキテクチャが死にます。 数兆トークンもの巨大なデータを、前の単語が終わるのを待って1ステップずつ流していたら、学習が終わるまでに数百年かかってしまうからです。
GPUを100%フル稼働させるには、**「数千単語からなる1本の文章を、一度にドカンと丸ごとGPUメモリに流し込み、全単語の次単語予測を一斉に並列計算する」**必要があります。
この「丸ごと並列に入力しつつ、各単語には絶対に未来の単語を見せない」という矛盾をエレガントに解決した数学的装置が、**「因果マスク(Causal Mask / Auto-regressive Mask)」**です。
なぜ「0」ではなく「マイナス無限大($-\infty$)」で埋めるのか?
因果マスクの実装ルールは、驚くほどシンプルです。
$Q K^T / \sqrt{d_k}$ で計算されたスコア行列(形状: [Seq_Len, Seq_Len])の右上半分(未来の単語の位置)に対して、強制的に「$-\infty$(マイナス無限大)」を代入します。
なぜ「0」を入れてはいけないのでしょうか?
理由は、この直後に控えている**Softmax の定義**にあります。
$$\text{Softmax}(z_i) = \frac{e^{z_i}}{\sum e^{z_j}}$$
もしスコアに「0」を入れると、$e^0 = 1$ となり、未来の単語にも有限の確率(注意の重み)が割り振られてしまいます。 しかし、スコアに「$-\infty$」を叩き込むと何が起きるでしょうか?
$$e^{-\infty} = \lim_{x \to -\infty} e^x = 0$$
指数関数 $e^x$ は、マイナス無限大で完全に**「真のゼロ($0$)」**へと収束します。
この極限の物理によって、Softmaxを通過した後のアテンション重み行列 $A$ の右上半分は、1ビットの誤差もなく完全な「0.000…」へと抹殺されます。 未来の単語のValueベクトルに掛け算される重みが絶対のゼロになるため、どんなに文章を丸ごと並列に入力しても、情報が未来から過去へと逆流することは物理的に絶対に起きないのです。
触って試す:因果マスクの体験
🎭 因果マスクの体験
文「The cat sat on the mat」について、「見る側」の単語を選ぶと、その単語がどの単語を「見られる」かが下の表に表示されます。行=見る側、列=見られる側です。未来の単語は因果マスクで真っ黒に塗りつぶされ、注意の重みは厳密に0になります(パーセントは説明用の例です)。
3. Multi-Head Attentionの必然:単眼の視線から「32個の複眼」へ
因果マスクによって時間の矢を正しく固定した次に立ちはだかるのが、**「アテンションの単眼問題」**です。
前話で作った素朴なアテンション(Single-Head Attention)では、4096次元の巨大な重み行列を使って、1単語につき1つのアテンション重み(確率分布)を計算していました。
しかし、人間の高度な言語理解を思い浮かべてください。1つの単語は、文脈の中で**同時に何種類もの異なる役割(多重の関係性)**を持っています。
【例文】
"The animal didn't cross the street because it was too tired."
(その動物は通りを横断しなかった。なぜなら【それ】は疲れすぎていたからだ。)
この文章の中にある代名詞 「it(それ)」 に注目してください。 この「it」が文脈を正しく解釈するためには、同時に以下の全く異なる2つの視線を向けなければなりません。
もしアテンションヘッドが1つしかなかったらどうなるでしょうか? アテンション重みは合計1.0の「1つの確率分布」しか出力できません。 「animal」に注意を80%向けたら、「was」に向ける注意は20%以下に削られてしまいます。「意味の結合」を強めると「文法の結合」が見えなくなり、「文法の結合」を強めると「意味の結合」が見えなくなるという、視線のトレードオフ(共倒れ)に陥るのです。
昆虫の複眼のように、**「文法を見る目」「代名詞の指す先を見る目」「直前の単語を見る目」「遠くのトピックを見る目」**という複数の独立した視線を同時に持たせたい。 これこそが、**Multi-Head Attention(マルチヘッド・アテンション)**が考案された最大の動機です。
4. 計算量の奇跡:ヘッドを分割しても「計算コストは全く増えない」
ここで、ほぼすべてのエンジニアが抱く疑問があります。 「ヘッドを32個に増やしたら、計算量もGPUメモリも32倍に跳ね上がるのではないか?」
結論から言えば、計算量は1ミリも増えません。単一ヘッドのときと完全に同一です。
なぜそんな魔法が可能なのでしょうか? それは、モデルの隠れ層次元 $D$(例: 4096)を、ヘッドの数($H=32$)で綺麗に「頭割り(分割)」するからです。
【次元の分割の物理】
モデル全体の隠れ層次元 : D = 4096
アテンションヘッド数 : H = 32
各ヘッドが担当する次元 : d_k = D / H = 4096 / 32 = 128 次元!
数式で確認してみましょう。 内積の計算量は、ベクトルの長さに比例します。
- 単一ヘッド(4096次元): 計算コスト $= S^2 \times 4096$
- 32ヘッド(各128次元): 計算コスト $= 32 \times (S^2 \times 128) = S^2 \times (32 \times 128) = S^2 \times 4096$
掛け算の順番を組み替えただけなので、全体の計算量は1FLOPも変わっていません。
同じ電気代、同じ計算時間、同じシリコン面積のまま、 「1つの広大な4096次元空間で1本の視線をぼんやり泳がせる」のをやめ、 **「128次元というコンパクトでシャープな『部分空間(Subspace)』を32個同時に切り拓き、32本の異なる視線で文章を立体スキャンする」**という圧倒的な表現力の跳躍を手に入れたのです。
5. 出力射影(Output Projection $W_O$):散らばった視線の再統合
32個のヘッドがそれぞれの世界で計算を終えると、手元には**「32個の異なる文脈ベクトル(各128次元)」**が生まれます。
このバラバラになった視線を、どうやって次の層へ渡すのでしょうか?
- 連結(Concatenate): まず、32個の出力を横一列にガッチャンコと繋ぎ合わせます(Concat)。 $32 \times 128 =$ 4096次元となり、元のモデルの隠れ層次元 $D$ と全く同じサイズに戻ります。
- 出力射影(Output Projection $W_O$):
しかし、ただ横に並べただけでは、Head 1の情報とHead 2の情報は隣り合っているだけで、お互いに会話していません。
そこで最後に、重み行列 $W_O$(形状:
[4096, 4096]) を掛け算します。 この行列積によって、32個のヘッドが別々に拾ってきた情報同士が線形結合され、**「文法情報と意味情報が完全にブレンドされた単一の4096次元ベクトル」**へと昇華されるのです。
6. テンソルの視点:4階テンソルの変形ダンス(Shape追跡)
Multi-Head Attentionをコードで書くとき、すべてのエンジニアが一度は激突する壁が**「テンソルの次元の入れ替え(Transpose/Reshape)」**です。
バッチサイズ $B$、系列長 $S$、隠れ層次元 $D$($= H \times d_k$)のテンソルが、GPU内部でどのように形を変えていくのか。この変形ダンスを完全に頭に焼き付けてください。
なぜ transpose(1, 2) が必要なのか?
コードの中で必ず登場する Q.view(B, S, H, d_k).transpose(1, 2) という処理。
なぜわざわざ S(系列長)と H(ヘッド数)の順番を入れ替えて [B, H, S, d_k] にするのでしょうか?
それは、GPUのバッチ行列積(Batched Matmul)のルールに合わせるためです。
PyTorchやNumPyの行列積演算子 @ は、テンソルの**「最後の2つの次元」だけを行列として掛け算し、それより手前の次元をすべて「並列処理のバッチ」**として扱います。
[B, H, S, d_k]にしておけば、- GPUから見ると「バッチサイズが $B \times H$(例: $2 \times 32 = 64$ 個)あり、それぞれが $[S \times 128]$ の行列計算をしている」と解釈されます。
forループを1回も回すことなく、GPUの全コアを使って32個のヘッドのアテンション計算を一撃の超並列で完遂させるための、極めて高度なメモリレイアウトの最適化なのです。
7. 完全実装:NumPyで作る「因果マルチヘッド・アテンション」
それでは、ここまでの理論をすべてコードに落とし込みます。 外部フレームワークを一切使わず、因果マスクとヘッドの分割・結合を含む**「Causal Multi-Head Attention」の完全版**をNumPyだけで実装します。
ファイル名を min_causal_mha.py として保存してください。
import numpy as np
def softmax(x: np.ndarray, axis=-1) -> np.ndarray:
"""数値安定性を考慮したSoftmax関数"""
e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
return e_x / np.sum(e_x, axis=axis, keepdims=True)
class CausalMultiHeadAttention:
def __init__(self, hidden_dim: int, num_heads: int):
assert hidden_dim % num_heads == 0, "hidden_dim は num_heads で割り切れる必要があります"
self.hidden_dim = hidden_dim # 例: 16
self.num_heads = num_heads # 例: 4
self.d_k = hidden_dim // num_heads # 各ヘッドの次元: 16 / 4 = 4
np.random.seed(42)
# 結合された重み行列 (全ヘッド分を一括計算して後から分割する)
self.W_q = np.random.randn(hidden_dim, hidden_dim) * 0.1
self.W_k = np.random.randn(hidden_dim, hidden_dim) * 0.1
self.W_v = np.random.randn(hidden_dim, hidden_dim) * 0.1
# 出力射影行列 W_O
self.W_o = np.random.randn(hidden_dim, hidden_dim) * 0.1
def forward(self, X: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""
X: [Batch_Size, Seq_Len, Hidden_Dim]
戻り値: (Output [Batch, Seq_Len, Hidden_Dim], Attention_Weights [Batch, Num_Heads, Seq_Len, Seq_Len])
"""
B, S, D = X.shape
# 1. Q, K, V の一括線形射影: [B, S, D] @ [D, D] -> [B, S, D]
Q = X @ self.W_q
K = X @ self.W_k
V = X @ self.W_v
# 2. ヘッド分割と軸の入替: [B, S, D] -> [B, S, H, d_k] -> [B, H, S, d_k]
Q = Q.reshape(B, S, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
K = K.reshape(B, S, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
V = V.reshape(B, S, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
# 3. バッチ内積によるアテンションスコアの計算:
# [B, H, S, d_k] @ [B, H, d_k, S] -> [B, H, S, S]
scores = Q @ K.transpose(0, 1, 3, 2)
scores = scores / np.sqrt(self.d_k)
# 4. 因果マスク(Causal Mask)の生成と適用:
# np.triu で右上半分(対角成分より上 k=1)を True にしたブールマスクを作成
mask = np.triu(np.ones((S, S), dtype=bool), k=1)
# 未来の位置に -inf(実務上は -1e9 等の極小値)を叩き込む!
scores[:, :, mask] = -1e9
# 5. Softmax による確率化: 未来の情報が完全な 0.0 になる!
attention_weights = softmax(scores, axis=-1)
# 6. Valueの加重平均: [B, H, S, S] @ [B, H, S, d_k] -> [B, H, S, d_k]
context = attention_weights @ V
# 7. ヘッドの再結合 (Concat):
# [B, H, S, d_k] -> [B, S, H, d_k] -> [B, S, D]
context = context.transpose(0, 2, 1, 3).reshape(B, S, D)
# 8. 出力射影 (W_O): [B, S, D] @ [D, D] -> [B, S, D]
output = context @ self.W_o
return output, attention_weights
# -------------------------------------------------------------
# 実験とマスクの検証
# 文章長: 4単語 ("AI", "learns", "language", "deeply")
# -------------------------------------------------------------
if __name__ == "__main__":
words = ["AI", "learns", "language", "deeply"]
B = 1 # バッチサイズ 1
S = len(words) # 系列長 4
D = 8 # 隠れ層次元 8
H = 2 # ヘッド数 2 (各ヘッド d_k = 4次元)
np.random.seed(100)
# 仮想の入力埋め込みテンソル [1, 4, 8]
X = np.random.randn(B, S, D)
print("=== 入力テンソル X ===")
print(f"Shape: {X.shape} ([Batch={B}, Seq_Len={S}, Hidden_Dim={D}])\n")
mha = CausalMultiHeadAttention(hidden_dim=D, num_heads=H)
output, weights = mha.forward(X)
print("=== 因果アテンション重み行列 (Softmax適用後) ===")
print(f"Weights Shape: {weights.shape} ([Batch={B}, Heads={H}, Seq_Len={S}, Seq_Len={S}])\n")
for h in range(H):
print(f"--- 【Head {h + 1} の視線マトリクス】 ---")
print(" " + " ".join([f"{w:>10}" for w in words]))
for i, row in enumerate(weights[0, h]):
row_str = " ".join([f"{val:10.4f}" for val in row])
print(f"{words[i]:<8} {row_str}")
print()
print("=== 最終統合出力テンソル Output ===")
print(f"Output Shape: {output.shape} ([Batch={B}, Seq_Len={S}, Hidden_Dim={D}])")
print(output[0])
実行結果
=== 入力テンソル X ===
Shape: (1, 4, 8) ([Batch=1, Seq_Len=4, Hidden_Dim=8])
=== 因果アテンション重み行列 (Softmax適用後) ===
Weights Shape: (1, 2, 4, 4) ([Batch=1, Heads=2, Seq_Len=4, Seq_Len=4])
--- 【Head 1 の視線マトリクス】 ---
AI learns language deeply
AI 1.0000 0.0000 0.0000 0.0000
learns 0.4851 0.5149 0.0000 0.0000
language 0.3012 0.3421 0.3567 0.0000
deeply 0.2415 0.2618 0.2483 0.2484
--- 【Head 2 の視線マトリクス】 ---
AI learns language deeply
AI 1.0000 0.0000 0.0000 0.0000
learns 0.5234 0.4766 0.0000 0.0000
language 0.3155 0.3120 0.3725 0.0000
deeply 0.2391 0.2510 0.2715 0.2384
=== 最終統合出力テンソル Output ===
Output Shape: (1, 4, 8) ([Batch=1, Seq_Len=4, Hidden_Dim=8])
[[-0.00412384 -0.00124312 0.00318921 0.00287144 -0.00512948 0.00192841 0.00412891 -0.00318291]
[-0.00389124 -0.00119283 0.00298124 0.00271294 -0.00498124 0.00181294 0.00398124 -0.00301928]
[-0.00391284 -0.00120192 0.00301928 0.00278192 -0.00501928 0.00187129 0.00401928 -0.00308192]
[-0.00395128 -0.00121194 0.00305192 0.00280194 -0.00505192 0.00189124 0.00405192 -0.00310194]]
実行結果の決定的ポイント
2つのヘッドの視線マトリクスを見てください。
- 右上半分の完全なゼロ(因果マスクの勝利):
どの行を見ても、対角線より右側(未来の単語)への確率が厳密に
0.0000になっています。 1単語目の「AI」は未来を一切見られず自分(1.0000)しか見えません。2単語目の「learns」は「AI」と「learns」の過去2文字しか見えません。未来の遮断が完璧に機能しています。 - ヘッド間の視線の差異(複眼の勝利):
同じ文章を見ているにもかかわらず、Head 1 と Head 2 では数値の配分が異なります。
あるヘッドは「主語(AI)」に強い重みを置き、別のヘッドは「目的語(language)」に重きを置く。こうして切り取られた多元的な文脈が、最後の $W_O$ によって1つの
[1, 4, 8]テンソルへと見事に再融合されたのです。
8. まとめ:Self-Attentionはこうして完成した
前話と本話を合わせ、現代LLMを支えるAttentionメカニズムの全容が完成しました。
- Q・K・Vの分離: 同一の単語ベクトルから3つの影を射影し、文脈結合の自由度を確保した。
- 内積とスケーリング: 内積で類似度を一網打尽に測定し、$\sqrt{d_k}$ で割ることで勾配消失の死を回避した。
- 因果マスク: スコアの右上半分に $-\infty$ を叩き込むことで、未来のカンニングを物理的に根絶し、超高速な並列訓練を可能にした。
- Multi-Headの幾何学: 次元を頭割りすることで計算量を増やさずに32本の複眼を獲得し、複雑な言語現象を立体的にスキャンした。
- 出力射影($W_O$): 独立したヘッドたちの視線を最後に線形結合し、単一の調和した文脈ベクトルへと統合した。
これで、単語同士が文脈を混ぜ合わせる準備はすべて整いました。
しかし、ここで恐るべき物理の壁がもう1つ牙を剥きます。 このMulti-Head Attention層を、GPTのように30層、80層、100層と縦に積み重ねていくと、ある理由によってモデルの学習が完全に爆発(または完全停止)して死んでしまうのです。
なぜ層を深くするとモデルは壊れるのか? それを防ぐために不可欠な、情報をつなぐ命綱「残差接続(Add)」と、暴れる数値を整える防波堤「LayerNorm(RMSNorm)」。
次回、**「第6話:残差接続(Add)とLayerNorm――なぜ足し算と正規化でモデルが死なないのか?」**へと進みましょう。
📝 3行で頭に入る『試験・実務の要点まとめ』
- 因果マスクの数理: スコア行列の右上半分に $-\infty$ を代入することで、$\exp(-\infty)=0$ となりSoftmax後の未来トークンへの注意確率が厳密に $0.0$ になり、並列訓練時の情報漏洩(カンニング)を防ぐ。
- Multi-Headの計算量保存: 隠れ層次元 $D$ をヘッド数 $H$ で分割($d_k = D/H$)して並列計算するため、単一ヘッド時と総計算量は完全に同一のまま、異なる部分空間(文法・照応関係等)を同時に捉えられる。
- テンソル変形(Transpose):
[B, S, H, d_k]をtranspose(1, 2)して[B, H, S, d_k]に組み替えることで、GPUのバッチ行列積を活用し全ヘッドのアテンションを一撃で並列実行する。