1. 先搞清楚它在解决什么问题

1.1 从二分类的“Yes/No”到多分类的“单选题”

很多同学第一次遇到 softmax 是在学图像分类的时候，模型最后一层接 10 个输出，然后就冒出这个函数，看起来像是“取指数再归一化”。我当初也是这么过来的，背了公式，也调通了训练脚本，但真正理解它其实花了挺久。关键不在于记住e^x / sum(e^x)这个形式，而是先明白它要解决的是一个什么样的问题。

分类任务本质上是在做一个单选题：输入一张图、一段文本、一串特征，模型需要判断它是哪个类别。模型内部最后通常会输出一组“评分”，也就是 logits。logits 的绝对值没有严格的概率含义，可以很大，也可以为负，甚至某些类别之间的关系只有相对大小才有意义。如果直接用原始评分去判断类别，虽然也能取最大的那个，但没法回答一个更基本的问题：这个预测到底有多大把握？

这时 softmax 登场了。它的任务就是把一组任意实数评分压缩成一组和为 1 的非负数，让输出可以像一个概率分布一样被解读。每个类别拿到的数值，代表模型认为“这一项是正确答案”的占比。听起来很简单，但这里有一个容易被忽略的细节：它不一定是最合理的归一化方式，为什么偏偏是它，后面会说。

1.2 一口气看懂数学定义

Softmax 的数学定义可以写成：对第 i 个类别的原始评分 z_i，先取指数，再除以所有类别指数的总和。

用公式表达就是 softmax(z_i) = exp(z_i) / Σ_j exp(z_j)。

其中 z 是模型直接输出的 logits，可以看成一个长度为类别数 C 的向量，i 和 j 都是类别索引。分母的求和范围是从第 1 类到第 C 类，所以最终输出的向量每一个分量都在 0 到 1 之间，并且所有分量加起来正好等于 1。

举个例子。假设三分类任务的 logits 是 [2.0, 1.0, 0.0]，那么：

• exp(2.0) 约等于 7.389
• exp(1.0) 约等于 2.718
• exp(0.0) 等于 1.0

分母总和约等于 11.107，所以 softmax 后的概率大约是 [0.665, 0.245, 0.090]。这里能看出第一类占明显优势，但并没有把概率完全压到 1 或 0，而是保留了其他类别的可能，方便后续做损失计算和不确定性估计。

真正的工程实现里，输入往往是一个带 batch 的矩阵，形状是 [batch_size, num_classes]。比如一批图片同时过模型，每张图给出一行 logits，softmax 需要沿着类别维度计算，而不会去动 batch 维度，否则概率会跨样本互相影响，导致每个样本自己的输出不再和为 1。这个维度问题看着简单，但写代码的时候特别容易翻车，后面会具体提。

2. 为什么偏偏是“指数”而不是别的函数

2.1 指数让分数差异变得足够敏感

很多人第一次看 softmax 都会问：直接除以所有分数的总和不也行吗？我当年也有这个疑惑。真实原因在于，原始 logits 可能包含负数，直接线性归一化会出现很多反直觉的结果。

假设 logits 是 [2.0, 1.0, -1.0]，如果直接除以总和，总和是 2.0，结果是 [1.0, 0.5, -0.5]，第三个概率变成负数，这完全不可接受。所以必须做一个变换，把负数也变成非负数，同时还要保持原始评分之间的相对关系。

指数函数天然适合这个任务。它有三个非常关键的性质：输入再小，输出也是正数；输入越大，输出增长越快；输入整体平移相同的数值，输出结果不会变。第三点尤其重要，因为 logits 的绝对大小经常受到网络初始化、归一化方式的影响，但真正重要的是类与类之间的差值，平移不变性保证了模型对评分的整体偏移不那么敏感。

另外，指数会放大类间差异。同样是 [2.0, 1.0, 0.0] 这组分数，彼此只相差 1，softmax 后已经能把第一类的概率推到 0.665。如果线性归一化，得到的大概是 [0.667, 0.333, 0.0]，表面看着差不多，但遇到更极端的情况差异就出来了。比如 logits 是 [10.0, 9.0, 8.0]，线性归一化后是 [0.370, 0.333, 0.296]，而 softmax 后大约是 [0.665, 0.245, 0.090]，高分段会被推得更陡。分类任务需要决策果断，指数带来的这种“放大效果”往往更符合需要。

2.2 温度参数：决定模型敢不敢下重注

Softmax 并不是只有一个固定形态，很多时候会在指数里加一个温度系数 T，变成 softmax(z_i / T)。T 越大，输出越平滑；T 越小，输出越尖锐。这个技术特别常见，比如知识蒸馏里，教师模型会把 T 调高，让学生模型看到更多“软标签”。

以 logits [2.0, 1.0, 0.0] 为例，如果 T = 2，那么指数里面是 [1.0, 0.5, 0.0]，softmax 后概率大约是 [0.465, 0.283, 0.251]，分布变平了很多，原先几乎被压掉的最低分类别也能有一定份额。如果 T = 0.5，指数里面是 [4.0, 2.0, 0.0]，概率大约是 [0.936, 0.055, 0.009]，几乎变成一次“下重注”的预测。

这个特性在分析模型是否过度自信时很有用。训练之后观察 softmax 输出的概率分布，如果大多数样本的最高概率都接近 1，但验证集准确率并没有完全跟上，说明模型可能过拟合或者校准不好。此时可以在推理阶段用温度缩放，通俗点讲就是调高 T，把概率往回拉一点。温度的选择不改变类别预测的排序，所以准确率不会变，但会改变概率的绝对大小，在很多需要阈值决策的场景里非常实用。

2.3 它是“可微的 argmax”

Softmax 还有一层身份：它是 argmax 的一种平滑近似。argmax 只能告诉你哪个最大，但它不可导，梯度没法回传。分类网络必须通过损失函数反向传播，做一个不可导的决策等于断掉了梯度链路。

Softmax 把“最大值”变成了一个近似 one-hot 的分布。当某一类的 logits 明显领先时，它的概率会非常接近 1，其余类别接近 0，行为上像 argmax；但由于整个公式是可导的，梯度可以顺畅地传回前面的层。简单说，它保留了 argmax 的决策能力，同时给了网络一个可以“慢慢调整”的平滑表面。

虽然有人会较真说 softmax 并不保证靠近 one-hot，确实在 logits 差距不大时它是一个较为温和的分布。但这正是训练时需要的样子，如果模型从一开始就百分之百确信某个类，反而不利于学习，梯度会很快消失。在分类模型训练过程中，softmax 承担着“给出决策，但留出学习空间”的角色。

3. 数值稳定性与工程实现里的坑

3.1 指数上溢：先减最大值再取指数

Softmax 公式看起来简洁，直接照着写却很容易出事。指数函数增长非常快，exp(1000) 在绝大多数普通浮点数体系里直接变成无穷大，带上overflow警告，NAN 顺着传播出去，整个 loss 就崩了。

我在一个模拟项目里遇到过这种问题，某次把 logits 的初始化范围改得太大，第一次前向传播就出现了 NAN。排查很久才发现不是网络结构的问题，而是 softmax 实现太“裸”，没有做数值稳定处理。解决办法不复杂：先对 logits 减去它这一行的最大值，再做指数和归一化。

原因是 softmax 的平移不变性。logits 减去同一个常数后，exp 的分子和分母会同时缩放相同的倍数，最终比值不变。所以先算 z_i - max(z)，最大值对应的那一项会变成 exp(0) = 1，其他项都小于或等于 1，不会再出现指数爆炸。

在代码里，我习惯写成这样：

def softmax(logits, dim=-1): max_logits = logits.max(dim=dim, keepdim=True) shifted = logits - max_logits exp_shifted = exp(shifted) return exp_shifted / exp_shifted.sum(dim=dim, keepdim=True)

注意keepdim要保持原有维度，否则 logits 的形状会被压缩，广播的时候报错。这个操作看似多了一步，但代价很小，换来的是稳定性，属于标准的工程习惯。

3.2 log 域救场：从 softmax 到 log-softmax

不仅要算 softmax，分类训练还得算交叉熵，而交叉熵里经常出现 log(p)。这里有一个隐藏风险：如果某个 p 特别小，log(p) 会趋近负无穷；如果因为浮点误差 p 恰好被算成 0，那 log(0) 直接 NAN。

更优雅的做法是把指数、归一化、取对数这些操作合并到一个函数里，直接计算 log-softmax。数学上可以变形：

log_softmax(z_i) = z_i - log(Σ_j exp(z_j))

这个写法本身也有数值风险，分母里的 log(Σ exp(z_j)) 可能溢出，所以继续使用减最大值的技巧，定义 logsumexp：

logsumexp(z) = max(z) + log(Σ_j exp(z_j - max(z)))

然后用 z_i - logsumexp(z) 得到 log-softmax。整套操作只需要一次指数求和和一次对数，既避免了概率下溢，也不会有大数溢出。真正训练时，大多数框架提供的交叉熵接口都已经把 softmax 和 log 合并成了这一步，而不是先算 softmax 再算 log，这样一来数值稳定性好很多。

3.3 维度方向：最容易踩的 tensor 形状问题

Softmax 实现里最容易出的问题，反而不是数学公式，而是dim参数选错。很多初学者对一个小批量输入 [batch_size, num_classes] 计算 softmax 时，可能图省事直接对整个矩阵做计算，导致概率不是按每个样本单独归一化的。

如果输入是 [batch_size, num_classes]，应该沿最后一个维度做 softmax，也就是dim=-1。遇到 [batch_size, num_classes, height, width] 这种分割任务输出时，类别维度多半在第二个位置，需要明确指定dim=1或dim=2，绝对不能依赖默认行为。

我踩过的一个坑是：网络输出形状是 [batch_size, 1, 28, 28]，类别只有两类，但我在推理时用了默认 softmax，结果把整个张量都归一化了，所有元素加起来才等于 1，预测自然错误。后来把所有 softmax 调用都改成显式传dim，并封装成函数，类似问题才彻底得到控制。遇到不确定形状时，可以先打印 logits 的 shape，再用一个小张量验证每行概率和是否为 1，这种检查习惯能省下大量调试时间。

4. 反向传播推导：Softmax 为什么总跟交叉熵一起用

4.1 单独求 Jacobian 会很麻烦

Softmax 虽然公式简单，但反向传播如果对着向量逐项求导，会得到一个矩阵，专业叫法是 Jacobian。因为输出 p_i 对某个输入 z_j 求导时，存在两种情况：i 等于 j 时是一种表达式，i 不等于 j 时是另一种表达式。分开写会很乱，容易绕晕。

不过实际训练中很少单独为 softmax 算梯度，因为它几乎总是和交叉熵损失绑在一起。如果把两者合并起来看，数学结论会简洁得让人惊讶：对第 j 类的 logits z_j 求交叉熵梯度，结果竟然是 p_j - y_j，也就是预测概率减去真实标签。这是一个非常干净、好用的结果，值得亲手推导一次。

4.2 亲手反推一遍梯度

假设一批样本里的某个样本，真实标签写成 one-hot 向量 y，其中正确类别 k 的位置是 1，其余位置是 0。交叉熵损失定义是 L = -Σ_i y_i * log(p_i)。

把 softmax 的表达式代进去，对正确类别 k 有 L = -z_k + log(Σ_j exp(z_j))，因为 one-hot 会把求和里无关项全部清零。对这个式子求 z_j 的偏导，分两种情况，但结果可以统一写成一个简洁形式：

∂L / ∂z_j = -y_j + exp(z_j) / Σ_t exp(z_t) = p_j - y_j

当 j 是正确类别时，y_j 等于 1，梯度就是 p_j - 1，因为 p_j 小于 1，所以通常是负数。也就是说，网络会发现“正确类别的概率还不够高”，于是往参数里传一个正向调整信号，把正确类别的 logits 往上推。当 j 是错误类别时，y_j 等于 0，梯度就是 p_j，一个正数，于是反向传播会把这个错误类别的 logits 往下压。这就是学习过程最核心的机制。

4.3 这个“巧合”背后的意义

Softmax 和交叉熵一起用时，梯度大小直接取决于模型预测和真实标签之间的差距。如果预测完全正确，p_j 等于 1 且错误类别概率为 0，梯度为 0，参数不再更新。如果预测明显错误，比如真实类别概率只有 0.1，梯度绝对值为 0.9，更新的力度就会很大。这种“错得越狠，改得越猛”的性质非常契合分类任务。

相比之下，如果单独用均方误差去拟合 one-hot 标签，因为 softmax 的输出被限制在 0 到 1 之间，均值误差的梯度在输出接近 0 或 1 时会容易被压得很小，造成训练缓慢。交叉熵没有这个问题，它天然地把驱动信号绑在概率残差上，也让 softmax 和交叉熵几乎成了分类任务里的固定组合。理解了这一层，再去调分类模型，看 loss 曲线的走向时会更有底气。

5. 一次完整的端到端实战

5.1 任务描述和数据准备

为了把概念落到地面上，我拿一个模拟的十分类图像任务来演示。假设输入是 28x28 的灰度图，类别是 0 到 9。模型内部网络可以很简单，几层卷积加全连接，最终输出 10 维 logits。

数据准备阶段有个容易被忽略的细节：交叉熵损失接口通常接收整数标签，而不是 one-hot 标签。如果你已经手动做了 one-hot，也可以，但要注意接口类型，有些实现会严格校验标签形状。我习惯用整数索引作为标签，代码更简单，内存也更省。

这里提一个常犯的错误：不要在数据 normalization 时把标签也归一化了。标签是索引或 one-hot 值，混合进浮点标准化会让后续损失计算完全乱套。准备数据时，把输入特征和标签分开处理，训练循环里根本不应该出现标签被“归一化”的代码。

5.2 关键实现：从原始 logits 到 loss

训练时最好不要分别调用 softmax 和 log，直接用合并版本的 log-softmax，再配合交叉熵。我习惯写一个简单的手工版本用于验证公式：

def log_softmax(logits, dim=-1): max_logits = logits.max(dim=dim, keepdim=True) shifted = logits - max_logits log_sum_exp = log(exp(shifted).sum(dim=dim, keepdim=True)) return shifted - log_sum_exp def cross_entropy(logits, labels): # logits: [batch_size, num_classes] # labels: [batch_size] 整数标签 log_probs = log_softmax(logits, dim=-1) batch_size = logits.shape[0] # 取出每个样本对应真实类别的 log 概率 selected = log_probs[range(batch_size), labels] return -selected.mean()

沿着正确标签索引出 log 概率后取平均，再取负号，就成了标准的交叉熵。实际训练每步更新时，还可以做一个更简短的写法，利用 logsumexp 化简：

def cross_entropy_short(logits, labels): # 每个样本的 loss = logsumexp(logits) - logits[实际类别] max_logits = logits.max(dim=-1, keepdim=True) shifted = logits - max_logits logsumexp = max_logits.squeeze(-1) + log(exp(shifted).sum(dim=-1)) batch_index = range(logits.shape[0]) return (logsumexp - logits[batch_index, labels]).mean()

这个写法在有 batch 的情况下尤其高效，它同时做了数值稳定和损失合并，比先 softmax 再取 log 再算 loss 更稳。

5.3 训练监控时我习惯看的三组数字

训练时除了盯 loss 和准确率，我还会额外记录三个东西：验证集 loss、验证集准确率、以及模型输出的平均最高概率。第三项常常被忽略，但它能反映模型的自信程度。

如果 loss 一直在降，准确率也还可以，但平均最高概率一跌到接近 0.5，说明模型可能没有把类别区分开。如果平均最高概率接近 0.99，而准确率只有 80%，那模型很可能只是“嘴上自信”，实际泛化不够，需要检查过拟合、数据增强或标签是否有噪声。

训练早期我还会观察 loss 的绝对数值。十分类随机初始化时，交叉熵大约等于 log(10) 约 2.3026。如果第一个 epoch 的 loss 离这个值太远，比如直接是 5 以上，往往说明预处理或权重初始化有问题，而不是单纯模型还没学会。判断训练是否正常，拿随机初始状态下理论交叉熵作基准非常高效。

6. 常见问题与排查实录

6.1 训练一开始 loss 就是 NaN：优先检查 logits 和标签

这类问题最直接的原因几乎都是数值溢出。出现 NaN 时，第一步不是改网络结构，而是检查模型最后一层的输出有没有巨大的值，很多情况下是因为初始化范围太大。如果使用默认的线性层初始化，一般不会出问题，但一旦自定义初始化，需要注意标准差控制在一个较小范围。

另外，如果标签索引超出了类别数，比如模型输出 10 类，但某一条数据的标签写成 10，索引越界也会产生奇怪的行为。我曾经在一个数据合并脚本里把标签从 1 到 10 而不是 0 到 9，前几个 batch 的 loss 出现异常大值，打印标签才发现是索引偏移。遇到 loss 异常，最有效的做法是把前三个 batch 的 logits、标签、mask 全部打印出来，逐项核对。

6.2 模型过于自信：标签平滑和温度缩放

很多分类模型训练末尾会出现输出概率极度尖锐的情况，预测正确时概率接近 1，错误时也硬给出一个接近 1 的错误类别。表面看准确率不差，但概率明显不够可靠。如果业务场景需要阈值判断，比如“置信度低于 0.8 的样本要交给人工”，这个问题就很要命。

标签平滑是一个常见解法：把 one-hot 标签的 1 改成 1 - ε，把 ε 摊到其他类别上。常用的 ε 是 0.1 或 0.05。这么做相当于告诉模型：不要对某个类过于绝对，保留一点点对其他类别的“怀疑”。它会带来一定的准确率下降，但往往能提升概率校准和泛化。

如果模型已经训练完了，不想重新训练，那么可以做温度缩放。做法是在验证集上找一个 T，让缩放后的 softmax 概率分布更贴近真实正确率，T 一般大于 1。温度缩放不改变最终预测的类别，只改变置信度大小，适合后处理阶段使用。

6.3 Softmax 输出的“概率”不等于真实不确定性

这是新手上路最容易产生的误解：softmax 给了一个介于 0 和 1 之间的数，就认为它是模型真实的置信度。实际上它会受训练数据分布、网络结构、标签质量等因素影响，是一个被校准过的“类概率”，但不是一个严格的统计概率。

比如一个训练得很充分的模型，对某些远离训练分布的输入，可能仍然给出很高的 softmax 概率。因为模型没见过类似数据，它只能按照学到的决策边界硬猜。想要真正衡量“未知”程度，需要用其他方法，比如对输入做扰动、做模型集成、或者用额外的不确定性估计模块。

我整理过一个排查参考表，遇到问题可以对照看一下：

这个表只覆盖了常见情况，具体问题还是得从打印中间结果开始排查。

7. 一些容易混淆的概念和我的实际体会

7.1 多标签分类别用 Softmax

Softmax 的前提是类别互斥，一个样本只能属于一个类别。如果业务是“一张图片里同时有猫和狗”，或者“一篇文档可以同时属于多个主题”，那就不是多分类，而是多标签分类，这时候 softmax 就不合适。

多标签任务应该对每个类别单独计算二分类概率，常用的是 sigmoid。因为每个类别可以独立为 0 或 1，类别之间不要求概率和为 1。选错函数会让模型被迫在多个正确标签之间做竞争，其中一个标签概率变大，另一个就被压小，效果自然受影响。

我自己接手一个模拟项目时，一开始没有仔细看任务定义，想当然用了 softmax 加交叉熵，模型整体表现很差。后来确认每条样本的标签不止一个，改成基于 sigmoid 的二元交叉熵之后，效果立刻正常。所以 start 之前第一步永远是确认类别互斥关系。

7.2 Softmax 也不止用在分类

Softmax 的应用范围远不止分类最后一层。注意力机制里的 attention 权重，本质上就是一个 softmax，是对一组相似度分数做归一化，让模型把注意力分配比例限制在 0 到 1 之间，并且总和为 1。这个思想和分类 softmax 完全一致，只是作用的对象从类别换成了不同的位置或 token。

另外，在一些生成模型或强化学习策略中，也会用带温度的 softmax 从动作分布中采样。温度高，采样更随机；温度低，采样更接近最优动作。理解 softmax 就不只是为了调分类模型，而是理解一大类“从分数到分布”的转换逻辑。

7.3 最后分享一个小技巧

给刚接触 softmax 的同学一个建议：遇到和它相关的报错或者反常现象，先别急着找调参策略，先在 10 分钟里手工跑一遍前向计算和反向梯度。拿一个三分类、两个样本的小例子，纸笔算一次，再用代码打印中间结果对一遍。这个方法能同时解决数值稳定性、维度错误、梯度方向三个大坑。

我自己每次写 softmax 相关实现时，也会习惯加一个单元测试，输入固定的 logits，断言输出是否等于手算预期值。看似笨重，但项目复杂以后，它帮我少踩了非常多回归性 bug。深度学习很多时候像在搭建积木，最底层的积木如果不够踏实，上层建得越高越容易塌。Softmax 就是那块值得多花一点时间磨实的积木。