第9章 Softmaxと重み付き和:文脈をベクトルに取り込む

こんにちは。ゆうせいです。

第8章までで、どのトークンがどのトークンにどれだけ注目すべきかを表すスコア行列が完成しました。しかし、この段階のスコアは、まだ確率ではありません。マイナスの値も含まれ、合計も1になっていません。第9章では、このスコアをSoftmax関数によって確率に変換し、その確率を使って実際にValueを合成する処理を扱います。ここが、自己注意という仕組みが完成する瞬間です。Transformer Explainerで、スコア行列の各行が、色の濃淡で表された確率の行に変わり、その先でValueと組み合わされる様子を追いながら読み進めてください。第5章から続いてきた自己注意の説明が、この章で一つにつながります。

Softmax関数とは何か

まず、Softmax関数の定義を確認します。

Softmax関数は、複数の数値を受け取り、それらをすべて0から1の範囲に収め、合計が1になるように変換する関数です。

softmax(x_i) = \frac{e^{x_i}}{\sum_{j} e^{x_j}}

分子は、その要素の指数関数です。分母は、すべての要素の指数関数の合計です。

計算の手順は、次の三段階です。

第一に、すべての値について、指数関数を計算します。

第二に、それらをすべて足し合わせます。

第三に、各値を、その合計で割ります。

具体的な数値で確認する

三つのトークンに対するスコアが、次の値だったとします。

[2.0, 1.0, -1.0]

第一段階として、指数関数を計算します。

e^{2.0} \approx 7.389

e^{1.0} \approx 2.718

e^{-1.0} \approx 0.368

第二段階として、合計を求めます。

7.389 + 2.718 + 0.368 = 10.475

第三段階として、それぞれを合計で割ります。

\frac{7.389}{10.475} \approx 0.705

\frac{2.718}{10.475} \approx 0.260

\frac{0.368}{10.475} \approx 0.035

結果は、次の通りです。

[0.705, 0.260, 0.035]

すべて0から1の範囲に収まり、合計は1になっています。

もとの値にはマイナスが含まれていましたが、変換後はすべてプラスになっている点に注目してください。指数関数は、どんな入力に対しても、必ず正の値を返すためです。

なぜ指数関数を使うのか

「合計が1になるようにしたいだけなら、単純に各値を合計で割ればよいのではないか」という疑問が浮かぶかもしれません。この疑問に、三つの観点から答えます。

理由1:マイナスの値を扱えるようにするため

単純に合計で割る方法では、マイナスの値がそのまま残ってしまいます。

先ほどの例で言えば、[2.0, 1.0, -1.0] の合計は2.0です。それぞれを2.0で割ると、[1.0, 0.5, -0.5] となります。

マイナス0.5という注目度は、意味を持ちません。「マイナスの割合で注目する」という状態が、何を表すのか説明できないのです。

さらに、合計がちょうど0になる場合には、ゼロ除算が発生して計算が破綻します。

指数関数を通すことで、すべての値が正になり、この問題が解消されます。

理由2:差を強調するため

指数関数は、入力の差を、出力では比として拡大します。

\frac{e^{2.0}}{e^{1.0}} = e^{1.0} \approx 2.718

入力の差が1であれば、出力の比は約2.7倍になります。入力の差が2であれば、出力の比は約7.4倍になります。

この性質により、スコアがわずかに高いトークンに、より強く注目が集まります。単純に合計で割る方法では、この差の強調が起きません。

第7章で扱ったスケーリングは、この強調が効きすぎるのを抑える処理でした。強調しすぎず、しなさすぎず、ちょうどよい程度に調整するために、両者が組み合わされているのです。

理由3:微分の計算が扱いやすいため

指数関数には、微分しても形が変わらないという性質があります。

\frac{d}{dx}e^x = e^x

この性質のおかげで、Softmax関数の微分が、比較的単純な形になります。

\frac{\partial softmax(x_i)}{\partial x_i} = softmax(x_i)(1 - softmax(x_i))

誤差逆伝播法で勾配を計算する際、この単純さは大きな利点です。

同時に、この式から、第7章で説明した飽和の問題も確認できます。softmaxの出力が1に近づくと、(1 - softmax(x_i)) が0に近づくため、微分値がほぼ0になります。出力が0に近づく場合も、前半の項が0に近づくため、やはり微分値がほぼ0になります。

出力が極端な値に張り付くと勾配が消える、という現象が、この式から読み取れます。

自己注意でのSoftmaxの適用方向

Softmaxをスコア行列に適用する際、重要な注意点があります。どの方向に適用するかです。

第6章で確認したとおり、スコア行列の行はQueryを出す側、列はKeyを提示する側です。

Softmaxは、行ごとに適用します。つまり、1つの行の中の値を、合計が1になるように変換します。

なぜ行方向なのでしょうか。

行は「1つのトークンが、他のすべてのトークンにどれだけ注目するか」を表しています。この注目の配分は、合計して100パーセントになるべきものです。「Aに70パーセント、Bに20パーセント、Cに10パーセント注目する」という形です。

一方、列方向に合計しても、意味のある量にはなりません。「あるトークンが、他のトークンたちから合計でどれだけ注目されているか」は、100パーセントに正規化すべき量ではないからです。人気の高いトークンは多くの注目を集め、そうでないトークンはあまり注目されない、という差があってよいのです。

実装上は、PyTorchであれば dim=-1 と指定します。この指定が、行方向(最後の次元方向)を意味します。

因果マスクとの関係

第8章で適用したマスクの効果が、ここで現れます。

マイナス無限大に設定された位置は、指数関数を通すと0になります。

e^{-\infty} = 0

分子が0になるため、その位置の確率は完全に0になります。

そして、Softmaxは残された要素だけで合計が1になるように配分するため、遮断された分の重みは、参照可能な位置に再配分されます。

具体例で確認します。3番目のトークンの行では、4番目以降がすべて遮断されています。したがって、1番目、2番目、3番目の三つだけで、合計が1になるように配分されます。

最初のトークンの行では、自分自身しか残っていないため、その位置の確率が1になります。

重み付き和:Valueを合成する

Softmaxによって得られた確率は、注意の重み(Attention Weight)と呼ばれます。この重みを使って、実際に情報を取り込む処理が、自己注意の最終段階です。

計算の仕組み

各トークンについて、すべてのトークンのValueを、注意の重みに応じて足し合わせます。

出力_i = \sum_j 重み_{ij} \times Value_j

具体例で確認します。3番目のトークンの注意の重みが、次の値だったとします。

1番目への重み:0.2

2番目への重み:0.5

3番目への重み:0.3

4番目以降:0(マスクにより遮断)

このとき、3番目のトークンの出力は、次のように計算されます。

出力_3 = 0.2 \times Value_1 + 0.5 \times Value_2 + 0.3 \times Value_3

各Valueは64次元のベクトルですから、この計算は、三つのベクトルを重み付きで足し合わせる操作になります。結果も、64次元のベクトルです。

この計算が意味すること

ここで、自己注意が何をしているかが明確になります。

3番目のトークンの新しい表現は、自分自身の情報を30パーセント、2番目のトークンの情報を50パーセント、1番目のトークンの情報を20パーセント含んだ、混合物になります。

第5章で挙げた「it」の例に戻ります。「it」というトークンが「animal」に強く注目していれば、「it」の出力ベクトルには、「animal」の意味情報が大きく取り込まれます。

こうして、単なる代名詞だった「it」のベクトルが、「animalを指している代名詞」という文脈を反映したベクトルへと更新されるのです。

第3章の最後で述べた「文脈による意味の違いを区別できない」という問題が、ここでようやく解決されます。

行列の掛け算として一括で計算する

実装上は、この計算も行列の掛け算で一括処理されます。

出力 = 重み行列 \times V

重み行列の形は、n行n列です。Vの形は、n行64列です。したがって、出力の形は、n行64列になります。

入力と同じ形のベクトルが、文脈の情報を取り込んだ状態で出てくる、という構造です。

自己注意の全体像を数式で確認する

第5章から第9章までの処理を、一つの数式にまとめると、次のようになります。

Attention(Q, K, V) = softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)V

この式を、これまでの章と対応させます。

Q、K、Vの生成が、第5章です。

QK^T が、第6章の内積です。

\sqrt{d_k} で割る部分が、第7章のスケーリングです。

式には明示されていませんが、Softmaxの直前にマスクが適用されるのが、第8章です。

softmax関数と、その結果にVを掛ける部分が、この第9章です。

一行の数式に見えますが、そこには五つの章にわたる処理が凝縮されています。数式を見て内容を思い出せるようになれば、自己注意の理解は十分です。

Transformer Explainerで観察できること

Transformer Explainerでは、Softmax適用後の重み行列が、色の濃淡を伴って表示されます。

次の点を確認してみてください。

一つ目は、各行の値をすべて足すと1になるという点です。色の濃い位置と薄い位置がありますが、1つの行の中で合計すると必ず1になります。

二つ目は、右上の三角形の領域が、完全に0になっているという点です。マスクの効果が、確率として現れています。

三つ目は、最初のトークンの行が、1つの位置だけ最も濃い色になっているという点です。参照できる先が自分自身しかないため、確率が1になっています。

四つ目は、重みが特定の位置に偏っている行と、比較的均等に分散している行があるという点です。文法的に強い結びつきを持つトークンでは偏りが大きく、そうでないトークンでは分散する傾向があります。

初学者がつまずきやすい点

つまずき1:Softmaxを適用する方向を間違える

行方向か列方向かは、実装で最も間違えやすい箇所です。

間違えても、エラーにはなりません。計算は正常に完了し、それらしい数値が出てきます。しかし、モデルはまったく学習しません。

デバッグの際は、重み行列の各行の合計が1になっているかを確認する習慣をつけてください。列の合計が1になっていれば、方向を間違えています。

つまずき2:注意の重みが「説明」だと考えすぎる

注意の重みは、モデルがどこを見ているかを示す指標として、可視化によく使われます。

しかし、注意の重みが高いことが、そのままモデルの判断根拠であるとは限りません。この点については、注意の重みは説明にならないという趣旨の研究も発表されており、解釈には慎重さが求められます。

注意の重みは、モデルの内部状態を観察する手がかりの一つではありますが、判断の理由そのものではない、と理解しておいてください。

つまずき3:Softmaxの数値的な安定性の問題を知らない

Softmaxの計算には、実装上の落とし穴があります。

入力の値が大きいと、指数関数の計算で、コンピュータが扱える数値の上限を超えてしまいます。これをオーバーフローと呼びます。

e^{1000}

この値は、通常の浮動小数点数では表現できません。

この問題を回避するため、実際の実装では、入力から最大値を引いてから計算するという工夫が行われます。

softmax(x_i) = \frac{e^{x_i - max(x)}}{\sum_j e^{x_j - max(x)}}

すべての要素から同じ値を引いても、結果は変わりません。分子と分母の両方が同じ定数倍されるため、割り算で打ち消されるからです。

この工夫により、指数関数の入力が必ず0以下になり、オーバーフローを避けられます。

PyTorchやTensorFlowのSoftmax関数には、この処理があらかじめ組み込まれています。自分で実装する場合は、忘れずに入れてください。

演習

演習1:Softmaxを手計算する

次のスコアに対して、Softmaxを計算してください。小数点以下3桁まで求めてください。

[1.0, 0.0, -1.0]

必要な値として、e^{1.0} \approx 2.718e^{0.0} = 1.000e^{-1.0} \approx 0.368 を使ってください。

演習2:重み付き和を手計算する

3番目のトークンの注意の重みが、次の通りだとします。

1番目への重み:0.5

2番目への重み:0.3

3番目への重み:0.2

各トークンのValueが、次の3次元ベクトルだとします。

Value_1 = [2, 0, 1]

Value_2 = [0, 4, 2]

Value_3 = [1, 1, 0]

3番目のトークンの出力ベクトルを計算してください。

演習3:自己注意を最初から最後まで実装する

import torch
import torch.nn.functional as F
import math

torch.manual_seed(0)

# 設定
seq_len = 4
d_k = 8

# Q、K、Vをランダムに用意(本来は入力から生成される)
Q = torch.randn(seq_len, d_k)
K = torch.randn(seq_len, d_k)
V = torch.randn(seq_len, d_k)

# 第6章:内積によるスコア計算
scores = Q @ K.T
print("=== 第6章: スコア行列 ===")
print(scores.numpy().round(3))

# 第7章:スケーリング
scores = scores / math.sqrt(d_k)
print("\n=== 第7章: スケーリング後 ===")
print(scores.numpy().round(3))

# 第8章:因果マスク
mask = torch.tril(torch.ones(seq_len, seq_len))
scores = scores.masked_fill(mask == 0, float('-inf'))
print("\n=== 第8章: マスク適用後 ===")
print(scores.numpy().round(3))

# 第9章:Softmax
weights = F.softmax(scores, dim=-1)
print("\n=== 第9章: Softmax後の注意の重み ===")
print(weights.numpy().round(3))
print(f"各行の合計: {weights.sum(dim=-1).tolist()}")

# 第9章:重み付き和
output = weights @ V
print("\n=== 第9章: 出力 ===")
print(f"出力の形: {output.shape}")
print(output.numpy().round(3))

# 検証:PyTorchの標準実装と比較
print("\n=== 標準実装との比較 ===")
official = F.scaled_dot_product_attention(Q, K, V, is_causal=True)
print(f"一致しているか: {torch.allclose(output, official, atol=1e-5)}")

演習の解答例と解説

演習1の解答は、次の通りです。

指数関数の合計は、次のようになります。

2.718 + 1.000 + 0.368 = 4.086

それぞれを割ります。

\frac{2.718}{4.086} \approx 0.665

\frac{1.000}{4.086} \approx 0.245

\frac{0.368}{4.086} \approx 0.090

結果は、[0.665, 0.245, 0.090] です。

もとのスコアの差は1ずつでしたが、確率では0.665、0.245、0.090と、比が拡大されています。指数関数による差の強調が確認できます。

演習2の解答は、次の通りです。

各次元ごとに計算します。

第1次元:0.5 \times 2 + 0.3 \times 0 + 0.2 \times 1 = 1.0 + 0 + 0.2 = 1.2

第2次元:0.5 \times 0 + 0.3 \times 4 + 0.2 \times 1 = 0 + 1.2 + 0.2 = 1.4

第3次元:0.5 \times 1 + 0.3 \times 2 + 0.2 \times 0 = 0.5 + 0.6 + 0 = 1.1

出力ベクトルは、[1.2, 1.4, 1.1] です。

この計算から、出力が三つのValueの混合物であることが、具体的に確認できます。1番目のValueに最も強く影響されているものの、他の二つの情報も含まれています。

演習3では、自分で組み立てた処理と、PyTorchの標準実装が一致することが確認できます。第6章から第9章までの内容が、正しく理解できているかを検証する演習です。

講師向けの補足

この章は、第5章から続いてきた自己注意の説明が完結する回です。時間があれば、第5章からの流れを、一枚の図にまとめて振り返る時間を設けてください。

受講者が最も感動するのは、演習3で、自分が組み立てた計算がPyTorchの標準実装と一致することを確認した瞬間です。ブラックボックスだと思っていたものが、自分の手で再現できたという体験は、その後の学習意欲に大きく影響します。この演習は、必ず実施することを推奨します。

また、「it」が「animal」を指すという例を、この章でもう一度取り上げてください。第5章では役割の説明にとどまっていたものが、この章で実際に「animalの情報が取り込まれる」という結果として完結します。同じ例を最初と最後で使うことで、一連の処理が何のためのものだったかが明確になります。

理解度を確認する問い

第9章の内容を、次の問いで確認してみてください。

Softmax関数で、単純に合計で割るのではなく指数関数を使う理由を、三つ挙げてください。

Softmaxをスコア行列に適用する際、行方向に適用する理由は何でしょうか。

注意の重みを使ってValueを重み付きで足し合わせることは、意味的には何をしていることになるでしょうか。

まとめ

Softmax関数は、スコアを0から1の範囲に収め、合計が1になる確率へと変換する関数です。指数関数を使うことで、マイナスの値を扱えるようにし、値の差を比として強調し、微分の計算を扱いやすくしています。

Softmaxは、スコア行列の行ごとに適用されます。行は、1つのトークンが他のトークンにどれだけ注目するかを表しており、その配分の合計が1になるべきだからです。

第8章でマイナス無限大に設定された位置は、指数関数を通すと0になり、確率が完全に0になります。遮断された分の重みは、参照可能な位置に再配分されます。

得られた確率を注意の重みと呼び、この重みを使って、すべてのトークンのValueを重み付きで足し合わせます。この操作により、各トークンのベクトルに、周囲のトークンの情報が、注目度に応じて取り込まれます。

これが、自己注意という仕組みの全体です。第5章から第9章までの処理をまとめると、次の一行の数式になります。

Attention(Q, K, V) = softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)V

次の第10章では、マルチヘッドアテンションを扱います。ここまで説明してきた自己注意を、なぜ一つではなく複数並列に動かすのか、その結果をどのように統合するのかを解説します。

第10章に進む前に、演習3のプログラムを実際に動かし、自己注意の一連の処理が、自分の手で再現できることを確認しておいてください。この確認ができていると、次章の並列化の話が、既知の処理の繰り返しとして、すんなり理解できるはずです。

セイ・コンサルティング・グループでは新人エンジニア研修のアシスタント講師を募集しています。

投稿者プロフィール

山崎講師
山崎講師代表取締役
セイ・コンサルティング・グループ株式会社代表取締役。
岐阜県出身。
海外放浪の末、2000年創業、2004年会社設立。
IT企業向け人材育成研修歴業界歴20年以上。
すべての無駄を省いた費用対効果の高い「筋肉質」な研修を提供します!
この記事に間違い等ありましたらぜひお知らせください。

学生時代は趣味と実益を兼ねてリゾートバイトにいそしむ。長野県白馬村に始まり、志賀高原でのスキーインストラクター、沖縄石垣島、北海道トマム。高じてオーストラリアのゴールドコーストでツアーガイドなど。現在は野菜作りにはまっている。