Attention中softmax的梯度消失及scaled原因

    在bert模型中的attention构造中

Q：[batch, 12, seq, dk]
K：[batch, 12, seq, dk]

softmax中的梯度消失

     $x=(a,a,2a,4a)$, $a>0$，softmax(x)更多的将权重分配到大数上面，$softmax(4a)$远远大于$softmax(a)$。基于这个原因，很大的输入数据，会使反传时梯度消失。
$$y=softmax(x)\in R^d$$
$$y=(y_1,y_2,...,y_d)$$
计算导数，由于有softmax函数，y关于x的导数的一般计算结果如下
$$\begin{gathered} \frac{\partial y_i}{\partial x_j}=-y_i{y}_j(i!=j) \\ \frac{\partial y_i}{\partial x_j}=y_i(1-y_i)(i==j) \end{gathered}$$
可以将上式写成矩阵形式，如下
$$\frac{\partial y}{\partial x}=diag(y)-y{y}^T$$
$$=\left\{\begin{array}{ccccc}{y}_1& 0& 0& \cdots & 0\\ 0& y_2& 0& \cdots & 0\\ ⋮& ⋮& ⋮& \cdots & ⋮\\ 0& 0& 0& \cdots & y_d\end{array}\right\}-\left\{\begin{array}{ccccc}{y}_1^2& y_1{y}_2& y_1{y}_3& \cdots & y_1{y}_d\\ y_2{y}_1& y_2^2& y_2{y}_3& \cdots & y_2{y}_d\\ ⋮& ⋮& ⋮& \cdots & ⋮\\ y_d{y}_1& y_d{y}_2& y_d{y}_3& ...& y_d^2\end{array}\right\}$$
当x中某个值非常大时，softmax会把权重更多放到该标签上，生成更接近1的数。y更接近于one-hot变量。举个例子，a非常大时，$x=(a,a,2a,4a)$，经softmax后，$y=(0,0,0,1{)}^T$，即a所处位置的数会非常小，趋于0。此时
$$\frac{\partial y}{\partial x}=\left\{\begin{array}{cccc}0& 0& 0& 0\\ 0& 0& 0& 0\\ 0& 0& 0& 0\\ 0& 0& 0& 1\end{array}\right\}-\left\{\begin{array}{cccc}0& 0& 0& 0\\ 0& 0& 0& 0\\ 0& 0& 0& 0\\ 0& 0& 0& 1\end{array}\right\}=0$$

增加scaled的原因

$Q=(q_1,q_2,...,q_{d_k})$,$q_i\in R^{[seq,1]}$，$K$类似。可以将$q_i$看做随机变量。假设$q_i$与$k_i$都服从标准正态分布$N(0,{\sigma }^2)$且独立，下面求期望与方差。
$$E(q_i{k}_i)=E{q}_{i}E{k}_i=0$$
$$\begin{array}{rl}D(q_i{k}_i)& =E(q_i^2{k}_i^2)-(E(q_i{k}_i){)}^2\\ & =E{q}_i^{2}E{k}_i^2\\ & =(D(q_i)+(E{q}_i{)}^2)(D(k_i)+(E{k}_i{)}^2)\\ & ={\sigma }^4=1\end{array}$$
那么有
$$E(Q{K}^T)=\sum _{i=1}^{d_k}E(q_i{k}_i)=0$$
$$D(Q{K}^T)=\sum _{i=1}^{d_k}D(q_i{k}_i)=d_k{\sigma }^2=d_k$$
为保证softmax后的方差不变，所以对方差进行缩放。

矩阵的期望与方差

矩阵本没有期望与方差的定义，查看上面步骤，$Q{K}^T$的维度为[batch, 32, seq, seq]，只考虑后两个维度，$Q{K}^T$为方阵，$Q$维度为[seq,dk]。计算期望与方差时 ，可以按照行，也可以按照列计算，看最后的实际用处。