ソフトマックス関数の定義・性質・数値安定な実装

多クラス分類の出力層で「各クラスの確率」を出力したいとき、モデルの生の出力(ロジット)をどう確率に変換すればよいでしょうか。単純に各ロジットを合計で割る正規化では、負のロジットがあると負の「確率」が出てしまいます。

この問題をエレガントに解決するのがソフトマックス関数(softmax function)です。ソフトマックスは任意の実数ベクトルを「全要素が正で合計が1」の確率ベクトルに変換します。

ソフトマックス関数は見た目以上に奥が深く、統計力学のボルツマン分布との関係や、argmaxの「滑らかな近似」としての側面など、多くの理論的背景を持っています。

ソフトマックス関数を理解すると、以下のような場面で適切に扱えるようになります。

  • 多クラス分類: ニューラルネットワークの出力層で確率分布を生成する
  • 注意機構(Attention): TransformerのQuery-Key類似度を確率に変換する
  • 強化学習: ボルツマン方策でアクションの選択確率を決める
  • 温度スケーリング: モデルの確信度を事後的に調整する

本記事の内容

  • ソフトマックス関数の定義と直感的な理解
  • 数学的性質(不変性・単調性・ヤコビアン)
  • 温度パラメータの役割とボルツマン分布との関係
  • 数値安定な実装のテクニック
  • 逆伝播の導出とPythonでの実装

前提知識

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

ソフトマックス関数とは

直感的な理解

ソフトマックス関数はargmaxの滑らかな(微分可能な)近似として理解できます。

argmaxは「最大値をとるインデックスに1を、他に0を割り当てる」操作ですが、これは微分不可能です。ソフトマックスはこの操作を「柔らかく」(soft)して、最大値に近い要素に大きな確率を、小さい要素には小さな確率を割り当てます。

例えば、ロジットが $\bm{z} = (3, 1, -1)$ のとき、argmaxは $(1, 0, 0)$ を返しますが、ソフトマックスは $(0.867, 0.117, 0.016)$ のように連続的な値を返します。最大のロジットに最も大きな確率が割り当てられますが、他のクラスにも0でない確率が残ります。

argmaxとsoftmaxの比較: スコアを滑らかに確率化する

上の図がこの対応です。左の生スコア(ロジット)を、中央のargmaxは勝者総取りの one-hot に潰してしまいますが、右のソフトマックスは大小関係を保ったまま滑らかな確率分布に変換します。この「滑らかさ」が勾配を流せる=学習できることの源です。

定義

$K$ 次元の実数ベクトル $\bm{z} = (z_1, z_2, \ldots, z_K) \in \mathbb{R}^K$ に対して、ソフトマックス関数は

$$ \begin{equation} \text{softmax}(\bm{z})_k = \frac{e^{z_k}}{\sum_{j=1}^K e^{z_j}}, \quad k = 1, 2, \ldots, K \end{equation} $$

と定義されます。出力を $\hat{p}_k = \text{softmax}(\bm{z})_k$ と書きます。

なぜ指数関数 $e^{z_k}$ を使うのかというと、指数関数は

  1. 常に正の値を返す($e^{z_k} > 0$)ため、全要素が正になる
  2. 単調増加関数であるため、大きなロジットにより大きな値が割り当てられる
  3. 差の比較を比に変換する($e^{z_1 – z_2} = e^{z_1}/e^{z_2}$)

という性質を持つからです。

二値分類での特殊ケース — シグモイド関数

$K = 2$ のとき、ソフトマックスはシグモイド関数に帰着します。

$$ \hat{p}_1 = \frac{e^{z_1}}{e^{z_1} + e^{z_2}} = \frac{1}{1 + e^{-(z_1 – z_2)}} = \sigma(z_1 – z_2) $$

つまり、シグモイド関数はソフトマックスの二値版であり、2つのロジットのだけで確率が決まります。これが二値分類ではスカラーのロジット1つで十分な理由です。

ソフトマックスの基本的な定義を理解したところで、次にその数学的性質を詳しく見ていきましょう。

2クラスのsoftmaxはシグモイドと一致する

図のとおり、2クラスのソフトマックスでクラス1の確率を計算すると、ロジットの差 $z_1 – z_2$ のシグモイド曲線に完全に一致します(赤い点が2クラスsoftmaxの計算値)。シグモイドは「差だけで決まる2クラス版ソフトマックス」だったわけです。

ソフトマックスの数学的性質

性質1: 確率分布の条件を満たす

ソフトマックスの出力は常に正規の確率分布を形成します。

$$ \hat{p}_k > 0 \quad \text{(正値性)}, \qquad \sum_{k=1}^K \hat{p}_k = 1 \quad \text{(正規性)} $$

正値性は $e^{z_k} > 0$ から直ちに従います。正規性は

$$ \sum_{k=1}^K \frac{e^{z_k}}{\sum_j e^{z_j}} = \frac{\sum_k e^{z_k}}{\sum_j e^{z_j}} = 1 $$

から確認できます。

性質2: 平行移動不変性

ソフトマックスは入力ベクトルに定数を加えても結果が変わりません。

$$ \text{softmax}(\bm{z} + c\bm{1}) = \text{softmax}(\bm{z}), \quad \forall c \in \mathbb{R} $$

これは分子と分母の両方に $e^c$ がかかり、約分されるためです。

$$ \frac{e^{z_k + c}}{\sum_j e^{z_j + c}} = \frac{e^c \cdot e^{z_k}}{e^c \cdot \sum_j e^{z_j}} = \frac{e^{z_k}}{\sum_j e^{z_j}} $$

この性質は数値安定性のために重要であり、$c = -\max_k z_k$ とすることでオーバーフローを防ぎます。

平行移動不変性: 定数を足しても出力は同じ

実際に $z=(3,1,-1)$ に $+5$ や $-100$ を足しても、出力は3つとも $(0.867, 0.117, 0.016)$ で寸分違わず同じです。「差だけが意味を持つ」というソフトマックスの体質がよく現れています。

性質3: 単調性の保存

ロジットの大小関係がそのまま確率の大小関係に保存されます。

$$ z_i > z_j \Leftrightarrow \hat{p}_i > \hat{p}_j $$

$e^x$ が単調増加関数であることから明らかです。したがって、ソフトマックスの出力での最大確率のインデックスは、ロジットの最大値のインデックスと必ず一致します。

性質4: ヤコビアン

ソフトマックスのヤコビ行列(各出力の各入力に対する偏微分)を計算しましょう。$\hat{p}_i = \text{softmax}(\bm{z})_i$ として

$$ \frac{\partial \hat{p}_i}{\partial z_j} = \hat{p}_i (\delta_{ij} – \hat{p}_j) $$

ここで $\delta_{ij}$ はクロネッカーのデルタです。

$i = j$ の場合を導出します。商の微分則を適用すると

$$ \frac{\partial \hat{p}_i}{\partial z_i} = \frac{e^{z_i} \cdot \sum_j e^{z_j} – e^{z_i} \cdot e^{z_i}}{(\sum_j e^{z_j})^2} $$

分子を整理すると $e^{z_i}(\sum_j e^{z_j} – e^{z_i})$ なので

$$ = \frac{e^{z_i}}{\sum_j e^{z_j}} \cdot \frac{\sum_j e^{z_j} – e^{z_i}}{\sum_j e^{z_j}} = \hat{p}_i(1 – \hat{p}_i) $$

$i \neq j$ の場合:

$$ \frac{\partial \hat{p}_i}{\partial z_j} = \frac{0 – e^{z_i} \cdot e^{z_j}}{(\sum_l e^{z_l})^2} = -\hat{p}_i \hat{p}_j $$

両方をまとめると $\hat{p}_i(\delta_{ij} – \hat{p}_j)$ となります。

行列形式では

$$ \begin{equation} \frac{\partial \hat{\bm{p}}}{\partial \bm{z}} = \text{diag}(\hat{\bm{p}}) – \hat{\bm{p}}\hat{\bm{p}}^T \end{equation} $$

ヤコビアンが出力 $\hat{\bm{p}}$ のみで表されるのは計算上便利な性質です。次に温度パラメータについて見ていきましょう。

ソフトマックスのヤコビアン: 対角が正・非対角が負

$z=(3,1,-1,0.5)$ でヤコビアンを数値化したのが上の図です。対角成分 $p_i(1-p_i)$ は正、非対角成分 $-p_i p_j$ は負——ある成分の確率を上げると他の成分から確率を「奪う」構造が、行列にそのまま刻まれています。各行・各列の和が0になっている(確率の総和1が保存される)ことも確認できます。

温度パラメータとボルツマン分布

温度付きソフトマックス

ソフトマックスに温度パラメータ $T > 0$ を導入すると

$$ \begin{equation} \text{softmax}(\bm{z}/T)_k = \frac{e^{z_k/T}}{\sum_j e^{z_j/T}} \end{equation} $$

温度 $T$ はソフトマックスの「鋭さ」を制御します。

  • $T \to 0$: 出力がone-hot(argmax)に近づく。最大ロジットのクラスにほぼ全ての確率が集中
  • $T = 1$: 通常のソフトマックス
  • $T \to \infty$: 出力が一様分布に近づく。全クラスにほぼ等しい確率

温度パラメータによる分布の鋭さの変化

同じロジット $z=(3,1,-1)$ でも、$T=0.2$ ではほぼ one-hot、$T=50$ ではほぼ一様と、温度ひとつで分布の顔つきが激変します。ソフトマックスの「柔らかさ」は固定ではなく、ダイヤルで調整できるのです。

統計力学との関係

温度付きソフトマックスは、統計力学のボルツマン分布(カノニカル分布)そのものです。

$$ p_k = \frac{e^{-E_k/(k_B T)}}{\sum_j e^{-E_j/(k_B T)}} $$

エネルギー $E_k$ を $-z_k$ に、ボルツマン定数 $k_B$ を1に対応させると、まさにソフトマックスです。低温では低エネルギー状態(高ロジット)に集中し、高温ではすべての状態がほぼ等確率になります。

知識蒸留での活用

Hintonらが提案した知識蒸留(knowledge distillation)では、大きなモデル(教師)の知識を小さなモデル(生徒)に転移する際に、高温のソフトマックス出力を使います。

$T > 1$ にすることで、教師モデルのソフトマックス出力が「柔らかく」なり、正解でないクラス間の相対的な類似度情報(ダークナレッジ)が保存されます。例えば画像分類でロジットが(猫: 6.0, 犬: 3.5, 車: 0.0)のとき、$T = 1$ では「猫: 0.922, 犬: 0.076, 車: 0.002」とほぼ正解一色ですが、$T = 4$ では「猫: 0.569, 犬: 0.304, 車: 0.127」となり、「犬は車より猫に近い」という情報が活かされます。

知識蒸留: 高温softmaxがダークナレッジを露出させる

高温にすると、$T=1$ では0.076しかなかった「犬」の確率が0.304まで持ち上がり、クラス間の類似構造が生徒モデルに見える形になります。これがダークナレッジの正体です。

import numpy as np
import matplotlib.pyplot as plt

# 温度パラメータの効果
logits = np.array([5.0, 2.0, 1.0, 0.5, -1.0])
K = len(logits)
class_names = ["Cat", "Dog", "Bird", "Fish", "Car"]

temperatures = [0.1, 0.5, 1.0, 2.0, 5.0, 20.0]

fig, axes = plt.subplots(2, 3, figsize=(16, 10))

for ax, T in zip(axes.flat, temperatures):
    z_scaled = logits / T
    z_shifted = z_scaled - np.max(z_scaled)
    probs = np.exp(z_shifted) / np.sum(np.exp(z_shifted))

    bars = ax.bar(class_names, probs, color="steelblue", alpha=0.8, edgecolor="gray")
    ax.set_title(f"T = {T}", fontsize=14)
    ax.set_ylabel("Probability", fontsize=11)
    ax.set_ylim(0, 1.05)
    ax.grid(True, alpha=0.3, axis="y")

    # 確率値を棒の上に表示
    for bar, p in zip(bars, probs):
        if p > 0.01:
            ax.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.02,
                   f"{p:.3f}", ha="center", va="bottom", fontsize=9)

    # エントロピーの計算
    entropy = -np.sum(probs * np.log(probs + 1e-15))
    ax.text(0.98, 0.92, f"H = {entropy:.3f}", transform=ax.transAxes,
            ha="right", fontsize=10,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow"))

plt.suptitle("Effect of Temperature on Softmax Output", fontsize=16, y=1.02)
plt.tight_layout()
plt.savefig("softmax_temperature.png", dpi=150, bbox_inches="tight")
plt.show()

このグラフから、温度パラメータの効果が直感的に理解できます。

  1. $T = 0.1$(極低温): ほぼone-hotに近く、最大ロジット「Cat」にほぼ100%の確率が集中しています。エントロピーはほぼ0で、分布は非常に「尖って」います

  2. $T = 1.0$(標準): 通常のソフトマックス出力です。Catが最大確率ですが、他のクラスにも多少の確率が割り当てられています

  3. $T = 5.0$(高温): 分布が「柔らかく」なり、クラス間の差が縮まっています。Cat(ロジット5.0)とDog(ロジット2.0)の確率差は小さくなりますが、ロジットの大小関係は保存されています

  4. $T = 20.0$(極高温): ほぼ一様分布に近づいています。エントロピーが $\ln 5 \approx 1.609$ の理論最大値に近くなっています

温度が高くなるほどエントロピーが増加する(分布が「フラット」になる)ことが、各パネルのH値から確認できます。

数値安定な実装

オーバーフロー問題

ソフトマックスの素朴な実装 exp(z) / sum(exp(z)) は、$z_k$ が大きいとき $e^{z_k}$ がオーバーフロー(inf)になります。例えば、float64 では $z > 709$ 程度でオーバーフローします。

平行移動不変性を利用した安定化

前述の平行移動不変性を利用して、$c = \max_k z_k$ を引きます。

$$ \text{softmax}(\bm{z})_k = \frac{e^{z_k – \max(\bm{z})}}{\sum_j e^{z_j – \max(\bm{z})}} $$

$z_k – \max(\bm{z}) \leq 0$ なので $e^{z_k – \max(\bm{z})} \leq 1$ となり、オーバーフローが防げます。

実装と検証

import numpy as np

def softmax_naive(z):
    """素朴な実装(オーバーフローの危険あり)"""
    exp_z = np.exp(z)
    return exp_z / np.sum(exp_z, axis=-1, keepdims=True)

def softmax_stable(z):
    """数値安定な実装"""
    z_shifted = z - np.max(z, axis=-1, keepdims=True)
    exp_z = np.exp(z_shifted)
    return exp_z / np.sum(exp_z, axis=-1, keepdims=True)

def log_softmax_stable(z):
    """log-softmaxの安定な実装"""
    c = np.max(z, axis=-1, keepdims=True)
    log_sum_exp = c + np.log(np.sum(np.exp(z - c), axis=-1, keepdims=True))
    return z - log_sum_exp

# テスト1: 通常のロジット
z_normal = np.array([2.0, 1.0, 0.1])
print("Normal logits:", z_normal)
print("Naive:  ", softmax_naive(z_normal))
print("Stable: ", softmax_stable(z_normal))
print()

# テスト2: 大きなロジット(オーバーフローのテスト)
z_large = np.array([1000.0, 999.0, 998.0])
print("Large logits:", z_large)
print("Naive:  ", softmax_naive(z_large))   # inf/inf = nan
print("Stable: ", softmax_stable(z_large))  # 正常に計算

# テスト3: 非常に小さなロジット
z_small = np.array([-1000.0, -999.0, -998.0])
print("\nSmall logits:", z_small)
print("Naive:  ", softmax_naive(z_small))   # 0/0 = nan
print("Stable: ", softmax_stable(z_small))  # 正常に計算

# テスト4: バッチ処理
z_batch = np.array([
    [2.0, 1.0, 0.1],
    [1000.0, 999.0, 998.0],
    [-1000.0, -999.0, -998.0]
])
print("\nBatch stable softmax:")
print(softmax_stable(z_batch))
print("Row sums:", softmax_stable(z_batch).sum(axis=1))

# テスト5: log-softmax
print("\nLog-softmax (stable):")
print(log_softmax_stable(z_large))
print("Verify exp:", np.exp(log_softmax_stable(z_large)))

数値安定化: 素朴な実装はnan、最大値を引けば正しく計算できる

この実装の検証結果から、数値安定性の重要さが確認できます。素朴な実装では $z = 1000$ のようなロジットで nan が出力されますが、安定な実装では正しい結果が得られます。また、$z = -1000$ のケースでもアンダーフロー問題を回避しています。バッチ処理でも各行の確率の合計がきちんと1になることが確認できます。

PyTorchでは torch.nn.functional.log_softmax が内部でこの安定な実装を行っているため、交差エントロピー損失と組み合わせるときは LogSoftmax + NLLLoss または CrossEntropyLoss(ロジット入力)を使うのが推奨されます。

逆伝播の導出

ソフトマックス + 交差エントロピーの勾配

ニューラルネットワークの学習では、損失のロジットに対する勾配が必要です。ソフトマックス + 交差エントロピーの組み合わせでは、既に導出したように

$$ \frac{\partial L}{\partial z_k} = \hat{p}_k – y_k $$

この結果を連鎖律(chain rule)から確認しましょう。

損失 $L = -\sum_i y_i \ln \hat{p}_i$ のロジット $z_k$ に関する偏微分を連鎖律で計算します。

$$ \frac{\partial L}{\partial z_k} = \sum_i \frac{\partial L}{\partial \hat{p}_i} \cdot \frac{\partial \hat{p}_i}{\partial z_k} $$

$\frac{\partial L}{\partial \hat{p}_i} = -\frac{y_i}{\hat{p}_i}$ と、ソフトマックスのヤコビアン $\frac{\partial \hat{p}_i}{\partial z_k} = \hat{p}_i(\delta_{ik} – \hat{p}_k)$ を代入すると

$$ = \sum_i \left(-\frac{y_i}{\hat{p}_i}\right) \hat{p}_i (\delta_{ik} – \hat{p}_k) $$

$\hat{p}_i$ が約分されて

$$ = -\sum_i y_i (\delta_{ik} – \hat{p}_k) = -y_k + \hat{p}_k \sum_i y_i = -y_k + \hat{p}_k = \hat{p}_k – y_k $$

最後の等号で $\sum_i y_i = 1$(one-hotベクトル)を使いました。

Python実装

import numpy as np

def softmax_cross_entropy_forward_backward(logits, y_onehot):
    """
    ソフトマックス + 交差エントロピーの順伝播と逆伝播

    Parameters:
        logits: (batch_size, K) ロジット
        y_onehot: (batch_size, K) one-hotラベル

    Returns:
        loss: スカラー、バッチ平均の交差エントロピー損失
        grad: (batch_size, K) ロジットに対する勾配
    """
    # 順伝播(数値安定なソフトマックス + CE)
    c = np.max(logits, axis=-1, keepdims=True)
    log_sum_exp = c + np.log(np.sum(np.exp(logits - c), axis=-1, keepdims=True))
    log_probs = logits - log_sum_exp  # log-softmax

    # 交差エントロピー損失
    loss = -np.sum(y_onehot * log_probs, axis=-1)  # 各サンプルの損失
    mean_loss = np.mean(loss)

    # 逆伝播(勾配)
    probs = np.exp(log_probs)
    grad = (probs - y_onehot) / logits.shape[0]  # バッチ平均の勾配

    return mean_loss, grad

# 数値勾配との比較で正しさを検証
np.random.seed(42)
batch_size = 4
K = 5
logits = np.random.randn(batch_size, K)
y_true = np.random.randint(0, K, batch_size)
y_onehot = np.eye(K)[y_true]

loss, grad = softmax_cross_entropy_forward_backward(logits, y_onehot)
print(f"Loss: {loss:.6f}")
print(f"Gradient shape: {grad.shape}")
print(f"Gradient:\n{grad}")

# 数値勾配で検証
eps = 1e-5
numerical_grad = np.zeros_like(logits)
for i in range(batch_size):
    for j in range(K):
        logits_plus = logits.copy()
        logits_plus[i, j] += eps
        loss_plus, _ = softmax_cross_entropy_forward_backward(logits_plus, y_onehot)

        logits_minus = logits.copy()
        logits_minus[i, j] -= eps
        loss_minus, _ = softmax_cross_entropy_forward_backward(logits_minus, y_onehot)

        numerical_grad[i, j] = (loss_plus - loss_minus) / (2 * eps)

print(f"\nMax difference (analytical vs numerical): {np.max(np.abs(grad - numerical_grad)):.2e}")

この実装では、解析的に計算した勾配 $\hat{p}_k – y_k$ と数値微分の結果を比較しています。両者の差が $10^{-10}$ オーダー以下であれば、導出と実装が正しいことが確認できます。実際の深層学習フレームワーク(PyTorch, TensorFlow)も内部で同様の計算を行っています。

ソフトマックスの拡張

Gumbel-Softmax

微分可能なサンプリングが必要な場合(VAEのカテゴリカル潜在変数など)、Gumbel-Softmax(Concrete distribution)が使われます。

$$ \hat{p}_k = \frac{e^{(\ln \pi_k + g_k)/T}}{\sum_j e^{(\ln \pi_j + g_j)/T}} $$

ここで $g_k \sim \text{Gumbel}(0, 1)$ はGumbel分布からのサンプルです。$T \to 0$ のとき、これはカテゴリカル分布からの離散サンプリングに収束します。

Sparse Softmax(sparsemax)

通常のソフトマックスは全ての要素に0でない確率を割り当てますが、sparsemaxは一部の要素をちょうど0にすることができます。

$$ \text{sparsemax}(\bm{z}) = \arg\min_{\bm{p} \in \Delta^{K-1}} \|\bm{p} – \bm{z}\|^2 $$

確率単体 $\Delta^{K-1}$ への射影として定義されます。注意機構で不要なトークンへのアテンション重みを完全に0にしたい場合などに有用です。

import numpy as np
import matplotlib.pyplot as plt

def softmax(z):
    z_shifted = z - np.max(z)
    exp_z = np.exp(z_shifted)
    return exp_z / np.sum(exp_z)

def sparsemax(z):
    """Sparsemax関数の実装"""
    K = len(z)
    sorted_z = np.sort(z)[::-1]
    cumsum = np.cumsum(sorted_z)
    # k* = max{k : 1 + k*z_(k) > sum_{j<=k} z_(j)}
    k_array = np.arange(1, K + 1)
    support = 1 + k_array * sorted_z > cumsum
    k_star = np.max(np.where(support)[0]) + 1
    tau = (cumsum[k_star - 1] - 1) / k_star
    return np.maximum(z - tau, 0)

# 比較: softmax vs sparsemax
z_values = [
    np.array([2.0, 1.0, 0.1, -0.5, -1.0]),
    np.array([3.0, 3.0, 0.0, -1.0, -2.0]),
    np.array([1.0, 1.0, 1.0, 1.0, 1.0]),
]

fig, axes = plt.subplots(1, 3, figsize=(16, 5))

for ax, z in zip(axes, z_values):
    p_softmax = softmax(z)
    p_sparsemax = sparsemax(z)

    x = np.arange(len(z))
    width = 0.35

    ax.bar(x - width/2, p_softmax, width, color="steelblue", alpha=0.8, label="Softmax")
    ax.bar(x + width/2, p_sparsemax, width, color="coral", alpha=0.8, label="Sparsemax")

    ax.set_xlabel("Class", fontsize=12)
    ax.set_ylabel("Probability / Weight", fontsize=12)
    ax.set_title(f"Logits: {z}", fontsize=12)
    ax.set_xticks(x)
    ax.legend(fontsize=10)
    ax.grid(True, alpha=0.3, axis="y")

    # スパース性を表示
    n_zero_softmax = np.sum(p_softmax < 1e-10)
    n_zero_sparsemax = np.sum(p_sparsemax < 1e-10)
    ax.text(0.98, 0.92,
            f"Softmax zeros: {n_zero_softmax}\nSparsemax zeros: {n_zero_sparsemax}",
            transform=ax.transAxes, ha="right", fontsize=9,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow"))

plt.suptitle("Softmax vs Sparsemax", fontsize=15, y=1.02)
plt.tight_layout()
plt.savefig("softmax_vs_sparsemax.png", dpi=150, bbox_inches="tight")
plt.show()

このグラフから、ソフトマックスとsparsemaxの違いが明確に読み取れます。

  1. 左図: ロジットに差がある場合、ソフトマックスは全クラスに非零の確率を割り当てますが、sparsemaxはロジットが低いクラス(-0.5, -1.0)に正確に0を割り当てます。確率が0のクラスが2つあり、スパースな出力が得られています

  2. 中央図: 上位2クラスのロジットが等しい場合、sparsemaxはその2クラスにのみ均等に確率を割り当て、残り3クラスは0です。一方、ソフトマックスは全クラスに非零の確率を割り当てます

  3. 右図: 全ロジットが等しい場合、ソフトマックスもsparsemaxも一様分布を返します。sparsemaxでも全要素が等しければスパース化は起きません

sparsemaxは注意機構で解釈性が重要な場合(どのトークンに注目しているかを明確にしたい場合)に特に有用です。

実践的な注意点

ソフトマックスの飽和

ロジットの値が非常に大きいとき、ソフトマックスの出力はほぼone-hotになり、勾配がほぼ0になります。これはソフトマックスの飽和(saturation)と呼ばれ、学習の停滞を引き起こすことがあります。

ただし、交差エントロピー損失と組み合わせた場合、ロジットに対する勾配 $\hat{p}_k – y_k$ には飽和の問題がないことは既に確認しました。飽和が問題になるのは、ソフトマックスの出力を中間層で使う場合(例: Attention weightの計算)です。

ソフトマックスの飽和: ロジット差が開くと勾配が消える

ロジットの差が開くにつれて最大クラスの確率は1に張り付き、ヤコビアンの最大成分(勾配の大きさ)は急落します。差が8を超えるとヤコビアン最大成分は約0.0007——実質的に勾配が流れない飽和領域です。attentionのスコアを $\sqrt{d_k}$ でスケーリングするのは、まさにこの領域を避けるためです。

ソフトマックスの温度スケーリングによるキャリブレーション

学習済みモデルの予測確率が実際の確率と一致しない(miscalibration)場合、温度スケーリングを使って事後的にキャリブレーションできます。

検証データ上で交差エントロピーを最小化する温度 $T^*$ を求め、テスト時に $\text{softmax}(\bm{z}/T^*)$ を使います。$T^* > 1$ ならばモデルはoverconfident(自信過剰)であり、$T^* < 1$ ならばunderconfidentです。

import numpy as np
from scipy.optimize import minimize_scalar

rng = np.random.default_rng(42)
n, K = 2000, 3
y_true = rng.integers(0, K, n)

# モデルの予測は70%しか当たらないのに、確信度だけ高い状況を作る(=自信過剰)
model_pred = y_true.copy()
flip = rng.random(n) < 0.30
model_pred[flip] = (y_true[flip] + rng.integers(1, K, flip.sum())) % K
logits = rng.standard_normal((n, K)) * 0.5
logits[np.arange(n), model_pred] += 3.0
y_onehot = np.eye(K)[y_true]

def nll_with_temperature(T, logits, y_onehot):
    z = logits / T
    c = np.max(z, axis=-1, keepdims=True)
    lse = c + np.log(np.sum(np.exp(z - c), axis=-1, keepdims=True))
    return -np.mean(np.sum(y_onehot * (z - lse), axis=-1))

res = minimize_scalar(lambda T: nll_with_temperature(T, logits, y_onehot),
                      bounds=(0.01, 10.0), method="bounded")
T_opt = res.x

def softmax_T(z, T=1.0):
    z = z / T
    e = np.exp(z - z.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

print(f"正解率 = {(logits.argmax(1) == y_true).mean():.3f}")
print(f"平均確信度 (T=1)  = {softmax_T(logits).max(1).mean():.3f}")
print(f"最適温度 T* = {T_opt:.2f}")
print(f"平均確信度 (T=T*) = {softmax_T(logits, T_opt).max(1).mean():.3f}")
# => 正解率 = 0.698
# => 平均確信度 (T=1)  = 0.893
# => 最適温度 T* = 2.08
# => 平均確信度 (T=T*) = 0.675

シミュレーション結果は、絵に描いたような自信過剰です。正解率は 0.698 しかないのに、モデルの平均確信度は 0.893——「9割の自信で7割しか当たらない」状態です。NLLを最小化する温度は $T^* = 2.08$ と1より大きくなり、ソフトマックスを高温側に補正して確信度を 0.675 まで下げることで、実際の正解率とほぼ釣り合う予測確率になりました。

温度スケーリングによるキャリブレーション: 信頼度ダイアグラムとNLLの温度依存性

左の信頼度ダイアグラムでは、補正前(赤)が対角線の下に大きく垂れ下がり(確信度0.9の予測の実際の正解率は0.7程度)、補正後(青)は対角線に近づいています。右のNLL曲線は $T=1$ ではなく $T^* = 2.08$ で最小——$T^* > 1$ はモデルが自信過剰であることの定量的な証拠です。

まとめ

本記事では、ソフトマックス関数の定義・性質・実装について包括的に解説しました。

  • ソフトマックス関数は任意の実数ベクトルを確率分布に変換する関数で、argmaxの「滑らかな近似」として理解できる
  • 平行移動不変性 $\text{softmax}(\bm{z} + c\bm{1}) = \text{softmax}(\bm{z})$ は数値安定な実装の鍵であり、$\max(\bm{z})$ を引くことでオーバーフローを防ぐ
  • ヤコビアンは $\hat{p}_i(\delta_{ij} – \hat{p}_j)$ で表され、出力のみで計算可能である
  • 温度パラメータ $T$ は出力の「鋭さ」を制御し、知識蒸留やキャリブレーションで活用される
  • 交差エントロピーとの組み合わせでロジットに対する勾配が $\hat{p}_k – y_k$ と非常にシンプルになり、勾配消失が起きない
  • sparsemaxはスパースな出力を返す拡張で、解釈性が重要な場面に適する

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