Tech カテゴリのイメージ
Tech

因果マスクとは？LLMが学習中に未来の答えを見ない仕組み

Self-Attentionは、ある位置のQueryと系列内のKeyを比べ、どのValueから情報を受け取るか決める。ところが文章生成モデルで全位置を自由に参照させると、学習中だけ未来の正解トークンが見えてしまう。

たとえば「猫 は 眠る 。」から次のトークンを学ぶとき、「猫」の位置が後ろの「は」や「眠る」を参照できれば、予測ではなく答えの盗み見になる。因果マスクは、各位置から未来方向へ伸びるAttentionをsoftmaxの前に無効化する仕組みだ。

見えてよい範囲は三角形になる

4トークンを左から0、1、2、3番目とする。位置0が参照できるのは位置0だけ、位置1は位置0と1、位置2は位置0から2までだ。参照してよい場所を1、隠す場所を0とすると次の行列になる。

[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]

左下だけが残るので、下三角行列と呼ばれる。対角成分を残すのは、各位置が自分自身を参照してよいからだ。

Transformer原論文のデコーダーでも、位置 ii の予測は ii より前の既知の出力だけに依存するようSelf-Attentionを変更している。入力と正解を1位置ずらす処理と、このマスクを組み合わせることで次トークン予測の条件を守る。

softmaxの前に未来を無効化する

マスク対象のAttentionスコアへ −∞-\infty を加えてからsoftmaxを計算する。指数関数では e−∞=0e^{-\infty}=0 になるため、未来位置の重みは0になる。

softmax⁡(S+M)\operatorname{softmax}(S + M)

SS は通常のAttentionスコア、MM は参照可能な場所が0、未来の場所が −∞-\infty の行列だ。softmax後の重みへ0を掛けるだけでは、残った重みの合計を1へ戻す処理が別途必要になる。スコア段階でマスクすれば、softmaxが参照可能な位置だけを使って正規化する。

4トークンで重みの変化を見る

次の例では、説明用にすでに計算済みのスコア行列を使う。実際にはQueryとKeyの積を次元数の平方根で割ってスコアを作る。

import numpy as np


scores = np.array([
    [2.0, 1.0, 0.0, -1.0],
    [1.0, 2.0, 1.0, 0.0],
    [0.0, 1.0, 2.0, 1.0],
    [-1.0, 0.0, 1.0, 2.0],
])


def softmax(x):
    shifted = x - np.max(x, axis=-1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=-1, keepdims=True)


plain = softmax(scores)


# 対角線より右上を未来として隠す
causal_mask = np.triu(np.ones_like(scores, dtype=bool), k=1)
masked_scores = np.where(causal_mask, -np.inf, scores)
causal = softmax(masked_scores)


np.set_printoptions(precision=3, suppress=True)
print("mask (1 means hidden):")
print(causal_mask.astype(int))
print("without mask:")
print(plain)
print("with causal mask:")
print(causal)
print("row sums:", causal.sum(axis=1))

出力:

mask (1 means hidden):
[[0 1 1 1]
 [0 0 1 1]
 [0 0 0 1]
 [0 0 0 0]]
without mask:
[[0.644 0.237 0.087 0.032]
 [0.197 0.534 0.197 0.072]
 [0.072 0.197 0.534 0.197]
 [0.032 0.087 0.237 0.644]]
with causal mask:
[[1.    0.    0.    0.   ]
 [0.269 0.731 0.    0.   ]
 [0.09  0.245 0.665 0.   ]
 [0.032 0.087 0.237 0.644]]
row sums: [1. 1. 1. 1.]

マスクなしの1行目は、まだ現れていない位置1、2、3にも重みを割り当てている。因果マスクを入れると未来側はすべて0になり、位置0は自分だけを見る。2行目は位置0と1だけで重みを配り直す。最後の位置には未来がないため、マスク前後の値は変わらない。

推論時にもマスクは必要か

1トークンずつ生成する単純な推論では、未来のKeyとValueがまだ存在しない。そのため、最後のQueryだけを計算する場面では未来を参照しようがない。一方、学習では複数位置を一度に並列計算するため、因果マスクがなければ後続トークンを見られる。

実装では学習だけに限らず、複数Queryをまとめて処理する推論やキャッシュを使う処理でも因果関係を正しく揃える必要がある。「生成時は常にマスク不要」と一般化せず、QueryとKeyの系列長や位置の対応を確認する。

パディングマスクとは目的が違う

因果マスクは未来を隠す。パディングマスクは、系列長を揃えるために追加した空の位置を隠す。両方を同時に使う場合もあるが、理由は別だ。

PyTorchではAPIによって真偽値マスクの意味が逆になる点にも注意がいる。torch.nn.functional.scaled_dot_product_attention の真偽値 attn_mask は True が参加可能な位置を示す。一方、nn.MultiheadAttention の key_padding_mask では True が無視する位置を示す。見た目が同じ真偽値行列でも、そのまま流用しないほうがよい。

scaled_dot_product_attention では is_causal=True を指定すると、正方形のAttention行列に下三角の因果マスクが適用される。自前の attn_mask と同時指定できないため、別の制約も重ねたいときは利用するAPIの仕様を確認する必要がある。

因果マスクは、モデルへ時系列の意味を新しく教える部品ではない。学習課題の条件を壊さないための情報制限だ。次は、マスクされた表現から語彙上の各トークンへlogitsを出し、softmaxで次の候補を比べる流れへ進む。

参照資料

生成AI・機械学習講座の一覧へ

関連記事