Softmax函数：为什么这个公式统治了神经网络的概率输出

在神经网络的输出层，我们经常看到这样一个公式：

$$\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$$

它看起来如此简单——指数、求和、归一化。但为什么是这个特定的形式？为什么不直接除以总和？为什么要有指数？为什么这个函数能统治从图像分类到大语言模型的几乎所有概率输出？

答案藏在三个关键词中：可微性、放大效应和概率解释。

一个简单的直觉问题

假设你在训练一个图像分类器，需要将图像分成猫、狗、鸟三类。神经网络的最后一层输出了三个数值：[2.1, 1.5, 0.8]。

如果直接归一化，每个类别的"概率"是 [0.49, 0.35, 0.19]。看起来合理，但问题来了：第二大的狗的得分只比猫少0.6，概率却只有猫的71%。这是否意味着模型对猫的信心是狗的1.4倍？

但现实是，模型的"信心"程度应该与得分差异呈非线性关系。一个得分明显高于其他的类别，模型应该更加"确信"它的预测。

Softmax通过指数函数解决了这个问题。同样的输入经过Softmax后变成 [0.63, 0.28, 0.09]。现在猫的概率是狗的2.25倍，而不是1.4倍。指数函数放大了差异。

flowchart LR
    subgraph Input["输入层"]
        A[原始得分]
    end
    
    subgraph Compare["对比分析"]
        B["简单归一化<br/>[0.49, 0.35, 0.19]<br/>比值: 1.4x"]
        C["Softmax<br/>[0.63, 0.28, 0.09]<br/>比值: 2.25x"]
    end
    
    A --> B
    A --> C
    
    style A fill:#e1f5ff
    style B fill:#ffcccb
    style C fill:#90EE90

指数函数的放大效应

为什么指数函数能放大差异？考虑一个简单例子。假设两个输入值分别是 2 和 1。

直接相除：$2/1 = 2$，比值是2。

经过指数：$e^2 / e^1 = e^{2-1} = e \approx 2.718$，比值变成了2.718。

指数函数将线性差异转化为指数差异。差距被放大了。

更重要的是，这种放大是"自适应"的。当所有输入值都很小时，比如 [0.1, 0.05, 0.02]，指数后的比值是 $e^{0.1} : e^{0.05} : e^{0.02} = 1.65 : 1.05 : 1$，分布相对均匀。当输入值差异巨大时，比如 [10, 1, 0.1]，比值变成 $e^{10} : e^1 : e^{0.1} \approx 22026 : 2.7 : 0.1$，分布极度集中。

这正是我们想要的：当模型"不确定"时，概率分布相对平坦；当模型"很确定"时，概率分布高度集中。

graph TD
    subgraph LowVariance["低差异输入 [0.1, 0.05, 0.02]"]
        A1["原始比值: 5:2.5:1"]
        A2["指数后比值: 1.65:1.05:1<br/>分布相对均匀"]
    end
    
    subgraph HighVariance["高差异输入 [10, 1, 0.1]"]
        B1["原始比值: 100:10:1"]
        B2["指数后比值: 22026:2.7:0.1<br/>分布高度集中"]
    end
    
    A1 --> A2
    B1 --> B2
    
    style A1 fill:#fff3e0
    style A2 fill:#e8f5e9
    style B1 fill:#fff3e0
    style B2 fill:#ffcccc

温度参数：从统计力学借来的智慧

Softmax有一个重要变体——带温度参数的Softmax：

$$\text{softmax}(z_i, T) = \frac{e^{z_i/T}}{\sum_{j=1}^{K} e^{z_j/T}}$$

这个T并非随意命名。它来自统计力学中的玻尔兹曼分布。

在统计力学中，一个系统处于某个状态的概率与该状态的能量成反比，与温度成正比。具体来说，粒子处于能量为 $E_i$ 状态的概率是：

$$P(E_i) = \frac{e^{-E_i/kT}}{Z}$$

其中 $k$ 是玻尔兹曼常数，$T$ 是温度，$Z$ 是配分函数（归一化因子）。

将能量 $E_i$ 替换为 $-z_i$（得分越高，“能量"越低，概率越大），就得到了Softmax的形式。温度 $T$ 控制分布的"尖锐程度”：

高温 (T → ∞)：分布趋于均匀。就像高温气体分子运动混乱，所有状态概率相近。
低温 (T → 0)：分布高度集中。就像低温下分子趋向最低能量状态，最大得分项概率趋向1。

在大语言模型中，温度参数直接影响输出的多样性。低温（如0.1）让模型"保守"，总是选择最可能的词；高温（如2.0）让模型"创意"，敢于选择概率较低的词。

flowchart LR
    A[输入 logits] --> B{温度 T 调节}
    B --> C{指数运算 e^z/T}
    C --> D{归一化}
    D --> E[概率分布]
    
    F["低温 T=0.1<br/>分布集中"] -.-> E
    G["高温 T=2.0<br/>分布均匀"] -.-> E
    
    style A fill:#e1f5ff
    style B fill:#f9f
    style C fill:#bbf
    style D fill:#9f9
    style E fill:#ffd700
    style F fill:#ffcccc
    style G fill:#e8f5e9

为什么是Softmax而不是简单归一化

假设我们用简单归一化替代Softmax：

$$\text{simple_norm}(z_i) = \frac{z_i}{\sum_{j} z_j}$$

这看起来也能得到概率分布。但问题在于：

问题一：负值处理

如果输入包含负值，简单归一化可能产生负概率，这在数学上没有意义。Softmax通过指数函数保证所有值都是正数。

问题二：梯度消失

这是更关键的问题。考虑一个二分类问题，真实标签是类别1，模型输出是 $[p_1, p_2]$。交叉熵损失是：

$$L = -\log(p_1)$$

对于简单归一化，$p_1 = \frac{z_1}{z_1 + z_2}$，损失对 $z_1$ 的梯度是：

$$\frac{\partial L}{\partial z_1} = -\frac{1}{p_1} \cdot \frac{z_2}{(z_1 + z_2)^2}$$

当 $z_1$ 和 $z_2$ 都很大时，梯度可能接近零。

但对于Softmax，$p_1 = \frac{e^{z_1}}{e^{z_1} + e^{z_2}}$，梯度是：

$$\frac{\partial L}{\partial z_1} = p_1 - 1$$

梯度形式极其简洁，且当预测正确时（$p_1 \approx 1$），梯度接近零；当预测错误时（$p_1 \approx 0$），梯度接近-1，能有效地纠正错误。

这个优雅的梯度形式并非巧合，而是数学上的必然结果。

flowchart TB
    subgraph SimpleNorm["简单归一化问题"]
        S1["梯度复杂"]
        S2["梯度可能消失"]
        S3["无法处理负值"]
    end
    
    subgraph Softmax["Softmax优势"]
        M1["梯度简洁: p - 1"]
        M2["梯度自适应"]
        M3["保证正值"]
    end
    
    SimpleNorm -.->|"对比"| Softmax
    
    style S1 fill:#ffcccc
    style S2 fill:#ffcccc
    style S3 fill:#ffcccc
    style M1 fill:#90EE90
    style M2 fill:#90EE90
    style M3 fill:#90EE90

与交叉熵的完美配合

Softmax与交叉熵损失的结合是深度学习中最优雅的数学设计之一。

设Softmax输出为 $\mathbf{p}$，真实标签的one-hot向量为 $\mathbf{y}$，交叉熵损失为：

$$L = -\sum_i y_i \log(p_i)$$

由于 $\mathbf{y}$ 是one-hot向量，只有真实类别 $k$ 对应的位置是1，其他为0，所以：

$$L = -\log(p_k)$$

反向传播时，我们需要计算 $\frac{\partial L}{\partial z_i}$。神奇的事情发生了：

$$\frac{\partial L}{\partial z_i} = p_i - y_i$$

这个结果简单得令人难以置信。梯度直接等于预测概率减去真实标签。对于真实类别 $k$，梯度是 $p_k - 1$；对于其他类别，梯度是 $p_i$。

这种简洁性的来源是Softmax的指数函数与交叉熵的对数函数的完美抵消：

$$\frac{\partial (-\log(\frac{e^{z_k}}{\sum_j e^{z_j}}))}{\partial z_i} = \frac{e^{z_i}}{\sum_j e^{z_j}} - \mathbb{1}_{i=k}$$

Log-Sum-Exp函数的梯度恰好是Softmax本身。

flowchart TD
    A[前向传播] --> B[计算 logits z]
    B --> C["Softmax: p = e^z / Σe^z"]
    C --> D["交叉熵损失: L = -log p_k"]
    D --> E[反向传播]
    E --> F["梯度: ∂L/∂z = p - y"]
    
    G["对数与指数完美抵消<br/>梯度形式极其简洁"] -.-> F
    
    style A fill:#e1f5ff
    style B fill:#fff3e0
    style C fill:#ff9999
    style D fill:#f9f0f0
    style E fill:#90EE90
    style F fill:#90EE90
    style G fill:#ffd700

这个优雅的梯度形式解释了为什么神经网络能高效训练：当预测错误时，梯度大，更新幅度大；当预测正确时，梯度小，更新幅度小。这是一种自适应学习率的机制。

在Transformer中的关键角色

Softmax在Transformer中有两个关键应用位置：注意力机制和输出层。

注意力机制中的Softmax

在自注意力中，Query和Key的点积经过Softmax得到注意力权重：

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

这里有一个关键细节：为什么要除以 $\sqrt{d_k}$？

假设 $d_k = 512$，Query和Key的每个维度是独立同分布的，均值为0，方差为1。那么点积 $q \cdot k$ 的期望是0，方差是 $d_k = 512$。

当输入到Softmax的值方差很大时，Softmax输出会高度集中在某个位置——几乎变成one-hot向量。这会导致梯度几乎只流向一个位置，其他位置的梯度接近零，造成训练困难。

除以 $\sqrt{d_k}$ 将方差归一化到1，使Softmax的输入分布在合理范围内。

输出层中的Softmax

在大语言模型中，输出层的Softmax将隐藏状态映射到词表上每个词的概率。对于GPT-3这样的模型，词表大小约5万，Softmax需要计算5万个指数并归一化。

flowchart LR
    subgraph Transformer["Transformer架构"]
        A[输入序列] --> B[Embedding]
        B --> C[多头注意力]
        C --> D[Feed Forward]
        D --> E[输出层]
    end
    
    subgraph Softmax位置["Softmax应用位置"]
        F["注意力权重计算<br/>softmax QK^T/√d_k"]
        G["词表概率分布<br/>softmax logits"]
    end
    
    C -.-> F
    E -.-> G
    
    style A fill:#e1f5ff
    style C fill:#ff9999
    style E fill:#ff9999
    style F fill:#ffd700
    style G fill:#ffd700

数值稳定性的隐形陷阱

Softmax有一个看似简单但危险的数值问题。

假设输入是 [1000, 1001, 1002]。$e^{1000} \approx 10^{434}$，这远超过计算机能表示的范围（float64最大约 $10^{308}$），直接计算会溢出。

解决方案基于一个数学性质：Softmax对平移不变。

$$\text{softmax}(z) = \text{softmax}(z - c)$$

证明：

$$\frac{e^{z_i}}{\sum_j e^{z_j}} = \frac{e^{z_i - c}}{\sum_j e^{z_j - c}} = \frac{e^{z_i}/e^c}{\sum_j e^{z_j}/e^c} = \frac{e^{z_i}}{\sum_j e^{z_j}}$$

因此，我们可以减去最大值：

def safe_softmax(z):
    z_max = max(z)
    exp_z = exp(z - z_max)
    return exp_z / sum(exp_z)

这个技巧称为Log-Sum-Exp Trick，在所有深度学习框架中都是标准实现。

flowchart TD
    A["原始输入 [1000, 1001, 1002]"] --> B{直接计算?}
    B -->|"会溢出"| C["溢出错误!"]
    B -->|"减去最大值"| D["平移后 [0, 1, 2]"]
    D --> E["计算指数 [1, e, e²]"]
    E --> F["归一化得到概率"]
    
    G["数学保证:<br/>softmax z = softmax z-c"] -.-> D
    
    style A fill:#fff3e0
    style C fill:#ffcccc
    style D fill:#e8f5e9
    style E fill:#e8f5e9
    style F fill:#90EE90
    style G fill:#ffd700

更进一步，当需要计算对数概率时（如交叉熵损失），直接使用Log-Sum-Exp：

$$\log(\text{softmax}(z_i)) = z_i - \log\sum_j e^{z_j} = z_i - \text{LSE}(z)$$

其中LSE是Log-Sum-Exp函数。这避免了中间的指数运算，更加稳定。

梯度的数学之美

Softmax的梯度有一个优雅的矩阵形式。

设Softmax输出向量为 $\mathbf{p}$，则对输入 $\mathbf{z}$ 的梯度是：

$$\frac{\partial \mathbf{p}}{\partial \mathbf{z}} = \text{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^T$$

其中 $\text{diag}(\mathbf{p})$ 是以 $\mathbf{p}$ 为对角线的矩阵，$\mathbf{p}\mathbf{p}^T$ 是外积矩阵。

具体来说：

对角线元素：$\frac{\partial p_i}{\partial z_i} = p_i(1-p_i)$
非对角线元素：$\frac{\partial p_i}{\partial z_j} = -p_i p_j$ ($i \neq j$)

这个雅可比矩阵揭示了一个重要性质：Softmax的每个输出都依赖于所有输入。改变任何一个输入，所有输出都会变化。这与ReLU等逐元素激活函数不同，后者每个输出只依赖于对应的输入。

这也是为什么Softmax被称为"软"——它没有硬性的选择，所有选项都有非零概率，只是概率大小不同。

graph TD
    subgraph Jacobian["雅可比矩阵结构"]
        A["对角线: p_i 1-p_i<br/>表示自身影响"]
        B["非对角线: -p_i * p_j<br/>表示相互抑制"]
    end
    
    subgraph Properties["重要性质"]
        C["每个输出依赖所有输入"]
        D["改变一个输入<br/>所有输出都变化"]
        E["全局归一化机制"]
    end
    
    Jacobian --> Properties
    
    style A fill:#e8f5e9
    style B fill:#fff3e0
    style C fill:#e1f5ff
    style D fill:#e1f5ff
    style E fill:#e1f5ff

替代方案与未来演进

Softmax并非唯一选择。研究者提出了多种替代方案：

Sparsemax：能产生真正的稀疏输出（部分概率严格为0），基于凸投影理论。但梯度在边界处不连续，优化更困难。

Sigmoid Attention：用Sigmoid替代Softmax计算注意力，避免全局归一化。2024年的研究表明其在样本效率上优于Softmax，但尚未在大规模模型中验证。

多项式替代：2024年的论文质疑是否必须产生概率分布，提出了多项式核替代指数函数，在某些任务上表现更好。

为什么Softmax仍然主导？因为它在三个方面取得了平衡：

数学优雅：与交叉熵损失完美配合，梯度简洁
计算可行：指数运算在现代GPU上有良好支持
经验有效：在数十年实践中证明有效

flowchart LR
    subgraph Alternatives["替代方案"]
        A[Sparsemax<br/>稀疏输出]
        B[Sigmoid Attention<br/>避免全局归一化]
        C[多项式核<br/>替代指数函数]
    end
    
    subgraph Softmax["Softmax优势"]
        D[数学优雅]
        E[计算可行]
        F[经验有效]
    end
    
    Alternatives -->|"挑战"| Softmax
    Softmax -->|"仍主导"| G[实际应用]
    
    style A fill:#e8f5e9
    style B fill:#e8f5e9
    style C fill:#e8f5e9
    style D fill:#ffd700
    style E fill:#ffd700
    style F fill:#ffd700
    style G fill:#90EE90

写在最后

Softmax的统治力来自三个层面的设计智慧：

数学层面：指数函数的非线性放大差异，对数与指数在梯度计算中的完美抵消，以及与统计力学的深刻联系。

工程层面：数值稳定的计算技巧，与交叉熵损失的优雅配合，在GPU上的高效实现。

哲学层面：它体现了机器学习的核心思想——将原始信号转化为概率，让模型能够表达不确定性而非做出硬性判断。

理解Softmax不仅仅是记忆一个公式，而是理解为什么这个特定的数学形式能同时满足概率解释、梯度优化和工程实现的多重需求。这种"为什么"的思考，是从使用工具到理解工具的跨越。

参考文献

Bridle, J. S. (1990). Probabilistic Interpretation of Feedforward Classification Network Outputs. NATO ASI Series.
Goodfellow, I., Bengio, Y., & Courville, A. (2016). Deep Learning. MIT Press. Chapter 6.2.2.3.
Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS.
Martins, A., & Astudillo, R. (2016). From Softmax to Sparsemax. ICML.
Boltzmann, L. (1868). Studien über das Gleichgewicht der lebendigen Kraft. Wiener Berichte.
Blanchard, M., et al. (2024). Polynomial Alternatives to Softmax in Transformers. arXiv.
Milakov, M., & Gimelshein, N. (2018). Online Normalizer Calculation for Softmax. arXiv.
Jaegle, A., et al. (2021). Perceiver IO: A General Architecture for Structured Inputs & Outputs. ICML.
Luce, R. D. (1959). Individual Choice Behavior. John Wiley & Sons.
Guo, D., et al. (2025). Sigmoid Self-Attention has Lower Sample Complexity than Softmax Self-Attention. arXiv.

一个简单的直觉问题#

指数函数的放大效应#

温度参数：从统计力学借来的智慧#

为什么是Softmax而不是简单归一化#

与交叉熵的完美配合#

在Transformer中的关键角色#

数值稳定性的隐形陷阱#

梯度的数学之美#

替代方案与未来演进#

写在最后#

参考文献#

更多精彩内容

Logits：神经网络输出的原始真相，从概念到实践的完整解析

万能逼近定理：为什么两层神经网络能逼近任意函数

Sigmoid与Softmax：多分类与多标签任务的激活函数选择逻辑

Transformer 的注意力机制究竟在计算什么？从 QKV 到多头的完整解析

Softmax的数值稳定性：为什么一行简单的代码能让训练崩溃