Softmax（Softmax） | Transformer 内部構造ガイド

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

何をする機能か

[B, T, V] の logits を、同じ形の確率分布へ変換します。

各値を指数関数に通し、合計で割るだけです。結果は次の性質を満たします。

すべての値が 0 より大きい
最後の次元（V 方向）で合計するとちょうど 1 になる
大小関係は変わらない（logit が大きいものは確率も大きい）

なぜ必要か

logits は範囲の制限がない実数で、そのままでは「どれくらい確からしいか」を比較できません。サンプリングで次のトークンを選ぶにも、複数の候補を確率として扱う必要があります。

指数関数を使うのは、次の理由からです。

必ず正の値になる（確率として扱える）
差が大きいほど比が急激に開く（自信のある予測が際立つ）
微分が扱いやすく、交差エントロピー損失と組み合わせたときの勾配が単純な形になる

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

図では Linear（語彙数へ変換）の直後、次トークンの確率の直前です。

ただし Softmax はここだけの関数ではありません。 Attention の内部（スコアを重みに変換するところ）でも同じ関数が使われています。違いは、正規化する方向と対象だけです。

使われる場所	入力	正規化する方向	出力の意味
Attention 内部	スコア `[B, H, S, S]`	Key 位置の方向	各トークンをどれだけ参照するかの重み
出力ヘッド	logits `[B, T, V]`	語彙の方向	次のトークンの確率

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

前段から受け取るもの

logits [B, T, V]。確率ではなく、範囲の制限がない実数です。

後段へ渡すもの

probabilities [B, T, V]。次トークンの確率のサンプリング処理へ渡されます。

例文での処理

共通例文

The keys on the table are lost.

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

The keys on the table まで確定していて、次のトークンを選ぶ場面です。

候補	logit	exp(logit)	確率
`are`	8.2	3641	0.94
`is`	5.1	164	0.042
`were`	4.3	74	0.019
`on`	1.4	4.1	0.001
その他（語彙の残り）	—	—	合計 0.001 未満

差だけが結果を決める

すべての logit に同じ値を足しても、確率は変わりません。

\mathrm{softmax}(z + c)_i = \mathrm{softmax}(z)_i

この性質は、実装上の数値安定化（最大値を引いてから計算する）にも利用されます。

温度パラメータ

logits を $\tau$ で割ってから Softmax を取ると、分布の鋭さを調整できます。

$\tau < 1$ : 分布が鋭くなる（確率の高い候補に集中し、出力が決定的になる）
$\tau = 1$ : そのまま
$\tau > 1$ : 分布が平らになる（多様だが、一貫性を欠きやすくなる）

これは生成時によく使う temperature の正体です。

EncoderとDecoderでの違い

出力ヘッドの Softmax は Decoder 側にのみあります。一方、Attention 内部の Softmax は Encoder・Decoder の両方にあります。

入力と出力

入力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`logits`	語彙ごとのスコア（確率ではない実数）	float32 Tensor	`[B, T, V]`	B = バッチサイズ（同時に処理する文の数） / T = 出力系列長（Decoder側のトークン数） / V = 語彙数（Tokenizer が扱うトークンの総数）	[1, 8, 32000]	指数関数と正規化
`temperature`	分布の鋭さを調整する値（生成時のみ）	float	—	—	0.7	logits を割る

出力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`probabilities`	次のトークンの確率分布（V 方向の合計が 1）	float32 Tensor	`[B, T, V]`	B = バッチサイズ（同時に処理する文の数） / T = 出力系列長（Decoder側のトークン数） / V = 語彙数（Tokenizer が扱うトークンの総数）	[1, 8, 32000]	サンプリング／貪欲法による選択

B: バッチサイズ（同時に処理する文の数）
T: 出力系列長（Decoder側のトークン数）
V: 語彙数（Tokenizer が扱うトークンの総数）

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

よくある誤解

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

数式

\mathrm{softmax}(z)_i = \frac{\exp(z_i)}{\sum_{j=1}^{V} \exp(z_j)}

温度 $\tau$ を入れる場合：

\mathrm{softmax}_\tau(z)_i = \frac{\exp(z_i / \tau)}{\sum_{j=1}^{V} \exp(z_j / \tau)}

実装では、オーバーフローを避けるため最大値を引きます（前述の平行移動不変性により結果は変わりません）。

\mathrm{softmax}(z)_i = \frac{\exp(z_i - \max_j z_j)}{\sum_{k}\exp(z_k - \max_j z_j)}

交差エントロピーとの関係

正解トークンを $y$ とすると、損失は次のようになります。

\mathcal{L} = -\log \mathrm{softmax}(z)_y = \log\!\left(\sum_j \exp(z_j)\right) - z_y

この形（log-sum-exp）で計算すると、確率を経由せずに済み、数値的に安定します。 cross_entropy が logits を受け取る理由です。

勾配は次のように単純な形になります。

\frac{\partial \mathcal{L}}{\partial z_i} = \mathrm{softmax}(z)_i - \mathbb{1}[i = y]

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

import torch
from torch import Tensor
 
 
def probabilities_from_logits(logits: Tensor, temperature: float = 1.0) -> Tensor:
    """logits [B, T, V] を確率 [B, T, V] へ変換する。
 
    Args:
        logits: [B, T, V]。範囲の制限がない実数。
        temperature: 分布の鋭さ。1.0 でそのまま、小さいほど鋭くなる。
 
    Returns:
        [B, T, V]。最後の次元の合計が 1 になる確率分布。
    """
    scaled = logits / temperature
    # 数値安定化: 最大値を引いても結果は変わらない
    shifted = scaled - scaled.max(dim=-1, keepdim=True).values
    exponentiated = shifted.exp()
    return exponentiated / exponentiated.sum(dim=-1, keepdim=True)
 
 
def loss_from_logits(logits: Tensor, labels: Tensor) -> Tensor:
    """学習時の損失。softmax を明示的に掛けず logits をそのまま渡す。
 
    Args:
        logits: [B, T, V]
        labels: [B, T] int64。各位置の正解トークンID。
 
    Returns:
        スカラーの損失値。
    """
    batch, seq_len, vocab = logits.shape
    return torch.nn.functional.cross_entropy(
        logits.reshape(batch * seq_len, vocab),  # [B*T, V]
        labels.reshape(batch * seq_len),         # [B*T]
    )
 
 
# 最小使用例
logits = torch.tensor([[[8.2, 5.1, 4.3, 1.4, 0.2, -3.7]]])  # [B=1, T=1, V=6]
probs = probabilities_from_logits(logits)
 
print(probs.shape)                    # torch.Size([1, 1, 6]) == [B, T, V]
print(probs.sum(dim=-1))              # tensor([[1.]])  合計は必ず 1
print(probs[0, 0, 0].item())          # 約 0.94（are に相当）
 
# 温度を上げると分布が平らになる
print(probabilities_from_logits(logits, temperature=2.0)[0, 0, 0].item())  # 約 0.75
 
# 全体に定数を足しても結果は変わらない
print(torch.allclose(probs, probabilities_from_logits(logits + 100.0)))    # True

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

標準の torch.softmax を使わず、計算内容がわかる形で書いています。実務では組み込み関数を使ってください。
top-k / top-p による絞り込みは next-token-prediction のページで扱います。
バッチ内のパディング位置の扱いは省略しています。

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

計算量

$O(B \cdot T \cdot V)$ です。行列積ではないため計算自体は軽いですが、 V が大きいため [B, T, V] の読み書きがメモリ帯域を使います。

推論時は最後の位置だけを処理すれば十分なので [B, 1, V] で済みます。

実装差と派生

log_softmax: 対数確率を直接求める関数。損失計算やビームサーチで使われます。
Attention 内部の Softmax: 同じ関数ですが、正規化する方向が Key の位置方向です。マスクされた位置は $-\infty$ のため、 $\exp(-\infty) = 0$ で厳密に除外されます。
Gumbel-Softmax: 離散選択を微分可能に近似する変種。通常の推論では使いません。
精度: logits を float32 に上げてから Softmax を取る実装が一般的です（bfloat16 のままだと分解能が足りず、確率が潰れることがあります）。

SoftmaxSoftmax