なぜAttentionはスケールドドット積なのか — √d_kで割る理由を徹底解説

Transformerのアテンションの式を初めて見たとき、多くの人が「なぜ内積をわざわざ $\sqrt{d_k}$ で割るのだろう?」と引っかかります。

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

分子の $QK^\top$(クエリとキーの内積)は分かります。softmax で確率にするのも分かります。でも、その間に挟まっている $\sqrt{d_k}$ での割り算は、パッと見では正体不明です。定数で割るだけなら、softmax の平行移動不変性みたいに「大した影響はない」のでは?と思うかもしれません。

ところが、この一手間こそがアテンションを「ちゃんと学習できるもの」にしている急所です。原論文 “Attention Is All You Need” でも、この項は “Scaled Dot-Product Attention”(スケールドドット積アテンション)と名付けられ、わざわざ節のタイトルになっています。割らないと、次元 $d_k$ が大きいときに softmax が飽和して勾配がほぼ0になり、学習が止まってしまう のです。

この $\sqrt{d_k}$ の意味を理解すると、次のような場面で役に立ちます。

  • モデル設計: ヘッドの次元 $d_k$ を変えたとき、なぜスコアの分散が変わるのか、初期化やスケーリングをどう合わせるべきかが分かる
  • 温度スケーリング: $\sqrt{d_k}$ で割ることは「温度 $T=\sqrt{d_k}$ のソフトマックス」と同じであり、知識蒸留やサンプリングの温度制御と同じ枠組みで理解できる

本記事の内容

  • なぜ内積 $q\cdot k$ の分散が $d_k$ になるのか(期待値・分散の性質から丁寧に導出)
  • 大きい内積が softmax を飽和させ、勾配を消す仕組み
  • $\sqrt{d_k}$ で割ると分散が1に戻り、softmax が適度な鋭さに保たれること
  • スケーリングを「温度」として見る視点
  • numpy による実測(分散・エントロピー・勾配の大きさ)とグラフでの確認

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

スケールドドット積アテンションの概念図(内積→√d_kで割る→softmax→重み)

上の概念図が、この記事で追いかける流れの全体像です。クエリ $q$ とキー $k$ の内積を取り、それを $\sqrt{d_k}$ で割ってから softmax に通してアテンション重みを作ります。ポイントは中央の割り算で、内積の「大きくばらつく」分散 $d_k$ を「適度な」分散 1 に戻す役割を担っています。この一手間が softmax の鋭さを整え、勾配が流れる状態を保ちます。

まず結論:内積は次元とともに「大きくばらつく」

細かい式に入る前に、直感を掴みましょう。

アテンションのスコアは、クエリベクトル $q$ とキーベクトル $k$ の内積 $q \cdot k = \sum_i q_i k_i$ です。これは「$q$ と $k$ がどれだけ似ているか」を測る量です。ここで大事なのは、内積が たくさんの成分の足し算 だということです。

いま $q$ と $k$ の各成分が、平均0・分散1くらいのランダムな値だとしましょう(学習初期はまさにこんな状態です)。次元 $d_k$ が4なら、4個の項を足します。$d_k$ が256なら、256個の項を足します。足す項が増えれば増えるほど、合計の値は大きく振れます。つまり 次元 $d_k$ が大きいほど、内積 $q\cdot k$ は大きな正・大きな負の値を取りやすくなる のです。

コインを4枚投げた合計と、256枚投げた合計を比べるようなものです。枚数が増えるほど、合計のばらつき(分散)は大きくなります。この「ばらつきの拡大」を数値で見てみましょう。

import numpy as np

rng = np.random.default_rng(0)

for dk in [4, 16, 64, 256]:
    n = 100000
    q = rng.standard_normal((n, dk))   # 各成分 平均0・分散1
    k = rng.standard_normal((n, dk))
    dots = np.sum(q * k, axis=1)        # 内積 q・k を n 回計算
    print(f"d_k={dk:4d}: 平均={dots.mean():7.3f}  分散={dots.var():8.2f}"
          f"  標準偏差={dots.std():6.2f} (理論 √d_k = {np.sqrt(dk):.2f})")

このコードを実行すると、次のような出力が得られます。

d_k=   4: 平均= -0.014  分散=    4.03  標準偏差=  2.01 (理論 √d_k = 2.00)
d_k=  16: 平均=  0.000  分散=   16.01  標準偏差=  4.00 (理論 √d_k = 4.00)
d_k=  64: 平均= -0.000  分散=   63.86  標準偏差=  7.99 (理論 √d_k = 8.00)
d_k= 256: 平均=  0.008  分散=  255.13  標準偏差= 15.97 (理論 √d_k = 16.00)

出力から、はっきりした規則が読み取れます。内積の 平均はどの次元でもほぼ0 ですが、分散は次元 $d_k$ とほぼ一致 します。$d_k=4$ なら分散は約4、$d_k=256$ なら約256です。標準偏差は $\sqrt{d_k}$ に一致し、$d_k=256$ では内積が平均的に $\pm16$ 程度も振れています。次元を上げただけで、スコアのスケールがこんなに膨らむのです。

同じことをヒストグラムで見ると、分布の広がり方が視覚的に分かります。

次元d_kが大きいほど内積の分布が広がるヒストグラム

この図では、$d_k$ を4から256まで変えたときの内積 $q\cdot k$ の分布を重ねています。どの分布も0を中心にした釣鐘型ですが、$d_k$ が大きくなるほど裾が横に広がっていきます。$d_k=4$ では大半が $\pm5$ に収まるのに対し、$d_k=256$ では $\pm40$ を超える値も珍しくありません。次元を上げるとスコアが「桁違いに大きくなり得る」ことが一目で分かります。

なぜ分散がちょうど $d_k$ になるのか。これは偶然ではなく、期待値と分散の性質からきっちり導けます。次にその導出を見ていきましょう。

なぜ内積の分散は d_k になるのか(導出)

ここでのゴールは、「$q$ と $k$ の各成分が独立に平均0・分散1のとき、内積 $q\cdot k$ の平均は0、分散は $d_k$ になる」ことを示すことです。

準備:各成分についての仮定

$q = (q_1, \dots, q_{d_k})$、$k = (k_1, \dots, k_{d_k})$ とし、すべての成分が互いに独立で、次を満たすとします。

$$ \mathbb{E}[q_i] = 0, \quad \mathbb{E}[k_i] = 0, \quad \mathrm{Var}(q_i) = 1, \quad \mathrm{Var}(k_i) = 1 $$

分散1・平均0という仮定は、期待値と分散の定義から次のように言い換えられます。あとで使うので確認しておきます。

$$ \mathbb{E}[q_i^2] = \mathrm{Var}(q_i) + (\mathbb{E}[q_i])^2 = 1 + 0 = 1 $$

同様に $\mathbb{E}[k_i^2] = 1$ です。

ステップ1:各項 q_i k_i の平均を求める

内積は $q\cdot k = \sum_{i=1}^{d_k} q_i k_i$ という項の和です。まず1つの項 $q_i k_i$ の平均を調べます。$q_i$ と $k_i$ は独立なので、積の期待値は期待値の積に分解できます。

$$ \mathbb{E}[q_i k_i] = \mathbb{E}[q_i]\,\mathbb{E}[k_i] = 0 \times 0 = 0 $$

つまり、各項は平均0です。

ステップ2:各項 q_i k_i の分散を求める

次に1つの項の分散です。分散の定義 $\mathrm{Var}(X) = \mathbb{E}[X^2] – (\mathbb{E}[X])^2$ を使います。いま $\mathbb{E}[q_i k_i] = 0$ なので、第2項は消えて $\mathrm{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2]$ です。

$(q_i k_i)^2 = q_i^2 k_i^2$ と分け、$q_i$ と $k_i$ の独立性から期待値を分解すると、

$$ \mathrm{Var}(q_i k_i) = \mathbb{E}[q_i^2 k_i^2] = \mathbb{E}[q_i^2]\,\mathbb{E}[k_i^2] $$

ここで先ほど準備した $\mathbb{E}[q_i^2] = 1$、$\mathbb{E}[k_i^2] = 1$ を代入すると、

$$ \mathrm{Var}(q_i k_i) = 1 \times 1 = 1 $$

各項は平均0・分散1になります。

ステップ3:和の分散を求める

いよいよ内積全体の分散です。$d_k$ 個の項 $q_1 k_1, \dots, q_{d_k} k_{d_k}$ は互いに独立です(異なる $i$ の成分はすべて独立だから)。独立な確率変数の和の分散は、各項の分散の 単純な足し算 になります。これが導出の核心です。

$$ \mathrm{Var}\!\left(\sum_{i=1}^{d_k} q_i k_i\right) = \sum_{i=1}^{d_k} \mathrm{Var}(q_i k_i) $$

なぜ和の分散が足し算になるのかを補足します。一般に $\mathrm{Var}(X+Y) = \mathrm{Var}(X) + \mathrm{Var}(Y) + 2\,\mathrm{Cov}(X, Y)$ ですが、独立なら共分散 $\mathrm{Cov}(X, Y) = 0$ なので 交差項がすべて消えます。だから分散だけが足し合わさるのです。ここに各項の分散1を代入すると、

$$ \mathrm{Var}(q \cdot k) = \underbrace{1 + 1 + \cdots + 1}_{d_k \text{ 個}} = d_k $$

平均については、期待値は常に線形なので(独立性すら不要で)、

$$ \mathbb{E}[q\cdot k] = \sum_{i=1}^{d_k} \mathbb{E}[q_i k_i] = 0 $$

以上で、内積 $q\cdot k$ は 平均0・分散 $d_k$、したがって 標準偏差 $\sqrt{d_k}$ であることが示せました。先ほどの numpy 実測($d_k=64$ で分散約64、標準偏差約8)とぴったり一致します。

この導出の流れを1枚にまとめたのが次の図です。

内積の分散がd_kになる導出の直感図

図の通り、内積は「平均0・分散1の項」を $d_k$ 個足したものです。独立性のおかげで交差項が消え、分散はただ $d_k$ 個の1を足すだけになります。だから標準偏差は $\sqrt{d_k}$。この $\sqrt{d_k}$ こそ、後で割り算に使う量そのものです。式のどこから $\sqrt{d_k}$ が出てくるのかが、これでクリアになりました。

分散が $d_k$ に一致することを、次元を横軸に取って改めて確かめておきます。

import numpy as np

rng = np.random.default_rng(1)
dks = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
for dk in dks:
    n = 40000
    q = rng.standard_normal((n, dk))
    k = rng.standard_normal((n, dk))
    dots = np.sum(q * k, axis=1)
    print(f"d_k={dk:4d}  実測分散={dots.var():8.2f}  理論 d_k={dk}")

内積の分散がd_kに一致することの実測

両対数軸で見ると、実測点(赤)が理論直線 $\mathrm{Var}=d_k$(破線)の上にきれいに乗っています。$d_k$ を1から512まで動かしても、内積の分散は常に $d_k$ に一致します。乱数の揺らぎで多少ずれても、傾き1の直線からは外れません。「分散 = 次元」という関係が、広い範囲で成り立つことが確認できました。

さて、内積のスケールが $\sqrt{d_k}$ に比例して膨らむことは分かりました。では、その大きなスコアを softmax に通すと何が起きるのでしょうか。

大きい内積は softmax を飽和させる

内積のスコアはそのまま使うわけではなく、softmax に通してアテンション重み(確率)に変換します。ここで、スコアが大きすぎると深刻な問題が起きます。

softmax は差が大きいと one-hot に潰れる

softmax は入力の に反応します。復習すると、要素 $z_i$ に対する出力は

$$ \mathrm{softmax}(\bm{z})_i = \frac{e^{z_i}}{\sum_j e^{z_j}} $$

でした。指数関数は差を比に変えるので、最大要素とそれ以外の差が大きいほど、最大要素の確率が1に近づきます。たとえば $(1, 0)$ の差は1で softmax は $(0.73, 0.27)$ ですが、$(10, 0)$ になると $(0.9999, 0.0001)$ とほぼ one-hot です。

内積の分散が $d_k$ ということは、スコアの差の大きさも $\sqrt{d_k}$ に比例して膨らむということです。$d_k=64$ なら、最大スコアとそれ以外の差は平均的に10前後にもなります。その結果、softmax は ほぼ one-hot に潰れてしまいます。1つのキーだけに重み1、他はほぼ0、という極端な分布です。

実際に、$d_k=64$・キー8本で、割らない場合と割る場合の softmax 出力を比べてみましょう。

import numpy as np

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

def entropy(p):
    return -np.sum(p * np.log(p + 1e-12))

rng = np.random.default_rng(2)
dk, n_keys = 64, 8
q = rng.standard_normal(dk)
K = rng.standard_normal((n_keys, dk))

raw = K @ q                       # 未スケールの内積スコア
scaled = raw / np.sqrt(dk)        # √d_k で割ったスコア

p_raw = softmax(raw)
p_scaled = softmax(scaled)

print(f"未スケール:  最大重み={p_raw.max():.3f}  エントロピー={entropy(p_raw):.3f}")
print(f"√d_kで割る:  最大重み={p_scaled.max():.3f}  エントロピー={entropy(p_scaled):.3f}")
print(f"一様分布のエントロピー = ln({n_keys}) = {np.log(n_keys):.3f}")

このコードの出力は次の通りです。

未スケール:  最大重み=0.906  エントロピー=0.314
√d_kで割る:  最大重み=0.328  エントロピー=1.812
一様分布のエントロピー = ln(8) = 2.079

出力から違いは明らかです。未スケールでは最大重みが0.906 で、1本のキーにほぼ全部の重みが集中しています。エントロピーは0.314しかなく、分布がとても尖っています。一方 $\sqrt{d_k}$ で割ると最大重みは0.328 に下がり、エントロピーは1.812(一様分布の上限 $\ln 8 \approx 2.079$ に近い)まで上がります。割るだけで、分布が「1点集中」から「適度に分散」へ変わるのです。

この差を棒グラフで見ると、直感的です。

スケーリング有無でのsoftmax出力の棒グラフ比較

左(スケーリング無し)は、キー0番の重みが0.9を超え、他はほぼ0。実質的に1本のキーしか見ていません。右($\sqrt{d_k}$ で割る)は、複数のキーに重みが分散していて、番号0・3を中心にしつつ他も無視していません。アテンションは「複数の情報を混ぜて使う」のが本来の役割なので、左のような1点集中は望ましくありません。

「1点集中でも、正しいキーを選べているなら良いのでは?」と思うかもしれません。しかし本当の問題は、この飽和状態では 学習が進まない ことにあります。次にそれを見ます。

飽和すると勾配が消える

深層学習のモデルは、勾配を使ってパラメータを更新します。勾配がほぼ0だと、その部分は学習が止まります。softmax が飽和すると、まさにこの勾配消失が起きます。

softmax のヤコビアンは飽和でほぼ0になる

softmax のヤコビアン(各出力の各入力に対する偏微分)は、ソフトマックス関数の記事で導いたように

$$ \frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij} – p_j), \qquad \frac{\partial \bm{p}}{\partial \bm{z}} = \mathrm{diag}(\bm{p}) – \bm{p}\bm{p}^\top $$

です。ここで出力 $\bm{p}$ がほぼ one-hot、つまりある成分が1で他が0に近いとどうなるか考えます。$p_i \approx 1$ のとき対角成分は $p_i(1-p_i) \approx 1\times 0 = 0$、$p_i \approx 0$ の成分も $p_i(\cdots) \approx 0$。どの成分もほぼ0 になり、ヤコビアン全体が0行列に近づきます。

これは1次元のイメージで見ると分かりやすいです。2要素の softmax は、差 $\delta = z_1 – z_2$ のシグモイド $p_1 = \sigma(\delta)$ で、その勾配は $p_1(1-p_1)$ です。

softmax飽和とロジット差の関係

青の実線が softmax 出力 $p_1$、赤の破線がその勾配 $p_1(1-p_1)$ です。差 $\delta$ が0付近では勾配が最大(0.25)ですが、差が $\pm5$ を超える灰色の「飽和域」に入ると、出力は0か1に張り付き、勾配はほぼ0に落ちます。内積の差が $\sqrt{d_k}$ に比例して大きくなると、まさにこの飽和域に押し込まれてしまうのです。

勾配の大きさを実測する

実際に、割らない場合と割る場合でヤコビアンのノルム(勾配の大きさの目安)を比べてみます。

import numpy as np

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

def jac_norm(logits):
    p = softmax(logits)
    J = np.diag(p) - np.outer(p, p)     # softmax のヤコビアン
    return np.linalg.norm(J)            # フロベニウスノルム

rng = np.random.default_rng(4)
dk, n_keys = 64, 8
trials = 5000
jr, js, er, es, mr, ms = [], [], [], [], [], []
for _ in range(trials):
    q = rng.standard_normal(dk)
    K = rng.standard_normal((n_keys, dk))
    raw = K @ q
    scaled = raw / np.sqrt(dk)
    jr.append(jac_norm(raw));    js.append(jac_norm(scaled))
    p_r, p_s = softmax(raw), softmax(scaled)
    mr.append(p_r.max());        ms.append(p_s.max())

print(f"未スケール:  平均最大重み={np.mean(mr):.3f}  平均ヤコビアンノルム={np.mean(jr):.4f}")
print(f"√d_kで割る:  平均最大重み={np.mean(ms):.3f}  平均ヤコビアンノルム={np.mean(js):.4f}")

5000試行の平均を取ると、次の結果になります。

未スケール:  平均最大重み=0.875  平均ヤコビアンノルム=0.1515
√d_kで割る:  平均最大重み=0.361  平均ヤコビアンノルム=0.3657

この結果から、飽和と勾配消失の関係がはっきりします。未スケールでは平均最大重みが0.875 と飽和気味で、そのときの平均ヤコビアンノルムは0.1515。一方 $\sqrt{d_k}$ で割ると最大重みは0.361 に下がり、ヤコビアンノルムは0.3657と 2倍以上大きく なります。割らないと勾配が半分以下に痩せてしまうわけです。学習初期にこれが起きると、アテンション層に十分な勾配が流れず、学習が停滞します。

スコアのスケールを連続的に変えて、勾配がどう変化するかも見ておきましょう。

スコアスケールが大きいと勾配が消失する様子

横軸はスコアのスケール係数で、係数1が「ちょうど $\sqrt{d_k}$ で割った状態」に対応します。係数が1(緑の破線)付近では勾配(ヤコビアンのノルム)がしっかり残っていますが、係数を2倍・3倍と大きくすると、勾配は急速に0へ向かって落ちていきます。スコアを大きくしすぎると、勾配が消えて学習が進まなくなることが定量的に見て取れます。

同じ横軸でエントロピーの変化も確認しておきます。

スコアスケールとsoftmaxエントロピーの関係

スケール係数が小さいうちはエントロピーが一様分布 $\ln 8$ に近く、分布が「フラット」です。係数を大きくするとエントロピーは単調に下がり、分布がどんどん尖っていきます。勾配の図とあわせて見ると、エントロピーが下がる(尖る)ほど勾配も小さくなる という対応が読み取れます。適度なエントロピーを保つ係数1、すなわち $\sqrt{d_k}$ で割る選択が、ちょうど良いバランス点になっているのです。

ここまでで「割らないと飽和して勾配が消える」ことが分かりました。では $\sqrt{d_k}$ で割ると、なぜちょうど良くなるのでしょうか。

√d_k で割ると分散が1に戻る

ここが $\sqrt{d_k}$ を選ぶ理由のクライマックスです。

内積 $q\cdot k$ の標準偏差は $\sqrt{d_k}$ でした。分散でいえば $d_k$ です。この量を $\sqrt{d_k}$ で割ると、分散はどうなるでしょうか。確率変数を定数 $c$ で割ると分散は $c^2$ で割られる、という基本性質を使います。

$$ \mathrm{Var}\!\left(\frac{q\cdot k}{\sqrt{d_k}}\right) = \frac{1}{(\sqrt{d_k})^2}\,\mathrm{Var}(q\cdot k) = \frac{1}{d_k}\cdot d_k = 1 $$

きれいに 分散が1に戻ります。次元 $d_k$ が4だろうと512だろうと、$\sqrt{d_k}$ で割ればスコアの分散は常に1。次元によらずスコアのスケールが揃うのです。

これがなぜ嬉しいかというと、softmax に入るスコアの「差の大きさ」が次元に依存しなくなるからです。分散1なら、スコアの差は平均的に1〜2程度に収まります。この範囲では softmax は one-hot に潰れず、適度に複数のキーへ重みを配ります。しかも先ほどの1次元の図で見た通り、差が0〜2の範囲は勾配がよく残る領域です。分散を1に保つことは、softmax を「適度な鋭さ」と「十分な勾配」の両方を満たす動作点に置くこと に等しいのです。

もし $\sqrt{d_k}$ ではなく $d_k$ で割ったら(分散が $1/d_k$ になる)、今度はスコアが小さくなりすぎて softmax がほぼ一様分布になり、どのキーも区別できなくなります。逆に割らなければ飽和します。$\sqrt{d_k}$ は、この2つの極端のちょうど中間、分散1という「黄金比」を実現する唯一の選択なのです。

多数のクエリ・キーで、割ることが最大重みの分布をどう変えるかも見ておきましょう。

最大アテンション重みの分布(スケール有無)

赤(スケール無し)は最大重みが1のすぐ手前に山を作っています。つまりほとんどのケースで飽和しています。緑($\sqrt{d_k}$ で割る)は最大重みが0.25前後に山があり、飽和から遠い健全な分布です。$d_k=128$ のような大きな次元では、割るか割らないかで挙動がまるで違うことが分かります。

アテンション行列全体で見ても、同じ傾向が現れます。

スケール前後のアテンション行列ヒートマップ

各行が1つのクエリの重み分布です。左(スケール無し)は各行でほぼ1マスだけが明るく、他は真っ黒、つまり各クエリが1つのキーしか見ていません。右($\sqrt{d_k}$ で割る)は各行の明るさが滑らかに分散していて、複数のキーの情報を混ぜて使えています。アテンションが「柔らかく注意を配る」という本来の働きをするには、このスケーリングが欠かせないのです。

分散を1に保つ、という視点で $\sqrt{d_k}$ の意味が腹落ちしました。実はこのスケーリングには、もう一つ美しい見方があります。「温度」としての解釈です。

スケーリングは「温度」の一種

$\sqrt{d_k}$ で割る操作は、softmax の 温度パラメータ と同じ枠組みで理解できます。

温度付き softmax は、スコアを温度 $T$ で割ってから softmax に通します。

$$ \mathrm{softmax}(\bm{z}/T)_i = \frac{e^{z_i/T}}{\sum_j e^{z_j/T}} $$

これと $\mathrm{softmax}(QK^\top/\sqrt{d_k})$ を見比べると、形が完全に一致しています。つまり アテンションのスケーリングは、温度 $T = \sqrt{d_k}$ のソフトマックス に他なりません。

温度の役割を思い出すと意味がよく分かります。

  • $T$ が小さい(割る数が小さい)→ スコアが相対的に大きくなり、分布が尖る(低温)
  • $T$ が大きい(割る数が大きい)→ スコアが相対的に小さくなり、分布がフラットになる(高温)

割らない($T=1$)と低温すぎて one-hot に凍りつき、$\sqrt{d_k}$ で割る($T=\sqrt{d_k}$、たとえば $d_k=64$ なら $T=8$)ことで適度な高温にして、分布を「溶かす」わけです。知識蒸留で $T$ を上げて教師の出力を柔らかくするのと、まったく同じ発想です。

同じスコアを異なる温度に通して、鋭さの変化を見てみましょう。

import numpy as np

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

logits = np.array([2.0, 1.2, 0.5, -0.3, -1.0, 0.8, -0.5, 1.5])
for T in [0.25, 1.0, 4.0]:
    p = softmax(logits / T)
    H = -np.sum(p * np.log(p + 1e-12))
    print(f"T={T:>4}:  最大重み={p.max():.3f}  エントロピー={H:.3f}")

出力は次の通りです。

T=0.25:  最大重み=0.842  エントロピー=0.556
T= 1.0:  最大重み=0.356  エントロピー=1.725
T= 4.0:  最大重み=0.175  エントロピー=2.050

温度を上げるほど最大重みが下がり、エントロピーが上がる(分布がフラットになる)ことが数値で確認できます。温度の効果を棒グラフでも見ておきます。

温度としてのスケーリングの見方

左(低温 $T=0.25$)は1つの要素に重みが集中し、まさに割らないアテンションの飽和状態に対応します。中央($T=1$)は適度に分散し、右(高温 $T=4$)はほぼ一様になります。$\sqrt{d_k}$ で割るとは、次元に応じて温度を自動で上げ、飽和を防いで中央のような適度な分布に保つ操作だと理解できます。次元が大きいほど飽和しやすいので、次元とともに温度 $\sqrt{d_k}$ を上げるのは理にかなっています。

温度の視点まで押さえれば、$\sqrt{d_k}$ の意味はもう完全に見えています。最後に要点を整理しましょう。

まとめ

本記事では、アテンションで内積を $\sqrt{d_k}$ で割る理由を、導出と実測の両面から解説しました。

  • アテンションのスコアは内積 $q\cdot k = \sum_i q_i k_i$ で、各成分が平均0・分散1・独立なら、平均0・分散 $d_k$(標準偏差 $\sqrt{d_k}$)になる。独立性のおかげで和の分散が各項の分散の足し算になり、$d_k$ 個の1を足して $d_k$ になる
  • 次元 $d_k$ が大きいとスコアの差が $\sqrt{d_k}$ に比例して膨らみ、softmax がほぼ one-hot に飽和 する
  • 飽和すると softmax のヤコビアン $\mathrm{diag}(\bm{p}) – \bm{p}\bm{p}^\top$ がほぼ0になり、勾配が消えて学習が止まる(実測でヤコビアンノルムが割らない場合の半分以下になった)
  • $\sqrt{d_k}$ で割ると分散が $1/(\sqrt{d_k})^2 \times d_k = 1$ に戻り、次元によらずスコアのスケールが揃う。softmax が「適度な鋭さ」と「十分な勾配」を両立する動作点に保たれる
  • このスケーリングは 温度 $T=\sqrt{d_k}$ のソフトマックス と同じで、次元とともに温度を上げて飽和を防いでいると解釈できる

$\sqrt{d_k}$ という一見地味な定数が、アテンションを学習可能にしている急所だった、というのが結論です。定数で割るだけと侮ってはいけません。

次のステップとして、以下の記事も参考にしてください。