softmaxの温度とattentionエントロピー — 注意の「集中・分散」を定量化する

同じアテンションスコアを見ても、あるヘッドは1つのトークンだけをじっと見つめ、別のヘッドは文全体を広く眺めています。この「注意の絞り」は、モデルのどこを見ているかを決める大事な性質です。ところが、スコアの数字を眺めているだけでは、いま注意が「集中モード」なのか「分散モード」なのかはよく分かりません。

そこで登場するのが2つの道具です。ひとつは softmaxの温度 $\tau$ — 同じスコアから、尖った分布も平らな分布も作り分けるダイヤルです。もうひとつは Shannonエントロピー $H$ — その分布がどれだけ「広がっているか」を1つの数字で測る温度計です。この2つをセットで理解すると、アテンションの振る舞いを定性的な印象ではなく、定量的な指標として扱えるようになります。

この視点は、次のような場面で効いてきます。

  • 学習の安定化: 訓練中にアテンションエントロピーが急落する「エントロピー崩壊」は、損失の振動や発散の前触れになります。エントロピーを監視すれば、不安定化を早期に察知できます(Zhaiら, ICML 2023)。
  • 長い文脈への対応: 系列が長くなると注意が薄く広がり、鋭さが失われます。YaRNなどの長文脈手法は、softmaxに温度補正を入れてこの希釈を防ぎます。エントロピーで「薄まり」を定量化できると、なぜ温度補正が要るのかが腹落ちします。

本記事の内容

  • 温度付きsoftmax $\mathrm{softmax}(z/\tau)$ の定義と、$\tau\to 0$ でone-hot・$\tau\to\infty$ で一様になる極限の導出
  • アテンションの $1/\sqrt{d_k}$ が「温度 $\sqrt{d_k}$ のsoftmax」に他ならないこと
  • Shannonエントロピー $H(p)$ をアテンション重みに当てはめ、$0$(one-hot)〜 $\log N$(一様)で注意の広がりを測ること、有効注意幅 $\exp(H)$
  • 温度とエントロピーの関係、次元 $d_k$ とエントロピー崩壊の実測
  • エントロピー崩壊と訓練不安定性、$\sigma$Reparam による対策
  • 系列長とエントロピー、長文脈での温度補正の動機

前提知識

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

温度によるattention分布の集中と分散のスペクトラム

上の概念図が、この記事の主題です。まったく同じスコア(キーへの関連度)でも、温度 $\tau$ を小さくすれば1点集中の尖った分布(左)に、大きくすれば広く浅い一様に近い分布(右)になります。左はスポットライトを絞った状態、右は部屋全体を照らした状態です。そして各パネルの右上にある $H$ の値が、その「絞り具合」を測る温度計 — エントロピーです。$H$ が小さいほど集中、大きいほど分散。以降、この温度 $\tau$ とエントロピー $H$ の関係を丁寧にほどいていきます。

温度付きsoftmax — 分布の鋭さを1つのダイヤルで操る

まず、注意重みを作る softmax に「温度」という調整つまみを付けます。イメージは料理の火加減です。弱火(低温)だと素材の差がくっきり残り、強火(高温)だと全部がなじんで均される — そんな感覚です。

温度付き softmax は、スコアベクトル $\bm{z}=(z_1,\dots,z_N)$ を温度 $\tau>0$ で割ってから softmax に通します。

$$ \mathrm{softmax}(\bm{z}/\tau)_i = \frac{e^{z_i/\tau}}{\sum_{j=1}^{N} e^{z_j/\tau}} $$

$\tau$ が分母にあることに注目してください。$\tau$ を小さくするとスコアが相対的に「拡大」され、差が強調されます。逆に $\tau$ を大きくするとスコアが「縮小」され、差がぼやけます。この一手間で、同じ $\bm{z}$ から尖った分布も平らな分布も作れます。

極限 τ→0:argmax に凍りつく(one-hot)

なぜ低温で「1点集中」になるのか、極限を丁寧に見てみましょう。最大のスコアを持つインデックスを $m=\arg\max_j z_j$ とし、その最大値を $z_m$ と書きます。$i$ 番目の重みを、分子分母を $e^{z_m/\tau}$ で割って整理します。

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

ここで指数の肩 $(z_i – z_m)/\tau$ に注目します。$z_i – z_m$ は、最大要素なら $0$、それ以外なら です。$\tau \to 0^+$ とすると、負の肩は $-\infty$ に飛ぶので $e^{(z_i-z_m)/\tau}\to 0$、最大要素だけが $e^0 = 1$ として残ります。したがって分母は $1$(最大が1つの場合)に収束し、

$$ \lim_{\tau \to 0^+} \mathrm{softmax}(\bm{z}/\tau)_i = \begin{cases} 1 & (i = m) \\ 0 & (i \neq m) \end{cases} $$

つまり温度を下げきると softmax は argmax、すなわち one-hot ベクトルに凍りつきます。「最も関連するキー1本だけを見る」極端な集中状態です。

極限 τ→∞:全部が均される(一様分布)

逆に高温側を見ます。$\tau \to \infty$ とすると、どのスコアも $z_i/\tau \to 0$ に潰れます。すると指数はすべて $e^0 = 1$ に近づくので、

$$ \lim_{\tau \to \infty} \mathrm{softmax}(\bm{z}/\tau)_i = \frac{1}{\sum_j 1} = \frac{1}{N} $$

どの要素も同じ重み $1/N$、つまり 一様分布 です。スコアの差はすべて無視され、「全キーを平等に見る」状態になります。数値でも確認しておきましょう。

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([3.0, 2.1, 1.5, 0.4, -0.2, 1.8, -0.8, 0.9])
for T in [1e-3, 1e3]:
    p = softmax(logits / T)
    print(f"T={T:>7g}:  p = {np.round(p, 3)}")

出力は次の通りです。

T=  0.001:  p = [1. 0. 0. 0. 0. 0. 0. 0.]
T=   1000:  p = [0.125 0.125 0.125 0.125 0.125 0.125 0.125 0.125]

極限の計算通り、$\tau=0.001$ では先頭(最大スコア)に重み1が集中した one-hot、$\tau=1000$ ではすべて $1/8=0.125$ の一様分布になっています。温度が両極端の間で分布の鋭さを連続的に動かすダイヤルであることが、これで見て取れます。

温度が分布の形を決めることは分かりました。実はアテンションの式には、この温度がすでにこっそり組み込まれています。次にそれを見ます。

1/√d_k は「温度の一種」である

スケールドドット積アテンションの式を思い出しましょう。

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

この $QK^\top/\sqrt{d_k}$ という部分を、温度付き softmax $\mathrm{softmax}(\bm{z}/\tau)$ と見比べてください。スコア $\bm{z}=QK^\top$ を定数 $\sqrt{d_k}$ で割っているのですから、これは 温度 $\tau = \sqrt{d_k}$ のsoftmax そのものです。$d_k=64$ なら $\tau=8$、$d_k=256$ なら $\tau=16$ に相当します。

なぜこの温度が必要なのでしょうか。クエリ $q$ とキー $k$ の各成分が独立に平均0・分散1だと、内積 $q\cdot k=\sum_{i=1}^{d_k} q_i k_i$ の分散は $d_k$、標準偏差は $\sqrt{d_k}$ になります(スケールドドット積の記事で導出しています)。つまり次元が大きいほどスコアが $\sqrt{d_k}$ に比例して大きく振れ、放っておくと softmax が低温状態(one-hot)に凍りついてしまう。そこで温度を $\sqrt{d_k}$ に上げてスコアの分散を1に戻し、適度な鋭さを保つわけです。

この記事では、その「適度な鋭さ」を エントロピー という物差しで測ります。$\sqrt{d_k}$ で割ることは「次元に応じて温度を上げ、アテンションエントロピーを一定の高さに保つ操作」だと言い換えられます。この言い換えが正しいことは、後ほど次元 $d_k$ を変えた実測(エントロピー崩壊のグラフ)ではっきり確かめます。

温度がアテンションに内蔵されていることが分かりました。では、その分布の「広がり」を測るエントロピーを定義しましょう。

エントロピー — 注意の広がりを1つの数字で測る

温度で分布の形を変えられるようになりましたが、「いまどれくらい集中/分散しているか」を人間が毎回グラフで見るのは大変です。分布の広がりを1つの数字に要約したい。それがShannonエントロピーです。

直感的には、エントロピーは「その分布がどれだけ迷っているか」の量です。1つのキーに決め打ちしていれば迷いは0、全キーを平等に見ていれば迷いは最大。この「迷いの量」を、確率 $p_i$ の対数の期待値として定義します。

$$ H(\bm{p}) = -\sum_{i=1}^{N} p_i \log p_i $$

各項 $-p_i\log p_i$ は、確率 $p_i$ の「驚き」$-\log p_i$ を、その確率で重み付けした量です。$p_i=1$(確実)なら驚きは $-\log 1 = 0$、$p_i$ が小さいほど驚きは大きくなります。それらを平均したのがエントロピーです。

エントロピーの取りうる範囲:0 から log N まで

アテンション重み $\bm{p}$ にこの $H$ を当てはめると、極端な2つのケースで値が決まります。

  • one-hot(1つのキーに重み1、他は0): ある $p_m=1$、他が0。$-1\cdot\log 1 = 0$、$0\cdot\log 0 = 0$(極限で0)なので $H=0$。最小値です。
  • 一様(全キーに $1/N$): $H = -\sum_{i=1}^N \frac{1}{N}\log\frac{1}{N} = -N\cdot\frac{1}{N}\log\frac{1}{N} = \log N$。最大値です。

つまりアテンションエントロピーは $0 \le H \le \log N$ の範囲を動き、$0$ が完全集中、$\log N$ が完全分散に対応します。2要素の場合でこの様子を描くと分かりやすいです。

2要素分布のエントロピーの取りうる範囲

横軸は片方の確率 $p$(もう片方は $1-p$)、縦軸がエントロピーです。両端 $p=0,1$(one-hot)で $H=0$、真ん中 $p=0.5$(一様)で最大値 $\log 2 \approx 0.693$ になります。上に凸のなめらかな山で、少しでも偏ると即座にエントロピーが下がることが分かります。実際の値も確かめましょう。

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))

dists = {
    "one-hot": np.array([1, 0, 0, 0, 0, 0, 0, 0.0]),
    "尖った":   np.array([0.7, 0.15, 0.05, 0.03, 0.02, 0.02, 0.02, 0.01]),
    "一様":     np.ones(8) / 8,
}
for name, p in dists.items():
    p = p / p.sum()
    print(f"{name:8s} H={entropy(p):.4f}  exp(H)={np.exp(entropy(p)):.3f}")
print(f"上限 log(8) = {np.log(8):.4f}")

出力は次の通りです。

one-hot  H=-0.0000  exp(H)=1.000
尖った    H=1.0700  exp(H)=2.915
一様      H=2.0794  exp(H)=8.000
上限 log(8) = 2.0794

3つの分布の性格がエントロピー1つで綺麗に区別できています。one-hot は $H=0$、一様は上限 $\log 8 = 2.079$、その中間の尖った分布は $H=1.07$。注意の「広がり具合」を、グラフを見ずに数値だけで語れるようになりました。

有効注意幅 exp(H):実質何本のキーを見ているか

出力に $\exp(H)$ という量も添えました。これは 有効注意幅(あるいはパープレキシティ的な解釈)と呼ばれ、「実質的に何本のキーに注意を配っているか」を表します。一様に $k$ 本を見ているなら $H=\log k$ なので $\exp(H)=k$ に一致します。先の出力でも、一様8本で $\exp(H)=8.0$、one-hot で $\exp(H)=1.0$ ときれいに「本数」になっています。

エントロピーは対数スケールなので直感が湧きにくいですが、$\exp(H)$ に直すと「だいたい何本分」という具体的な本数になり、頭に入りやすくなります。温度を変えたときの有効注意幅を見てみましょう。

有効注意幅exp(H)の直感

左は12本のキーに対し、温度を3段階に変えたときの重み(大きい順)です。低温 $\tau=0.3$ では実質1.6本しか見ておらず、高温 $\tau=5$ では11.7本とほぼ全部に配っています。右は温度を連続的に動かしたときの $\exp(H)$ で、$\tau$ を上げると有効注意幅が1本から上限12本へ滑らかに増えていきます。温度が「見るキーの本数」を直接コントロールしていることが、この図から読み取れます。

エントロピーで注意の広がりを測れるようになりました。次に、温度とエントロピーがどう連動するかを定量的に確かめます。

温度を上げるとエントロピーは単調に増える

温度が分布を平らにし、エントロピーが平らさを測るのですから、両者は連動するはずです。実際に固定したスコアへ様々な温度をかけて、最大重みとエントロピーを追ってみましょう。

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))

logits = np.array([3.0, 2.1, 1.5, 0.4, -0.2, 1.8, -0.8, 0.9])
for T in [0.2, 0.5, 1.0, 2.0, 5.0, 20.0]:
    p = softmax(logits / T)
    print(f"T={T:>4}:  最大重み={p.max():.3f}  H={entropy(p):.3f}  exp(H)={np.exp(entropy(p)):.2f}")
print(f"一様の上限 log(8) = {np.log(8):.3f}")

出力は次の通りです。

T= 0.2:  最大重み=0.986  H=0.082  exp(H)=1.09
T= 0.5:  最大重み=0.753  H=0.864  exp(H)=2.37
T= 1.0:  最大重み=0.456  H=1.573  exp(H)=4.82
T= 2.0:  最大重み=0.275  H=1.923  exp(H)=6.84
T= 5.0:  最大重み=0.178  H=2.052  exp(H)=7.79
T=20.0:  最大重み=0.137  H=2.078  exp(H)=7.99

温度を $0.2$ から $20$ へ上げるにつれ、最大重みは $0.986$(ほぼ集中)から $0.137$(ほぼ一様の $1/8=0.125$)へ単調に下がり、エントロピーは $0.082$ から $2.078$(上限 $\log 8=2.079$ 直下)へ単調に上がります。有効注意幅も1本弱から8本へ。温度とエントロピーが一対一で連動していることが数値で確認できます。連続的に描くと、この単調な関係がなめらかな曲線として現れます。

温度とエントロピーの単調増加の関係

対数軸の温度に対し、エントロピーは下限0(one-hot)から上限 $\log N$(一様)へS字を描いて単調に増えます。$\tau=1$ 付近が、両極端の間のちょうど中間的な鋭さに対応しています。分布そのものの変化も棒グラフで見ておきましょう。

温度スイープでのsoftmax分布の形状変化

同じロジットから、$\tau=0.25$ では1本に尖り、$\tau=1$ で適度に分散、$\tau=4$ では点線(一様 $1/N$)にほぼ張り付くまで平らになります。エントロピー $H$ の値(各パネル上)が、この「尖り→平ら」の度合いをきちんと追っていることが分かります。

温度でエントロピーが動く関係が分かりました。ここで冒頭に予告した「$1/\sqrt{d_k}$ は温度」の主張を、エントロピーの目で検証します。

次元 d_k とエントロピー崩壊:√d_k補正の意味

$1/\sqrt{d_k}$ が温度 $\sqrt{d_k}$ に相当するなら、補正しなければ次元とともにエントロピーが崩壊し、補正すれば一定に保たれる はずです。これを実測で確かめます。クエリとキーをランダムに生成し、生の内積 softmax と $\sqrt{d_k}$ で割った 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(6)
n_keys = 16
for dk in [4, 16, 64, 256]:
    hr, hs = [], []
    for _ in range(4000):
        q = rng.standard_normal(dk)
        K = rng.standard_normal((n_keys, dk))
        raw = K @ q
        hr.append(entropy(softmax(raw)))            # 補正なし
        hs.append(entropy(softmax(raw / np.sqrt(dk))))  # √d_k で割る
    print(f"d_k={dk:4d}: 補正なし H={np.mean(hr):.3f}  √d_k補正 H={np.mean(hs):.3f}")
print(f"一様の上限 log(16) = {np.log(16):.3f}")

出力は次の通りです。

d_k=   4: 補正なし H=1.760  √d_k補正 H=2.378
d_k=  16: 補正なし H=0.903  √d_k補正 H=2.367
d_k=  64: 補正なし H=0.409  √d_k補正 H=2.356
d_k= 256: 補正なし H=0.188  √d_k補正 H=2.354
一様の上限 log(16) = 2.773

結果ははっきりしています。補正なし ではエントロピーが $d_k=4$ の $1.76$ から $d_k=256$ の $0.19$ へと崩壊していきます。次元を上げるほどスコアが大きく振れ、softmax が one-hot に凍りつくためです。一方 $\sqrt{d_k}$ 補正 では、次元が64倍になってもエントロピーは $2.35$ 前後でほぼ一定。次元によらず注意の広がりが揃っています。これを次元の関数として描くと、崩壊と安定の対比が一目瞭然です。

次元d_kとアテンションエントロピーの関係

赤(補正なし)は次元とともにエントロピーが0へ向かって崩れ落ち、青($\sqrt{d_k}$ 補正)は一様上限 $\log 16$ の少し下で水平に保たれます。「$\sqrt{d_k}$ で割る=次元に応じて温度を上げ、アテンションエントロピーを一定に保つ」という冒頭の主張が、これで定量的に裏付けられました。

ここまでは「1回の softmax」の話でした。ここから先は、訓練の途中でエントロピーがどう動くか、そして急落するとなぜ困るのかを見ます。

エントロピー崩壊:訓練が不安定になる前触れ

これまで見たエントロピー崩壊は、次元設定による静的なものでした。しかし厄介なのは、$\sqrt{d_k}$ で正しく割っていても、訓練が進むうちに アテンションが特定トークンへ過度に集中し、エントロピーが急落することがある点です。ZhaiらはこれをTransformerの訓練不安定性の中心的なメカニズムとして指摘しました。

論文タイトルはそのものずばり “Stabilizing Transformer Training by Preventing Attention Entropy Collapse”(ICML 2023)です。彼らは各アテンションヘッドのエントロピーを訓練を通じて追跡し、低いアテンションエントロピーが訓練の不安定性(損失の振動や発散)と密接に結びつく ことを発見しました。この病的に低いエントロピーの状態を エントロピー崩壊(entropy collapse) と名付けています。

Zhaiらのentropy collapse図(学習率2倍で崩壊)

出典: Zhai et al., “Stabilizing Transformer Training by Preventing Attention Entropy Collapse”, ICML 2023, Fig.1

左が標準のVision Transformer、右は学習率を2倍にしたものです。左では訓練損失(上)が順調に下がり、各層のアテンションエントロピー(下)は一度落ち込んでも回復しています。ところが右では、学習率を上げただけでアテンションエントロピーが全層ほぼ0まで崩壊し(下)、それと同時に訓練損失が下がらず発散してしまっています(上)。エントロピー崩壊と訓練失敗が同時に起きていることが、この図から読み取れます。

対策:σReparam(スペクトル正規化)

なぜ崩壊が起きるのか。アテンションのロジット $QK^\top$ が大きくなりすぎると softmax が飽和し、エントロピーが落ちます。ロジットの大きさは、クエリ・キーを作る線形層の重み $W$ の スペクトルノルム(最大特異値)$\sigma(W)$ に効きます。重みが大きく育つとロジットが膨らみ、崩壊へ向かうのです。

Zhaiらの提案 $\sigma$Reparam は、この暴走を抑える簡潔な再パラメータ化です。各線形層の重み $W$ を、スペクトルノルムで割って正規化し、学習可能なスカラー $\gamma$ を掛け直します。

$$ \widehat{W} = \frac{\gamma}{\sigma(W)}\, W $$

$\sigma(W)$ で割ることで重みの「伸び」を抑え、$\gamma$ で必要なスケールだけを学習で取り戻します。これによりロジットの暴走が防がれ、アテンションエントロピーが高い(=適度に分散した)領域に保たれます。結果として、warmupなしや大きな学習率といった過酷な設定でも訓練が安定します。エントロピーという1つの指標が、不安定化の診断(何が起きているか)と処方(どこを抑えるか)の両方を与えてくれるわけです。

論文の現象を、手元の小さなTransformerでも部分的に再現できます。合成タスク(3つ前のトークンをコピーする長距離依存)で小型GPTを訓練し、全層平均のアテンションエントロピーを追ってみました。

訓練ステップでのアテンションエントロピー推移

訓練の初期はエントロピーが高い(注意がまだ広く散らばっている)状態から始まり、学習が進むにつれてタスクに必要なトークンへ注意が集中し、エントロピーが下がっていきます。これは正常な「学習による集中」ですが、下がりすぎると崩壊の危険域に入ります。ヘッドごとに見ると、集中の度合いには大きなばらつきがあります。

訓練後の層ヘッドごとのアテンションエントロピー

各マスが1つのヘッドの平均エントロピーです。第2層のヘッドはエントロピーが $0.2$〜$0.4$ と低く、特定トークン(この場合は3つ前のトークン)を鋭く拾う「コピーヘッド」に育っています。一方、第1層や第3層のヘッドは $1$ 以上と高めで、広く文脈を眺めています。エントロピーを層・ヘッド単位で見ると、モデルが役割分担していることが定量的に分かります。監視すべきは、この分布の一部が突然0へ崩れ落ちる瞬間です。

集中しすぎ(低エントロピー)の弊害を見ました。では逆に、分散しすぎ(高エントロピー)は問題ないのでしょうか。次にその裏面を見ます。

逆の極限:一様すぎる注意と長系列の希釈

エントロピーが高い、つまり注意が一様に近い状態は「崩壊しない安全側」に見えますが、これはこれで問題です。全キーを平等に見るということは、どれも選べていない ということ。アテンションの役割は「関連するものを選んで情報を混ぜる」ことなので、一様注意では何も抽出できません。各トークンの表現がどれも似通ってしまう rank collapse 的な状態にもつながります。集中しすぎと分散しすぎの間に、ちょうど良い動作点があるわけです。

この「分散しすぎ」が構造的に起きやすいのが 長い系列 です。系列長 $N$ が大きいほど一様分布の上限 $\log N$ が上がり、softmax は多数のトークンへ薄く広がりやすくなります。スコアの鋭さ(分散)が同じままだと、1本あたりの注意が希釈されるのです。

系列長とエントロピーおよび注意の希釈

分散1の固定スコアで系列長を変えると、エントロピー(青)は上限 $\log N$(緑)に沿って増え続けます。同時に、最大注意重み(赤、右軸)は $N=8$ の $0.36$ から $N=1024$ の $0.02$ 近くまで落ちていきます。系列が長くなるほど、一番見たいトークンにすら十分な重みを割けなくなる — これが長文脈での注意の希釈です。

短い文脈で訓練したモデルを、そのまま長い文脈へ広げると、この希釈でアテンションが眠くなります。そこで YaRN などの長文脈拡張手法は、softmaxのロジットに温度補正(スケーリング係数)を掛けて、系列長が伸びてもアテンションの鋭さ=エントロピーが一定に保たれるよう調整します。発想は $\sqrt{d_k}$ とまったく同じで、「分布が薄まる要因(ここでは系列長)に応じて温度を動かし、エントロピーを狙った高さに固定する」ことです。温度とエントロピーの言葉を持っていると、こうした手法の狙いが一言で理解できます。

最後に、訓練済みモデルで有効注意幅が実際どの程度かを見て、締めくくりましょう。

有効注意幅で見る「実質いくつ見ているか」

エントロピーそのものより、有効注意幅 $\exp(H)$ の方が実感を伴います。先ほどの小型Transformerの最終層ヘッドについて、クエリ位置ごとの有効注意幅を測ってみました。

有効注意幅の位置別推移

灰色の破線は、そのクエリ位置で見られるキーの最大数(因果マスクにより位置 $t$ では $t+1$ 本)です。オレンジの実線が実際の有効注意幅 $\exp(H)$ で、位置が進んでも最大数よりずっと下、実質数本〜十数本に留まっています。系列が長くなっても、このヘッドは全部を平等に見るのではなく、少数の関連トークンへ選択的に注意を配っている — アテンションが本来の「選ぶ」働きをしていることが、有効注意幅から読み取れます。エントロピーを $\exp$ に通すだけで、抽象的な情報量が「何本分」という具体的な本数になり、モデルの振る舞いが手触りを持って見えてきます。

まとめ

本記事では、softmaxの温度とアテンションエントロピーを使って、注意の「集中・分散」を定量化しました。

  • 温度付きsoftmax $\mathrm{softmax}(\bm{z}/\tau)$ は分布の鋭さを操るダイヤル。$\tau\to 0$ で argmax(one-hot)、$\tau\to\infty$ で一様分布になることを極限計算で確認した
  • アテンションの $1/\sqrt{d_k}$ は 温度 $\sqrt{d_k}$ のsoftmax に他ならない。次元に応じて温度を上げ、スコアの分散を1に戻す操作である
  • Shannonエントロピー $H(\bm{p})=-\sum p_i\log p_i$ は注意の広がりを1数字で測る。$0$(one-hot)〜 $\log N$(一様)を動き、$\exp(H)$ は「実質何本のキーを見ているか」を表す
  • 温度を上げるとエントロピーは単調増加。補正なしは次元とともにエントロピーが崩壊し、$\sqrt{d_k}$ 補正は一定に保つことを実測で確認した
  • 訓練中の エントロピー崩壊 は損失の振動・発散と結びつく。$\sigma$Reparam(重みのスペクトル正規化)でロジットの暴走を抑えると安定する(Zhaiら, ICML 2023)
  • 逆に一様すぎる注意は「何も選べない」状態。長系列では $\log N$ が増えて注意が希釈されるため、YaRN等は温度補正でエントロピーを一定に保つ

温度とエントロピーは、アテンションを「なんとなく集中している/分散している」という印象から、監視・診断・制御できる定量的な対象へ引き上げてくれます。モデルの内部で何が起きているかを覗く、強力な物差しです。

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

画像なし
なぜAttentionはスケールドドット積なのか — √d_kで割る理由を徹底解説
内積の分散がd_kになる導出から、softmaxの飽和・勾配消失、温度としての見方までnumpy実測で解説。
画像なし
ソフトマックス関数の定義・性質・数値安定な実装
ソフトマックスの定義・数学的性質・温度パラメータの役割と、数値安定な実装・逆伝播をPythonで。
画像なし
Attention Sink(注意の吸い込み)とは
softmaxの合計1制約が生む「余った注意の捨て場」がなぜ先頭トークンに集まるのかを解説。
画像なし
アテンションを直感で理解する
難しい数式の前に『関連するものに注目して情報を集める』イメージを検索アナロジーと図で。