softmax函数，又称softargmax或归一化指数函数，是一种将实数向量转换为概率分布的数学函数。它是逻辑函数在多维空间中的推广，广泛应用于机器学习，尤其在深度学习中用于多类分类问题。
定义
给定输入向量 \( \mathbf{z} = (z_1, \dots, z_K) \in \mathbb{R}^K \)，其中 \( K > 1 \)，softmax函数 \( \sigma: \mathbb{R}^K \to (0,1)^K \) 对每个分量 \( i \) 定义为：
\[ \sigma(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} \]
输出是一个正值向量，其元素之和为1，因此可解释为类别概率分布。指数运算放大了输入之间的差异，因此较大的输入值会产生不成比例地更大的概率。例如，对向量 \( (1, 2, 8) \) 应用softmax，结果约为 \( (0.001, 0.002, 0.997) \)，将几乎所有的概率质量分配给最大元素。
性质
softmax函数具有以下几个重要性质：
- 输出范围：每个分量严格介于0和1之间。
- 和为1：所有输出分量之和等于1。
- 顺序保持：如果 \( z_i > z_j \)，则 \( \sigma(\mathbf{z})_i > \sigma(\mathbf{z})_j \)。
- 对常数平移的不变性：对所有输入加上常数 \( c \) 不会改变输出，因为 \( e^{z_i + c} / \sum_j e^{z_j + c} = e^{z_i} / \sum_j e^{z_j} \)。这一性质常用于数值稳定性，通过在计算前减去最大输入值来实现。
温度参数
softmax函数的一种变体引入了温度参数 \( \beta \)（或其倒数 \( T = 1/\beta \)）：
\[ \sigma(\mathbf{z})_i = \frac{e^{\beta z_i}}{\sum_{j=1}^{K} e^{\beta z_j}} \]
当 \( \beta > 1 \) 时，分布变得更加集中（“更尖锐”）于最大输入附近，而 \( 0 < \beta < 1 \) 则产生更均匀的分布。该参数常用于top-k采样和温度缩放，以控制生成式AI模型的输出随机性。
应用
softmax是许多神经网络架构中的核心组件：
- 分类：作为最终激活层，将原始logits转换为类别概率，从而支持使用交叉熵损失进行训练。
- 注意力机制：在Transformer (architecture)模型中，softmax用于计算序列位置上的注意力权重，如多头注意力所示。
- 强化学习：将动作偏好转换为随机策略。
- 混合模型：用于计算高斯混合模型等模型中的混合权重。
数值稳定性
直接计算指数可能导致大输入时溢出。常见的解决方法是先减去最大输入值再进行指数运算：
\[ \sigma(\mathbf{z})_i = \frac{e^{z_i - \max(\mathbf{z})}}{\sum_{j=1}^{K} e^{z_j - \max(\mathbf{z})}} \]
由于平移不变性，此变换不会改变结果，但确保最大指数为零，从而防止溢出。
相关概念
softmax函数与逻辑函数密切相关，后者是其在 \( K = 2 \) 时的特例。它还出现在损失函数（如交叉熵）的上下文中，以及用于序列生成的束搜索解码中。在大语言模型推理中，softmax应用于logits以获得token概率，通常通过温度调整来平衡创造性和连贯性。