self-attention 为什么要除以 \(\sqrt{d_k}\)？
为什么不除以 \(d_k\)？
先弄清楚：每个概念是什么
- Attention score \(q\cdot k\)
- query 是"我想找什么"，key 是"我这里有什么"，点积就是匹配分，过 softmax 变成权重。例：\(q = (1, 2)\)，\(k = (3, 1)\)，\(q\cdot k = 5\)。
- 点积的 variance 随 \(d_k\) 变大
- 点积是 \(d_k\) 个随机小项相加，像连掷 \(d_k\) 次 ±1 的硬币，总和的典型幅度是 \(\sqrt{d_k}\)。例：\(d_k = 64\) 时 std = 8。
- Softmax saturation
- 输入差距一大，softmax 几乎变成 one-hot，gradient 接近 0，\(W_Q\)、\(W_K\) 收不到信号。
30 秒回答
Because the variance of the dot product grows with d_k. If q and k have independent components with mean 0 and variance 1, q·k has variance d_k and standard deviation √d_k. With d_k = 64 the scores are about eight times larger, softmax saturates to near one-hot, and its gradient p_i(δ_ij − p_j) is almost zero, so the query and key projections barely learn. Dividing by √d_k brings the variance back to 1. Dividing by d_k would overshoot: the variance becomes 1/d_k and attention becomes nearly uniform.
推导 / 关键公式
variance：为什么是 \(\sqrt{d_k}\)
假设 \(q\)、\(k\) 各分量独立，mean 0、variance 1（原论文 footnote 4）：
\[ \mathrm{Var}(q_i k_i) = \mathbb E[q_i^2]\,\mathbb E[k_i^2] = 1 \] \[ \mathrm{Var}(q\cdot k) = \sum_{i=1}^{d_k}\mathrm{Var}(q_i k_i) = d_k \] \[ \mathrm{Var}\big(q\cdot k/\sqrt{d_k}\big) = 1 \]乘以 \(c\) 后 variance 是 \(c^2 d_k\)，要它等于 1 只能取 \(c = 1/\sqrt{d_k}\)。
为什么 score 大了学不动
\[ \frac{\partial p_i}{\partial s_j} = p_i(\delta_{ij} - p_j) \]\(p\) 接近 one-hot 时每个 \(p_i\) 不是 ≈ 0 就是 ≈ 1，这一项都 ≈ 0。
手算例子：\(d_k = 1\) vs \(d_k = 64\)
两个 key 的分数分别比平均高、低一个 std。\(d_k = 1\) 时是 (1, −1)，\(d_k = 64\) 时 std = 8，变成 (8, −8)。
| \(d_k = 64\) 时 | score | softmax | \(p_1(1-p_1)\) |
|---|---|---|---|
| 不缩放 | (8, −8) | (≈1, \(1.1\times10^{-7}\)) | \(1.1\times10^{-7}\)：saturate |
| 除以 \(\sqrt{d_k} = 8\) | (1, −1) | (0.88, 0.12) | 0.105：和 \(d_k = 1\) 一样 |
| 除以 \(d_k = 64\) | (0.125, −0.125) | (0.56, 0.44) | 0.246：接近 uniform |
必答点
- 假设 \(q\)、\(k\) 分量独立、mean 0、variance 1 → \(\mathrm{Var}(q\cdot k) = d_k\)
- \(d_k\) 大 → score 大 → softmax saturate → gradient \(p_i(\delta_{ij}-p_j) \approx 0\)
- 除以 \(\sqrt{d_k}\)：variance 回到 1，与 \(d_k\) 无关
- 除以 \(d_k\) 过头：variance 变 \(1/d_k\)，attention 接近 uniform
追问
multi-head 里的 \(d_k\) 指哪个维度？
每个 head 的维度 \(d_{\text{model}}/h\)。原版 512/8 = 64，所以除以 8，不是 \(\sqrt{512}\)。
训练中 \(q\)、\(k\) 不再是 unit variance 怎么办？
假设只在初始化附近成立。ViT-22B 在 8B 规模遇到 attention logits 变大、loss 发散，用 QK-norm 解决，即先对 \(q\)、\(k\) 做 LayerNorm 再点积。
和 softmax temperature 什么关系？
除以 \(\sqrt{d_k}\) 就是 temperature \(T = \sqrt{d_k}\) 的 softmax，只是这里 \(T\) 固定不调。
为什么不用 additive attention？
原论文说 \(d_k\) 大时 additive 比不缩放的 dot-product 好。缩放后的 dot-product 能用矩阵乘法实现，更快、更省内存。
来源Vaswani et al. 2017 §3.2.1 及 footnote 4（p.4）· Dive into Deep Learning §11.3.1 eq. 11.3.2 · Dehghani et al. 2023（ViT-22B） §2（p.2，QK normalization）