稳定Transformer中的注意力分数 | AI生成和翻译
在缩放点积注意力机制中（如公式 \(\text{缩放分数} = \frac{Q K^T}{\sqrt{d_k}}\) 所示），将点积结果除以键向量维度平方根 \(\sqrt{d_k}\) 是经过深思熟虑的设计，旨在稳定梯度并防止数值过大，特别是在Transformer架构中。以下详细说明为何选择 \(\sqrt{d_k}\) 而非直接使用 \(d_k\)：
- 点积方差特性：
- 点积运算 \( Q K^T \) 用于计算查询向量（\( Q \)）与键向量（\( K \)）的相似度，每个向量的维度均为 \( d_k \)。若假设 \( Q \) 和 \( K \) 的元素相互独立且均值为0、方差为1（常见于初始化或归一化后），则单对查询键向量的点积 \( Q_i \cdot K_j \) 的方差为 \( d_k )。这是因为 \( d_k \) 个独立标准正态变量乘积之和的方差与 \( d_k \) 呈线性关系。
- 若不进行缩放，\( Q K^T \) 的数值量级会随 \( d_k \) 增大而增长（Transformer中 \( d_k \) 通常为64、128或更大）。注意力分数过大会导致后续通过softmax函数时出现问题。
- Softmax稳定性：
- 缩放后的注意力分数 \( \frac{Q K^T}{\sqrt{d_k}} \) 将输入softmax函数以计算注意力权重。若分数过大（未缩放或缩放不足时），softmax会产生尖锐的分布——某个元素接近1，其余元素趋近0。这会导致大多数元素的梯度消失，影响模型学习效果。
- 除以 \(\sqrt{d_k}\) 可确保缩放后分数的方差约等于1，使分数保持在soft函数的理想工作区间，从而产生更均衡的注意力权重和稳定的梯度。
- 为何不选用 \( d_k \)：
- 若直接除以 \( d_k \) 而非 \(\sqrt{d_k}\)，会导致过度缩放，使分数方差降至 \( \frac{1}{d_k} \)。当 \( d_k \) 较大时，分数会过小，导致softmax输出接近均匀分布（因为softmax在输入值较小时会输出近似 \( \frac{1}{n} \) 的结果）。这将削弱注意力机制聚焦关键信息的能力，因为分数差异被过度压缩。
- 使用 \( d_k \) 进行过度缩放还可能引发数值稳定性问题，因为过小的数值在浮点运算中可能难以精确处理。
- 选择 \(\sqrt{d_k}\) 的数学原理：
- 除以 \(\sqrt{d_k}\) 可将点积方差归一化至约等于1：\( \text{Var}\left(\frac{Q K^T}{\sqrt{d_k}}\right) = \frac{\text{Var}(Q K^T)}{d_k} = \frac{d_k}{d_k} = 1 )。该操作使分数保持在合理范围内，确保softmax既能生成有意义的注意力权重，又不会过度尖锐或均匀。
- 平方根的选择源于点积方差的数学特性，在防止数值爆炸（无缩放时）与避免分数过度抑制（使用 \( d_k \) 时）之间取得平衡。
- 实证有效性：
- \(\sqrt{d_k}\) 的缩放策略由原始Transformer论文（Vaswani等人在2017年发表的《Attention is All You Need》）提出，并已在众多基于Transformer的模型中经过实证验证。该设计能适应不同 \( d_k \) 取值，确保训练稳定性与模型性能。
综上所述，除以 \(\sqrt{d_k}\) 是基于数学原理的设计决策，通过稳定注意力分数的方差，确保softmax函数处于理想工作状态，同时避免梯度爆炸或消失问题。若改为除以 \( d_k \)，则会过度压缩分数分布，导致注意力权重区分度下降，最终损害模型性能。