マルチヘッド自己注意（Multi-Head Self-Attention）

Level 1概要役割と処理の流れ

何をする機能か

系列内の各トークンが、同じ系列内のすべてのトークンを見渡し、関連が強い相手ほど強く情報を取り込んで、自分のベクトルを更新します。

各トークンは 3 つの役割のベクトルに変換されます。

Query（問い合わせ）: 「自分はどんな情報を探しているか」
Key（見出し）: 「自分はどんな情報を提供できるか」
Value（中身）: 「実際に渡す情報」

Query と Key の内積で関連の強さを測り、その重みで Value を加重平均します。検索システムに例えると、Query で全件に問い合わせ、Key との一致度でスコアを付け、スコアに応じて Value を混ぜて返す処理です。ただし 1 件を選ぶのではなく、全件を重み付きで混ぜます。

「マルチヘッド」は、この処理を H 組（原典では 8 組）並行して行い、最後に結合することを指します。

なぜ必要か

トークンの意味は、単独では決まりません。 keys が「鍵」なのか「暗号鍵」なのかは、周囲の語によって変わります。また、are がどの語に対応する動詞なのかは、文中の関係を見ないと判断できません。

Self-Attention は、この「どのトークンを見て自分を解釈すべきか」を、距離ではなく内容に基づいて決めます。RNN や CNN と違い、どれだけ離れた語でも 1 ステップで参照できる点が大きな違いです。

全体の中でいつ呼ばれるか

Encoder の各層の最初のサブ層です。位置情報が加算されたあと、残差接続＋LayerNorm の直前に実行されます。これが N 層ぶん繰り返されます。

Level 2内部処理Tensor形状・内部ステップ・設計上の意味

前段から受け取るもの

hidden_states [B, S, D]。 1 層目は位置情報付加後の出力、 2 層目以降は前の層の出力です。

後段へ渡すもの

同じ形の hidden_states [B, S, D] と、可視化・分析用の attention_weights [B, H, S, S]（本流の計算には使いません）。

渡し先は残差接続＋LayerNorm です。このとき、Attention に入る前の値も残差として一緒に渡されます。

内部の5ステップ

投影: hidden_states [B, S, D] を、学習済みの重み $W^Q, W^K, W^V$ で Query・Key・Value [B, S, D] にそれぞれ変換する
Head分割: [B, S, D] を [B, H, S, Dk] に並べ替える（D = H × Dk）
スコア計算: Query と Key の内積を取り、 $\sqrt{D_k}$ で割る → [B, H, S, S]
正規化: 各行に Softmax をかけ、合計 1 の重みにする → [B, H, S, S]
加重和と結合: 重みで Value を加重平均し [B, H, S, Dv]、 Head を結合して [B, S, D]、最後に出力用の $W^O$ を掛ける

[B, H, S, S] は「各 Head が、各 Query 位置から、各 Key 位置をどれだけ見たか」を表す行列です。 S の 2 乗に比例するため、長い系列でメモリを最も使う部分になります。

例文での処理

共通例文

The keys on the table are lost.

概念上のトークン列（S = 8）: [The, keys, on, the, table, are, lost, .]

The
keys
on
the
table
are
lost
.

位置5の are が Query。位置1の keys と位置4の table が候補です。

are が自分の表現を作るとき、直前の table（位置4）は距離としては最も近い名詞です。しかし文法上は、主語である keys（位置1）との対応が重要です。

are の Query から見た重みの例

The0.04
keys0.38
on0.05
the0.04
table0.22
are0.18
lost0.07
.0.02

説明用の数値です

ここに示す数値は、仕組みを説明するために用意した仮の値です。学習済みモデルの実測値ではありません。また、Attention の重みが特定の文法関係（主語と動詞の対応など）を必ず表すとは限りません。学習済みモデルでは、そのような関係に対応する Attention が観察されることがある、という程度に捉えてください。

この例では、近くにある table だけでなく、離れた keys にも強い重みが置かれています。Self-Attention は距離ではなく内容で相手を選べるため、このような参照が可能です。ただし、実際の学習済みモデルでどの Head がどの関係を担うかは決まっておらず、Head ごとに異なる傾向が観察されます。

Head が複数あることで、「主語と動詞の対応を見る Head」「隣接語をまとめて句を作る Head」のように、異なる観点を同時に扱えます。 on the table を 1 つのまとまりとして扱う Head があれば、 keys の表現に「テーブルの上にある」という情報が混ざります。

Self-Attention 前

位置5の are は「be動詞の複数形」という語そのものの情報だけを持つ

Self-Attention 後

「主語 keys に対応する動詞」という文脈の情報が混ざった [B, S, D] になる

EncoderとDecoderでの違い

	Encoder の Self-Attention	Decoder の Masked Self-Attention	Cross Attention
Query の出どころ	自分の系列	自分の系列	Decoder の系列
Key・Value の出どころ	自分の系列	自分の系列	Encoder の出力
参照範囲	全位置（前後とも）	自分以前の位置のみ	Encoder の全位置
スコア行列	`[B, H, S, S]`	`[B, H, T, T]`	`[B, H, T, S]`

Encoder のこのページの処理では未来方向のマスクをかけません。入力文は最初から全部わかっているため、The が lost を見ても問題ないからです。

入力と出力

入力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`hidden_states`	各位置のベクトル列（Query・Key・Value すべての元になる）	float32 Tensor	`[B, S, D]`	B = バッチサイズ（同時に処理する文の数） / S = 入力系列長（Encoder側のトークン数） / D = モデル内部のベクトル次元（d_model）	[1, 8, 512]	Q/K/V への線形変換
`attention_mask`	パディング位置を参照させないためのマスク（省略可）	bool / float32 Tensor	`[B, S]`	B = バッチサイズ（同時に処理する文の数） / S = 入力系列長（Encoder側のトークン数）	[[1, 1, 1, 1, 1, 1, 1, 1]]	スコアへの加算（マスク位置を -inf に）

出力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`hidden_states`	周囲の情報が混ざったベクトル列	float32 Tensor	`[B, S, D]`	B = バッチサイズ（同時に処理する文の数） / S = 入力系列長（Encoder側のトークン数） / D = モデル内部のベクトル次元（d_model）	[1, 8, 512]（形は変わらない）	残差接続＋ LayerNorm

補助出力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`attention_weights`	各Headの、Query位置からKey位置への重み（行ごとの合計が1）	float32 Tensor	`[B, H, S, S]`	B = バッチサイズ（同時に処理する文の数） / H = Attention Head 数 / S = 入力系列長（Encoder側のトークン数）	[1, 8, 8, 8]	可視化・分析（本流の計算には使わない）

B: バッチサイズ（同時に処理する文の数）
S: 入力系列長（Encoder側のトークン数）
D: モデル内部のベクトル次元（d_model）
H: Attention Head 数

このサイトでは Tensor 形状を PyTorch の `batch_first=True` 相当、つまり `[B, S, D]`（バッチ次元が先頭）で統一しています。実装によっては `[S, B, D]` を採用するものもあります。

よくある誤解

Level 3数式・Python数式・疑似実装・計算量・実装差

数式

1 つの Head の計算（Scaled Dot-Product Attention）は次の通りです。

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{D_k}}\right)V

$\sqrt{D_k}$ で割るのは、 $D_k$ が大きいほど内積の分散が大きくなり、 Softmax が極端に尖って勾配が消えるのを防ぐためです。

マルチヘッドは次のようにまとめられます。

\mathrm{head}_h = \mathrm{Attention}(XW^Q_h,\; XW^K_h,\; XW^V_h)

\mathrm{MultiHead}(X) = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_H)\,W^O

ここで $X \in \mathbb{R}^{B \times S \times D}$ 、 $W^O \in \mathbb{R}^{D \times D}$ です。

Pythonで表したインターフェース

from dataclasses import dataclass
from typing import Optional
import math
import torch
from torch import Tensor, nn
 
 
@dataclass
class AttentionOutput:
    hidden_states: Tensor      # [B, S, D]
    attention_weights: Tensor  # [B, H, S, S]
 
 
class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int) -> None:
        super().__init__()
        assert d_model % num_heads == 0, "D は H で割り切れる必要がある"
        self.num_heads = num_heads
        self.d_k = d_model // num_heads          # Dk = D / H
        self.w_q = nn.Linear(d_model, d_model)   # [D, D]
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)
 
    def _split_heads(self, x: Tensor) -> Tensor:
        """[B, S, D] -> [B, H, S, Dk]"""
        batch, seq_len, _ = x.shape
        return x.view(batch, seq_len, self.num_heads, self.d_k).transpose(1, 2)
 
    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: Optional[Tensor] = None,
    ) -> AttentionOutput:
        """各トークンを同じ系列内のトークンと関連付ける。
 
        Args:
            hidden_states: [B, S, D]。Query/Key/Value すべての元になる。
            attention_mask: [B, S]。1 が実トークン、0 がパディング。
 
        Returns:
            hidden_states [B, S, D] と attention_weights [B, H, S, S]。
        """
        q = self._split_heads(self.w_q(hidden_states))  # [B, H, S, Dk]
        k = self._split_heads(self.w_k(hidden_states))  # [B, H, S, Dk]
        v = self._split_heads(self.w_v(hidden_states))  # [B, H, S, Dv]
 
        # [B, H, S, Dk] @ [B, H, Dk, S] -> [B, H, S, S]
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)
 
        if attention_mask is not None:
            # [B, S] -> [B, 1, 1, S] にしてブロードキャスト
            mask = attention_mask[:, None, None, :].bool()
            scores = scores.masked_fill(~mask, float("-inf"))
 
        weights = scores.softmax(dim=-1)                # [B, H, S, S]
        context = weights @ v                           # [B, H, S, Dv]
 
        # [B, H, S, Dv] -> [B, S, D]
        batch, _, seq_len, _ = context.shape
        merged = context.transpose(1, 2).reshape(batch, seq_len, -1)
 
        return AttentionOutput(
            hidden_states=self.w_o(merged),             # [B, S, D]
            attention_weights=weights,                  # [B, H, S, S]
        )
 
 
def multi_head_self_attention(
    hidden_states: Tensor,
    layer: MultiHeadSelfAttention,
    attention_mask: Optional[Tensor] = None,
) -> AttentionOutput:
    """各トークンを同じ系列内のトークンと関連付ける。"""
    return layer(hidden_states, attention_mask)
 
 
# 最小使用例
B, S, D, H = 1, 8, 512, 8
layer = MultiHeadSelfAttention(d_model=D, num_heads=H)
 
hidden_states = torch.randn(B, S, D)                  # [B, S, D]
out = multi_head_self_attention(hidden_states, layer)
 
print(out.hidden_states.shape)                        # torch.Size([1, 8, 512])    == [B, S, D]
print(out.attention_weights.shape)                    # torch.Size([1, 8, 8, 8])   == [B, H, S, S]
print(out.attention_weights[0, 0, 5].sum())           # tensor(1.0000)  各行の合計は1

この例で単純化している点

Dropout、バイアス項、重みの初期化を省略しています。
マスクは [B, S] のパディングマスクのみを扱い、因果マスクは Masked Self-Attention のページで扱います。
FlashAttention のようなメモリ効率化は行わず、[B, H, S, S] を実体化しています。

型と Tensor 形状の見え方を掴むための擬似実装です。そのまま学習・推論に使えるものではありません。

計算量とメモリ

項目	オーダー	例文（S=8, D=512, H=8）での目安
Q/K/V 投影	$O(B \cdot S \cdot D^2)$	主にここが支配的（短い系列の場合）
スコア計算	$O(B \cdot H \cdot S^2 \cdot D_k) = O(B \cdot S^2 \cdot D)$	S が伸びると急増
重み行列のメモリ	$O(B \cdot H \cdot S^2)$	S=4096 なら 1 層あたり数百 MB になりうる

系列長が短いうちは投影（ $D^2$ の項）が支配的ですが、 $S > D$ になるとスコア計算の $S^2$ が効いてきます。長文脈モデルで FlashAttention のようなメモリ効率化が重要になるのはこのためです。

実装差と派生

Grouped Query Attention (GQA): Key と Value の Head 数を減らして共有する方式。推論時の KV Cache のメモリを削減できます。Llama 2 70B 以降などで採用。
Multi-Query Attention (MQA): Key と Value の Head を 1 組だけにする、GQA の極端な形。
FlashAttention: [B, H, S, S] を実体化せずにブロック単位で計算し、メモリを削減します。数学的な結果は同じです。
scaled_dot_product_attention: PyTorch 2 以降の組み込み関数を使うと、上記の最適化が自動的に選択されます。実務ではこちらを使います。

マルチヘッド自己注意Multi-Head Self-Attention