学习笔记|自注意力机制(self-attention)——考虑全局又聚焦重点

1. 自注意力机制概述

有时候我们期望网络能够看到全局，但是又要聚焦到重点信息上。比如在在做自然语言处理时，句子中的一个词往往不是独立的，和它上下文相关，但是和上下文中不同的词的相关性是不同的，所以我们在处理这个词时，在看到它的上下文的同时也要更加聚焦与它相关性更高的词，这就要用到常说的自注意力机制。比如下面这两幅图，通过自注意力机制处理后，计算出了词间的相关性，可以看到第一个图的it与animal的相关性很强，第二个图it与street的相关性很强。那么如何实现自注意力机制呢？

2. 向量相关性计算

自注意力机制的核心是捕捉向量之间的相关性。比如下面这幅图，输出一个向量$b^1$不只看$a^1$本身，还要看$a^2$、$a^3$、$a^4$，但是看它们的程度不一样。这就需要分别计算$a^1$与$a^2$、$a^3$、$a^4$之间的相关性$\alpha$,$\alpha$越大，相关性越高，给予的重视程度就越高。那么如何让网络自动计算出两个向量之间的相关性呢？

计算两个向量之间的相关性的常见方法是求点积（dot-product），如下图所示。具体的做法是左边的向量乘以一个变换矩阵$W^q$得到向量$q$，右边的向量乘以一个变换矩阵$W^k$得到向量$k$，然后将向量$q$和向量$k$点积就可以得到相关性$\alpha$。由点积的性质可知，两个向量的相似度越高，点积的值就会越大。当然，计算向量相关性的方法不只点积这一种，也有其他方式，但是点积这种是最常见的。

基于点积计算，我们就可以向量两两之间的关联性了，比如首先分别计算$a^1$与$a^2$、$a^3$、$a^4$之间的相关性。我们首先将$a^1$乘以变换矩阵$W^q$得到向量$q^1$，这里的$q^1$向量有个专门的名字，叫做 “query” 。然后将$a^1$、$a^2$、$a^3$、$a^4$分别乘以变换矩阵$W^k$得到向量$k^1,k^2,k^3,k^4$，这里的$k^i$向量也有个专门的名字，叫做 “key”。然后将$q^1$和这四个key分别做点积，就得到四个相关性数值${\alpha }_{1,1},{\alpha }_{1,2},{\alpha }_{1,3},{\alpha }_{1,4}$。求出这四个相关性的值后，然后通过一个Soft-max层进行归一化，得到${\alpha }_{1,1}^{{}^{\prime }},{\alpha }_{1,2}^{{}^{\prime }},{\alpha }_{1,3}^{{}^{\prime }},{\alpha }_{1,4}^{{}^{\prime }}$，这是最后输出的相关性值，我们将这些值又称为**“注意力分数”**。现在我们得到了$a^1$对$a^1$、$a^2$、$a^3$、$a^4$之间的注意力分数，那么如何做到考虑全局又聚焦重点呢？

3. 基于注意力分数抽取向量信息

通过上面计算出的注意力分数${\alpha }_{1,1}^{{}^{\prime }},{\alpha }_{1,2}^{{}^{\prime }},{\alpha }_{1,3}^{{}^{\prime }},{\alpha }_{1,4}^{{}^{\prime }}$，我们已经知道$a^1$要给予$a^1$、$a^2$、$a^3$、$a^4$的关注程度了，接下来我们抽取这些向量中重要的信息以输出$b^1$了。具体的做法如下图所示。首先我们再将$a^1$、$a^2$、$a^3$、$a^4$乘以一个新的变换矩阵$W^v$得到向量$v^1,v^2,v^3,v^4$，这里的$v^i$向量也有个专门的名字，叫做 “value”。然后将向量$v^1,v^2,v^3,v^4$分别乘以对应的注意力分数${\alpha }_{1,1}^{{}^{\prime }},{\alpha }_{1,2}^{{}^{\prime }},{\alpha }_{1,3}^{{}^{\prime }},{\alpha }_{1,4}^{{}^{\prime }}$，并进行求和，输出向量$b^1$。从这里可以看出，所有向量都有参与计算，这样就做到了看全局。但是各向量参与计算的程度不一样，${\alpha }_{1,i}^{{}^{\prime }}$就相当权重值，权重值越大的，对应向量参与计算的程度就越大，最后得到的输出向量$b^1$就和该向量越相似。这样就做到了看全局又聚焦重点。通过上述同样的计算方式，也可以计算得到$b^2,b^3,b^4$，而且$b^1,b^2,b^3,b^4$是可以并行计算的。以上就是自注意力机制的全部了，但是对自注意力机制的解析并没有结束，下面从矩阵计算的角度来看自注意力机制。

4. 自注意力机制中的矩阵计算

4.1 计算矩阵$K、V、Q$

前面提到将$a^1$、$a^2$、$a^3$、$a^4$分别乘以变换矩阵$W^k$得到向量$k^1,k^2,k^3,k^4$。我们将输入向量$a^1$、$a^2$、$a^3$、$a^4$拼在一起，得到一个矩阵用$I$表示，即，$$I=[a^1{a}^2{a}^3{a}^4]$$将key向量$k^1$、$k^2$、$k^3$、$k^4$拼在一起得到一个矩阵用$K$表示，即，$$K=[k^1{k}^2{k}^3{k}^4]$$用矩阵相乘表示$K$矩阵的计算过程即，$$K=W^{k}I$$同理，query向量拼成的矩阵$Q$等于，$$Q=W^{q}I$$value向量拼成的矩阵$V$等于,$$V=W^{v}I$$。下图展示了上述计算过程。

4.2 计算注意力分数矩阵$A^{\prime }$

前面提到将$q^1$和四个key向量$k^1,k^2,k^3,k^4$分别做点积，得到四个相关性数值${\alpha }_{1,1},{\alpha }_{1,2},{\alpha }_{1,3},{\alpha }_{1,4}$。注意这里的向量都是列向量，所以点积可以写成，$${\alpha }_{1,1}=q^1\cdot k^1=(k^1{)}^T{q}^1$$ $${\alpha }_{1,2}=q^1\cdot k^2=(k^2{)}^T{q}^1$$ $${\alpha }_{1,3}=q^1\cdot k^3=(k^3{)}^T{q}^1$$ $${\alpha }_{1,4}=q^1\cdot k^4=(k^4{)}^T{q}^1$$
用矩阵计算表示上述计算过程为 $$\left[\begin{array}{c}{\alpha }_{1,1}\\ {\alpha }_{1,2}\\ {\alpha }_{1,3}\\ {\alpha }_{1,4}\end{array}\right]=\left[\begin{array}{c}(k^1{)}^T\\ (k^2{)}^T\\ (k^3{)}^T\\ (k^4{)}^T\end{array}\right]q^1=K^T{q}^1$$将$K^T$与$q^2、q^3、q^4$相乘可以得到相似的结果，即，$$A=\left[\begin{array}{cccc}{\alpha }_{1,1}& {\alpha }_{2,1}& {\alpha }_{3,1}& {\alpha }_{4,1}\\ {\alpha }_{1,2}& {\alpha }_{2,2}& {\alpha }_{3,2}& {\alpha }_{4,2}\\ {\alpha }_{1,3}& {\alpha }_{2,3}& {\alpha }_{3,3}& {\alpha }_{4,3}\\ {\alpha }_{1,4}& {\alpha }_{2,4}& {\alpha }_{3,4}& {\alpha }_{4,4}\end{array}\right]=\left[\begin{array}{c}(k^1{)}^T\\ (k^2{)}^T\\ (k^3{)}^T\\ (k^4{)}^T\end{array}\right][q^1{q}^2{q}^3{q}^4]=K^{T}Q$$ $A$矩阵通过softmax层归一化后得到$A^{\prime }$ 。上述计算过程如下图所示。

4.3 计算输出矩阵$O$

前面讲到将向量$v^1,v^2,v^3,v^4$分别乘以对应的注意力分数${\alpha }_{1,1}^{{}^{\prime }},{\alpha }_{1,2}^{{}^{\prime }},{\alpha }_{1,3}^{{}^{\prime }},{\alpha }_{1,4}^{{}^{\prime }}$，并进行求和，输出向量$b^1$，这个过程用矩阵计算可表示为，$$b^1=[v^1{v}^2{v}^3{v}^4]\left[\begin{array}{c}{\alpha }_{1,1}^{{}^{\prime }}\\ {\alpha }_{1,2}^{{}^{\prime }}\\ {\alpha }_{1,3}^{{}^{\prime }}\\ {\alpha }_{1,4}^{{}^{\prime }}\end{array}\right]=V\left[\begin{array}{c}{\alpha }_{1,1}^{{}^{\prime }}\\ {\alpha }_{1,2}^{{}^{\prime }}\\ {\alpha }_{1,3}^{{}^{\prime }}\\ {\alpha }_{1,4}^{{}^{\prime }}\end{array}\right]$$
通过相似的计算，也可以得到$b^2、b^3、b^4$,即，$$O=[b^1{b}^2{b}^3{b}^4]=[v^1{v}^2{v}^3{v}^4]\left[\begin{array}{cccc}{\alpha }_{1,1}^{{}^{\prime }}& {\alpha }_{2,1}^{{}^{\prime }}& {\alpha }_{3,1}^{{}^{\prime }}& {\alpha }_{4,1}^{{}^{\prime }}\\ {\alpha }_{1,2}^{{}^{\prime }}& {\alpha }_{2,2}^{{}^{\prime }}& {\alpha }_{3,2}^{{}^{\prime }}& {\alpha }_{4,2}^{{}^{\prime }}\\ {\alpha }_{1,3}^{{}^{\prime }}& {\alpha }_{2,3}^{{}^{\prime }}& {\alpha }_{3,3}^{{}^{\prime }}& {\alpha }_{4,3}^{{}^{\prime }}\\ {\alpha }_{1,4}^{{}^{\prime }}& {\alpha }_{2,4}^{{}^{\prime }}& {\alpha }_{3,4}^{{}^{\prime }}& {\alpha }_{4,4}^{{}^{\prime }}\end{array}\right]=V{A}^{\prime }$$

4.4 自注意力矩阵计算总结

综上，自注意力机制的计算过程可总结为,
（1）计算$Q、K、V$矩阵$$Q=W^{q}I$$ $$K=W^{k}I$$ $$V=W^{v}I$$
（2）计算注意力分数矩阵$A^{\prime }$$$A=K^{T}Q$$ $$A^{\prime }=\text{softmax}(A)$$
（3）计算输出矩阵$O$$$O=V{A}^{\prime }$$

可以看出，自注意力机制看起来比较复杂，其实计算过程并不复杂，需要学习的参数只有$W^q、W^k、W^v$。

5. 多头自注意力机制

自注意力机制还有一个进阶版，叫多头自注意力机制（multi-head self-attention）。为什么要多头呢？自注意力机制实质上是用过$q$向量去找相关的$k$向量，但是相关性可能有多种，一个$q$只能找到一种相关的$k$向量，因此就要引入多个$q$向量和$k$向量来捕捉多种相关性。多头自注意力机制很简单，设置多组矩阵$W^{q,i}、W^{k,i}、W^{v,i}$，每一组$W^{q,i}、W^{k,i}、W^{v,i}$只进行内部计算，得到相应的输出$O^i$，如下图所示。

在得到不同的输出$O^i$后，再将其拼到一起，形成一个大的矩阵。如果是2头，就将这2个输出直接拼到一起。然后通过一个转换矩阵$W^o$将拼接的矩阵转换成原输出的长度的向量，即，$$O=W^o\left[\begin{array}{c}{O}^1\\ O^2\end{array}\right]$$
因此，多头注意力机制要多一个参数矩阵，即$W^o$。