注意力机制中除以Dk的方差归一化

要理解注意力机制中除以 $\sqrt{d_k}$ 的数学本质，需从点积的方差问题切入 —— 这是避免 softmax 输出极端化、保证模型可学习性的关键。下面将分 “核心问题推导”“为什么是 $\sqrt{d_k}$”“为什么不是 $d_k$ 或 $d_k^2$” 三部分解释，并结合具体数学例子验证。

一、核心问题：未缩放的点积会导致方差爆炸

注意力机制（以缩放点积注意力为例）的核心计算是查询向量 $\mathbf{q}$（维度 $d_k$）与键向量 $\mathbf{k}$（维度 $d_k$）的点积：

$\mathbf{q}\cdot \mathbf{k}={\sum }_{i=1}^{d_k}{q}_i{k}_i$

为了分析其数值特性，我们做合理假设（符合模型初始化惯例）：

• $\mathbf{q}$ 和 $\mathbf{k}$ 的每个元素 $q_i,k_i$ 是独立同分布（i.i.d.） 的随机变量；

• 均值为 $0$（消除偏移影响），方差为 $1$（标准化初始化，如 Xavier 初始化）；

• $\mathbf{q}$ 和 $\mathbf{k}$ 相互独立（查询与键来自不同分支）。

1.1 点积的方差计算

根据方差的性质：

• 独立变量之和的方差 = 各变量方差之和（即 $\text{Var}(a+b)=\text{Var}(a)+\text{Var}(b)$，当 $a,b$ 独立时）；

• 独立变量乘积的方差：$\text{Var}(q_i{k}_i)=\mathbb{E}[(q_i{k}_i{)}^2]-(\mathbb{E}[q_i{k}_i]{)}^2$。由于 $\mathbb{E}[q_i]=\mathbb{E}[k_i]=0$，且 $q_i,k_i$ 独立，故 $\mathbb{E}[q_i{k}_i]=\mathbb{E}[q_i]\mathbb{E}[k_i]=0$，且 $\mathbb{E}[(q_i{k}_i{)}^2]=\mathbb{E}[q_i^2]\mathbb{E}[k_i^2]$。

又因为 $\text{Var}(q_i)=\mathbb{E}[q_i^2]-(\mathbb{E}[q_i]{)}^2=\mathbb{E}[q_i^2]=1$（同理 $\mathbb{E}[k_i^2]=1$），因此：

$\text{Var}(q_i{k}_i)=1\times 1=1$

最终，点积 $\mathbf{q}\cdot \mathbf{k}$ 是 $d_k$ 个独立变量 $q_i{k}_i$ 之和，其方差为：

$\text{Var}(\mathbf{q}\cdot \mathbf{k})={\sum }_{i=1}^{d_k}\text{Var}(q_i{k}_i)=d_k\times 1=d_k$

1.2 方差爆炸的危害

当 $d_k$ 增大时（如 Transformer 中 $d_k=64$ 或更大），点积的方差会随 **** **** 线性增长，导致点积的数值范围急剧扩大（例如 $d_k=64$ 时，方差 = 64，3σ 范围为 $[-24,24]$）。

此时经过 softmax 函数（$\text{softmax}(x)=\frac{e^x}{\sum e^x}$）会出现严重问题：

• 若点积数值极大，$e^x$ 会指数级增大，导致该位置的概率接近 $1$，其他位置接近 $0$（分布 “尖锐”）；

• 若点积数值极小，$e^x$ 接近 $0$，概率也接近 $0$。

这种极端分布会导致 梯度消失（softmax 在概率接近 $0$ 或 $1$ 时梯度趋近于 $0$），模型无法学习到有效的注意力权重。

二、为什么除以 $\sqrt{d_k}$：方差归一化

我们的目标是将点积的方差归一化到常数（与 $d_k$ 无关），确保 softmax 输出分布平缓、梯度正常。

对於点积进行缩放：$\frac{\mathbf{q}\cdot \mathbf{k}}{\sqrt{d_k}}$，计算其方差：

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

代入 $\text{Var}(\mathbf{q}\cdot \mathbf{k})=d_k$，可得：

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

无论 $d_k$ 取何值，缩放后的点积方差始终为 $1$，数值范围稳定在 $[-3,3]$（3σ 范围）。此时 softmax 能输出合理的概率分布，既可以区分不同键的重要性，又不会导致梯度消失。

三、为什么不是 $d_k$ 或 $d_k^2$：方差过小导致注意力失效

若选择其他缩放因子，会破坏方差平衡，导致注意力机制失去作用。我们通过方差计算和具体例子验证：

3.1 若除以 $d_k$：方差随 $d_k$ 减小

缩放后的方差为：

$\text{Var}\left(\frac{\mathbf{q}\cdot \mathbf{k}}{d_k}\right)=\frac{1}{d_k^2}\times d_k=\frac{1}{d_k}$

当 $d_k$ 增大时，方差会随 **** **** 反比例减小（例如 $d_k=64$ 时，方差 = 1/64，3σ 范围为 $[-0.9375,0.9375]$）。

此时点积数值会集中在 $0$ 附近，softmax 后所有位置的概率接近 $\frac{1}{n}$（$n$ 为键的数量），形成 “均匀分布”—— 注意力无法突出重要键，完全失效。

3.2 若除以 $d_k^2$：方差急剧减小

缩放后的方差为：

$\text{Var}\left(\frac{\mathbf{q}\cdot \mathbf{k}}{d_k^2}\right)=\frac{1}{d_k^4}\times d_k=\frac{1}{d_k^3}$

当 $d_k=64$ 时，方差 = 1/(64³)≈3.8×10⁻⁵，3σ 范围仅为 $[-0.0058,0.0058]$。

点积数值几乎全部集中在 $0$，softmax 后概率完全均匀（如 3 个键的概率接近 $0.333$），注意力机制彻底失效。

四、数学例子：直观验证不同缩放因子的影响

假设 $d_k=16$（常见维度），$\mathbf{q}$ 和 $\mathbf{k}$ 的元素均服从 $\mathcal{N}(0,1)$（均值 0，方差 1），且有 3 个键（即 3 个点积结果）。

步骤 1：生成未缩放的点积

假设 3 个点积（未缩放）为：$[12,8,10]$（符合方差 = 16，数值范围 $[-12,12]$）。

步骤 2：不同缩放因子的结果对比

缩放因子	缩放后的点积	softmax 概率分布	分布特点
不缩放	$[12,8,10]$	$[0.867,0.016,0.117]$	极端尖锐
除以 $\sqrt{16}=4$	$[3,2,2.5]$	$[0.506,0.186,0.307]$	平缓且有区分度
除以 $16$	$[0.75,0.5,0.625]$	$[0.376,0.293,0.332]$	接近均匀
除以 $16^2=256$	$[0.0469,0.03125,0.0391]$	$[0.336,0.331,0.333]$	完全均匀

结论

• 不缩放：分布极端，梯度消失；

• 除以 $\sqrt{d_k}$：分布合理，注意力有效；

• 除以 $d_k$ 或 $d_k^2$：分布均匀，注意力失效。

总结

注意力机制除以 $\sqrt{d_k}$ 的本质是方差归一化：通过该操作将点积的方差固定为 1，避免因 $d_k$ 增大导致的数值极端化，确保 softmax 输出平缓、梯度正常。若选择 $d_k$ 或 $d_k^2$ 作为缩放因子，会导致方差过小，注意力失去区分重要信息的能力，最终失效。