Transformer 中缩放点积注意力机制探讨：除以根号 dk 理由及其影响

`Transformer 中缩放点积注意力机制的探讨`

1. 引言

自2017年Transformer模型被提出以来，它迅速成为自然语言处理（NLP）领域的主流架构，并在各种任务中取得了卓越的表现。其核心组件之一是注意力机制，尤其是缩放点积注意力（Scaled Dot-Product Attention）。本文将深入探讨为什么在计算注意力分数时要除以 $\sqrt{d_k}$ ，以及如果使用不同的缩放因子会带来什么后果。

在这里插入图片描述

2. 缩放点积注意力机制简介

缩放点积注意力机制是一种高效的注意力计算方法，它通过计算查询向量（Query）、键向量（Key）和值向量（Value）之间的点积来衡量相关性。公式如下：

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

其中， $Q$ 是查询矩阵， $K$ 是键矩阵， $V$ 是值矩阵，而 $d_k$ 是键向量的维度。分母中的 $\sqrt{d_k}$ 是一个关键的缩放因子。

3. 为什么要除以 $\sqrt{d_k}$

3.1 防止激活函数饱和

当 $d_k$ 较大时，查询向量和键向量之间的点积可能会变得非常大，导致softmax函数的输入值过大，输出接近于0或1，使得梯度变得很小。通过除以 $\sqrt{d_k}$ ，可以将点积的结果缩小到一个合理的范围，避免激活函数进入饱和区。

3.2 保持数值稳定性

高维空间中，点积的结果会迅速增大，可能导致指数运算中的溢出或下溢问题。通过除以 $\sqrt{d_k}$ ，可以在一定程度上缓解这些问题，维持数值的稳定性。

3.3 维度不变性

较小的 $d_k$ 可能会导致点积值较低，而较大的 $d_k$ 则可能导致点积值较高。通过引入 $\sqrt{d_k}$ 的缩放因子，可以使不同维度下的点积具有相似的尺度，保证了模型性能的一致性和可预测性。

3.4 理论依据

根据统计理论，当两个随机向量是独立同分布（i.i.d.）的标准正态分布时，它们的点积期望值为0，方差为 $d_k$ 。为了使点积的方差保持恒定，需要对点积结果进行 $\sqrt{d_k}$ 的缩放，确保随着 $d_k$ 的增加，点积不会线性增长，而是维持在一个相对稳定的水平。

4. 不同缩放因子的影响

4.1 如果除以的数比 $\sqrt{d_k}$ 大

4.1.1 梯度消失的情况

思考过程：
- 当除以的数大于 $\sqrt{d_k}$ 时，实际上是在进一步缩小点积的结果。这意味着输入到softmax函数的值会变得更小。
- 对于softmax函数来说，当输入值变得非常小的时候，输出的概率分布将趋于均匀分布，即每个位置的概率都接近 $\frac{1}{n}$ ，其中 $n$ 是候选的数量。
- 在这种情况下，即使某些键向量与查询向量之间存在较强的相关性，它们在最终的注意力权重中也不会显著突出，导致模型难以学习到有效的特征表示。
- 由于softmax后的梯度是基于概率分布计算的，过于平滑的概率分布会导致梯度变得很小，进而引起梯度消失的问题。
结果：
- 如果除以的数过大，可能会导致梯度消失，使得模型难以收敛或收敛速度极慢，影响训练效果。

4.2 如果除以的数比 $\sqrt{d_k}$ 小

4.2.1 梯度爆炸的情况

思考过程：
- 当除以的数小于 $\sqrt{d_k}$ 时，实际上是放大了点积的结果。这会导致输入到softmax函数的值变得更大。
- 对于softmax函数而言，当输入值非常大的时候，少数几个最大值对应的输出概率会趋近于1，而其他位置的概率则趋近于0。
- 这种极端的概率分布会导致模型对这些最大值对应的位置产生过强的依赖，忽略了其他潜在的重要信息。
- 此外，softmax函数的导数在输入值很大的情况下也会变得非常小（对于非最大值），但对于最大值处的梯度却可能变得很大，这容易引发梯度爆炸现象，特别是在反向传播过程中。
结果：
- 如果除以的数过小，可能会导致梯度爆炸，使得模型训练不稳定，参数更新幅度剧烈，甚至可能导致数值溢出或模型无法正常训练。