Softmax関数とは

Softmax / ソフトマックス関数

実数値のベクトルを、合計1の確率分布へ変換する活性化関数

ひとことで言うと

AIが出したいくつかの候補の点数を、合計100%になる「確率」に変換してくれる計算。

概要

Softmax関数は、実数値のベクトル(ロジット)を受け取り、各要素を指数関数によって変換したのち総和が1になるよう正規化し、確率分布として解釈できるベクトルへ変換する関数。 多クラス分類モデルの出力層で、各クラスに対応するスコアを「そのクラスである確率」として扱うために使われることが最も一般的。 LLMのテキスト生成では、次のトークン候補それぞれのスコアをSoftmaxで確率分布に変換し、その分布からTemperatureやTop-p、Top-kなどのサンプリング手法を使って実際に出力するトークンを選ぶ。 TransformerのSelf-Attentionでも、各トークン間の関連度スコアをSoftmaxで正規化し、注意の重みとして使う。

背景

ニューラルネットワークの出力層はそのままだと任意の実数値(ロジット)を出すため、確率として解釈したり複数クラスを比較したりするには、値を非負かつ合計1に正規化する変換が必要になる。 Softmax関数はこの要件を満たしつつ、指数関数を使うことで最大値を持つ要素をより強調する性質を持つ。

アーキテクチャ

入力ベクトル z = (z_1, ..., z_n) の各要素に対し exp(z_i) を計算し、その総和を使って正規化することで出力 y_i = exp(z_i) / Σ_j exp(z_j) を得る。 出力ベクトルの各要素は0以上1以下となり、全要素の合計は常に1になるため、確率分布として扱える。 数値計算上のオーバーフローを避けるため、実装では各z_iから最大値を引いてから指数関数を計算する工夫が一般的に行われる。

コード例

import torch
import torch.nn.functional as F

logits = torch.tensor([2.0, 1.0, 0.1])
probs = F.softmax(logits, dim=-1)
# tensor([0.6590, 0.2424, 0.0986])

利点

• 出力を確率分布として解釈でき、交差エントロピー損失など確率ベースの損失関数と組み合わせやすい
• 指数関数を用いることで、最大値を持つ要素の重みを相対的に強調できる
• 多クラス分類やAttentionの重み計算など、幅広い場面で標準的に利用できる

欠点

• クラス数が非常に多いと、計算コスト(特に学習時の正規化項の計算)は大きくなりやすい
• 外れ値的に大きいロジットがあると、他の候補の確率がほぼ0に潰れてしまう
• 確率分布として出力されるが、モデルの実際の予測の「自信度」を正確に表しているとは限らない

比較

• Attention — Attentionでは各トークン間の関連度スコアをSoftmaxで正規化し、注意の重みとして使う
• Temperature — Temperatureは、Softmax適用前のロジットをスケーリングすることで確率分布の鋭さを調整する
• ロジスティック回帰 — ロジスティック回帰の2クラス版シグモイド関数を、多クラスに一般化したものがSoftmax関数

関連用語

よくある質問

関連Zenn記事

• Softmax関数とエントロピー正則化
• 生成AIの推論状態を分解する―― Grokking から見る logit、softmax、KV cache の役割
• 生成AIはなぜメジャー技術に強く見えるのか── Logit と Softmax で眺める“AIの迷いやすさ”
• 生成AIの挙動を支配する「2つの数理法則」：Softmax CrowdingとSemantic Drift