为什么Transformer的注意力要除以√dₖ：从方差到梯度消失的完整数学解析

2017年，Vaswani等人在《Attention Is All You Need》论文中提出了Transformer架构。在缩放点积注意力的公式中，有一个看似不起眼的细节：点积结果要除以$\sqrt{d_k}$。这个操作在代码中只是一行，但它背后隐藏着深度学习中最核心的问题之一——梯度消失。

一个简单的问题

先看缩放点积注意力的公式：

其中$Q$是查询矩阵，$K$是键矩阵，$V$是值矩阵，$d_k$是键向量的维度。

问题来了：为什么要除以$\sqrt{d_k}$？为什么不直接用$QK^T$？为什么是平方根而不是其他？

答案要从两个角度理解：点积的统计特性，以及Softmax的梯度行为。

点积在高维空间的"膨胀"

假设$Q$和$K$的每个元素都是独立同分布的随机变量，均值为0，方差为1。考虑它们的点积：

由于$q_i$和$k_i$相互独立，根据期望的线性性质：

方差呢？对于独立随机变量的乘积之和：

当$E[q_i] = E[k_i] = 0$且$\text{Var}[q_i] = \text{Var}[k_i] = 1$时：

因此：

这就是关键：点积的方差随维度$d_k$线性增长。

具体数值意味着什么

假设$d_k = 64$（Transformer原论文的设置），那么点积的标准差是$\sqrt{64} = 8$。如果$d_k = 512$，标准差就变成了$\sqrt{512} \approx 22.6$。

这不是小数字。考虑一个序列长度为10的注意力计算，每对$(q_i, k_j)$的点积都是一个随机变量。如果$d_k = 512$，这些点积的值可能在$-60$到$60$之间波动（大约3倍标准差范围内）。

这会带来什么问题？答案在Softmax。

Softmax的"尖峰陷阱"

Softmax函数将向量转换为概率分布：

Softmax有一个重要特性：输入差异越大，输出分布越"尖"。

数值示例

假设三个注意力分数是$[1, 2, 3]$：

softmax([1, 2, 3]) = [0.090, 0.245, 0.665]

分布相对均匀，三个位置都有一定权重。

现在把分数放大10倍$[10, 20, 30]$：

softmax([10, 20, 30]) = [0.000, 0.000, 1.000]

几乎变成了one-hot向量。第一个和第二个位置的权重趋近于0。

这就是"尖峰"问题：当输入差异过大时，Softmax会输出一个极其尖锐的分布，几乎所有的概率质量都集中在最大值上。

梯度消失的根源

问题不仅仅在于输出的分布形状，更在于梯度。

Softmax的梯度公式是：

当某个$p_k \approx 1$，其他$p_i \approx 0$时：

• 对于$i \neq k$：$\frac{\partial p_i}{\partial x_i} = p_i(1-p_i) \approx 0$
• 对于$i = k$：$\frac{\partial p_i}{\partial x_j} = -p_i p_j \approx 0$

梯度几乎全部消失。

这就是为什么缩放如此重要：如果不控制点积的量级，Softmax的输入会非常大，输出会极其尖峰，梯度会消失，模型几乎无法学习。

为什么偏偏是√dₖ

现在问题变成：为什么除以$\sqrt{d_k}$就能解决问题？

回到方差的计算。如果点积的方差是$d_k$，那么标准差就是$\sqrt{d_k}$。

统计学中，将随机变量除以其标准差，就能得到单位方差：

这就是"缩放"的本质：将点积的方差归一化为1，无论$d_k$是多少。

为什么不直接除以dₖ

有人可能会问：为什么不除以$d_k$，这样方差不就变成$1/d_k$了吗？

问题在于，除以$d_k$会过度压缩。当$d_k$很大时，所有点积都会变得很小，Softmax的输入都接近0，输出会变成均匀分布：

softmax([0.001, 0.002, 0.003]) ≈ [0.333, 0.333, 0.334]

这同样有问题：注意力机制失去了区分能力，无法聚焦到关键位置。

除以$\sqrt{d_k}$是一个精妙的平衡：既避免了方差过大导致的梯度消失，又保留了足够的区分度让模型学习。

与Xavier初始化的深层联系

这种"除以$\sqrt{d}$“的思想并非Transformer首创，它深深植根于深度学习的初始化理论。

Xavier初始化

2010年，Xavier Glorot和Yoshua Bengio提出了一种权重初始化方法。核心思想是：在正向传播中，保持每一层输出的方差稳定；在反向传播中，保持梯度的方差稳定。

对于权重矩阵$W \in \mathbb{R}^{n_{out} \times n_{in}}$，Xavier初始化建议：

当输入$x$的方差为$\sigma^2$时，输出$y = Wx$的方差为：

注意力的缩放因子$\sqrt{d_k}$遵循同样的逻辑：让输出的方差保持稳定。

这并非巧合。深度学习的核心挑战之一就是控制信号在网络中的传播。无论是初始化权重还是设计注意力机制，都需要确保数值的量级不会失控。

加性注意力为什么不需要缩放

在Transformer之前，Bahdanau等人在2015年提出了加性注意力（Additive Attention），计算公式是：

为什么加性注意力不需要缩放？

关键在于$\tanh$函数。$\tanh$的输出范围是$[-1, 1]$，无论输入多大，输出都被限制在这个区间内。即使$W_q q + W_k k$的值很大，经过$\tanh$后仍然是有界的。

点积则不同。理论上，点积的值域是$(-\infty, +\infty)$，没有任何约束。当维度增加时，点积的量级会不断增长。

Transformer选择点积而非加性注意力的原因是效率：点积可以通过高度优化的矩阵乘法实现，而加性注意力需要额外的神经网络层。

不使用缩放会发生什么

为了理解缩放的重要性，可以做一个简单的实验对比。

假设$d_k = 512$，序列长度为128。不使用缩放时，点积的值可能达到$\pm 50$甚至更大。经过Softmax后：

import numpy as np

# 模拟不缩放的注意力分数
d_k = 512
seq_len = 128
np.random.seed(42)
Q = np.random.randn(seq_len, d_k)
K = np.random.randn(seq_len, d_k)

# 不缩放
scores_no_scale = Q @ K.T
print(f"无缩放分数范围: [{scores_no_scale.min():.1f}, {scores_no_scale.max():.1f}]")
# 输出: 无缩放分数范围: [-51.3, 49.8]

# 缩放后
scores_scaled = scores_no_scale / np.sqrt(d_k)
print(f"缩放后分数范围: [{scores_scaled.min():.1f}, {scores_scaled.max():.1f}]")
# 输出: 缩放后分数范围: [-2.3, 2.2]

无缩放的分数范围是$[-51, 50]$，缩放后变成$[-2.3, 2.2]$。

看看Softmax的效果：

def softmax(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

# 不缩放的Softmax
attn_no_scale = softmax(scores_no_scale[0])
print(f"无缩放注意力分布: 最大值={attn_no_scale.max():.6f}, 最小值={attn_no_scale.min():.10f}")
# 输出: 无缩放注意力分布: 最大值=0.999999, 最小值=0.0000000000

# 缩放后的Softmax
attn_scaled = softmax(scores_scaled[0])
print(f"缩放后注意力分布: 最大值={attn_scaled.max():.4f}, 最小值={attn_scaled.min():.4f}")
# 输出: 缩放后注意力分布: 最大值=0.1564, 最小值=0.0042

无缩放时，注意力权重几乎是one-hot：最大值接近1，最小值接近0。这意味着模型只能关注一个位置，无法学习到分布式的表示。

缩放后，注意力分布更加均匀，最小值也有$0.42\%$的权重，模型可以同时关注多个相关位置。

多头注意力中的dₖ选择

在多头注意力中，$d_k$的选择直接影响模型的表达能力。

原始Transformer的设置是$d_{model} = 512$，$h = 8$个头，因此每个头的维度：

这意味着每个头的缩放因子是$\sqrt{64} = 8$。

不同模型的选择：

注意到大多数模型保持$d_k = 64$或$128$，即使总维度很大。这是因为$d_k$不仅影响缩放，还影响每个头的表达能力。$d_k$太小，单个头无法学习复杂的表示；$d_k$太大，计算成本增加，且缩放因子的校正作用减弱。

Softmax的数值稳定性

还有一个与缩放相关的技术细节：Softmax本身的数值稳定性。

直接计算$e^{x_i}$可能会导致数值溢出。常见的解决方案是减去最大值：

这个技巧保证了数值稳定性，因为指数的输入总是非正的。

但要注意：减去最大值和除以$\sqrt{d_k}$是两回事。

减去最大值只是平移数值，不改变相对差异：

[100, 102, 103] - 103 = [-3, -1, 0]
差异仍然是2和1

除以$\sqrt{d_k}$则是缩放，会压缩相对差异：

[100, 102, 103] / √100 = [10, 10.2, 10.3]
差异变成0.2和0.1

前者解决的是数值溢出问题，后者解决的是方差和梯度问题。两者都需要。

从工程角度理解

整个流程是一个精巧的工程设计：

1. 问题定义：点积的方差随维度增长
2. 后果分析：Softmax饱和，梯度消失
3. 解决方案：除以标准差，归一化方差
4. 验证方法：确保梯度可以正常反向传播

这个设计体现了深度学习的核心方法论：理解数学本质，然后用简单的方法解决复杂的问题。

总结

Transformer的缩放因子$\sqrt{d_k}$并非凭空出现，它是对深度学习基本问题的精准回答：

问题：点积的方差随维度线性增长，导致Softmax输入过大、输出尖峰、梯度消失。

方案：除以标准差$\sqrt{d_k}$，将方差归一化为1，保持数值稳定。

洞察：这与Xavier初始化的思想一脉相承——在深度网络中控制信号的量级，是成功训练的关键。

看似简单的一个除法操作，背后是对概率论、优化理论和数值计算的深刻理解。这也是Transformer能够取代RNN成为主流架构的原因之一：它不仅在架构上创新，更在每一个细节上都经过深思熟虑。

参考文献

1. Vaswani, A., et al. (2017). Attention Is All You Need. arXiv:1706.03762.
2. Bahdanau, D., Cho, K., & Bengio, Y. (2014). Neural Machine Translation by Jointly Learning to Align and Translate. arXiv:1409.0473.
3. Glorot, X., & Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. AISTATS.
4. He, K., et al. (2015). Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. ICCV.
5. Luong, M. T., et al. (2015). Effective Approaches to Attention-based Neural Machine Translation. EMNLP.
6. Devlin, J., et al. (2018). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv:1810.04805.
7. Brown, T., et al. (2020). Language Models are Few-Shot Learners. NeurIPS.
8. Touvron, H., et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288.