前言

这篇文章是本站《Transformer 详解：从注意力机制到代码实战》的数学补充。它的定位非常具体：只讲"服务 Attention 的数学"，不重复基础线性代数。

我们假设你已经读过本站的《矩阵微积分如何服务反向传播》，理解 Jacobian 与向量-Jacobian 乘积（VJP）这两个概念。在此基础上，本文回答三个问题：

  1. softmax 的梯度长什么样，为什么是这个形式？
  2. 为什么 softmax 与交叉熵组合在一起，梯度会变得极简且数值稳定？
  3. 缩放点积注意力中的 $\sqrt{d_k}$ 究竟从哪里来，为什么非除不可？

Softmax 定义与直觉

给定一个实数向量 $z \in \mathbb{R}^n$（常称为 logits，即未归一化的分数），softmax 把它映射成一个概率分布：

$$ \text{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}} $$

三个值得记住的性质：

  • 输出非负且和为 1：天然是概率分布。
  • 放大差异：指数运算放大了 logits 之间的差距。$z = [1, 2]$ 时，softmax 输出约 $[0.27, 0.73]$，而不是线性的 $[0.33, 0.67]$——赢家更突出。
  • 平移不变：$\text{softmax}(z + c) = \text{softmax}(z)$（对任意常数 $c$）。分子分母同乘 $e^c$，抵消了。这个性质是后面 log-sum-exp 技巧的基础。

还有一个常被提及的变体：温度参数 $T$，写作 $\text{softmax}(z / T)$。$T < 1$ 时分布更"尖锐"（注意力更集中），$T > 1$ 时更"平滑"。在 Attention 里，softmax 的温度直接控制"注意力有多集中"。

Softmax 的梯度

softmax 的输入输出都是向量，所以它的"导数"是一个 Jacobian 矩阵。记 $s_i = \text{softmax}(z)_i$，并令分母 $Z = \sum_k e^{z_k}$。

当 $i = j$ 时，用商法则：

$$ \frac{\partial s_i}{\partial z_i} = \frac{e^{z_i} Z - e^{z_i} \cdot e^{z_i}}{Z^2} = \frac{e^{z_i}}{Z} - \left(\frac{e^{z_i}}{Z}\right)^2 = s_i(1 - s_i) $$

当 $i \neq j$ 时，分母 $Z$ 是常数，只需对分子求导：

$$ \frac{\partial s_i}{\partial z_j} = \frac{-e^{z_i} \cdot e^{z_j}}{Z^2} = -s_i s_j $$

把两种情况合并，用克罗内克记号 $\delta_{ij}$（$i=j$ 时为 1，否则为 0）：

$$ \frac{\partial s_i}{\partial z_j} = s_i(\delta_{ij} - s_j) $$

写成完整的 Jacobian 矩阵：

$$ \frac{\partial s}{\partial z} = \begin{bmatrix} s_1 & & \\ & \ddots & \\ & & s_n \end{bmatrix} - s s^\top = \text{diag}(s) - ss^\top $$

直觉：softmax 对第 $i$ 个输入的偏导是"自己乘以（单位阵 − 对方）"。这带来一个微妙的效果——提高 $z_i$ 不仅会增大 $s_i$，还会抑制所有其他 $s_j$（因为 $i \neq j$ 时偏导为负）。这正是概率分布的归一化特性：总和固定为 1，此消彼长。

Softmax 与交叉熵的组合梯度

在分类任务里，softmax 后面通常接交叉熵损失。设 $y$ 是 one-hot 标签（$y_i \in \{0, 1\}$，$\sum_i y_i = 1$）：

$$ L = -\sum_i y_i \log s_i $$

我们要的是 $\partial L / \partial z_j$，可以直接用链式法则穿过 softmax：

$$ \frac{\partial L}{\partial z_j} = \sum_i \frac{\partial L}{\partial s_i} \frac{\partial s_i}{\partial z_j} = \sum_i \left(-\frac{y_i}{s_i}\right) s_i(\delta_{ij} - s_j) = \sum_i -y_i(\delta_{ij} - s_j) $$

展开：$\sum_i -y_i \delta_{ij} = -y_j$，而 $\sum_i y_i s_j = s_j \sum_i y_i = s_j$（因为 $\sum_i y_i = 1$）。于是：

$$ \frac{\partial L}{\partial z} = s - y $$

一个极其优雅的结论：softmax + 交叉熵的组合梯度就是"预测概率减标签"——$s_j > y_j$ 时梯度为正，说明模型对第 $j$ 类的置信度偏高，需要下调。这一结论是工程上把 softmax 与交叉熵合并实现（PyTorch 的 F.cross_entropy 直接接收 logits，而非先 softmax 再算 NLL）的根本原因：合并后梯度计算完全绕开了中间的指数项，数值上干净利落。

数值稳定性：log-sum-exp 技巧

现在考虑数值实现。softmax 要算 $e^{z_i}$，如果 $z_i$ 很大（比如 1000），$e^{1000}$ 会直接溢出为 inf。利用 softmax 的平移不变性，先把所有 logits 减去最大值 $m = \max_j z_j$：

$$ \text{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}} = \frac{e^{z_i - m}}{\sum_j e^{z_j - m}} $$

减去 $m$ 后，最大的指数项是 $e^{0} = 1$，其余都小于 1，绝不会溢出。数学上完全等价，因为分子分母同时乘了 $e^{-m}$。这就是 log-sum-exp 技巧，几乎所有深度学习框架的 softmax 都是这么实现的。

顺带一提，减最大值也解释了为什么需要先做缩放再进 softmax：如果 $QK^\top$ 里出现很大的值，即便不溢出，$e^{z_i}$ 也会把其它项压成 0，导致下一节讲的问题。

核心：QKᵀ 为何除以 √dₖ

这是 Attention 里最容易"知其然而不知其所以然"的一个细节。原论文的说明很短——“点积的幅度随 $d_k$ 增大而变大，会把 softmax 推向梯度极小的区域，因此除以 $\sqrt{d_k}$"。我们把这句话的数学补全。

假设 $q, k \in \mathbb{R}^{d_k}$ 的各分量独立同分布，且归一化后均值为 0、方差为 1（这是经过合理初始化/归一化后的常见状态）。点积展开：

$$ q \cdot k = \sum_{i=1}^{d_k} q_i k_i $$

先看期望。由于 $q_i$ 与 $k_i$ 独立且均值都为 0：

$$ \mathbb{E}[q_i k_i] = \mathbb{E}[q_i] \cdot \mathbb{E}[k_i] = 0 $$

再看每一项的方差。利用"独立且均值 0 时，乘积方差 = 方差之积”（$\mathbb{E}[q_i^2] = \mathrm{Var}(q_i) = 1$）：

$$ \mathrm{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2] - \underbrace{\mathbb{E}[q_i k_i]^2}_{=0} = \mathbb{E}[q_i^2] \mathbb{E}[k_i^2] = 1 \cdot 1 = 1 $$

点积是 $d_k$ 个独立项之和，方差直接相加：

$$ \mathrm{Var}(q \cdot k) = d_k, \qquad \text{标准差} = \sqrt{d_k} $$

结论：点积的标准差随 $\sqrt{d_k}$ 线性增长。$d_k = 512$ 时，点积的标准差是 22.6 左右，数值普遍偏大。

为什么不缩放会出问题？回看 softmax 的梯度 $\frac{\partial s_i}{\partial z_j} = s_i(\delta_{ij} - s_j)$：当 logits 数值很大时，softmax 输出的分布趋于极端——最大的 $s_i$ 接近 1，其余接近 0。此时：

  • 对 $i \neq j$，梯度 $-s_i s_j$ 中至少一个因子接近 0，梯度趋近 0；
  • 对 $i = j$，梯度 $s_i(1 - s_i)$ 也趋近 0。

这就是饱和区：softmax 越"自信"，梯度越接近零，Attention 学不动。而除以 $\sqrt{d_k}$ 后，点积的方差被压回 1，softmax 的输入落在温和区间，梯度健康。

把这个结论和本站的 VJP 文章连起来看：Attention 层反向传播时，上游梯度要乘上 softmax 的 Jacobian。缩放与否直接决定了这个 Jacobian 的尺度，进而影响整个网络梯度流的稳定性——这又是一个"数值设计决定训练行为"的例子。

总结

Transformer 的 Attention 公式 $\text{softmax}\left(QK^\top / \sqrt{d_k}\right)V$ 看似简单，背后其实是三个数学工具的协作：

  1. softmax：把相似度得分变成概率分布，并放大"谁更相关"的差异；
  2. log-sum-exp 技巧：减去最大值，保证指数运算永不溢出；
  3. 除以 $\sqrt{d_k}$：把点积方差压回 1，避免 softmax 进入梯度消失的饱和区。

这三者共同保证了 Attention 在数值上可控、在梯度上健康。回到主文，就能更自然地理解为什么《Transformer 详解》里，缩放因子 $\sqrt{d_k}$ 会出现在公式最显眼的位置。

参考