在最初的 transformer 架构中，完整的 attention 公式为：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$

这篇文章推导公式中的 $\sqrt{d_k}$ 是如何得到的。

值得说明一点，这里的 $\sqrt{d_k}$ 在代码中常被称为 norm_factor，它的取值可能根据模型架构、attention 变体以及长度外推方法不同而有所改变。

推导过程

$d_k$ 表示每个 attention head 中 query/key 向量的维度，也就是通常所说的 head_dim。

假设某个 query 向量和 key 向量分别为：

$$q = [q_1, q_2, \ldots, q_{d_k}], \quad k = [k_1, k_2, \ldots, k_{d_k}]$$

它们的点积为：

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

也就是：

$$q \cdot k = q_1 k_1 + q_2 k_2 + \cdots + q_{d_k} k_{d_k}$$

为了分析这个点积的数值尺度，假设 $q_i$ 和 $k_i$ 是相互独立的随机变量，并且满足：

$$E[q_i] = 0, \quad E[k_i] = 0, \quad \text{Var}(q_i) = 1, \quad \text{Var}(k_i) = 1$$

令：

$$X_i = q_i k_i$$

则：

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

首先计算单项 $X_i$ 的期望：

$$E[X_i] = E[q_i k_i]$$

由于 $q_i$ 和 $k_i$ 独立：

$$E[q_i k_i] = E[q_i] E[k_i]$$

又因为 $E[q_i] = 0$，$E[k_i] = 0$，所以：

$$E[X_i] = 0$$

接着计算 $X_i$ 的方差：

$$\text{Var}(X_i) = E[X_i^2] - E[X_i]^2$$

因为 $E[X_i] = 0$，所以：

$$\text{Var}(X_i) = E[X_i^2]$$

代入 $X_i = q_i k_i$：

$$\text{Var}(X_i) = E[(q_i k_i)^2] = E[q_i^2 k_i^2]$$

由于 $q_i$ 和 $k_i$ 独立：

$$E[q_i^2 k_i^2] = E[q_i^2] E[k_i^2]$$

又因为 $\text{Var}(q_i) = E[q_i^2] - E[q_i]^2 = 1$，且 $E[q_i] = 0$，所以 $E[q_i^2] = 1$。同理 $E[k_i^2] = 1$。

因此：

$$\text{Var}(X_i) = E[q_i^2] E[k_i^2] = 1$$

也就是说，每一项乘积 $q_i k_i$ 的方差约为 1。

由于点积 $q \cdot k$ 是 $d_k$ 个这样的项相加：

$$q \cdot k = X_1 + X_2 + \cdots + X_{d_k}$$

在各项近似独立的情况下，协方差为0，故和的方差等于方差之和：

$$\text{Var}(q \cdot k) = \text{Var}\left(\sum_{i=1}^{d_k} X_i\right) = \sum_{i=1}^{d_k} \text{Var}(X_i) = d_k$$

因此：

$$\text{Var}(q \cdot k) = d_k$$

对应的标准差为：

$$\text{Std}(q \cdot k) = \sqrt{\text{Var}(q \cdot k)} = \sqrt{d_k}$$

这说明，当 $d_k$ 变大时，未缩放的点积 $q \cdot k$ 的典型数值幅度会随着 $\sqrt{d_k}$ 增大。

而 attention score 后面会进入 softmax：

$$\text{softmax}(QK^T)$$

如果 $QK^T$ 的数值过大，softmax 会变得非常尖锐。例如：

$$\text{softmax}([20, 1, -3, 0]) \approx [1, 0, 0, 0]$$

这会导致 attention 分布过早饱和，梯度变小，训练变得不稳定。

因此，为了让进入 softmax 的 logits 保持在相对稳定的尺度，需要对点积结果做归一化缩放：

$$\frac{q \cdot k}{\sqrt{d_k}}$$

缩放之后：

$$\text{Var}\left(\frac{q \cdot k}{\sqrt{d_k}}\right) = \frac{\text{Var}(q \cdot k)}{d_k} = \frac{d_k}{d_k} = 1$$

对应标准差为：

$$\text{Std}\left(\frac{q \cdot k}{\sqrt{d_k}}\right) = 1$$

也就是说，除以 $\sqrt{d_k}$ 后，attention logits 的尺度大致回到 $O(1)$ 的量级，不会随着 $\text{head\_dim}$ 增大而持续变大。

因此，在原始 Transformer 中：

$$\text{norm\_factor} = \sqrt{d_k}$$

其作用是控制 $QK^T$ 的数值尺度，使 softmax 不至于过度饱和，从而提升训练稳定性。