Attention 05：softmax 为什么会变成注意力权重 • Jerry's Blog

01 > 02 > 03 > 04 > [05] > 06 | 07 > 08 > 09 > 10 > 11 > 12

“softmax 不是在证明谁最重要，而是在给当前 token 生成一个对所有候选位置的信息分配比例。“

这一讲要回答什么#

上一讲我们知道 $QK^\top$ 产生了原始分数矩阵——每一行是一个 Query 对所有 Key 的匹配分数。但这些分数可以是任意实数：正的、负的、大的、小的。

它们不能直接当”注意力权重”用，因为：

负分数怎么理解？注意力为 -0.3 是什么意思？
分数没有统一量纲，不同行之间无法比较
不满足”分配比例”的语义——权重应该加起来等于 1

所以需要一个函数，把任意实数变成合法的概率分布。这个函数就是 softmax：

\alpha_j = \text{softmax}(s_j) = \frac{e^{s_j}}{\sum_{k=1}^{n} e^{s_k}}

这一讲要回答：为什么偏偏是 softmax？它的三个性质为什么恰好满足”注意力权重”的要求？

先手算一次#

拿 3 个分数 $[2.0, \; 1.0, \; 0.5]$ 做一次完整的 softmax。

交互演示：实时 softmax 计算器

修改下面的分数，实时查看 softmax 的每一步计算和输出分布

s₁

s₂

s₃

s₄

手算过程（以默认值 $[2.0, 1.0, 0.5, -1.0]$ 为例）：

步骤	Token 1	Token 2	Token 3	Token 4
原始分数 $s$	2.0	1.0	0.5	-1.0
$e^s$	7.389	2.718	1.649	0.368
$\alpha = e^s / \sum e^s$	0.609	0.224	0.136	0.030

验证： $0.609 + 0.224 + 0.136 + 0.030 = 1.000$ ✓

softmax 的三个关键性质#

性质 1：非负#

e^{s_j} > 0 \quad \Longrightarrow \quad \alpha_j > 0

不管原始分数多小、多负，指数函数的输出永远是正数。这保证了每个位置的权重 ≥ 0——“关注程度”不会是负数。

性质 2：归一化#

\sum_{j=1}^{n} \alpha_j = \sum_{j=1}^{n} \frac{e^{s_j}}{\sum_k e^{s_k}} = 1

所有权重加起来恰好等于 1，可以直接解读为概率分布或信息分配比例。“位置 $j$ 获得 30% 的注意力”比”位置 $j$ 的分数是 1.8”有意义得多。

性质 3：保序且放大差异#

s_i > s_j \quad \Longrightarrow \quad \alpha_i > \alpha_j

分数大的位置权重一定更大（保序），但差距会被指数函数放大。原始分数差 1.0，经过 $e^x$ 后差距变成 $e \approx 2.718$ 倍。

交互演示：指数放大效应

拖动滑块改变分数差距，观察 softmax 如何放大微小的分数差异

最高分与次高分的差距 Δ = 1.0

关键观察：

Δ = 0 时，所有位置权重相等（均匀分布）
Δ 逐渐增大时，最高分位置越来越”独占”
但即使 Δ 很大，其他位置的权重永远不会变成 0——softmax 是”软”的

这就是 softmax 名字的由来：它放大差异，但不完全消灭低分位置，总是保留一点”后路”。

温度参数：控制 softmax 的”锐利度”#

在实际应用中，有时需要调节 softmax 的尖锐程度。方法很简单——在分数上除以一个温度参数 $T$ ：

\alpha_j = \frac{e^{s_j / T}}{\sum_k e^{s_k / T}}

交互演示：温度如何改变注意力分布

拖动滑块调节温度 T，观察 softmax 输出从均匀到尖锐的变化

温度 T = 1.0

温度	行为	类比
$T \to 0^+$	趋近 one-hot（hardmax）	“只看最相关的一个位置”
$T = 1$	标准 softmax	Attention 的默认行为
$T \to \infty$	趋近均匀分布	”所有位置一视同仁”

在 Attention 中， $\sqrt{d_k}$ 就起到了类似温度的作用——第 3 讲说过，除以 $\sqrt{d_k}$ 防止分数过大导致 softmax 退化成 one-hot。

对应关系： $\text{softmax}(QK^\top / \sqrt{d_k})$ 中的 $\sqrt{d_k}$ 就是温度 $T$ 。 $d_k$ 越大，“温度”越高，分布越平滑。

为什么不用别的归一化方式#

候选 1：直接除以总和#

\alpha_j = \frac{s_j}{\sum_k s_k}

交互对比：softmax vs 直接归一化 vs hardmax

点击切换，观察同一组分数在不同归一化方式下的输出差异

直接除以总和有两个致命问题：

负分数怎么办？ $s = [3, -1, 2]$ ，总和 = 4，但 $\alpha_2 = -1/4 = -0.25$ 是负权重——含义不明
总和为 0 怎么办？ $s = [1, -1]$ ，除以 0 直接崩溃

softmax 用 $e^x$ 把所有分数先映射到正数，从根源上避免了这两个问题。

候选 2：hardmax（只取最大值）#

\alpha_j = \begin{cases} 1 & \text{if } j = \arg\max_k s_k \\ 0 & \text{otherwise} \end{cases}

hardmax 的问题第 2 讲已经讲过：不可微、丢信息。这里补充一个更根本的问题——梯度为零。hardmax 输出是阶跃函数，几乎处处梯度为 0，反向传播时什么也学不到。

候选 3：sparsemax#

sparsemax 是 softmax 的”稀疏”变体：部分位置的权重会恰好等于 0（不只是接近 0）。

\text{sparsemax}(\mathbf{s}) = \arg\min_{\mathbf{p} \in \Delta^n} \|\mathbf{p} - \mathbf{s}\|^2

sparsemax 在一些场景下有优势（注意力更可解释），但标准 Transformer 不用它，因为：

计算比 softmax 复杂（需要排序）
GPU 上缺乏高效实现
实验证明对 Transformer 性能提升不大

总结#

方式	非负	归一化	可微	保留低分信息	计算效率
直接归一化	✗	✓	✓	✓	快
hardmax	✓	✓	✗	✗	快
sparsemax	✓	✓	✓	部分	慢
softmax	✓	✓	✓	✓	快

softmax 是唯一同时满足全部四个要求的函数。

softmax 的梯度：为什么它能训练#

softmax 可微这件事非常重要。它的 Jacobian 矩阵是：

\frac{\partial \alpha_i}{\partial s_j} = \alpha_i (\delta_{ij} - \alpha_j)

其中 $\delta_{ij}$ 是 Kronecker delta。这意味着：

$i = j$ （自身）： $\frac{\partial \alpha_i}{\partial s_i} = \alpha_i(1 - \alpha_i)$ — 权重越接近 0 或 1，梯度越小
$i \neq j$ （交叉）： $\frac{\partial \alpha_i}{\partial s_j} = -\alpha_i \alpha_j$ — 提高一个位置的分数会降低其他位置的权重

交互演示：softmax 的梯度

点击某个位置，查看提高它的分数如何影响所有位置的权重

梯度消失的风险：当某个 $\alpha_i \to 1$ （softmax 退化成 one-hot），自身梯度 $\alpha_i(1-\alpha_i) \to 0$ 。这就是为什么要除以 $\sqrt{d_k}$ ——防止分数过大导致 softmax 过于尖锐，从而导致梯度消失。

数值稳定性：一个工程细节#

直接计算 $e^{s_j}$ 有溢出风险。如果 $s_j = 1000$ ， $e^{1000}$ 超出 float64 范围。

解决方法很简单——先减去最大值：

\text{softmax}(s_j) = \frac{e^{s_j}}{\sum_k e^{s_k}} = \frac{e^{s_j - \max(s)}}{\sum_k e^{s_k - \max(s)}}

交互演示：数值稳定性

拖动滑块增大分数量级，观察减去最大值如何避免溢出

分数缩放因子 = 1×

数学上完全等价，但数值上安全得多——减去最大值后所有指数的参数 ≤ 0， $e^x \leq 1$ ，不可能溢出。

PyTorch 的 F.softmax 和 torch.softmax 内部已经做了这个优化，不需要手写。

完整的 PyTorch 实验#

import torch
import torch.nn.functional as F
import math

scores = torch.tensor([2.0, 1.0, 0.5, -1.0])

# ---- 手动实现 softmax ----
def manual_softmax(s):
    s_shifted = s - s.max()        # 数值稳定：减去最大值
    exps = torch.exp(s_shifted)    # e^(s - max)
    return exps / exps.sum()       # 归一化

print("手动 softmax:", manual_softmax(scores))
print("PyTorch softmax:", F.softmax(scores, dim=0))
# 结果完全一致

# ---- 性质验证 ----
alpha = F.softmax(scores, dim=0)
print(f"\n非负: 所有 α > 0? {(alpha > 0).all()}")         # True
print(f"归一化: Σα = {alpha.sum().item():.6f}")             # 1.000000
print(f"保序: scores 排序 = α 排序? "
      f"{(scores.argsort() == alpha.argsort()).all()}")       # True

# ---- 温度实验 ----
print("\n温度实验:")
for T in [0.1, 0.5, 1.0, 2.0, 5.0]:
    a = F.softmax(scores / T, dim=0)
    entropy = -(a * a.log()).sum().item()
    print(f"  T={T:.1f}: α={a.numpy().round(3)}, 熵={entropy:.3f}")

# ---- 梯度 ----
scores_g = scores.clone().requires_grad_(True)
alpha_g = F.softmax(scores_g, dim=0)
alpha_g[0].backward()  # 对第一个位置的权重求梯度
print(f"\n提高 s₁ 对各位置权重的影响:")
print(f"  ∂α/∂s = {scores_g.grad.numpy().round(4)}")
# 第一个位置正梯度（提高自己），其余负梯度（降低别人）

# ---- 数值稳定性 ----
big_scores = scores * 500  # 放大 500 倍
try:
    naive = torch.exp(big_scores) / torch.exp(big_scores).sum()
    print(f"\n朴素计算: {naive}")  # 可能出现 nan/inf
except:
    print("\n朴素计算: 溢出!")
print(f"稳定计算: {F.softmax(big_scores, dim=0)}")  # 正常

python

把 softmax 放回 Attention 里#

现在我们可以完整理解注意力权重矩阵 $\alpha$ 的含义：

\alpha_{ij} = \frac{e^{S_{ij}/\sqrt{d_k}}}{\sum_{k=1}^{n} e^{S_{ik}/\sqrt{d_k}}}

交互演示：注意力权重热力图

点击任意行（Query 位置），查看该位置的完整注意力分布和 softmax 计算

逐行阅读这个矩阵：

第 $i$ 行：位置 $i$ 把注意力分配给所有位置的比例
每行之和 = 1（softmax 保证）
颜色越深 = 权重越大 = 从该位置取越多信息

softmax 把”分数”变成了”比例”——这就是为什么 Attention 的输出 $\alpha V$ 是一个加权平均，而不是某种奇怪的线性组合。

这一讲的三个核心结论#

softmax 满足注意力权重的所有要求：非负、归一化、可微、保序。是唯一同时满足这四个条件且计算高效的函数。
指数函数放大差异但不消灭低分。分数高的位置获得更多权重，但其他位置永远不会完全归零——这是”soft”的核心含义。
$\sqrt{d_k}$ 就是温度。它控制 softmax 的锐利度：防止分数过大导致 one-hot 退化（梯度消失），也防止分数过小导致均匀分布（不区分重要性）。

试着想一想#

如果所有分数都相等（ $s_1 = s_2 = \cdots = s_n$ ），softmax 输出什么？这对应 Attention 的什么行为？
温度 $T$ 和 $\sqrt{d_k}$ 的区别是什么？一个是超参数，另一个是……？
softmax 的梯度公式 $\alpha_i(1 - \alpha_i)$ 和 sigmoid 的梯度 $\sigma(1 - \sigma)$ 长得一样。这是巧合吗？（提示：softmax 是 sigmoid 的多类推广。）
如果用 $\text{ReLU}(s_j) / \sum_k \text{ReLU}(s_k)$ 替代 softmax，会出什么问题？
为什么 Flash Attention 的论文要花大量篇幅讨论 softmax 的计算？它和矩阵分块有什么冲突？（提示：softmax 的分母需要看到一整行。）

下一讲#

softmax 产出了权重，下一步是 $\alpha V$ ——加权求和。但为什么是求和？为什么不直接取权重最大的那个 Value？

第 6 讲：为什么是加权求和而不是取最大值