この記事で決められること
読み終えたあと、num_heads をいくつにするかと、nn.MultiheadAttention を結果を壊さずに呼ぶための3つの設定を、自分のコードで決められるようになります。
対象読者
- PyTorchで nn.Linear や nn.Module を書いたことがある人
- Self-Attention という名前は知っているが、ヘッドを分ける意味を説明できない人
この記事で得られること
- num_heads を選ぶ基準と、パラメータ数がヘッド数で変わらないという事実
- エラーが出ない3つの罠の再現コードと見分け方
- nn.MultiheadAttention を呼ぶ前のチェックリストと、エラーメッセージ別の対処表
この記事で扱わないこと
- Transformer 全体の構造(残差接続、LayerNorm、FFN)
- 速度の比較と、SDPA による自前実装や GQA への拡張(続編で扱います)
なぜこれを書くのか
これまで Self-Attention や Flash Attention の記事を書いてきましたが、その間をつなぐ「頭を分ける」という一段を、きちんと書いていませんでした。PyTorch 2.14 の公式ドキュメントを読み返して驚いたのは、nn.MultiheadAttention が自分自身を「基礎を理解するための参照実装」と説明していたことです。
便利な既製モジュールほど、既定値が自分の前提とずれていても、エラーが出ない限り気付けません。この記事では、既製品の中身を一度手で組み直して確かめる価値を、一緒に考えてもらえたらうれしいです。検証は GPU なしの Xeon 2コア・メモリ8GB の Linux 環境で、Python 3.13 と PyTorch 2.14.0 を CPU 実行しています。
1. MHAを一言でいうと:同じ原稿を8人の編集者が読む
最初に、記事全体で使うたとえを決めておきます。用語は末尾の用語集にもまとめています。
このセクションで分かること:MHA(Multi-Head Attention、マルチヘッド注意)が何をしていて、なぜ1人ではなく複数人で読むのか。
Attention は、文の各単語が「ほかのどの単語をどれだけ参考にするか」を決める仕組みです。これを、1人の編集者が原稿を読む作業にたとえます。
- 編集者が持つ「何を探しているか」のメモが Query(問い)です
- 各単語に付いた「自分は何の話か」の見出し札が Key(見出し)です
- 見出し札の裏に書かれた単語の中身が Value(中身)です
編集者は問いと見出し札を照らし合わせ、合いそうな単語ほど赤ペンを濃く引き、濃さに応じて中身を集めます。
MHAはこの編集者を8人に増やす仕組みです。1人は文法の係り受けを、1人は代名詞の指す先を、というように、それぞれ違う色眼鏡で同じ原稿を読みます。最後に編集長が8人分のメモを1枚にまとめます。
なぜ1人ではだめなのか。原論文 "Attention Is All You Need" はこう書いています。
原文:"With a single attention head, averaging inhibits this."
和訳:注意ヘッドが1つだけだと、平均化がこれ(複数の表現部分空間の情報に同時に注意を向けること)を妨げる。
出典:Vaswani et al. 2017, §3.2.2
1人の編集者が係り受けと代名詞の両方を気にすると、赤ペンの濃さが2つの目的の平均になり、どちらも中途半端になります。だから担当を分けます。
2. 数式で見るMHA:3つの式と、たとえとの対応表
前節で、MHAとは8人の編集者が別々の色眼鏡で同じ原稿を読むことだと決めました。たとえだけではコードに落とせないので、同じ内容を式に置き換えます。
このセクションで分かること:MHAの3つの式が、編集者のたとえのどこに当たるか。
記号は使う直前に1つずつ導入します。
- $E$(イー、埋め込み次元、整数):1単語を表すベクトルの長さ。たとえでは原稿1語あたりの情報の幅。引数名は embed_dim
- $h$(エイチ、ヘッド数、整数):編集者の人数。引数名は num_heads
- $d$(ディー、ヘッド次元、整数):1人の編集者が担当する幅で、$d = E / h$
まず、1人分の注意の式です。$Q$, $K$, $V$(行列、行数は単語数、列数は $d$)は、問いのメモ、見出し札、中身を単語ぶん縦に並べたものです。
\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right) V
この式が言っているのは、問いと見出し札の内積で相性を測り、softmax で赤ペンの濃さ(合計1)に直し、その濃さで中身を混ぜる、ということです。$\sqrt{d}$ で割るのは、幅 $d$ が大きいと内積が大きくなりすぎて softmax の勾配が極端に小さくなるのを防ぐためだと、原論文 §3.2.1 は説明しています。
次に、$i$ 番目の編集者です。$W_i^Q$, $W_i^K$, $W_i^V$(行列、$E \times d$)は編集者 $i$ の色眼鏡です。
\mathrm{head}_i = \mathrm{Attention}(Q W_i^Q,\; K W_i^K,\; V W_i^V)
この式が言っているのは、同じ原稿を、編集者ごとに違う色眼鏡で幅 $E$ から幅 $d$ に絞って読む、ということです。
最後に、編集長のまとめです。$W^O$(行列、$E \times E$)は編集長のまとめ方です。
\mathrm{MultiHead}(Q, K, V) = \mathrm{Concat}(\mathrm{head}_1, \dots, \mathrm{head}_h)\, W^O
この式が言っているのは、$h$ 人分の幅 $d$ のメモを横に並べて幅 $E$ に戻し、編集長が1枚の要約に書き直す、ということです。
| 記号 | PyTorchでの名前 | たとえでの役割 |
|---|---|---|
| $E$ | embed_dim | 原稿1語あたりの情報の幅 |
| $h$ | num_heads | 編集者の人数 |
| $d = E/h$ | embed_dim // num_heads | 1人が担当する幅 |
| $Q$, $K$, $V$ | query, key, value | 問いのメモ、見出し札、中身 |
| $W_i^Q, W_i^K, W_i^V$ | in_proj_weight を3分割したもの | 編集者 $i$ の色眼鏡 |
| softmax の出力 | attn_output_weights | 赤ペンの濃さ |
| $W^O$ | out_proj | 編集長のまとめ方 |
3. 本当に式どおりか:手で組んで nn.MultiheadAttention と照合する
前節で、MHAは「射影、分割、注意、連結、出力射影」の5段だと式で確認しました。式が正しくても、PyTorch の実装が同じことをしているかは別問題です。照合しておけば、後半の罠でどこが壊れたかを手組みと比べて特定できます。
このセクションで分かること:nn.MultiheadAttention が式どおりに動いていること、そしてヘッド数を変えてもパラメータ数が1つも変わらないこと。
まず、後の実験でも使う小道具をまとめます。
import torch
import torch.nn as nn
import torch.nn.functional as F
def project_qkv(x_q, x_k, x_v, mha: nn.MultiheadAttention):
"""MHAの重みでQ/K/Vを射影する。(N, L, E) -> (N, L, E)"""
Wq, Wk, Wv = mha.in_proj_weight.chunk(3)
bq, bk, bv = mha.in_proj_bias.chunk(3)
return F.linear(x_q, Wq, bq), F.linear(x_k, Wk, bk), F.linear(x_v, Wv, bv)
def split_heads(t, num_heads):
"""(N, L, E) -> (N, H, L, d)"""
N, L, E = t.shape
return t.view(N, L, num_heads, E // num_heads).transpose(1, 2)
def merge_heads(t):
"""(N, H, L, d) -> (N, L, H*d)"""
N, H, L, d = t.shape
return t.transpose(1, 2).reshape(N, L, H * d)
def max_diff(a, b):
return (a - b).abs().max().item()
次に、式どおりに手で組んだMHAと、nn.MultiheadAttention の出力を比べます。
import torch
import torch.nn as nn
from mha_utils import project_qkv, split_heads, merge_heads, max_diff
torch.manual_seed(0)
E, H, N, L = 64, 8, 2, 5
mha = nn.MultiheadAttention(E, H, batch_first=True).eval()
x = torch.randn(N, L, E)
def manual_mha(x, mha):
d = mha.embed_dim // mha.num_heads
q, k, v = (split_heads(t, mha.num_heads) for t in project_qkv(x, x, x, mha))
w = torch.softmax(q @ k.transpose(-2, -1) / d**0.5, dim=-1) # (N,H,L,L)
return mha.out_proj(merge_heads(w @ v)), w
with torch.no_grad():
y_ref, w_ref = mha(x, x, x, average_attn_weights=False)
y_man, w_man = manual_mha(x, mha)
print(f"出力の最大誤差 : {max_diff(y_ref, y_man):.2e}")
print(f"重みの最大誤差 : {max_diff(w_ref, w_man):.2e}")
print(f"重みのshape : {tuple(w_ref.shape)}")
print("\nE=512 で num_heads を変えたときのパラメータ数")
for h in (1, 2, 4, 8, 16):
m = nn.MultiheadAttention(512, h)
n_params = sum(p.numel() for p in m.parameters())
print(f" heads={h:2d} head_dim={512 // h:3d} params={n_params:,}")
print(f" 理論値 4E^2+4E = {4 * 512**2 + 4 * 512:,}")
出力の最大誤差は 1.19e-07、重みの最大誤差は 5.96e-08 でした。float32 の丸め誤差の範囲なので、nn.MultiheadAttention は前節の式そのものです。パラメータ数は次のとおりです。
| num_heads | head_dim | パラメータ数 |
|---|---|---|
| 1 | 512 | 1,050,624 |
| 2 | 256 | 1,050,624 |
| 4 | 128 | 1,050,624 |
| 8 | 64 | 1,050,624 |
| 16 | 32 | 1,050,624 |
ヘッド数を1から16まで変えても、パラメータ数は1つも変わりません。色眼鏡 $W^Q, W^K, W^V$ と編集長の $W^O$ はどれも $E \times E$ で、バイアスが $4E$ あるので合計 $4E^2 + 4E$ です。ヘッド数は、この総額を何人で分けるかを決めているだけです。
つまり num_heads はモデルの大きさではなく1人あたりの幅を決めるつまみです。原論文は $E = 512$, $h = 8$, $d = 64$ を使っています。選び方の出発点は次の2つです。
- $d$ が64前後になる $h$ を選ぶ(原論文の設定に合わせる)
- $h$ は偶数にする。インストールされたソース(torch/nn/modules/activation.py)を読むと、推論用の高速経路である fastpath の判定に、ドキュメントの条件一覧には載っていない「num_heads が偶数であること」が入っているためです
4. nn.MultiheadAttention の引数を読む
前節で、中身が式どおりであることを確かめました。次に問題になるのは、その中身にどんな形でデータを渡すかです。この後で扱う3つの罠は、どれも計算ではなく入口の形で起きるので、先に引数を整理します。
このセクションで分かること:よく使う引数と、入出力テンソルの形。
| 引数 | 既定値 | 意味 |
|---|---|---|
| embed_dim | なし | モデル全体の次元 $E$ |
| num_heads | なし | ヘッド数 $h$。$E$ は $h$ で割り切れる必要がある |
| dropout | 0.0 | 注意の重みにかけるドロップアウト率 |
| kdim, vdim | None | key, value の特徴量の次元。None なら embed_dim |
| batch_first | False | True なら (N, L, E)、False なら (L, N, E) |
forward の主な引数です。$N$ はバッチサイズ、$L$ は問い側の系列長、$S$ は参照側の系列長です。
| 引数 | 形 | 意味 |
|---|---|---|
| query | (N, L, E) | 問いのメモ |
| key, value | (N, S, E) | 見出し札と中身 |
| key_padding_mask | (N, S) | True の位置を無視する(パディング用) |
| attn_mask | (L, S) または (N·h, L, S) | True の位置への注意を禁止する |
| need_weights | True | 注意の重みも返すか |
| average_attn_weights | True | 重みをヘッド間で平均して返すか |
戻り値は (attn_output, attn_output_weights) のタプルです。重みは既定でヘッド平均の (N, L, S) です。ヘッドごとの赤ペンを見たいときは average_attn_weights=False にすると (N, h, L, S) になります。
位置づけについては、ドキュメントに次の一文があります。
原文:"The intent of this layer is as a reference implementation for foundational understanding"
和訳:この層の意図は、基礎を理解するための参照実装である。
出典:PyTorch 2.14 torch.nn.MultiheadAttention
新しい構成には部品から組むことも勧めていますが、学習や実験の読みやすさでは既製品が便利です。この記事では既製品を安全に使う側に絞ります。
ここまでのまとめ
- MHAは、同じ原稿を $h$ 人の編集者が幅 $d = E/h$ の色眼鏡で読み、編集長がまとめる仕組み
- nn.MultiheadAttention は式どおりで、手組みとの差は 1e-7 程度
- ヘッド数を変えてもパラメータ数は $4E^2 + 4E$ のまま。num_heads は $d$ を決めるつまみで、偶数にしておく
- 既定は batch_first=False で、key_padding_mask は True が無視
後半は、この既定値とマスクの向きが引き起こす罠を、実際に壊して確かめます。
5. エラーにならずに壊れる3つの罠
前節の最後で、batch_first の既定が False であること、マスクは True が無視であることを確認しました。どちらも形さえ合えばエラーは出ません。エラーが出ない設定ミスは、学習が進まない、精度が少し低い、という症状でしか現れないので、気付くのが最も遅れます。
このセクションで分かること:黙って結果を壊す3つの設定ミスと、それぞれを見分ける方法。
罠1:batch_first を付け忘れると、別の原稿どうしが混ざる
(N, L, E) のデータを batch_first=False のまま渡すと、PyTorch は (L, N, E) だと解釈します。たとえで言えば、違う原稿の同じ位置の単語どうしを見比べさせることになります。
import torch
import torch.nn as nn
from mha_utils import max_diff
torch.manual_seed(0)
E, H, N, L = 64, 8, 4, 10
x = torch.randn(N, L, E) # 自分では (N, L, E) のつもり
x2 = x.clone()
x2[1] += 1.0 # サンプル1だけを書き換える
def leak_into_sample0(batch_first):
torch.manual_seed(1) # 2つの設定で同じ重みにする
mha = nn.MultiheadAttention(E, H, batch_first=batch_first).eval()
with torch.no_grad():
y1, w = mha(x, x, x)
y2, _ = mha(x2, x2, x2)
return max_diff(y1[0], y2[0]), tuple(w.shape)
for bf in (False, True):
leak, w_shape = leak_into_sample0(bf)
print(f"batch_first={bf!s:5} サンプル0の出力変化={leak:.3e} 重みshape={w_shape}")
| batch_first | サンプル1だけ変えたときの、サンプル0の出力変化 | 重みの形 |
|---|---|---|
| False(既定) | 4.804e-01 | (10, 4, 4) |
| True | 0.000e+00 | (4, 10, 10) |
サンプル1しか変えていないのに、batch_first=False ではサンプル0の出力が 0.48 も動きました。バッチ内の別の原稿の情報が漏れ込んでいる証拠で、エラーは一度も出ていません。正直、ここまで黙って動くとは思っていませんでした。
見分け方は重みの形です。$N = 4$, $L = 10$ なら正しくは (4, 10, 10) で、(10, 4, 4) が返ってきたらバッチと系列が入れ替わっています。
罠2:SDPAへ移すとき、マスクの向きが逆になる
注意の計算だけを行う最適化済みの関数 F.scaled_dot_product_attention(以下 SDPA)は、boolマスクの True の意味が nn.MultiheadAttention と逆です。ドキュメントにも注意書きがあります。
原文:"If migrating from MHA, ensure you invert your boolean mask"
和訳:MHAから移行する場合は、boolマスクを必ず反転すること。
出典:PyTorch 2.14 torch.nn.functional.scaled_dot_product_attention
nn.MultiheadAttention の key_padding_mask は True が無視、SDPA の attn_mask は True が参加です。たとえで言えば、片方は「読むな」の付箋、もう片方は「読め」の付箋です。同じマスクを渡して比べます。
import torch
import torch.nn as nn
import torch.nn.functional as F
from mha_utils import project_qkv, split_heads, merge_heads, max_diff
torch.manual_seed(0)
E, H, N, L = 64, 8, 2, 5
mha = nn.MultiheadAttention(E, H, batch_first=True).eval()
x = torch.randn(N, L, E)
pad = torch.tensor([[False, False, False, True, True], # True = 無視したい位置
[False, False, False, False, False]])
def sdpa_mha(x, bool_mask):
q, k, v = (split_heads(t, H) for t in project_qkv(x, x, x, mha))
o = F.scaled_dot_product_attention(q, k, v, attn_mask=bool_mask[:, None, None, :])
return mha.out_proj(merge_heads(o))
with torch.no_grad():
y_ref, _ = mha(x, x, x, key_padding_mask=pad)
print(f"同じマスクをSDPAへ : 誤差 {max_diff(y_ref, sdpa_mha(x, pad)):.2e}")
print(f"~maskで反転して渡す: 誤差 {max_diff(y_ref, sdpa_mha(x, ~pad)):.2e}")
all_pad = torch.ones(N, L, dtype=torch.bool) # 全位置を無視
for train in (True, False):
y, _ = mha.train(train)(x, x, x, key_padding_mask=all_pad)
print(f"全部マスク train={train!s:5}: NaNを含む={torch.isnan(y).any().item()}")
| 条件 | nn.MultiheadAttention との最大誤差 |
|---|---|
| 同じマスクをそのまま SDPA へ渡す | 1.35e+00 |
| ~mask で反転して渡す | 1.19e-07 |
反転を忘れると、読むべき単語を読まずパディングだけを読む計算になります。形は合っているので、ここでもエラーは出ません。
罠3:全部マスクしたサンプルは NaN になる
上のコードの後半は、あるサンプルの全キーを無視させたときの挙動です。
| モード | 出力に NaN を含むか |
|---|---|
| train | True |
| eval | True |
すべての見出し札に「読むな」の付箋が貼られると、softmax に入る値がすべて負の無限大になり、0割る0の形になって赤ペンの濃さが定義できなくなります。空のサンプルが1つ混ざるだけで、そのバッチの損失が NaN になり得ます。データ側で空のサンプルを除くか、少なくとも1トークンは残します。
6. 呼ぶ前のチェックリスト
前節までで、num_heads の選び方と、黙って壊れる3つの罠を実測で確かめました。これらを、呼ぶ前に上から順に確かめる形にまとめます。
このセクションで分かること:nn.MultiheadAttention を書く前と、結果がおかしいときに確かめる順番。
最後の箱が一番安く効きます。1バッチだけ流して attn_output_weights の形を print する。それだけで罠1は確実に捕まえられます。
7. トラブルシューティング
前節のチェックリストは、エラーが出ない罠を防ぐためのものでした。ここでは逆に、エラーとして表に出るものを集めます。実際にエラーを起こして、メッセージを確かめたコードです。
import torch
import torch.nn as nn
E, H = 64, 8
x = torch.randn(2, 5, E)
def try_run(label, fn):
try:
fn()
print(f"[通過] {label}")
except Exception as e: # 記事に載せるため、型とメッセージだけ拾う
print(f"[{type(e).__name__}] {label}\n {str(e).splitlines()[0][:110]}")
try_run("embed_dim=64, num_heads=6",
lambda: nn.MultiheadAttention(64, 6))
try_run("is_causal=True だけ渡す",
lambda: nn.MultiheadAttention(E, H, batch_first=True)(x, x, x, is_causal=True))
try_run("kdim違いのkeyをkdim指定なしで渡す",
lambda: nn.MultiheadAttention(E, H, batch_first=True)(x, torch.randn(2, 7, 32), torch.randn(2, 7, 32)))
try_run("kdim=32, vdim=32 を指定して渡す",
lambda: nn.MultiheadAttention(E, H, kdim=32, vdim=32, batch_first=True)(
x, torch.randn(2, 7, 32), torch.randn(2, 7, 32)))
| 症状 | 原因 | 対処 |
|---|---|---|
| AssertionError: embed_dim must be divisible by num_heads | $E$ が $h$ で割り切れない | $E$ を割り切る $h$ を選ぶ |
| RuntimeError: Need attn_mask if specifying the is_causal hint | is_causal=True だけを渡した | nn.Transformer.generate_square_subsequent_mask で作ったマスクを attn_mask にも渡す |
| RuntimeError: mat1 and mat2 shapes cannot be multiplied | key, value の次元が embed_dim と違う | kdim, vdim を指定する(指定すれば通る) |
| エラーなしで精度が出ない。重みの形が (L, N, N) | batch_first=False のまま (N, L, E) を渡した | batch_first=True を付ける |
| SDPA に移したら結果がずれた | boolマスクの意味が逆 | ~mask で反転する |
| 損失が突然 NaN になる | 全トークンがマスクされたサンプルがある | 空サンプルを除く、または1トークン残す |
is_causal の行は、引数説明の1文目だけを読むと「指定すれば因果マスクを適用する」と読めます。実際には、同じ説明の警告にあるとおり is_causal は「渡した attn_mask は因果マスクです」というヒントで、attn_mask を一緒に渡さないとエラーになります。
まとめ
MHAは、同じ原稿を $h$ 人の編集者が別々の色眼鏡で読み、編集長がまとめる仕組みでした。nn.MultiheadAttention は式どおりに動き、ヘッド数を変えてもパラメータ数は変わりません。num_heads は1人あたりの幅 $d$ を決めるつまみで、$d$ が64前後、かつ偶数を出発点にします。
一方で、batch_first の既定値、マスクの向き、全マスク行の NaN という3つの罠は、どれもエラーを出さずに結果を壊します。編集部が黙って原稿を取り違えても、編集長の要約はそれらしく仕上がってしまう。だから 1バッチ目で重みの形を見る習慣 が、一番安い保険になります。
個人的には、ドキュメントにない偶数条件をソースで見つけたことが一番の収穫でした。既製品の中身は、一度は手で組み直して確かめる価値があります。
続編:fastpath と SDPA の速さを CPU で実測し、SDPA で自前の MHA と GQA を組む記事を準備中です(公開後にリンクを追加します)。
学習ロードマップ
次に読むなら、注意を速く計算する Flash Attention、生成時のメモリを決める KVキャッシュ、注意を使わない Mamba-2 の順がおすすめです。
用語集
| 用語 | 意味 | たとえでの対応 |
|---|---|---|
| MHA | 注意を複数のヘッドに分けて並列に計算し、最後にまとめる仕組み | 8人の編集者と1人の編集長 |
| head_dim | $d = E/h$ | 1人が担当する幅 |
| Query / Key / Value | 問い、照合用の見出し、集める中身 | 問いのメモ、見出し札、札の裏の中身 |
| key_padding_mask | 無視するキーの位置。True が無視 | 「読むな」の付箋 |
| SDPA | F.scaled_dot_product_attention。True が参加 | 「読め」の付箋で動く別の編集部 |
| fastpath | 条件を満たした推論時だけ使われる高速経路 | 条件付きの急ぎ便 |
参考文献
- PyTorch 2.14 ドキュメント torch.nn.MultiheadAttention(引数、fastpath の条件、参照実装という位置づけが書かれた一次情報)
- PyTorch 2.14 ドキュメント torch.nn.functional.scaled_dot_product_attention(マスクの向きの注意書き)
- Vaswani et al., "Attention Is All You Need", arXiv:1706.03762, 2017(邦題訳:必要なのは注意だけ。MHA を提案した原論文。§3.2.2 に $h=8$, $d=64$ の設定)
関連用語記事
最後まで読んでいただきありがとうございました。X でも更新情報を流しています。