Softmax

Softmax 是一种函数，可将称为 logits 的原始分数向量转换为概率分布：每个输出都在 0 到 1 之间，且所有输出之和恰好为 1。对于向量 z，它计算 softmax(z)ᵢ = e^(zᵢ) / Σⱼ e^(zⱼ)，即对每个分数求指数，再除以所有指数值之和。Softmax 是从多个类别中选出一个类别的神经网络的标准输出层；它会在 Transformer 的注意力机制中将相似度分数转换为权重，还会把语言模型对其词表的分数转换为下一个词元的概率。它的名称体现了它是 argmax 的平滑、可微版本；argmax 会将全部概率分配给分数最高的单项。

Softmax 的工作原理#

指数函数会使所有输出都为正数，并放大分数之间的差异，因此最大 logit 会占据不成比例的概率份额。Softmax 会保留 logit 的顺序，所以最高分始终对应最高概率。只有 logit 之间的差值才有影响：给每个 logit 加上相同的常数，输出保持不变。库会利用这一性质来保证数值稳定性：在进行指数运算之前先减去最大 logit，因为 e¹⁰⁰⁰ 会导致 32 位浮点数溢出。

温度 T 控制分布的尖锐程度：softmax(z / T)。T 小于 1 时，分布会变得更尖锐，偏向最高类别；T 大于 1 时，分布会变得更平坦，趋近均匀分布。在统计力学中，同一公式以玻尔兹曼分布的形式出现，其中 T 就是字面意义上的温度。下面的 PyTorch代码片段展示了这两种效果：

import torch

logits = torch.tensor([2.0, 1.0, 0.1])

# Divide logits by a temperature before softmax: low sharpens, high flattens
for temperature in [0.5, 1.0, 2.0]:
    probs = torch.softmax(logits / temperature, dim=0)
    print(f"T={temperature}: {[round(p, 3) for p in probs.tolist()]}")

# Adding a constant leaves softmax unchanged, so libraries subtract the max to avoid overflow
big = logits + 1000
naive = big.exp() / big.exp().sum()
print(f"Naive: {naive.tolist()}, torch.softmax: {[round(p, 3) for p in torch.softmax(big, dim=0).tolist()]}")

当 T = 1 时，logits 2.0、1.0 和 0.1 会变成 [0.659, 0.242, 0.099]。将温度减半后，分布会变得更尖锐，成为 [0.864, 0.117, 0.019]；将温度加倍后，分布会变得更平坦，成为 [0.502, 0.304, 0.194]。给每个 logit 加上 1,000 后，朴素公式会溢出并输出 [nan, nan, nan]，而 torch.softmax 仍会返回 [0.659, 0.242, 0.099]。

训练中的 Softmax#

分类器通过结合 softmax 和交叉熵损失进行训练。损失为 −log(pᵧ)，即分配给正确类别 y 的概率的负对数；相对于 logits 的梯度很简单，就是 p − y，也就是预测概率减去 one-hot 标签。因此，框架会将两者融合：PyTorch 的 CrossEntropyLoss接收原始 logits，并在内部计算 log-softmax；与分别计算概率再取对数相比，这种方式数值更稳定。以这种方式训练的模型会输出 logits，只有在推理时需要概率时才应用 softmax。

Softmax 的应用场景#

• 图像分类：分类器的最后一层为每个类别生成一个 logit，softmax 会将它们转换为用户看到的置信度分数。
• 注意力：Transformer会对查询与每个键之间的相似度分数计算 softmax(QKᵀ / √dₖ)，因此每个词元的注意力权重之和为 1。除以 √dₖ 可以防止分数随维度增大而增长，导致 softmax 饱和。
• 文本生成：大语言模型（LLM）在每一步都会对包含数万或数十万个词元的词表应用 softmax，上文的 T 就是采样温度。核采样随后会从概率总和达到某个阈值的最小词元集合中进行采样。
• 强化学习：对动作价值应用 softmax 会得到一种策略：通常选择最佳动作，同时仍会探索其他动作；温度则用于权衡探索与利用。

YOLO26分类模型会在推理时于分类头中应用 softmax，每个结果的 probs属性包含 1,000 个 ImageNet类别的完整概率分布：

from ultralytics import YOLO

# Load a YOLO26 classification model trained on ImageNet
model = YOLO("yolo26n-cls.pt")
results = model("https://ultralytics.com/images/bus.jpg")

# probs holds the softmax output of the classification head
probs = results[0].probs
for index, conf in zip(probs.top5, probs.top5conf.tolist()):
    print(f"{results[0].names[index]}: {conf:.3f}")
print(f"Sum over all 1,000 classes: {probs.data.sum().item():.3f}")

它会输出概率最高的五个类别及其概率，而 1,000 个类别的概率总和为 1.000，因为分类头返回的是 softmax 分布。

Softmax 与相关函数的比较#

• Sigmoid：独立地将每个分数压缩到 0 到 1，因此输出之和不一定为 1。Sigmoid 适用于二分类和多标签问题，例如图像中同时包含一个人和一个背包的情况；目标检测器也常用它对每个类别进行评分。对于两个类别，softmax 恰好等价于对两个 logits 之差应用 sigmoid。
• Argmax：返回最大分数的索引，是一种没有梯度的硬决策。Softmax 保持相同的排序，但仍可微，因此可以用于训练；对 softmax 概率取 argmax，结果与对 logits 取 argmax 相同。
• Log-softmax：直接计算以保证稳定性的 softmax 对数，供交叉熵损失使用。
• 温度缩放：一种校准方法，在训练后用验证数据拟合单个温度，使 softmax 概率与模型实际预测正确的频率相符。由于将每个 logit 除以同一个正数 T 会保留它们的顺序，因此预测类别不会改变。