注意力机制与 dk 由来

Softmax函数

$max$ 函数输出多个值中最大的那个, 而 $softmax$ 函数则将数值转化为概率分布.

$$Softmax(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{n}e^{x_j}}$$

如某一向量$\bm{a}=(0,2)$, 则 $Softmax(\bm{a})=(\frac{1}{1+e^2},\frac{e^2}{1+e^2})$. 其作用是让原本较小的值也能够被取到. 然而该函数用到了指数, 这会扩大原本的差距, 可能会出现梯度消失现象, 即某些较小的数值变得非常接近 $0$.

余弦相似度

想要介绍注意力机制, 需要从余弦相似度说起. 人们通常用余弦来评判两个向量的相似度, 因为向量间夹角越小则越相似, 相似度公式如下:

$$cos(\theta) = \frac{\bm{a} \cdot \bm{b}}{||\bm{a}||\times||\bm{b}||}$$

如果有两个矩阵 $Q=(\bm{q}_1,\bm{q}_2,\cdots \bm{q}_n)^\top \in \mathbb{R}^{n \times d}, K=(\bm{k}_1,\bm{k}_2,\cdots \bm{k}_n)^\top \in \mathbb{R}^{n \times d}$, 其中 $n$ 为向量数量, $d$ 为向量维数, 想要计算两个矩阵之间的相似度, 我们只需要把它们相乘:

$$S = QK^\top$$

这样, $S_{i,j}$ 就表示 $\bm{q}_i$ 和 $\bm{k}_j$ 之间的相似度. 为了消去数值的影响, 我们将其转化为概率分布:

$$S = Softmax(QK^\top)$$

然而我们还应该除以一个数, 显然矩阵是没有模长的, 我们考虑 $Q,K$ 矩阵中具体的值.

假设矩阵 $S$ 中的某个元素为 $s$, 则

$$s = \bm{q} \cdot \bm{k} = \sum_{i=1}^{d}\bm{q}_i \cdot \bm{k}_i$$

再假设 $q,k$ (未加粗)为向量 $\bm{q}, \bm{k}$ 中的对应元素. 我们可以将其看作一种随机变量, 显然这些随机变量是独立同分布的, 通过适当的数学方法, 我们能够将其转化为标准正态分布:

$$q, k \sim \mathcal{N}(0, 1)$$

因而它们的协方差为 $0$:

$$\begin{array}{ll} & Cov(q, k) \\ =& \mathbb{E}(q \cdot k) - \mathbb{E}(q)\mathbb{E}(k) \\ =& 0 \end{array}$$

所以能够推出 $\mathbb{E}(q \cdot k)$ 也为 $0$, 因而

$$\mathbb{E}(s) = \sum_{i=1}^{d}\mathbb{E}(\bm{q}_i \cdot \bm{k}_i) = 0$$

对于方差, 直接套用公式

$$\begin{array}{ll} & \mathbb{D}(q \cdot k) \\ =& \mathbb{E}(q^2 \cdot k^2) - \mathbb{E}^2(q \cdot k) \\ =& \mathbb{E}(q^2) \mathbb{E}(k^2) - 0 \\ =& [\mathbb{E}^2(q) + \mathbb{D}(q)][\mathbb{E}^2(k) + \mathbb{D}(k)] \\ =& 1 \end{array}$$

因此

$$\mathbb{D}(s) = \sum_{i=1}^{d}\mathbb{D}(\bm{q}_i \cdot \bm{k}_i) = d$$

从上面的推导中, 我们可以发现, $s$ 的方差为 $d$, 更大的方差意味着数值间差距更大, 则套用了 $Softmax$ 后的梯度消失现象会更明显, 故我们要消去这个方差, 显然只需要除以标准差就行了.

$$\mathbb{D}(\frac{s}{\sqrt{d}}) = 1$$

由于 $Q,K$ 的维数是相同的,所以原作者以 $K$ 为准, 最终相似度矩阵为:

$$S = Softmax(\frac{QK^\top}{\sqrt{d_k}})$$