小白也能看懂的self-attention教程

1. 引言

自从2017年论文《Attention Is All You Need》问世以来，基于Attention的Transformer网络在NLP领域和计算视觉CV领域都得到了广泛的应用。

本文旨在用简单易懂的语句来告诉大家self-attention的数学原理，并配以尽可能直观的图片来进行讲解说明，希望大家可以实现自己的自注意力模块！

闲话少说，我们直接开始吧！

2. 定义

在原始Transformer论文中，self-attention主要用于为单词生成新的向量表示，这些向量表示可以帮助机器理解句子中不同单词之间的中重要程度。它的输入为用单词嵌入向量（token embeddings）用来代替文本语句，token embedding实质上就是单词的多维的数字表示。而经过self-attention处理后，将这个输入序列转换成一个新的表示，同时可以保留原始单词输入语序，使其成为Transformer等强大人工智能模型的重要组成部分。

在心理学中，注意力是有选择地专注于一件或几件事情而忽略其他不太重要的事情的认知过程。神经网络被认为是以简化方式来模仿人脑的动作。注意力机制也是类似，通过它可以让神经网络有选择地专注于一些重要的输入，而忽略不太重要的输入。

我们来打个比方，假设我们正在查看一张幼儿园毕业照的合影。通常，会有好几排老师学生组成的队列，而老师一般坐在多个孩子的中间。现在，如果有人问问题，“一共有多少人？”，你会如何回答？

只需数人头，对吧？大家无需考虑照片中的任何其他内容。现在，如果有人问一个不同的问题，“照片中的老师是谁？”，你的大脑确切地知道该怎么做。它只会开始寻找照片中中间位置具有成年人特征的人，其余的孩子将被忽略。这就是我们的大脑非常擅长实现的“注意力”。

3. 准备输入

接着，我们通过一个具体的示例，来对自注意力机制进行讲解。假设我们有三个四维的token embedding作为self-attention的输入，换句话说就是我们有三个words作为输入，我们将每个word用一个四维的向量进行代替，如下所示：

Input 1: [1, 0, 1, 0] 
Input 2: [0, 2, 0, 2]
Input 3: [1, 1, 1, 1]

图示如下：

4. 初始化权重

接着，我们需要将每个input和三个权重矩阵分别相乘，进而得到我们的key(橙色)，query(红色)，value(紫色)。如下：

由于我们的input的shape为1×4，假设key、query、value的shape为1×3，因此可以推出与input相乘的三个权重矩阵的shape均为4×3。观察上图，每个输入（绿色）都乘以一组keys的权重、一组querys的权重和一组values的权重，然后得到对应的key(橙色)，query(红色)，value(紫色)。

在我们的例子中，我们将三组权重初始化如下:

# init weight for key
w_key=[[0, 0, 1],
 [1, 1, 0],
 [0, 1, 0],
 [1, 1, 0]]
 
# init weight for query
 w_query=[[1, 0, 1],
 [1, 0, 0],
 [0, 0, 1],
 [0, 1, 1]]

 # init weight for value
 w_value=[[0, 2, 0],
 [0, 3, 0],
 [1, 0, 3],
 [1, 1, 0]]

在神经网络初始化权重设置中，这些权重通常是很小的浮点数，一般使用随机分布（如Gaussian、Xavier）来随机初始化。权重初始化放在训练之前。

5. 计算key

经过以上准备工作，现在我们有了三组权重，让我们获得每个输入下的key、query和value的向量表示。

首先来针对input1计算key:

               [0, 0, 1]
[1, 0, 1, 0] x [1, 1, 0] = [0, 1, 1]
               [0, 1, 0]
               [1, 1, 0]

接着使用同样的weight来为输入input2计算对应的key:

               [0, 0, 1]
[0, 2, 0, 2] x [1, 1, 0] = [4, 4, 0]
               [0, 1, 0]
               [1, 1, 0] 

接着使用同样的weight来为输入input3计算对应的key:

               [0, 0, 1]
[1, 1, 1, 1] x [1, 1, 0] = [2, 3, 1]
               [0, 1, 0]
               [1, 1, 0]

更快的方法是对合并上述操作进行矩阵运算：

               [0, 0, 1]
[1, 0, 1, 0]   [1, 1, 0]   [0, 1, 1]
[0, 2, 0, 2] x [0, 1, 0] = [4, 4, 0]
[1, 1, 1, 1]   [1, 1, 0]   [2, 3, 1]

图示如下：

6. 计算value和query

我们使用同样的方法来获取每个输入对应的value，如下：

               [0, 2, 0]
[1, 0, 1, 0]   [0, 3, 0]   [1, 2, 3] 
[0, 2, 0, 2] x [1, 0, 3] = [2, 8, 0]
[1, 1, 1, 1]   [1, 1, 0]   [2, 6, 3]

图示如下：

对于query，也进行类似操作，如下：

               [1, 0, 1]
[1, 0, 1, 0]   [1, 0, 0]   [1, 0, 2]
[0, 2, 0, 2] x [0, 0, 1] = [2, 2, 2]
[1, 1, 1, 1]   [0, 1, 1]   [2, 1, 3]

图示如下：

7. 计算attention scores

为了获得注意力得分(attention scores)，我们用input1的query（红色）和input1，input2，input3的key（橙色）的转置分别计算点积。因此针对input1我们获得了3个注意力得分（蓝色）,分别表示input1与input1,intpu2,intpu3的注意力得分。

            [0, 4, 2]
[1, 0, 2] x [1, 4, 3] = [2, 4, 4]
            [1, 0, 1]

图示如下：

请注意，我们只使用input1中的query。我们针对其他输入input的query可以重复同样的步骤，不在累述。

8. 计算softmax

我们在得到了针对input1中query的attention scores后，我们可以使用softmax函数来执行归一化操作，使得各项相加后和为1，如下：

softmax([2, 4, 4]) = [0.0, 0.5, 0.5]

图示如下：

请注意，为了便于阅读，我们针对softmax输出进行四舍五入到小数点后1位。

9. 与values相乘

接着将每个输入的softmax之后的注意力得分（蓝色）乘以其相应的值（紫色）。这将产生3个alignment vectors（黄色）。在本文中，我们将把它们称为weighted values。如下：

1: 0.0 * [1, 2, 3] = [0.0, 0.0, 0.0]
2: 0.5 * [2, 8, 0] = [1.0, 4.0, 0.0]
3: 0.5 * [2, 6, 3] = [1.0, 3.0, 1.5]

图示如下：

10. 对weighted values求和得到output

在上一步骤中，我们针对input1的query(红色)产生了三个weighted values(黄色)，我们将这三个weighted values 按元素相加，得到我们的输出output1（深绿色），表示针对input1中的query与所有输入的key（包含input1）之间的atten score作为对value求和时的权重。

  [0.0, 0.0, 0.0]
+ [1.0, 4.0, 0.0]
+ [1.0, 3.0, 1.5]
-----------------
= [2.0, 7.0, 1.5]

图示如下：

11. 针对其他输入重复即可

现在我们已经完成了针对output1的计算可视化过程，那么对于output2和output3我们只需要重复上述步骤4到7即可。我相信大家可以自己动身计算咯！

注意：由于计算atten score时使用的时点积操作，所以query和value的维度必须始终相同。但是，value的维度可能与query和value不同。由此产生的输出将和value的维度保持一致。

12. 总结

本文重点介绍了自注意力机制的相关原理，transformer正是利用注意力机制的力量在自然语言处理行业取得前所未有的成果。希望大家可以通过本文，对自注意力机制有直观的理解！

您学废了嘛？

13. 参考

本文参考链接如下：

链接1：戳我
链接2：戳我