考察点
这道题出自字节跳动 LLM 算法岗实习二面，是 Transformer 基础题里最高频的一道，但面试官要的不是背出「防止梯度消失」这六个字，而是你能不能把「为什么方差会随 d_k 涨」「softmax 饱和之后梯度具体怎么消失」讲清楚。追问通常会往两个方向走：为什么偏偏是根号 d_k 而不是 d_k，以及实际工程里还有哪些和缩放相关的 trick（QK-Norm、初始化、温度系数）。
参考答案
点积的方差随维度线性增长
Attention 的打分是 query 和 key 的点积：
$$score = q \cdot k = \sum_{i=1}^{d_k} q_i k_i$$
假设 q 和 k 的每个分量都是独立随机变量，均值为 0、方差为 1（这正是网络初始化后、LayerNorm 之后大致满足的状态），那么每一项 $q_i k_i$ 的均值是 0、方差是 1，$d_k$ 项独立求和之后，点积的均值仍为 0，但方差变成了 d_k，标准差是 $\sqrt{d_k}$。
拿实际数字感受一下：GPT 类模型里每个 attention head 的 d_k 常见是 64 或 128。不做缩放时，点积的标准差就是 8 到 11 左右，打分落到 ±30 的区间很常见；如果 d_k 更大或者分量方差比 1 大一点，打分的绝对值会更大。而 softmax 对输入的尺度非常敏感。
softmax 饱和与梯度消失
softmax 是个「赢家通吃」倾向很强的函数。输入分值的绝对值一大，输出就会急剧向 one-hot 靠拢：最大分值对应的权重逼近 1，其余全部压到接近 0。看反向传播就明白问题在哪——softmax 的雅可比矩阵对角元素是 $p_i(1-p_i)$，非对角是 $-p_i p_j$。当输出接近 one-hot 时，$p_i$ 不是接近 1 就是接近 0，所有这些项都趋近于 0，梯度几乎传不回去。
后果是：Q、K 两个投影矩阵拿到的梯度信号极弱，学得很慢甚至学不动；attention 的权重分布被「锁死」在偶然形成的硬选择上，模型很难再去调整注意力模式。训练初期这个问题尤其致命，因为初始化后的点积本来就是随机噪声放大版。
为什么恰好除以根号 dk
目标很朴素：让 softmax 的输入方差回到 1 这个量级，和它分量的方差一致。点积方差是 d_k，所以除以 $\sqrt{d_k}$ 之后，方差变成 $d_k / d_k = 1$，标准差回到 1，softmax 工作在输入绝对值个位数的「温和区间」——这个区间里输出分布平滑、梯度健康，每个 key 都有被训练竞争到的机会。
之所以不是除以 d_k：那会把方差压到 1/d_k，矫枉过正，打分全部挤在 0 附近，softmax 输出接近均匀分布，注意力学不出区分度。$\sqrt{d_k}$ 是「方差归一」这个原则下唯一自然的系数，不是调参调出来的玄学数字。
工程里的延伸认知
实际模型里围绕这个缩放还有不少值得知道的细节：
- 大部分框架的 attention 实现允许自定义 scale 参数（PyTorch 的
scaled_dot_product_attention就有scale入参），默认就是 $1/\sqrt{d_k}$。 - 一些新架构（如部分 Qwen、Gemma 变体）引入 QK-Norm，在点积前对 q 和 k 做 RMSNorm，从输入端控制点积的量级，缓解极深模型里 attention logits 爆炸的问题——这和除以根号 d_k 解决的是同一类问题，只是下手的位置不同。
- 推理侧的 temperature 参数本质也是在缩放 logits 影响 softmax 的尖锐程度，思想上和这里是相通的。
- FlashAttention 的 CUDA kernel 里，缩放系数被融合进累加过程，配合 online softmax 的分块 rescale，保证数值稳定的同时不牺牲显存效率——面试提到这层会加分。
回答收尾时可以点一句：这个缩放是「保持初始化下方差可控」的经典案例，和 Xavier/He 初始化、残差连接的方差分析是同一套思维，面试官往往就是想听你把点连成线。
可能的追问
- 为什么除以根号 d_k 而不是 d_k？ 答：目标是让点积方差归一到 1。除以 d_k 会把方差压到 1/d_k，打分全挤在零附近，softmax 输出趋近均匀分布，注意力失去区分度。
- d_k 很大时除了缩放还有什么办法？ 答：QK-Norm 在点积前归一化 q/k 的模长；或者对 attention logits 加 soft-cap（如 Gemma2 的 logit softcapping），都是控制打分量级的手段。
- 只缩放 q 或只缩放 k 行不行？ 答：数学上等价，都是对点积整体除以根号 d_k；工程上习惯在 q 上缩放一次，kernel 里常融合到 softmax 之前统一处理。