qk相乘得到attention矩阵后, 为什么要进行scale

维度大小与点积结果的方差

对于Attention矩阵中的每个值, 都是由某个位置的$\mathbf{q}$向量和对应位置的$\mathbf{k}$向量点积得到. 假设$\mathbf{q}$向量和$\mathbf{k}$向量的各分量都是相互独立的随机变量, 且均值是0, 方差是1, 那么点积$\mathbf{q} \cdot \mathbf{k}$的均值为0, 方差为$d_k$, 其中$d_k$是向量的长度.

推导过程为: 对于$\forall i=1, \cdots, d_{k}$, $q_i$和$k_i$是独立的随机变量, 记这两个随便变量为: $X=q_i$, $Y=k_i$. 则有$E(X)=E(Y)=0$, $D(X)=D(Y)=1$.

则:

$$E(XY)=E(X)E(Y) = 0$$

$$\begin{aligned} D(X Y) &=E\left(X^{2} \cdot Y^{2}\right)-[E(X Y)]^{2} \\ &=E\left(X^{2}\right) E\left(Y^{2}\right)-[E(X) E(Y)]^{2} \\ &=E\left(X^{2}-0^{2}\right) E\left(Y^{2}-0^{2}\right)-[E(X) E(Y)]^{2} \\ &=E\left(X^{2}-[E(X)]^{2}\right) E\left(Y^{2}-[E(Y)]^{2}\right)-[E(X) E(Y)]^{2} \\ &=D(X) D(Y)-[E(X) E(Y)]^{2} \\ &=1 \times 1-(0 \times 0)^{2} \\ &=1 \end{aligned}$$

那么遍历$\forall i=1, \cdots, d_{k}$, $q_i$和$k_i$的乘积都符合均值为0, 方差为1, 而且不同$i$之间是相互独立的, 把$q_i \cdot k_i$这个随机变量记为$Z_i$. 则有:

$E\left(\sum_{i} Z_{i}\right)=\sum_{i} E\left(Z_{i}\right)$

$D\left(\sum_{i} Z_{i}\right)=\sum_{i} D\left(Z_{i}\right)$

因此$\mathbf{q} \cdot \mathbf{k}$点积的分布符合: $E(\mathbf{q} \cdot \mathbf{k})=0$, $D(\mathbf{q} \cdot \mathbf{k})=d_k$. 即维度越高, 点积结果分布的方差越大.

Attention矩阵的基础是对所有的$\mathbf{q}$和$\mathbf{k}$向量两两点积计算. 而方差越大, 点积得到的结果值较大的概率也会越大, 即矩阵中每个位置取值越大, 则Attention矩阵中行向量的模长也会越大.

向量模长与Softmax梯度

对于一个输入向量$\mathbf{x} \in \mathbb{R}^{d}$, softmax函数将其归一化到一个分布$\hat{\mathbf{y}} \in \mathbb{R}^{d}$下. softmax计算过程中, 使用自然底数$e$将输入原始的差距先拉大, 然后再归一化为一个分布.

假设输入$\mathbf{x}$向量中最大的元素对应的下标为$k$, 经过softmax转换后对应的$\hat{y}_k$对应的概率也是最大的. 但如果输入向量的模长增加, 即输入向量中各个元素等比例扩大, 在每个输入元素都很大的情况下, 这时的$\hat{y}_k$会非常接近1.

这时由于softmax的特性决定的, 对于向量$\mathbf{x}=[a, a, 2 a]^{\top}$, softmax得到的$\hat{y}_3$值随自变量$a$的变化如下图所示, 横坐标是$a$的值, 纵坐标是$\hat{y}_3$的值:

可以看到, 输入向量的模长对softmax得到的分布影响非常大. 在模长较长时, softmax归一化的结果, 几乎将所有概率都分配给输入最大值的位置, 其余位置的结果基本为0.

这会导致一个问题: 反向传播时, softmax的梯度几乎为0, 即发生了梯度消失, 造成参数更新困难. softmax函数的导数, 以及输出接近one-hot向量时导数的结果参考transformer中的attention为什么scaled? - TniL的回答.

解决方法

因此在Transformer结构中, 在向量两两点积获得点积结果矩阵后, 需要先将点积结果除以$\sqrt{d_k}$, 然后再对行向量进行softmax, 以避免梯度消失. 因为对点积结果除以$\sqrt{d_k}$, 会使得结果的分布方差回归到1:

$$D\left(\frac{q \cdot k}{\sqrt{d}_{k}}\right)=\frac{d_{k}}{\left(\sqrt{d}_{k}\right)^{2}}=1$$

消除了模长过大的问题, 进而消除了梯度消失的情况.

参考资料

• transformer中的attention为什么scaled? - TniL的回答

• transformer中的attention为什么scaled? - 小莲子的回答