Scaled Dot-Product：那个根号 d_k 是怎么来的'本文从零推导注意力机制点积方差的来源，解释缩放因子如

序言：一个除号背后藏着的整门数学课

很多人第一次读 Vaswani 2017 的公式时，都会卡在那一个 $\sqrt{d_k}$ 上。

公式本身写得简洁：

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

但那个分母上的 $\sqrt{d_k}$ 看起来像是凭空冒出来的常数。

「为什么是 $\sqrt{d_k}$ 不是 $d_k$ ？」

「为什么不是 $\sqrt{d_{\mathrm{model}}}$ ？」

「为什么不是其他什么数？」

如果你只读论文的那一句话——"to counteract the effect of large dot products"——你会觉得这是一个 经验技巧。

事实远不止此。

那个 $\sqrt{d_k}$ 是一个 从概率论第一性原理推出来的、几乎别无选择的数字。

它涉及：随机变量的方差加法、softmax 的饱和性、链式法则下的梯度衰减、训练动力学的稳定性——你能想到的关于「为什么神经网络能优化」的核心问题，全都串在这一个除号上。

本文要做的事情，就是把这个除号拆开——一步一步、不跳逻辑——告诉你：为什么是 $\sqrt{d_k}$ ？它到底拯救了什么？以及，到 2026 年，对这个 $\sqrt{d_k}$ 的现代理解（包括 NTK 视角、FlashAttention 数值稳定性、Muon 优化器对 attention 的影响）有哪些新维度。

读完之后，你应该能在被人问到「为什么除以 $\sqrt{d_k}$ 」的时候，给出一个 5 分钟版本、一个 30 分钟版本、和一个「我可以为你推一遍」版本。

原文链接

一、问题缘起：先看不除会发生什么

1.1 复盘公式

第 13 篇我们看到：

\operatorname{attention}(q, k, v) = \operatorname{softmax}(qK^{\mathsf{T}}) V

第 14 篇我们看到：self-attention 让每个 token 同时扮演 q、k、v。

但实际工程里写的、Vaswani 论文里写的、所有 PyTorch 实现里写的，都是：

\operatorname{attention}(Q, K, V) = \operatorname{softmax}\left(\frac{QK^{\mathsf{T}}}{\sqrt{d_k}}\right) V

那个 $\sqrt{d_k}$ 是什么？为什么必须有？

我们做一个思想实验：把 $\sqrt{d_k}$ 拿掉，看会发生什么。

1.2 一个具体的数值实验

设 $d_k = 64$ 。

$q$ 和 $k$ 是两个 $d_k$ 维向量，每一维独立、均值 $0$ 、方差 $1$ （就当它们是从标准正态采样）。

我们想知道： $q \cdot k$ 的分布是什么？

q \cdot k = \sum_i q_i k_i

每一项 $q_i k_i$ 是两个独立标准正态的乘积——均值是 $0$ ，方差是 $1 \times 1 = 1$ 。因为

\operatorname{Var}(XY) = \mathbb{E}[X^2]\mathbb{E}[Y^2] - \mathbb{E}[X]^2 \mathbb{E}[Y]^2 = 1 \cdot 1 - 0 = 1,

对独立零均值变量成立。

这 $64$ 个独立项相加：

\mathbb{E}[q \cdot k] = 0, \qquad \operatorname{Var}(q \cdot k) = 64, \qquad \sigma = 8.

所以 $q \cdot k$ 的取值范围大概在 $\pm 24$ （ $3\sigma$ ）以内浮动。

1.3 把 q·k = 24 喂进 softmax

假设我们有 $8$ 个 key，对应 $8$ 个点积，碰巧最大那个是 $24$ ，其它都接近 $0$ 。

\operatorname{softmax}([24, 0, 0, 0, 0, 0, 0, 0]) = ?

e^{24} \approx 2.6 \times 10^{10}, \qquad e^0 = 1.

归一化后： $[\approx 1, \approx 0, \approx 0, \ldots, \approx 0]$ ——几乎纯 one-hot。

这看起来不是好事吗？模型「下定决心」选了一个 token？

恰恰相反——这是一场灾难。

1.4 灾难的来源：梯度消失

来看 softmax 的 Jacobian。

设 $p = \operatorname{softmax}(s)$ ，那么：

\frac{\partial p_i}{\partial s_j} = p_i (\delta_{ij} - p_j)

其中 $\delta_{ij}$ 是 Kronecker delta（ $i = j$ 时为 $1$ ，否则为 $0$ ）。

如果 $p$ 接近 one-hot，比如 $p_1 \approx 1$ ，其它 $\approx 0$ ，那么：

\frac{\partial p_1}{\partial s_1} = 1 \cdot (1 - 1) = 0, \qquad \frac{\partial p_1}{\partial s_j} = 1 \cdot (0 - 0) = 0 \; (j \neq 1), \qquad \frac{\partial p_i}{\partial s_j} \approx 0 \; (i \neq 1).

整个 Jacobian 几乎为零矩阵。

这意味着：通过 softmax 反向传播的梯度被掐死了。

下游的 loss 想告诉 attention 层「你应该多看看 token 5」，但这个信号被 softmax 饱和性吃掉了——logits $s$ 几乎不会被更新。

1.5 这就是 $\sqrt{d_k}$ 要解决的问题

如果 logits 的方差是 $d_k$ ，那把它们除以 $\sqrt{d_k}$ ，方差就变成 $1$ 。

logits 不再因为维度放大而漂到饱和区，softmax 输出保持在「有梯度的工作点」。

训练就能进行。

这是 $\sqrt{d_k}$ 的全部直觉——剩下的都是细节。

二、点积方差的严格推导

2.1 假设

我们在以下假设下推 $\operatorname{Var}(q \cdot k) = d_k$ ：

$q$ 和 $k$ 是 $d_k$ 维向量
$q$ 的每一维 $q_i$ 之间独立
$k$ 的每一维 $k_j$ 之间独立
$q$ 和 $k$ 之间也独立
所有 $q_i$ 和 $k_j$ 都是均值 $0$ 、方差 $1$

后面我们会讨论这些假设在真实模型中成立到什么程度。

2.2 推导

设

X = q \cdot k = \sum_i q_i k_i.

第一步：均值。

\mathbb{E}[X] = \sum_i \mathbb{E}[q_i k_i] = \sum_i \mathbb{E}[q_i] \mathbb{E}[k_i] = 0.

第二步：方差。

\operatorname{Var}(X) = \mathbb{E}[X^2] - \mathbb{E}[X]^2 = \mathbb{E}[X^2].

\mathbb{E}[X^2] = \mathbb{E}\left[\left(\sum_i q_i k_i\right)^2\right] = \sum_i \sum_j \mathbb{E}[q_i k_i q_j k_j].

对于 $i \neq j$ ：

\mathbb{E}[q_i k_i q_j k_j] = \mathbb{E}[q_i] \mathbb{E}[k_i] \mathbb{E}[q_j] \mathbb{E}[k_j] = 0.

对于 $i = j$ ：

\mathbb{E}[q_i^2 k_i^2] = \mathbb{E}[q_i^2] \mathbb{E}[k_i^2] = 1 \times 1 = 1.

\mathbb{E}[X^2] = \sum_i 1 = d_k.

\operatorname{Var}(X) = d_k.

2.3 标准差

\sigma = \sqrt{d_k}.

这就是 $\sqrt{d_k}$ 的来源——不是任何「经验试出来的常数」，而是 $\operatorname{Var}$ 加法的直接结果。

2.4 缩放后的分布

定义

X' = \frac{X}{\sqrt{d_k}}.

\operatorname{Var}(X') = \frac{\operatorname{Var}(X)}{d_k} = 1.

不管 $d_k$ 是 $64$ 、 $512$ 、还是 $4096$ ， $X'$ 的方差永远是 $1$ 。

logits 的尺度被「归一化」到了一个不依赖于维度的水平。

2.5 为什么是 $\sqrt{d_k}$ 不是 $d_k$

有人问：「除以 $d_k$ 不是更彻底吗？」

不行。

如果除以 $d_k$ ，那么

\operatorname{Var}\left(\frac{X}{d_k}\right) = \frac{d_k}{d_k^2} = \frac{1}{d_k} \to 0

当 $d_k$ 很大时就会发生这个退化。

logits 全部接近 0，softmax 输出接近均匀分布——attention 失去了「选择性」。

我们要的是「方差 $= 1$ 」（既不太尖锐也不太平），所以分母必须是 $\sqrt{d_k}$ 。

这是一个临界点，不是一个「随便挑的数」。

2.6 有人问：为什么不除以 $\sqrt{2 d_k}$ 之类

也可以，只要常数因子合理（比如让方差 $= 0.5$ 而不是 $1$ ）。

但这只会让 softmax 略偏平缓——本质和 $\sqrt{d_k}$ 没区别。

Vaswani 选 $\sqrt{d_k}$ 是因为它最自然——把方差归一化到 $1$ ，保留了「方差为 $1$ 」这个统计学最常用的标准化。

后续工作（比如 RoFormer、LLaMA）也都沿用这个选择。

2.7 一个数值表

$d_k$	$\sigma = \sqrt{d_k}$	logits 范围（ $3\sigma$ ）
8	2.83	±8.5
32	5.66	±17
64	8	±24
128	11.3	±34
256	16	±48
512	22.6	±68

如果不缩放， $512$ 维的点积可能跑到 $\pm 68$ ——softmax 看到这种 logits，对应的 $e^{68} \approx 10^{29}$ ——任何对手 logit 都被压成 $0$ 。

缩放后 logits 永远在 $\pm 3$ 左右——softmax 仍能区分大小，但梯度不会断流。

三、softmax 饱和性的可视化

softmax 尖锐度转存失败，建议直接上传图片文件

3.1 直观图景

右侧子图（unscaled）：一个 logit 比其它大很多——softmax 集中在那一个点上。

左侧子图（scaled）：logits 接近，softmax 平缓，多个 token 都有可见权重。

不是说「平缓就一定好」——实际训练中，模型最终会学到「需要 sharp 时 sharp」的能力。

但训练初期必须从「平缓 + 有梯度」的状态出发，否则模型一开始就被卡在饱和区出不来。

3.2 一个比喻

把 softmax 想象成一根弹簧。

弹簧未饱和时，你拉它它会动，反馈给你力——你能学到「拉的方向」。

弹簧饱和时（拉到极限），你怎么拉它都不动——你什么也学不到。

logits 越大，softmax 越饱和；scaled dot-product 就是把弹簧从「饱和区」拉回到「线性区」，让训练能进行。

3.3 数学复盘

\operatorname{softmax}(s)_i = \frac{e^{s_i}}{\sum_j e^{s_j}}.

求导：

\frac{\partial \operatorname{softmax}(s)_i}{\partial s_j} = p_i (\delta_{ij} - p_j)

最大值的对角项： $p_i (1 - p_i)$ ，当 $p_i \to 1$ 时为 $0$ ；当 $p_i \to 0$ 时也为 $0$ ；最大在 $p_i = 0.5$ 。

非对角项： $-p_i p_j$ ，仅当两个都不是 $0$ 也不是 $1$ 时才有效。

所以「梯度最大」的工作点是 $p_i \in [0.1, 0.9]$ ——这正是 logits 适中（ $\sigma \approx 1$ ）时的状态。

3.4 与温度参数的关系

很多人熟悉「softmax 温度」 $T$ ：

\operatorname{softmax}_T(s)_i = \frac{e^{s_i / T}}{\sum_j e^{s_j / T}}

$T$ 大，输出平缓； $T$ 小，输出 sharp。

scaled dot-product 中的 $\sqrt{d_k}$ 就扮演温度的角色——具体来说 $T = \sqrt{d_k}$ 。

但与 $T$ 不同， $\sqrt{d_k}$ 不是调参，而是定参——它的值由维度决定，不是用户选择。

3.5 一个常见的混淆

「我可以学一个 temperature 替代 $\sqrt{d_k}$ 吗？」

理论上可以，但实践中很少这么做。

因为 $\sqrt{d_k}$ 已经把 logits 归一化到 $\sigma \approx 1$ ，再学一个 temperature 等于多此一举——除非你想做「shaped attention」之类的研究。

LLaMA、GPT、PaLM 等都没有学习 temperature，全用 $\sqrt{d_k}$ 。

但有一些工作（如 NormFormer、QK-norm）提出在 $Q$ 和 $K$ 上做 LayerNorm，再不除 $\sqrt{d_k}$ ——效果近似但实现略有不同。

到 2026 年，QK-norm 方案在大模型训练中越来越常见。

四、为什么这件事到 d_k = 64 才显著

4.1 一个有趣的现象

在最早的 attention 工作（Bahdanau 2014）中，用的是加性注意力—— $\mathrm{score} = v^{\mathsf{T}} \tanh(W_q q + W_k k)$ ——根本没有 $\sqrt{d_k}$ 的除法。

为什么 Bahdanau 不需要？

因为 Bahdanau 用的是 RNN 的 hidden state（典型 $d = 256$ 但走 $\tanh$ ，输出落在 $[-1, 1]$ ）+ 学习的 $v$ ——score 永远在一个有限的 bounded 区间，不会因为维度爆炸。

dot-product attention（Luong 2015）开始有这个问题——因为 $q \cdot k$ 没有 $\tanh$ 包住，方差直接随 $d_k$ 增长。

但 Luong 的实验里 d 不大，问题不严重。

到 Vaswani 2017 multi-head 时代， $d_k = 64$ （每 head 的维度）， $Q$ 、 $K$ 的来源是线性投影后的向量——方差大约是 $1$ （因为初始化 + LayerNorm）——这时候 $q \cdot k$ 的方差就接近 $64$ ，问题就显现出来了。

4.2 d_k 越大，问题越严重

到 GPT-3： $d_k = 128$ （每 head），问题更严重。

到 PaLM： $d_k = 256$ （每 head），不缩放训练直接发散。

Vaswani 的论文里有一段话：「我们怀疑对于大的 d_k 值，dot products 在量级上变大，从而把 softmax 推到具有极小梯度的区域。」

这是一句经验观察——他们看到了「不缩放训练崩」，做了缩放，发现「训练好了」。

后来的理论分析（Xiong 2020 "On Layer Normalization in the Transformer Architecture"）才把这件事讲透。

4.3 为什么加性注意力没这个问题

$v^{\mathsf{T}} \tanh(W_q q + W_k k)$ 中， $\tanh$ 的输出落在 $[-1, 1]$ 。

随后 $v^{\mathsf{T}}(\cdot)$ 是一个 $d$ 维点积，这一步也会有方差放大。

但因为 tanh 已经把每一维 bound 住了，方差不会无界放大。

所以加性注意力天然「自带稳定性」，但代价是计算更慢（多一次矩阵乘 + 非线性）。

dot-product attention 要更快——因为它就是一个 matmul——但代价是必须手工加 $\sqrt{d_k}$ 来保稳定。

4.4 一个关于 norm 的细节

Vaswani 假设 $q$ 和 $k$ 每一维方差是 $1$ 。

实际模型里，这通过 LayerNorm 大致成立——LayerNorm 把每一层输出的 mean 归零、std 归一。

但有些层（比如 attention 输出）是 LayerNorm 之前还是之后？这就涉及 Pre-LN vs Post-LN 的选择。

Pre-LN（LayerNorm 在 sublayer 之前）让 $q$ 、 $k$ 在进入 attention 时严格 normalized—— $\sqrt{d_k}$ 的假设最契合。

Post-LN（LayerNorm 在 sublayer 之后）让 $q$ 、 $k$ 在进入 attention 时未必 normalized——可能需要 warmup 来稳住训练。

到 2026 年，Pre-LN 是主流（GPT、LLaMA 都用 Pre-LN）。

4.5 那 Q、K 不是单位方差怎么办

如果 $W_q$ 、 $W_k$ 初始化合理（比如 Xavier 或 Kaiming），且输入 $X$ 经过 LayerNorm，那么 $q = W_q x$ 的方差大致就是 $1$ 。

但如果你不做 LayerNorm、用奇怪初始化、或者训练到某一步参数漂移——方差就不是 1 了。

QK-norm（在 q、k 上做 LayerNorm）就是把这个假设显式强制——不再依靠「希望 LayerNorm 保住」。

五、点积方差的可视化

点积分布转存失败，建议直接上传图片文件

5.1 三个直方图

$d_k = 8$ 时分布窄， $\sigma \approx 2.83$ 。

$d_k = 64$ 时分布宽， $\sigma = 8$ 。

$d_k = 512$ 时分布很宽， $\sigma \approx 22.6$ 。

直方图的横轴是点积值——纵轴是出现频率。

5.2 为什么这个图重要

看到这张图，你应该马上意识到：

不缩放时，点积尺度完全由维度决定——你换一个模型规模，点积尺度就变了——你的训练超参（学习率、初始化等）就要重调。

缩放后，点积尺度永远是 1——超参可以跨规模迁移。

这是 scaling laws 能成立的一个隐性前提：架构内的统计尺度必须不依赖于规模。

5.3 Chinchilla scaling 的隐含条件

Hoffmann 2022 给出 Chinchilla 定律：参数 $N_p$ 与 token 数 $D$ 的最优比例 $N_p \approx D/20$ 。

这条定律的成立依赖于「同样的架构、同样的训练超参在不同规模下都能稳定训练」。

如果你不缩放点积，训练在 $d_k = 64$ 时还稳定，到 $d_k = 512$ 时就发散——scaling laws 整个就不成立。

$\sqrt{d_k}$ 是 scaling laws 的「隐性基础设施」之一。

六、训练曲线对比

训练曲线转存失败，建议直接上传图片文件

6.1 定性差异

红线（unscaled）：早期 loss 下降慢，很快卡在某个高位——softmax 饱和导致的优化困难。

绿线（scaled）：稳定下降。

初值相同（损失大约是 $\log(N)$ 那个均匀分布的 cross-entropy）。

6.2 一个真实的定量例子

Vaswani 2017 §3.2.1 没给具体的对比训练曲线（论文较老），但后续工作（Xiong 2020）做过实验。

在 $d_k = 64$ 的 Transformer-base 上，去掉 $\sqrt{d_k}$ ：

不调任何其它超参：loss 卡在 6.x，几乎不动。
把学习率降低 $10\times$ ：训练能进行，但 BLEU 显著低于带 $\sqrt{d_k}$ 的版本。

也就是说，没有 $\sqrt{d_k}$ 不是「完全不能训练」，而是「需要付出极大的超参代价、且最终质量更差」。

加上 $\sqrt{d_k}$ 等价于一个免费的、零计算开销的稳定性优化——为什么不用呢？

6.3 一个反直觉发现

有人发现：如果模型足够小（ $d_k = 16$ 之类），不缩放也能训。

这与第二节的方差分析一致—— $\sigma = 4$ ，logits 不会饱和。

但你不可能因为「小模型不需要」就在大模型里也省掉它——大模型里这个除号是必需品。

七、缩放与梯度下降稳定性

7.1 学习率与梯度的关系

如果不缩放，attention 的 logits 在 $\pm \sqrt{d_k}$ 量级——softmax 输出近似 one-hot——梯度近似 $0$ ——参数几乎不更新。

但偶尔某个 batch 有「比较平的 logits」，梯度突然爆发——参数大跳——loss 飞涨。

这是「饱和 + 偶尔不饱和」的混合模式——非常不稳定。

7.2 缩放后的 Lipschitz 性质

缩放后，softmax 的输入永远在 $[-3\sigma, 3\sigma] = [-3, 3]$ 左右——softmax 在这个区间内是 Lipschitz 连续的，导数有 bounded 上限。

这意味着「同样大小的输入扰动 $\to$ 同样大小的输出扰动」——训练动力学是稳定的。

八、参考资料

Vaswani 2017: Ashish Vaswani et al., "Attention Is All You Need" (首次提出 Transformer 与 Scaled Dot-Product).
Bahdanau 2014: Dzmitry Bahdanau et al., "Neural Machine Translation by Jointly Learning to Align and Translate" (提出 Additive Attention, $\tanh$ 限制方差).
Luong 2015: Minh-Thang Luong et al., "Effective Approaches to Attention-based Neural Machine Translation" (提出 Dot-Product Attention).
Xiong 2020: Ruibin Xiong et al., "On Layer Normalization in the Transformer Architecture" (分析缩放机制与 LayerNorm 对训练稳定性的深度影响).
Hoffmann 2022: Jordan Hoffmann et al., "Training Compute-Optimal Large Language Models" (Chinchilla 定律，隐含架构缩放的统计不变性假设).

Lipschitz 常数大致是 $1$ （用 $\ell_{\infty}$ 范数估计）。

7.3 与梯度裁剪的关系

很多 Transformer 训练里都有「gradient clipping」（梯度裁剪），把过大的梯度截断到 $\lVert g \rVert \le c$ 。

为什么需要梯度裁剪？因为偶尔会有「outlier batch」让某些参数的梯度爆掉——比如某个 batch 里所有 token 都是同一个。

scaled dot-product 让这种 outlier 的破坏力降低——但不能完全消除——所以梯度裁剪仍是必需。

7.4 与 warmup 的关系

Transformer 训练几乎都用 learning rate warmup（前若干步学习率从 0 线性涨到峰值）。

为什么？因为训练初期参数随机，logits 分布可能极不平衡——warmup 给模型时间「找到稳定区域」再放学习率。

scaled dot-product 让初期 logits 不那么大——warmup 期可以更短——但不能省略。

八、与 NTK / 无限宽神经网络理论的联系

8.1 NTK 是什么

NTK（Neural Tangent Kernel，Jacot 2018）是一种刻画「无限宽网络在小学习率下的训练动力学」的理论。

核心结论是：在某些假设下，无限宽网络的训练等价于一个线性化模型 + 核回归，其中的「核」就叫 NTK。

NTK 给我们一个工具：预测网络在不同初始化、不同尺度下的行为。

8.2 $\sqrt{d_k}$ 与 NTK

NTK 理论强调一个 principle：网络中每一层的输入与输出的统计尺度必须一致——否则梯度传播会失衡。

scaled dot-product 正是这个 principle 在 attention 层的体现——把点积归一化到 $\sigma = 1$ ，让 attention 层的「输入尺度」与「输出尺度」一致。

如果不缩放，attention 层把方差从 $1$ 放大到 $\sqrt{d_k}$ ——下一层 LayerNorm 又把它拉回 $1$ ——但中间这一段不稳定。

8.3 muP（Maximal Update Parametrization）

Yang & Hu 2021 的 muP 是 NTK 思想在工程上的实现：通过精心设计每层的 init scale 和 LR scale，让模型在改变宽度时超参不变。

muP 框架下，attention 的 $\sqrt{d_k}$ 是一个特殊处理——它不是 muP 自动推出的，而是早就独立存在的设计——但它与 muP 的精神高度契合。

到 2026 年，muP（特别是 mup-transfer 思路）成为大模型训练前调超参的重要工具——基础假设之一就是 attention 已经被 $\sqrt{d_k}$ 正则化过了。

8.4 NTK 视角的 attention

在 NTK 视角下：

$\operatorname{attention}(Q, K, V) = \operatorname{softmax}\left(\frac{QK^{\mathsf{T}}}{\sqrt{d_k}}\right) V$ 是一个 bilinear 算子。
bilinear 部分 $QK^{\mathsf{T}}$ 把维度从 $d_k$ 映射到 $N \times N$ 的相似度矩阵。
softmax 是一个非线性归一化。
与 $V$ 的乘法把 $N \times N$ 映射回 $N \times d_v$ 。

每一步的统计尺度都需要被控制—— $\sqrt{d_k}$ 是控制 $QK^{\mathsf{T}}$ 这一步的尺度的工具。

V 那边没有显式缩放，因为 softmax 输出已经是概率（行和为 1）——V 的均值和方差只取决于 V 自己的统计——这一步通常不需要额外正则化。

九、Vaswani 论文里的原话

9.1 § 3.2.1 的关键段落

原文（NeurIPS 2017）：

"We suspect that for large values of d_k, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by 1/√d_k."

翻译：「我们怀疑对于大的 $d_k$ ，点积量级变大，从而把 softmax 推入梯度极小的区域。为抵消这个效应，我们用 $1 / \sqrt{d_k}$ 缩放点积。」

9.2 这段话其实没给完整证明

注意「我们怀疑」（we suspect）——Vaswani 没有给出 $\operatorname{Var}(q \cdot k) = d_k$ 的形式化推导，也没有大规模消融实验来证明。

后来的工作（Xiong 2020, On Layer Normalization in the Transformer Architecture）才把这件事详细分析。

但工程上，Vaswani 的「直觉 + 简单理论」已经够用——大家用了 $\sqrt{d_k}$ ，模型能训，事情就成立了。

这是科学研究里很常见的模式：实践先于理论——直觉推动实验，实验验证后再被理论补全。

9.3 注释：dot product vs scaled dot product 的对比实验

Vaswani 论文里的 Table 3 提到：

"While for small values of d_k the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of d_k. We suspect that..."

也就是说，Vaswani 团队做过对照实验——在 d_k 大时不缩放的 dot product 比加性 attention 差——所以加了缩放。

这是 $\sqrt{d_k}$ 设计的直接动机。

十、对应的 PyTorch 实现

10.1 最朴素版

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5)  # ← 这里就是 √d_k
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    attn = F.softmax(scores, dim=-1)
    return attn @ V, attn

注意 d_k ** 0.5 就是 $\sqrt{d_k}$ 。

Q.size(-1) 自动取最后一维——所以这段代码不需要传 d_k 参数。

10.2 数值稳定版（log-sum-exp trick）

朴素 softmax 在大 logits 时可能溢出（exp(700) = inf）。

实践中 PyTorch 的 F.softmax 已经内置了 log-sum-exp trick——把所有 logits 减去 max 再 exp，结果不变但数值稳定。

这一点对 scaled dot-product 也有意义——因为缩放后 logits 仍可能在 $\pm 10$ 量级（在某些 head 学到 sharp pattern 时），log-sum-exp 仍是必要的。

10.3 PyTorch 2.0+ 的内置实现

out = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask)

这一行调用底层 CUDA / Metal 实现——可能是 FlashAttention，也可能是 Memory-Efficient Attention，由 backend 自动选择。

但「除以 $\sqrt{d_k}$ 」这件事仍然在背后发生——只是你不用手写。

10.4 FlashAttention 中的 $\sqrt{d_k}$

FlashAttention 的核心是「tile-by-tile 计算 softmax」——在 SRAM 内做 streaming softmax。

$\sqrt{d_k}$ 缩放发生在每个 tile 计算 $QK^{\mathsf{T}}$ 的瞬间——和朴素实现没有本质区别。

工程上的难点是「数值稳定的 streaming softmax」（要保持 running max 和 running sum_exp）——但 $\sqrt{d_k}$ 这一步是简单乘法，不影响 FlashAttention 的核心算法。

10.5 一个常见 bug：缩放放在哪里

有些实现会写：

Q = Q / (d_k ** 0.5)  # 提前缩放 Q
scores = Q @ K.transpose(-2, -1)

这等价于在 $QK^{\mathsf{T}}$ 上除以 $\sqrt{d_k}$ ——结果一样，但计算更高效（少一次矩阵元素级除法）。

但要注意：如果 K 有特殊处理（比如 RoPE），缩放放在哪里可能影响 RoPE 的正确性——一般推荐放在 score 上，最稳。

十一、几个常见的变体与争议

11.1 $\sqrt{d_k}$ vs $\sqrt{d_{\mathrm{model}}}$

有人混淆：

$d_{\mathrm{model}}$ 是 token 嵌入维度（比如 $512$ ）。

$d_k$ 是每个 head 的 $Q/K$ 维度（比如 $64$ ，如果有 $8$ 个 head）。

scaled dot-product 用的是 $\sqrt{d_k}$ ，不是 $\sqrt{d_{\mathrm{model}}}$ 。

为什么？因为方差推导中，加和的项数是 $d_k$ ——每个 head 的点积只涉及 $d_k$ 维。

如果你写错成 $\sqrt{d_{\mathrm{model}}}$ （除得太多），attention 会过于平缓——softmax 输出近似均匀——模型失去选择性。

11.2 $1/d_k$ vs $1/\sqrt{d_k}$

如果你看到某些代码或论文写「 $\div d_k$ 」（不是 $\sqrt{d_k}$ ），那是错的——除非他们定义 $d_k$ 是 $\sqrt{\text{原 } d_k}$ 。

把方差推导记牢： $\sigma = \sqrt{d_k}$ ，所以分母是 $\sqrt{d_k}$ 而不是 $d_k$ 。

11.3 学习的温度参数 vs 固定的 $\sqrt{d_k}$

有些工作（Shaped Attention、Stable Attention）把 $\sqrt{d_k}$ 替换成可学习的 $\tau$ ，让模型自适应温度。

这通常需要额外的稳定化（比如把 $\tau$ clamp 到 $[\sqrt{d_k}/2, 2\sqrt{d_k}]$ ），否则 $\tau$ 容易学到 $0$ 或 $\infty$ 。

主流大模型仍然用固定 $\sqrt{d_k}$ ——因为它已经够好了，省掉一个超参。

11.4 logit-cap：另一个稳定性技巧

Gemini 和某些 Anthropic 模型用「logit cap」技巧：

\mathrm{scores} = c \cdot \tanh\left(\frac{QK^{\mathsf{T}}}{\sqrt{d_k} \, c}\right)

其中 $c$ 是某个 cap 值（比如 $50$ ）。

这把 logits 强行 clip 到 $[-c, c]$ 区间，防止极端 outlier。

这是 $\sqrt{d_k}$ 之后的进一步加强——不替代它，而是补充它。

11.5 query 缩放还是 score 缩放

一些工程实现把缩放放在 $Q$ 上：

Q := \frac{Q}{d_k^{0.25}}, \qquad K := \frac{K}{d_k^{0.25}}

两个 $0.25$ 次方相乘恰好得 $\sqrt{d_k}$ 。

这种「分散到 Q 和 K」的写法在某些硬件上更高效，但数学上完全等价。

PyTorch 默认实现把缩放放在 score 上。

十二、与训练分布假设的关系

12.1 「q, k 是单位方差」这个假设有多严格

我们推导 $\operatorname{Var}(q \cdot k) = d_k$ 的关键假设是：每一维独立、零均值、单位方差。

实际中：

独立：不严格，但近似成立（参数学到的 Q、K 投影把不同维度去相关到一定程度）。
零均值：通过 LayerNorm 严格成立。
单位方差：通过 LayerNorm 严格成立。

整体上，「点积方差 $\approx d_k$ 」是一个近似——但近似得相当好。

12.2 训练后期的偏移

训练后期， $Q$ 、 $K$ 的分布可能偏离单位方差（特别是某些 head 学到 sharp pattern 时）——logits 的实际方差可能比 $d_k$ 小很多（因为 $Q$ 和 $K$ 学到了对齐方向，使 $q \cdot k$ 偏正）。

这时 $\sqrt{d_k}$ 给出的「过度缩放」让 softmax 仍然平缓——模型需要学习一个更大的 $W_q$ 或 $W_k$ 来「弥补」缩放。

QK-norm 的提案就是为了解决这个：让 $q$ 、 $k$ 在训练全程都保持单位方差， $\sqrt{d_k}$ 缩放始终精确。

12.3 一个反直觉发现

Su 2024 等人发现：训练初期 logits 近似高斯，但训练到收敛时，logits 分布严重偏离高斯——出现一些「极端 outlier」（某些位置 logit 突变到 $\pm 50$ 以上）。

这种 outlier 对训练稳定性是灾难——logit-cap 就是为了 cap 这些 outlier。

$\sqrt{d_k}$ 的「单位方差」假设在训练稳定期成立，但在收敛附近可能开始失效——这是一个开放研究方向。

12.4 高斯假设到底有多重要

有人会问：如果 Q、K 的分布不是高斯（而是 t 分布、混合高斯、甚至离散），方差推导还成立吗？

成立。

$\operatorname{Var}(X+Y) = \operatorname{Var}(X) + \operatorname{Var}(Y)$ 对任何独立随机变量都成立——和分布无关。

我们推 $\operatorname{Var}(q \cdot k) = d_k$ 时也只用了「独立 + 零均值 + 单位方差」——没有用到高斯假设。

但软预测（softmax 输出形状）要用 Central Limit Theorem—— $q \cdot k$ 是 $d_k$ 个项之和， $d_k$ 大时近似高斯。

到 $d_k = 64$ 已经足够 CLT 生效——分布看起来很高斯。

12.5 重尾分布（heavy tail）的影响

如果 $Q$ 、 $K$ 不是高斯而是 heavy tail（比如 $t$ 分布、Cauchy），那方差推导可能不成立——或者方差为 $\infty$ 。

实际上深度网络的中间表示确实会有 heavy tail（参考 Martin & Mahoney 2018 关于深度网络中间表示的 heavy tail spectrum）。

但 LayerNorm + 标准初始化让 Q、K 的尾部不至于失控——这是工程上的「经验救场」。

到 2026 年，理解 $Q$ 、 $K$ 的真实分布、以及 $\sqrt{d_k}$ 在 heavy tail 下的有效性，仍是开放问题。

十三、当 $\sqrt{d_k}$ 不够用时

13.1 long-context 的 logits 暴涨

当上下文长度 N 很大时（比如 100k），同一个 query 要 attend 100k 个 key——每个 key 都贡献 logits 候选。

即使每个 logits 都是 $\sigma = 1$ （缩放后），最大值 $\max_i s_i$ 在 $N$ 大时按 $\sqrt{2 \ln N}$ 增长（极值统计）——logits 仍然漂移到大值。

但 $\operatorname{softmax}(s)$ 中真正起作用的是 $s_{\max} - s$ ——这个差仍然是 $\sigma \approx 1$ 的量级——所以 attention 仍然分布合理。

13.2 attention sink

Xiao 2023 (StreamingLLM) 发现：在 long-context 中，第一个 token（BOS）会「吸走」大量 attention 权重——这是 softmax 归一化的副产物。

具体机制：当所有 logits 都接近 0 时（没有特别匹配的 key），softmax 趋向均匀——但每个 token 都倾向把「无信息」的 attention 转移到「最早出现的 token」（BOS）。

$\sqrt{d_k}$ 与 attention sink 没有直接关系——但在 long-context 中， $\sqrt{d_k}$ 提供了基础稳定性，attention sink 现象在此基础上才能被研究。

13.3 ALiBi 与 $\sqrt{d_k}$ 的相互作用

ALiBi（Press 2021）在 logits 上加一个负的距离偏置：

s_{ij} = \frac{q_i \cdot k_j}{\sqrt{d_k}} - m \cdot |i - j|

$m$ 是一个固定的负向斜率。

ALiBi 与 $\sqrt{d_k}$ 是叠加关系—— $\sqrt{d_k}$ 控制方差，ALiBi 控制位置偏置——两者各司其职。

13.4 RoPE 与 $\sqrt{d_k}$ 的相互作用

RoPE（Su 2021）在 Q、K 上做旋转编码：

q' = R(\theta) q, \qquad k' = R(\theta) k

其中 $R(\theta)$ 是旋转矩阵——保持向量长度不变。

因此 $q' \cdot k'$ 的方差与 $q \cdot k$ 一致—— $\sqrt{d_k}$ 缩放仍然有效。

RoPE 是「不影响 $\sqrt{d_k}$ 假设」的位置编码——这是它能广泛应用的一个隐性原因。

十四、Muon 优化器与 $\sqrt{d_k}$ 的现代视角

14.1 Muon 是什么

Muon（2024）是一个新型优化器，专为 Transformer 设计——它对 attention 的 $W_q$ 、 $W_k$ 矩阵做特殊正交化。

核心思想： $W_q$ 、 $W_k$ 在训练中容易变得「不正交」——这让 $q \cdot k$ 的统计性质偏离原始假设——Muon 强制周期性正交化。

14.2 Muon 与 $\sqrt{d_k}$ 的关系

Muon 维持 $W_q$ 、 $W_k$ 的正交性 $\to$ 维持 $q$ 、 $k$ 的单位方差 $\to$ 维持 $\sqrt{d_k}$ 的精确性。

也就是说，Muon 让「 $\sqrt{d_k}$ 假设」在训练全程都接近精确——这反过来让 attention 训练更稳定。

到 2026 年，Muon 在某些大模型预训练中开始被采用（比如 Kimi 的 K2 模型）——这印证了「保护 $\sqrt{d_k}$ 假设」的工程价值。

14.3 一种联合视角

把 $\sqrt{d_k}$ 、QK-norm、Muon、logit-cap 放到一起，你会发现一条主线：

保护 attention logits 的统计性质，让 softmax 始终在「有效梯度区」工作。

每一项技术都是这条主线的一个工具—— $\sqrt{d_k}$ 是最基础的、最便宜的、必须有的——其它都是渐进改进。

十五、一个完整的数值小例子

15.1 设置

设 $d_k = 4$ 。 $q$ 和 $k$ 都是 $4$ 维。

q = [1, 0.5, -0.5, 1]

k_1 = [0.5, 1, 0, -1]

k_2 = [1, 0.5, 0.5, 0]

k_3 = [-1, 0, 1, 0.5]

15.2 不缩放的 logits

s_1 = q \cdot k_1 = 0.5 + 0.5 + 0 - 1 = 0

s_2 = q \cdot k_2 = 1 + 0.25 - 0.25 + 0 = 1

s_3 = q \cdot k_3 = -1 + 0 - 0.5 + 0.5 = -1

\mathrm{logits} = [0, 1, -1]

\operatorname{softmax}([0, 1, -1]) \approx [0.244, 0.665, 0.090]

15.3 缩放的 logits

\sqrt{d_k} = 2

s'_1 = 0 / 2 = 0

s'_2 = 1 / 2 = 0.5

s'_3 = -1 / 2 = -0.5

\mathrm{logits}' = [0, 0.5, -0.5]

\operatorname{softmax}([0, 0.5, -0.5]) \approx [0.295, 0.487, 0.218]

15.4 比较

不缩放： $[0.244, 0.665, 0.090]$ ——中间项更突出。

缩放： $[0.295, 0.487, 0.218]$ ——分布更平。

在 $d_k = 4$ 这种小尺度下，差别不大——缩放只让 softmax 略缓和。

但当 $d_k = 64$ 时，原始 logits 范围会扩大 $4$ 倍（ $\sqrt{64}/\sqrt{4} = 4$ ），不缩放 logits 是 $[0, 4, -4]$ ， $\operatorname{softmax}([0, 4, -4]) \approx [0.018, 0.964, 0.000]$ ——几乎 one-hot！

而缩放后仍然是 $[0, 0.5, -0.5]$ ——分布合理。

这就是「维度越大，越需要 $\sqrt{d_k}$ 」的直观体现。

15.5 backward 梯度差异

设 loss 对 attention 输出有梯度信号 $\partial L / \partial \mathrm{out}$ 。

对未缩放 softmax：

\frac{\partial L}{\partial s} = J_{\mathrm{softmax}} \cdot \frac{\partial L}{\partial \mathrm{out}}

$J_{\mathrm{softmax}}$ 在「one-hot」状态下接近零—— $\partial L / \partial s \approx 0$ —— $s = QK^{\mathsf{T}}$ 的梯度也接近零—— $Q$ 、 $K$ 的更新极缓慢。

对缩放 softmax：

\frac{\partial L}{\partial s'} = J_{\mathrm{softmax}} \cdot \frac{\partial L}{\partial \mathrm{out}}

这里 $J_{\mathrm{softmax}}$ 不饱和，传播正常。

\frac{\partial L}{\partial s} = \frac{\partial L}{\partial s'} \cdot \frac{1}{\sqrt{d_k}}

注意还多了一个 $1/\sqrt{d_k}$ ——但这只让梯度等比缩小，不让它消失。

整体梯度量级仍然合理，训练能进行。

十六、关键概念回顾

点积方差： $q \cdot k$ 的方差等于 $d_k$ （在标准假设下）。
softmax 饱和：当 logits 量级远大于 1 时，softmax 输出近似 one-hot——梯度近似为零。
$\sqrt{d_k}$ 缩放：把 logits 方差归一化到 $1$ ，避免 softmax 饱和。
临界点：分母必须是 $\sqrt{d_k}$ ——除以 $d_k$ 太多，除以更小则不够。
scaling laws 隐性基础： $\sqrt{d_k}$ 让 attention 在不同维度下有相同的「统计工作点」——这是 Chinchilla scaling 能成立的前提之一。
NTK 视角： $\sqrt{d_k}$ 是「保持每一层输入输出尺度一致」的具体实现。
现代变体：QK-norm、logit-cap、Muon 都是「保护 $\sqrt{d_k}$ 假设」的延伸。
训练稳定性： $\sqrt{d_k}$ 不能完全替代 LayerNorm、warmup、gradient clipping——但它是这些手段的基础。
Vaswani 原文：只是一句「we suspect」+ 简单实验——后续工作才补全理论。
PyTorch 实现：F.scaled_dot_product_attention 已内置 $\sqrt{d_k}$ ——但理解原理仍然重要。

十七、常见误解

17.1 $\sqrt{d_k}$ 是经验技巧

错。这是一个有严格概率论推导的设计——不是「随便选一个数」。

17.4 $\sqrt{d_k}$ 是 attention 唯一的稳定化机制

错。LayerNorm、warmup、gradient clipping、初始化都是稳定性的一部分—— $\sqrt{d_k}$ 只是其中一项。

17.5 缩放 = 退化

错。缩放后 attention 仍然能学到 sharp pattern——只是初始时不卡饱和——训练完成后该 sharp 还是 sharp。

17.12 缩放只为加速训练

不仅如此。

缩放更核心的目的是「让训练能进行」——而不是「让训练加快」。

不缩放时，训练在大模型上根本无法成功——加上缩放后训练才能稳定走完。

这是「质」的区别，不是「量」的区别。

17.14 $\sqrt{d_k}$ 缩放破坏了 attention 的「概率含义」

不严格。

attention 输出仍然是「key 上的概率分布」（softmax 输出每行和为 $1$ ）—— $\sqrt{d_k}$ 只影响这个分布有多 sharp，不影响它是不是概率。

事实上， $\sqrt{d_k}$ 让初始的概率分布「更接近均匀」——这反而是更好的概率初始化。

十八、下一步

到这里，我们已经把 attention 机制原理这一块的核心讲完了。

下一篇（第 16 篇）开始进入【Part 3：Transformer 架构】——讨论完整的 Transformer 块如何把 attention、FFN、residual、LayerNorm 串起来。

我们会从「2017 原始 Transformer」讲起，逐步看到「现代 LLaMA-style Transformer」演化的每一个改动是为什么——Pre-LN vs Post-LN、SwiGLU vs ReLU、RMSNorm vs LayerNorm 等。

如果你已经掌握了：

第 11 篇的「attention 是什么」直觉
第 12 篇的 Bahdanau 加性注意力
第 13 篇的 Q/K/V 三件套
第 14 篇的 self-attention 概念
第 15 篇的 $\sqrt{d_k}$ 缩放原理

那你已经有了进入 Transformer 架构层的所有理论基础——下一篇就把这些拼起来。

十九、参考文献

下面按相关度排序列出本篇直接引用与延伸阅读，每条附一句话提示其在本篇中的角色。

阅读建议：1、2、3、12 是核心，其余是延伸。

Vaswani, A. et al. "Attention Is All You Need." NeurIPS 2017. §3.2.1 给出 √d_k 的最早动机。
Xiong, R. et al. "On Layer Normalization in the Transformer Architecture." ICML 2020. 形式化分析 √d_k 与 Pre-LN 的关系。
Luong, M.-T. et al. "Effective Approaches to Attention-based Neural Machine Translation." EMNLP 2015. dot-product attention 的经典工作，没有 √d_k——展示了不缩放的问题。
Bahdanau, D. et al. "Neural Machine Translation by Jointly Learning to Align and Translate." ICLR 2015. 加性 attention 没有 √d_k 问题，因为 tanh bound。
Jacot, A., Gabriel, F., Hongler, C. "Neural Tangent Kernel: Convergence and Generalization in Neural Networks." NeurIPS 2018. NTK 理论的源头。
Yang, G., Hu, E. J. "Tensor Programs IV: Feature Learning in Infinite-Width Neural Networks." ICML 2021. muP 的理论基础。
Yang, G. et al. "Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer." NeurIPS 2021. muP 的实操版，与 √d_k 互补。
Hoffmann, J. et al. "Training Compute-Optimal Large Language Models." NeurIPS 2022. Chinchilla scaling laws，隐性依赖架构稳定性。
Su, J. et al. "RoFormer: Enhanced Transformer with Rotary Position Embedding." Neurocomputing 2024. RoPE 不破坏 √d_k 假设。
Press, O., Smith, N. A., Lewis, M. "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation." ICLR 2022. ALiBi，与 √d_k 叠加使用。
Xiao, G. et al. "Efficient Streaming Language Models with Attention Sinks." ICLR 2024. attention sink 现象。
Dao, T. et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." NeurIPS 2022. 工程实现里 √d_k 的位置。
Henry, A. et al. "Query-Key Normalization for Transformers." EMNLP Findings 2020. QK-norm 提案。
Shazeer, N. "Fast Transformer Decoding: One Write-Head is All You Need." arXiv 2019. MQA。
Ainslie, J. et al. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints." EMNLP 2023. GQA。
Martin, C. H., Mahoney, M. W. "Implicit Self-Regularization in Deep Neural Networks." JMLR 2021. 深度网络中间表示的 heavy tail 现象。
Jordan, K. et al. "Muon: An Optimizer for the Hidden Layers of Neural Networks." 2024 blog/preprint. Muon 的提出。

← 上一篇：14｜Self-Attention　|　下一篇：16｜Multi-Head Attention →

Scaled Dot-Product：那个根号 d_k 是怎么来的'