[NLP]——Transformer中的attention为什么要做scale？

前言

说起Transformer的self-attention，很容易想到下面的公式：
$$Attention(Q,K,V)=softmax(Q{K}^T/\sqrt{d_k})V$$

假设X是输入，分别经过$W_Q$、$W_K$和$W_V$映射得到$Q$、$K$和$V$，【$d_k$=Q.size(-1)=K.size(-1)】，通过$Q$和$K$的dot-product来计算$weight$，经过softmax得到归一化后的$weight$，再使用这个$weight$去对$V$做加权求和。那么，公式$Q{K}^T/\sqrt{d_k}$中的$\sqrt{d_k}$是用来干嘛的呢？

一句话概括就是：如果不对softmax的输入做缩放，那么万一输入的数量级很大，softmax的梯度就会趋向于0，导致梯度消失。

本文的思路如下:

• softmax是怎样求导的？
• 通过对softmax求导，我们可以知道，softmax的输入的数量级越大，求导的梯度越会趋向于0
• 如何控制softmax的输入的数量级（也就是$\sqrt{d_k}$的作用）

softmax的求导过程

这部分主要参考:详解softmax函数以及相关求导过程 - 忆臻的文章 - 知乎
https://zhuanlan.zhihu.com/p/25723112

假设softmax的输入是$x=(x_1,x_2,...x_n)$，则$y=softmax(x)=(y_1,y_2,...y_n)$，其中，$y_i$ = $e^{x_i}/{\sum }_{k=1..n}{e}^{x_k}$。求导过程如下（打公式好累，有个错别字“妨”）：

为softmax的输入的数量级越大，求导的梯度越会趋向于0

这部分主要参考transformer中的attention为什么scaled? - TniL的回答 - 知乎
https://www.zhihu.com/question/339723385/answer/782509914

首先，对于输入的$x=(x_1,x_2,...x_n)$，softmax中的max体现在：通过一个自然底数e来将输入中最大的元素更大，softmax中的soft体现在：不忽略输入中的最小的元素，依然给它们一定的权重。总之，softmax函数能够将输入中的元素间差距拉大，然后归一化为一个分布。

假设输入的$x=(x_1,x_2,...x_n)$中最大元素为$x_k$，其对应的概率输出为$y_k$，将会呈现“x的数量级越大，$y_k$越趋向于1”的趋势，具体举例如下:

在这种情况下，如果输入$x$的数量级很大，而假设它的最大值是$x_0$，则经过softmax计算得到的$y$中，只有$y_0$趋向于1，其它概率元素全都趋向于0。进一步结合第一部分的求导结果，将会出现下图的情况：

即，softmax的梯度趋向于0

如何控制softmax的输入的数量级

那么如何控制softmax的输入，也就是$x=(x_1,x_2,...x_n)$的数量级呢？首先明确一点，在Transformer中，如果没有$\sqrt{d_k}$，则，softmax的输入是$Q{K}^T$。宏观的，我们需要保证$Q{K}^T$，也就是一个(batch_size x)?sent_num x sent_n的矩阵，其中的每一个元素的数量级都不要很大。那么取出其中一个元素，它由$q\cdot k$计算得来，接下来证明$q\cdot k$的数量级与$\sqrt{d_k}$的关系。

下面的推导设计到概率论与数理统计的知识，还没有复习到相关内容，所以依然先参考：transformer中的attention为什么scaled? - TniL的回答 - 知乎
https://www.zhihu.com/question/339723385/answer/782509914

经过上述推导，得到结论:

参考

transformer中的attention为什么scaled? - TniL的回答 - 知乎
https://www.zhihu.com/question/339723385/answer/782509914

详解softmax函数以及相关求导过程 - 忆臻的文章 - 知乎
https://zhuanlan.zhihu.com/p/25723112