Softmax

Softmaxは、ロジットと呼ばれる生のスコアのベクトルを確率分布に変換する関数です。出力はすべて0から1の間にあり、その合計はちょうど1になります。ベクトルzに対して、softmax(z)ᵢ = e^(zᵢ) / Σⱼ e^(zⱼ)を計算します。各スコアを指数関数で変換し、すべての指数値の合計で割ります。Softmaxは、多数のクラスから1つを選ぶニューラルネットワークの標準的な出力層です。また、Transformerのアテンション機構内で類似度スコアを重みに変換し、言語モデルの語彙に対するスコアを次トークンの確率に変換します。この名前は、最大のスコアにすべての確率を割り当てるargmaxを、滑らかで微分可能にしたものを表しています。

Softmaxの仕組み#

指数関数によってすべての出力が正になり、スコア間の差が拡大するため、最大のlogitが確率の大部分を占めるようになります。Softmaxはlogitの順序を保つため、最も高いスコアが常に最も高い確率を得ます。重要なのはlogit間の差だけです。すべてのlogitに同じ定数を加えても、出力は変わりません。ライブラリは数値安定性のためにこの性質を利用し、指数関数を計算する前に最大のlogitを差し引きます。e¹⁰⁰⁰は32ビット浮動小数点数の範囲を超えるためです。

温度Tは分布の鋭さを調整します。softmax(z / T)と表されます。Tが1未満の場合は分布が最上位のクラスに集中し、1を超える場合は一様分布に近づきます。同じ式は統計力学ではボルツマン分布として現れ、Tは文字どおり温度を表します。以下のPyTorchのコード例では、両方の効果を確認できます。

import torch

logits = torch.tensor([2.0, 1.0, 0.1])

# Divide logits by a temperature before softmax: low sharpens, high flattens
for temperature in [0.5, 1.0, 2.0]:
    probs = torch.softmax(logits / temperature, dim=0)
    print(f"T={temperature}: {[round(p, 3) for p in probs.tolist()]}")

# Adding a constant leaves softmax unchanged, so libraries subtract the max to avoid overflow
big = logits + 1000
naive = big.exp() / big.exp().sum()
print(f"Naive: {naive.tolist()}, torch.softmax: {[round(p, 3) for p in torch.softmax(big, dim=0).tolist()]}")

T = 1の場合、ロジット2.0、1.0、0.1は[0.659, 0.242, 0.099]になります。温度を半分にすると分布は[0.864, 0.117, 0.019]のように鋭くなり、2倍にすると[0.502, 0.304, 0.194]のように平坦になります。すべてのロジットに1,000を加えると、単純な式はオーバーフローして[nan, nan, nan]を出力しますが、torch.softmaxは引き続き[0.659, 0.242, 0.099]を返します。

学習におけるSoftmax#

分類器はsoftmaxと交差エントロピー損失を組み合わせて学習します。損失は−log(pᵧ)で、正解クラスyに割り当てられた確率の負の対数です。ロジットに対する勾配は単純にp − yとなり、予測確率からone-hotラベルを引いた値です。そのため、フレームワークはこの2つを融合します。PyTorchのCrossEntropyLossは生のロジットを受け取り、内部でlog-softmaxを計算するため、個別に計算した確率の対数を取るよりも安定しています。この方法で学習したモデルはロジットを出力し、推論時に確率が必要な場合にのみsoftmaxを適用します。

Softmaxの用途#

• 画像分類：分類器の最終層はクラスごとに1つのロジットを出力し、softmaxはそれらをユーザーに表示される信頼度スコアに変換します。
• アテンション：Transformerは、クエリとすべてのキーとの類似度スコアに対してsoftmax(QKᵀ / √dₖ)を計算するため、各トークンのアテンション重みの合計は1になります。√dₖで割ることで、次元数に応じたスコアの増大とsoftmaxの飽和を防ぎます。
• テキスト生成：大規模言語モデル（LLM）は、各ステップで数万から数十万のトークンからなる語彙に対してsoftmaxを適用します。サンプリング温度は上記のTです。続いて核サンプリングでは、確率の合計がしきい値に達する最小のトークン集合からサンプリングします。
• 強化学習：行動価値にsoftmaxを適用すると、通常は最善の行動を選びながら、ほかの行動も探索する方策が得られます。温度によって探索と活用のバランスを調整します。

YOLO26の分類モデルは、推論時に分類ヘッドでsoftmaxを適用します。各結果のprobs属性には、1,000のImageNetクラスに対する完全な分布が格納されます。

from ultralytics import YOLO

# Load a YOLO26 classification model trained on ImageNet
model = YOLO("yolo26n-cls.pt")
results = model("https://ultralytics.com/images/bus.jpg")

# probs holds the softmax output of the classification head
probs = results[0].probs
for index, conf in zip(probs.top5, probs.top5conf.tolist()):
    print(f"{results[0].names[index]}: {conf:.3f}")
print(f"Sum over all 1,000 classes: {probs.data.sum().item():.3f}")

確率とともに可能性の高い上位5クラスを出力し、ヘッドがsoftmax分布を返すため、1,000クラスすべての合計は1.000と表示されます。

Softmaxと関連関数の比較#

• Sigmoid：各スコアを個別に0から1の範囲に押し込むため、出力の合計は1になるとは限りません。Sigmoidは、画像に人とバックパックの両方が写っている場合などの二値分類やマルチラベル分類に適しており、物体検出器も一般に各クラスのスコア計算に使用します。2クラスの場合、softmaxは2つのロジットの差に対するsigmoidと厳密に等しくなります。
• Argmax：最大スコアのインデックスを返す、勾配のないハードな判定です。Softmaxは同じ順位を保ちながら微分可能であるため、学習に利用できます。また、softmax確率のargmaxを取ると、ロジットのargmaxと同じ結果になります。
• Log-softmax：安定性のために直接計算されるsoftmaxの対数で、交差エントロピー損失の計算に使用されます。
• 温度スケーリング：学習後に検証データを使って単一の温度を調整するキャリブレーション手法です。これにより、softmax確率がモデルの実際の正解率に合うようになります。すべてのロジットを同じ正のTで割っても順序は保たれるため、予測クラスは変わりません。