PyTorch——自注意力（self-attention）机制实现（代码详解）

参考链接

1. https://www.bilibili.com/video/BV1JE411g7XF?p=54
2. https://arxiv.org/abs/1706.03762
3. https://blog.csdn.net/qq_36653505/article/details/83375160

简述自注意力机制（self-attention）

self-attention可以视为一个特征提取层，给定输入特征$a^1,a^2,\cdot \cdot \cdot a^n$，经过self-attention layer，融合每个输入特征，得到新的特征$b^1,b^2,\cdot \cdot \cdot b^n$。具体如下：

设输入特征为$I$，分别将其乘以三个矩阵$W^q$、$W^k$和$W^v$得到$Q$（query）、$K$（key）和$V$（value）三个矩阵；接下来使用矩阵$Q$和$K$的乘积得到注意力矩阵$A$，归一化得到$\hat{A}$；最后，将归一化后的注意力矩阵$\hat{A}$乘上$V$，得到最后的输出特征$O$。

多头自注意力机制（multi-head self-attention）

上述的self-attention中，每个输入特征$a^i$乘上矩阵$W^q$、$W^k$和$W^v$后，分别得到一个向量$q^i$、$k^i$和$v^i$，称为单头自注意力机制。如果将这些向量