マスク付き自己注意（Masked (Causal) Self-Attention）

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

何をする機能か

計算の中身はマルチヘッド自己注意と同じです。違いは 1 点だけで、各位置が自分より後ろの位置を参照できないようにマスクをかけることです。

Attention スコア [B, H, T, T] のうち、「Query 位置より右にある Key 位置」に対応する要素を $-\infty$ に置き換えてから Softmax を取ります。 $-\infty$ は Softmax を通すと確率 0 になるため、その位置の Value は一切混ざりません。

なぜ必要か

Decoder の仕事は「次のトークンを予測する」ことです。予測すべき答えそのもの（および、それより後ろ）が見えていては学習になりません。

一方、学習を効率よく行うには、全位置の予測を 1 回の順伝播でまとめて計算したいところです。そこで、計算は並列に行いつつ、参照範囲だけを位置ごとに制限するためにマスクを使います。

これにより、学習時と推論時で各位置が見る範囲が一致します。推論時は、そもそも未来のトークンがまだ存在しないため、自然に同じ条件になります。

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

Decoder 各層の最初のサブ層です。そのあと残差接続＋ LayerNorm を経て、 Cross Attention へ進みます。

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

前段から受け取るもの

hidden_states [B, T, D]。1 層目は右シフトした出力を埋め込み、位置情報を加えたものです。

後段へ渡すもの

hidden_states [B, T, D] と attention_weights [B, H, T, T]。渡し先は Add & Norm です。

マスクの形

因果マスク（causal mask）は [T, T] の下三角行列です。行が Query 位置、列が Key 位置を表します。

        Key →  0    1    2    3    4    5    6    7
Query 0        ○    ×    ×    ×    ×    ×    ×    ×
      1        ○    ○    ×    ×    ×    ×    ×    ×
      2        ○    ○    ○    ×    ×    ×    ×    ×
      3        ○    ○    ○    ○    ×    ×    ×    ×
      4        ○    ○    ○    ○    ○    ×    ×    ×
      5        ○    ○    ○    ○    ○    ○    ×    ×
      6        ○    ○    ○    ○    ○    ○    ○    ×
      7        ○    ○    ○    ○    ○    ○    ○    ○

○ が参照可能、× が参照不可です。対角線（自分自身）は参照できます。バッチ内のパディングを隠すマスク [B, T] とは別物で、実装では両者を組み合わせます。

例文での処理

共通例文

The keys on the table are lost.

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

are を予測する時点を考えます。右シフトの対応関係から、これは Decoder 入力の位置 5（table）での予測です。

入力済み: The keys on the table
予測対象: are
参照可能: The / keys / on / the / table
参照不可: are / lost / .

The
keys
on
the
table
are
lost
.

緑が参照可能なトークン、取り消し線が参照不可のトークンです。Query は位置4の table。

参照可能な範囲の中で、モデルは主語を探す必要があります。

table（are を予測する位置） の Query から見た重みの例

The0.06
keys0.47
on0.08
the0.05
table0.34

EncoderとDecoderでの違い

	Encoder の Self-Attention	Decoder の Masked Self-Attention
参照範囲	全位置（前後とも）	自分以前の位置のみ
マスク	パディングマスクのみ	パディングマスク＋因果マスク
対象の系列	入力文 `[B, S, D]`	出力文 `[B, T, D]`
理由	入力文は最初から全部わかっている	未来の答えを見てはいけない

Decoder-only の LLM では、プロンプトを含む 1 本の系列すべてにこの因果マスクが適用されます。つまり Decoder-only モデルの Attention は、基本的にすべてこの形です。

入力と出力

入力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`hidden_states`	Decoder 側の各位置のベクトル列	float32 Tensor	`[B, T, D]`	B = バッチサイズ（同時に処理する文の数） / T = 出力系列長（Decoder側のトークン数） / D = モデル内部のベクトル次元（d_model）	[1, 8, 512]	Q/K/V への線形変換
`causal_mask`	自分より後ろを参照させない下三角マスク	bool Tensor	`[T, T]`	T = 出力系列長（Decoder側のトークン数）	[8, 8] の下三角行列	スコアへ適用（False の位置を -inf に）

出力

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

補助出力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`attention_weights`	上三角部分が必ず 0 になる重み行列	float32 Tensor	`[B, H, T, T]`	B = バッチサイズ（同時に処理する文の数） / H = Attention Head 数 / T = 出力系列長（Decoder側のトークン数）	[1, 8, 8, 8]	可視化・分析

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

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

よくある誤解

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

数式

マスク行列 $M \in \mathbb{R}^{T \times T}$ を次のように定義します。

M_{ij} = \begin{cases} 0 & (j \le i) \\ -\infty & (j > i) \end{cases}

これをスコアに加算してから Softmax を取ります。

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

$\exp(-\infty) = 0$ なので、 $j > i$ の重みは 0 になります。

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

from dataclasses import dataclass
from typing import Optional
import math
import torch
from torch import Tensor
 
 
@dataclass
class AttentionOutput:
    hidden_states: Tensor      # [B, T, D]
    attention_weights: Tensor  # [B, H, T, T]
 
 
def build_causal_mask(seq_len: int, device: torch.device | None = None) -> Tensor:
    """自分より後ろを参照させない下三角マスクを作る。
 
    Returns:
        [T, T] の bool。True が参照可能。
    """
    return torch.ones(seq_len, seq_len, dtype=torch.bool, device=device).tril()
 
 
def masked_self_attention(
    query: Tensor,               # [B, H, T, Dk]
    key: Tensor,                 # [B, H, T, Dk]
    value: Tensor,               # [B, H, T, Dv]
    padding_mask: Optional[Tensor] = None,  # [B, T] 1=実トークン
) -> AttentionOutput:
    """未来のトークンを参照せずに Self-Attention を計算する。
 
    Args:
        query, key, value: Head 分割済みの [B, H, T, *]。
        padding_mask: [B, T]。パディング位置を 0 にする。
 
    Returns:
        hidden_states [B, H, T, Dv] と attention_weights [B, H, T, T]。
    """
    d_k = query.shape[-1]
    scores = query @ key.transpose(-2, -1) / math.sqrt(d_k)   # [B, H, T, T]
 
    seq_len = query.shape[2]
    causal = build_causal_mask(seq_len, scores.device)        # [T, T]
    allowed = causal[None, None, :, :]                        # [1, 1, T, T]
 
    if padding_mask is not None:
        allowed = allowed & padding_mask[:, None, None, :].bool()
 
    scores = scores.masked_fill(~allowed, float("-inf"))
    weights = scores.softmax(dim=-1)                          # [B, H, T, T]
    context = weights @ value                                 # [B, H, T, Dv]
 
    return AttentionOutput(hidden_states=context, attention_weights=weights)
 
 
# 最小使用例
B, H, T, DK = 1, 8, 8, 64
q = torch.randn(B, H, T, DK)
k = torch.randn(B, H, T, DK)
v = torch.randn(B, H, T, DK)
 
out = masked_self_attention(q, k, v)
 
print(out.attention_weights.shape)          # torch.Size([1, 8, 8, 8]) == [B, H, T, T]
# 位置4（table）から見た重み。位置5以降は厳密に0になる
print(out.attention_weights[0, 0, 4, 5:])   # tensor([0., 0., 0.])
print(out.attention_weights[0, 0, 4].sum()) # tensor(1.0000)

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

Multi-Head の分割・結合は multi-head-self-attention のページと同じなので簡略化しています。
KV Cache（推論時の高速化）は含めていません。
実務では torch.nn.functional.scaled_dot_product_attention(is_causal=True) を使います。

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

推論時の計算：KV Cache

推論では 1 トークンずつ生成するため、素朴に実装すると毎ステップで系列全体の Key・Value を計算し直すことになります。

因果マスクがあるおかげで、過去の位置の Key・Value はステップが進んでも変化しません。そこで計算済みの Key・Value を保存しておき、新しい 1 トークンぶんだけを追加します。これが KV Cache です。

	Cache なし	Cache あり
1 ステップの計算量	$O(T^2 D)$	$O(T D)$
追加メモリ	なし	$O(B \cdot H \cdot T \cdot D_k)$ × 2（K と V）× 層数

長文脈の推論では、この KV Cache のメモリがボトルネックになります。 Grouped Query Attention（GQA）は、Key・Value の Head 数を減らしてこのメモリを削る手法です。

実装差

マスクの持ち方: bool マスクと、加算用の float マスク（0 / -inf）の 2 通りがあります。
is_causal フラグ: PyTorch の scaled_dot_product_attention は is_causal=True を渡すとマスクを実体化せずに処理でき、メモリを節約できます。
Sliding Window Attention: 直近 $w$ トークンだけを参照する変種（Mistral など）。マスクを帯行列にしたものと考えられます。
Prefix LM: プロンプト部分は双方向、生成部分は因果、という混合マスクを使う方式もあります。

マスク付き自己注意Masked (Causal) Self-Attention