Softmax

Softmax（归一化指数函数）将任意实数向量转换为合法的概率分布，使所有输出值介于 0 与 1 之间且总和为 1。它是多分类、注意力权重计算与语言模型输出采样的核心运算之一。

Softmax 也称「归一化指数函数」，是逻辑斯谛函数（sigmoid）在多维情形下的推广。

• 输入：K 维实数向量（logits），每个分量代表某类别的「原始得分」
• 输出：K 维概率向量，满足各分量非负且总和严格等于 1
• 公式：第 i 个输出 = exp(z_i) / Σ exp(z_j)，求和遍历所有 K 个类别
• 单调性：输入值越大，对应输出概率越大，但类别相对排名与直接比较 logits 完全一致
• argmax 近似：当 logits 差异极大时，softmax 输出近似于「赢者通吃」的 argmax

直接计算 exp(z_i) 在 z_i 较大时会导致浮点溢出，工程实现需做等价变换。

• 减去最大值：令 m = max(z)，则 softmax(z_i) = exp(z_i − m) / Σ exp(z_j − m)，数学等价但不会溢出
• 下溢处理：减去最大值后，其余分量为负数，exp 结果趋近于 0 而非无穷，可安全计算
• log-softmax：先计算 log 概率再做后续运算，与 NLLLoss 配合使用时更稳定，常用于交叉熵损失
• 混合精度风险：在半精度（FP16）训练中尤其需要数值稳定技巧，否则极易出现 NaN 或 Inf

在标准 softmax 公式中引入温度超参数 T，可控制输出分布的「尖锐程度」。

• 公式变形：第 i 个输出 = exp(z_i / T) / Σ exp(z_j / T)
• T → 0：分布趋近于 argmax，模型输出极为确定（贪心选择）
• T = 1：还原为标准 softmax
• T > 1：分布趋近于均匀分布，采样更多样，常用于语言模型的创意输出
• 知识蒸馏：Hinton 等人 2015 年的蒸馏论文中使用高温 softmax，让「软标签」携带类间相似度信息

Softmax 在深度学习的多个模块中扮演关键角色。

• 多分类输出层：神经网络末层使用 softmax，将 logits 转为类别概率，配合交叉熵损失端到端训练
• 注意力机制：Transformer 的 Self-Attention 中，对 Q·K^T / √d_k 应用 softmax 得到注意力权重，再加权聚合 V
• 语言模型采样：GPT 等自回归模型在推理时对词表 logits 做 softmax，结合温度 / top-p 等策略采样下一个 token
• 强化学习策略：策略网络输出层常用 softmax 表示离散动作的概率分布
• 多标签场景区分：多标签分类（各类互不排斥）通常改用 sigmoid，而非 softmax

Softmax 的数学形式可追溯到统计物理，但在机器学习中的命名与推广有明确时间线。

• 1902：吉布斯（Josiah Willard Gibbs）在统计力学中形式化了玻尔兹曼分布，其数学结构与 softmax 完全一致
• 1959：卢斯（R. Duncan Luce）在选择公理（Luce's choice axiom）中使用相同数学形式描述理性选择行为
• 1989/1990：John S. Bridle 在神经网络领域正式命名「softmax」，并将其解释为多分类的概率输出
• 2014：Bahdanau 等人将 softmax 引入注意力机制，用于机器翻译对齐权重
• 2017：Vaswani 等人在「Attention Is All You Need」中，softmax 成为 Transformer 架构的核心运算

标准 softmax 在大规模与极端场景下存在若干不足，研究者提出了多种改进方案。

• 大词表瓶颈：词表规模达数万至数十万时，每步均需对全词表计算 exp，代价高；常用层次 softmax（hierarchical softmax）或负采样缓解
• 注意力熵崩溃：在 Transformer 中，softmax 可能将注意力过度集中于少数 token，影响训练稳定性
• Sparsemax：一种稀疏替代方案，部分输出恰好为 0，使注意力更具可解释性
• 过度自信：对分布外（OOD）输入，softmax 仍会强行给出接近 1 的概率，是神经网络过度自信问题的根源之一