全网最独特解析:self Attention为何除根号dk?
一、假设条件:查询向量和键向量服从正态分布
假设查询向量
q
i
q_i
qi和键向量
k
j
k_j
kj的每个分量均为独立同分布的随机变量,且服从标准正态分布,即:
q
i
(
m
)
,
k
j
(
m
)
∼
N
(
0
,
1
)
(
m
=
1
,
2
,
…
,
d
k
)
q_i^{(m)}, k_j^{(m)} \sim \mathcal{N}(0,1) \quad (m=1,2,\dots,d_k)
qi(m),kj(m)∼N(0,1)(m=1,2,…,dk)
此时,每个分量的均值为0,方差为1。
二、点积的统计特性分析
查询向量
q
i
q_i
qi和键向量
k
j
k_j
kj的点积为:
q
i
⋅
k
j
=
∑
m
=
1
d
k
q
i
(
m
)
k
j
(
m
)
q_i \cdot k_j = \sum_{m=1}^{d_k} q_i^{(m)} k_j^{(m)}
qi⋅kj=m=1∑dkqi(m)kj(m)
根据独立随机变量和的方差性质,点积的方差为:
Var
(
q
i
⋅
k
j
)
=
∑
m
=
1
d
k
Var
(
q
i
(
m
)
k
j
(
m
)
)
\text{Var}(q_i \cdot k_j) = \sum_{m=1}^{d_k} \text{Var}(q_i^{(m)} k_j^{(m)})
Var(qi⋅kj)=m=1∑dkVar(qi(m)kj(m))
由于
q
i
(
m
)
q_i^{(m)}
qi(m)和
k
j
(
m
)
k_j^{(m)}
kj(m)独立且均服从
N
(
0
,
1
)
\mathcal{N}(0,1)
N(0,1),乘积的方差为:
Var
(
q
i
(
m
)
k
j
(
m
)
)
=
Var
(
q
i
(
m
)
)
⋅
Var
(
k
j
(
m
)
)
+
[
E
(
q
i
(
m
)
)
]
2
⋅
Var
(
k
j
(
m
)
)
+
[
E
(
k
j
(
m
)
)
]
2
⋅
Var
(
q
i
(
m
)
)
=
1
\text{Var}(q_i^{(m)} k_j^{(m)}) = \text{Var}(q_i^{(m)}) \cdot \text{Var}(k_j^{(m)}) + [E(q_i^{(m)})]^2 \cdot \text{Var}(k_j^{(m)}) + [E(k_j^{(m)})]^2 \cdot \text{Var}(q_i^{(m)}) = 1
Var(qi(m)kj(m))=Var(qi(m))⋅Var(kj(m))+[E(qi(m))]2⋅Var(kj(m))+[E(kj(m))]2⋅Var(qi(m))=1
因此,点积的方差为
d
k
d_k
dk,标准差为
d
k
\sqrt{d_k}
dk。
三、缩放的必要性:Softmax的输入敏感性
Softmax函数对输入值的量级极其敏感:
- 数值溢出问题:若点积的绝对值随 d k d_k dk增大而显著增大(例如 d k = 64 d_k=64 dk=64时标准差为8),输入Softmax的值可能超出浮点数表示范围。
- 梯度消失问题:当某些点积值远大于其他值时,Softmax输出接近独热分布(Hard Attention),导致梯度趋近于零,阻碍参数更新。
- 分布退化问题:未经缩放的输入会使注意力权重集中在极少数位置,失去“软性关注”的优势。
四、除以 d k \sqrt{d_k} dk的数学解释
通过将点积除以
d
k
\sqrt{d_k}
dk,可以将点积的标准差从
d
k
\sqrt{d_k}
dk缩放至1,即:
Var
(
q
i
⋅
k
j
d
k
)
=
Var
(
q
i
⋅
k
j
)
d
k
=
1
\text{Var}\left( \frac{q_i \cdot k_j}{\sqrt{d_k}} \right) = \frac{\text{Var}(q_i \cdot k_j)}{d_k} = 1
Var(dkqi⋅kj)=dkVar(qi⋅kj)=1
此时,点积的分布被标准化为
N
(
0
,
1
)
\mathcal{N}(0,1)
N(0,1),实现了以下效果:
- 数值稳定性:Softmax输入的均值为0、方差为1,避免极端值。
- 梯度均衡性:Softmax输出的概率分布更平缓,梯度更新更稳定。
- 模型鲁棒性:注意力权重在多位置间合理分配,保留软性关注能力。
五、为何不是其他缩放因子?
若采用其他缩放因子(如
d
k
d_k
dk或
2
d
k
2\sqrt{d_k}
2dk):
• 除以
d
k
d_k
dk:方差将缩小为
1
/
d
k
1/d_k
1/dk,导致Softmax输入过小,注意力权重趋于均匀分布,失去区分性。
• 除以
2
d
k
2\sqrt{d_k}
2dk:方差将缩小为
1
/
4
1/4
1/4,输入量级过小,同样影响注意力权重的有效性。
总结
从正态分布的角度看,除以 d k \sqrt{d_k} dk的本质是通过方差归一化,将点积的统计特性控制在合理范围内。