注意力机制 (Attention)：是人工神经网络中一种模仿认知注意力的技术。这种机制可以增强神经网络输入数据中某些部分的权重，同时减弱其他部分的权重，以此将网络的关注点聚焦于数据中最重要的一小部分。数据中哪些部分比其他部分更重要取决于上下文。

本篇文章为论文 Attention Is All You Need 的笔记，因此本文的注意力机制架构遵照原论文。

1 总览

首先，我们通过非常简略的数学方式来总览一下注意力机制的操作过程，感受一下总体框架后，下文再进行详细展开。

对于一段文字：

$$ \text{The quick brown fox jumps over a lazy dog.} $$

首先将其切分为 Token（假设按词切分）：

$$ \text{The| quick| brown| fox| jumps| over| a| lazy| dog.} $$

然后通过 Embedding 将单词转为 $d_{\text{model}}$ 维的词向量：

$$ \vec{\mathrm{E}_1},\vec{\mathrm{E}_2},\dots,\vec{\mathrm{E}_9} $$

然后给词向量乘上 Query 矩阵 $\boldsymbol{W}_Q$，生成每个单词的 $d_q$ 维的查询向量 $\vec{\mathrm{Q}_i}=\boldsymbol{W}_Q\vec{\mathrm{E}_i}$：

$$ \vec{\mathrm{Q}_1},\vec{\mathrm{Q}_2},\dots,\vec{\mathrm{Q}_9} $$

然后给词向量乘上 Key 矩阵 $\boldsymbol{W}_K$，生成每个单词的 $d_k$ 维的键向量 $\vec{\mathrm{K}_i}=\boldsymbol{W}_K\vec{\mathrm{E}_i}$：

$$ \vec{\mathrm{K}_1},\vec{\mathrm{K}_2},\dots,\vec{\mathrm{K}_9} $$

然后将查询向量和键向量一一做向量点积，得到一个实数，这些值组成一个新的矩阵：

$$ \begin{bmatrix} \vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_1}&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_1}&\dots&\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_1}\\ \vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_2}&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_2}&\dots&\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_2}\\ \vdots&\vdots&\ddots&\vdots\\ \vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_9}&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_9}&\dots&\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_9} \end{bmatrix} $$

然后对这些值进行列 SoftMax 操作：

$$ \begin{bmatrix} \text{Softmax}(\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_1}) & \text{Softmax}(\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_1}) & \dots & \text{Softmax}(\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_1}) \\ \text{Softmax}(\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_2}) & \text{Softmax}(\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_2}) & \dots & \text{Softmax}(\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_2}) \\ \vdots & \vdots & \ddots & \vdots \\ \text{Softmax}(\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_9}) & \text{Softmax}(\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_9}) & \dots & \text{Softmax}(\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_9}) \end{bmatrix} $$

然后给每一项乘上 Value 矩阵 $\boldsymbol{W}_V$：

$$ M=\begin{bmatrix} \text{Softmax}(\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_1}) \cdot \boldsymbol{W}_V & \text{Softmax}(\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_1}) \cdot \boldsymbol{W}_V & \dots & \text{Softmax}(\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_1}) \cdot \boldsymbol{W}_V \\ \text{Softmax}(\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_2}) \cdot \boldsymbol{W}_V & \text{Softmax}(\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_2}) \cdot \boldsymbol{W}_V & \dots & \text{Softmax}(\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_2}) \cdot \boldsymbol{W}_V \\ \vdots & \vdots & \ddots & \vdots \\ \text{Softmax}(\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_9}) \cdot \boldsymbol{W}_V & \text{Softmax}(\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_9}) \cdot \boldsymbol{W}_V & \dots & \text{Softmax}(\vec{\mathrm{Q}_9}\cdot\vec{\mathrm{K}_9}) \cdot \boldsymbol{W}_V \end{bmatrix} $$

然后计算每列的更新值 $\Delta\vec{\mathrm{E}_i}$：

$$ \Delta\vec{\mathrm{E}_i}=\sum_{k=1}^{n}M_{ki}\vec{\mathrm{E}_i} $$

这便是原始词向量的更新值，也能称为 Attention：

$$ \vec{\mathrm{E}_i}'=\vec{\mathrm{E}_i}+\Delta\vec{\mathrm{E}_i} $$

上述过程的结构可以以下图方式展现：

2 Query 查询

从 Token 嵌入得来的词向量是没有上下文语意信息的，例如 Transformer 到底指变形金刚、变压器还是一种机器学习模型，词向量中是没有相关信息的。为了获取这个 Token 具体的含义，就得从上下文 Token 中获取具体的信息。同时，有些上下文与该 Token 关联不大，比如 The，但有些词语对该 Token 关联非常大，例如 Model. 注意力机制便是让模型自己学会上下文之间的关联程度。

获取上下文关联程度的第一步是是求得查询向量，可以将查询向量理解为该 Token 对上下文提出的“问题”，接下来根据上下文 Token 的”回应“来确认他们之间的关联程度。

对于每个 Token 的词向量 $\vec{\mathrm{E}_i}$，让它与矩阵 $\boldsymbol{W}_Q$ 相乘获得对应的查询向量 $\vec{\mathrm{Q}_i}=\boldsymbol{W}_Q\vec{\mathrm{E}_i}$.

其中，$\vec{\mathrm{E}_i}$ 的维度为 $d_{\text{model}}$，$\vec{\mathrm{Q}_i}$ 的维度为 $d_Q$，那么显然 $\boldsymbol{W}_Q$ 就是一个 $d_q\times d_{\text{model}}$ 的矩阵。需要注意，这个矩阵是学习得来的参数，因此 Query 操作的参数量就是 $d_q\times d_{\text{model}}$.

3 Key 键

Key 向量便是对 Token 提出”问题“的”回应“。

对于每个 Token 的词向量 $\vec{\mathrm{E}_i}$，让它与矩阵 $\boldsymbol{W}_K$ 相乘获得对应的查询向量 $\vec{\mathrm{K}_i}=\boldsymbol{W}_K\vec{\mathrm{E}_i}$.

其中，$\vec{\mathrm{E}_i}$ 的维度为 $d_{\text{model}}$，$\vec{\mathrm{K}_i}$ 的维度为 $d_k$，原论文中 $d_k=d_q$，那么显然 $\boldsymbol{W}_K$ 就是一个 $d_q\times d_{\text{model}}$ 的矩阵。需要注意，这个矩阵是学习得来的参数，因此 Key 操作的参数量也是 $d_q\times d_{\text{model}}$.

获得了 Key 值后，就要通过 Compatibility Function (评分函数) 来获得关联程度了。评分函数可以是向量点积：

$$ \text{Compatibility Function}(\vec{\mathrm{Q}_i},\vec{\mathrm{K}_i})=\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_i} $$

上面我们说到原论文中 $d_k=d_q$，那么如果 $d_k\neq d_q$，那么显然不能靠向量点积的评分函数了，而是需要加性评分函数。加性评分函数使用一个前馈神经网络来计算 Query 向量和 Key 向量的关联程度，这个神经网络只有一个隐藏层。

计算得到关联程度形成矩阵：

$$ \begin{bmatrix} \vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_1}&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_1}&\dots&\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_1}\\ \vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_2}&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_2}&\dots&\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_2}\\ \vdots&\vdots&\ddots&\vdots\\ \vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_i}&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_i}&\dots&\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_i} \end{bmatrix} $$

接下来有一个可选步骤——Mask (掩码)。对于大语言模型，训练时是通过上文来生成下文，因此上文是不可以从下文获取信息的，否则这就泄露了要预测的信息，造成干扰。但对于文本翻译，就不存在这种问题，因此掩码就是不必要的。

为了防止上文获取下文信息，可以直接将这些关联程度项设为 $-\infty$，这一步就叫掩码：

$$ \begin{bmatrix} \vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_1}&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_1}&\dots&\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_1}\\ -\infty&\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_2}&\dots&\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_2}\\ \vdots&\vdots&\ddots&\vdots\\ -\infty&-\infty&\dots&\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_i} \end{bmatrix} $$

还需要注意的一点是，经过点积评分函数后会扩大原数据的方差，导致 Softmax 过于关注小梯度：若 $q,k$ 为独立的随机变量，那么它的均值为 $0$ 方差为 $1$，但是 $q\cdot k$ 的均值为 $0$，方差为 $d_k$. 为了抵消这个问题，给每一项缩小 $\sqrt{d_k}$：

$$ \begin{bmatrix} \frac{\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_1}}{\sqrt{d_k}} & \frac{\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_1}}{\sqrt{d_k}} & \dots & \frac{\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_1}}{\sqrt{d_k}} \\ -\infty & \frac{\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_2}}{\sqrt{d_k}} & \dots & \frac{\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_2}}{\sqrt{d_k}} \\ \vdots & \vdots & \ddots & \vdots \\ -\infty & -\infty & \dots & \frac{\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_i}}{\sqrt{d_k}} \end{bmatrix} $$

再进行列 SoftMax，这些被置为 $-\infty$ 的项就变成 $0$ 了：

$$ M=\begin{bmatrix} \text{Softmax}\left(\frac{\vec{\mathrm{Q}_1}\cdot\vec{\mathrm{K}_1}}{\sqrt{d_k}}\right) & \text{Softmax}\left(\frac{\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_1}}{\sqrt{d_k}}\right) & \dots & \text{Softmax}\left(\frac{\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_1}}{\sqrt{d_k}}\right) \\ & \text{Softmax}\left(\frac{\vec{\mathrm{Q}_2}\cdot\vec{\mathrm{K}_2}}{\sqrt{d_k}}\right) & \dots & \text{Softmax}\left(\frac{\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_2}}{\sqrt{d_k}}\right) \\ & & \ddots & \vdots \\ 0 & & & \text{Softmax}\left(\frac{\vec{\mathrm{Q}_i}\cdot\vec{\mathrm{K}_i}}{\sqrt{d_k}}\right) \end{bmatrix} $$

4 Value 值

计算了关联程度，下一步就是与 Value 矩阵相乘获得最终的更新，记上述矩阵 $p$ 行 $q$ 列元素为 $M_{pq}$，那么：

$$ \Delta\vec{\mathrm{E}_i}=\sum_{k=1}^{n}M_{ki}\boldsymbol{W}_V\vec{\mathrm{E}_i} $$

其中，$\Delta\vec{\mathrm{E}_i}$ 的维度为 $d_{\text{model}}$，$\vec{\mathrm{E}_i}$ 的维度为 $d_{\text{model}}$，那么显然 $\boldsymbol{W}_V$ 就是一个 $d_{\text{model}}\times d_{\text{model}}$ 的方阵。需要注意，这个矩阵是学习得来的参数，因此 Value 操作的参数量也是 $d_{\text{model}}\times d_{\text{model}}$​.

综上，注意力头总参数量为：$d_{\text{model}}(2d_q+d_{\text{model}})$

上述过程的结构可以以下图方式展现：

用公式来表示就是：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

5 多头注意力

实际情况中，$d_{\text{model}}$ 往往非常大，可以达到 $10^4$ 数量级，这就导致 $\boldsymbol{W}_V$ 这个 $d_{\text{model}}\times d_{\text{model}}$ 的方阵参数量非常大，占据了模型的大部分参数。

多头注意力的思想就是，与其用一个完整维度的大参数单头注意力，不如将其拆分为多个低维度的小参数注意力，称为多头注意力。这样的效果要好于单头注意力。

我们可以用低秩分解的思想，将 $\boldsymbol{W}_V$ 拆分为 $d_{\text{model}}\times d_v$ 和 $d_v\times d_{\text{model}}$ 的两个矩阵：$\boldsymbol{W}_{V\uparrow}$ 和 $\boldsymbol{W}_{V\downarrow}$. 那么，原来的 Value 计算步骤就变成了：

$$ \Delta\vec{\mathrm{E}_i}=\sum_{k=1}^{n}M_{ki}\boldsymbol{W}_{V\uparrow}\boldsymbol{W}_{V\downarrow}\vec{\mathrm{E}_i} $$

低秩拆分是一种有损压缩技术，可以在一定程度上保留数据的重要特征，但无法完全保留原始数据的所有信息。另外，我们也可以将这个操作理解为做了两次线性映射：先将 $\vec{\mathrm{E}_i}$ 从 $d_{\text{model}}$ 维空间线性映射到低维的 $d_v$ 空间，再从低维的 $d_v$ 空间重新线性映射回 $d_{\text{model}}$ 维空间（原论文中，拆分后的 $\boldsymbol{W}_{V\uparrow}$ 和 $\boldsymbol{W}_{V\downarrow}$ 就直接被记为 Linear 了）.

矩阵乘法可以看作线性映射，例如有一个 $n$ 维向量 $\vec{p}$，将它与 $m\times n$ 的矩阵 $\boldsymbol{M}$ 相乘可以得到一个 $m$ 维向量 $\vec{q}$：

$$ \boldsymbol{M}_{m\times n}\vec{p}_{n\times1}=\vec{q}_{m\times1} $$

这个矩阵 $\boldsymbol{M}$ 就实现了将 $n$ 维空间的向量 $\vec{p}$ 映射到 $m$ 维空间里，变成向量 $\vec{q}$​.

上述拆分后仍然是单头注意力，接下来就是叠加注意力头了。最简单想到的一种方式是，直接将单头注意力复制 $h$ 份，每个注意力头有独立的参数（如第 $i$ 个注意力头的参数 $\boldsymbol{W}^{Q}_i,\boldsymbol{W}^{K}_i,\boldsymbol{W}^{V\downarrow}_i,\boldsymbol{W}^{V\uparrow}_i$），最后将 $h$ 个注意力头的结果取平均。

但是原论文并没有用这个方式。在原论文中，将单头注意力复制 $h$ 份后，$\boldsymbol{W}_{V\uparrow}$ 并不是独立的参数，具体来说：

1. 每个注意力头拥有独立的参数 $\boldsymbol{W}^{Q}_i,\boldsymbol{W}^{K}_i,\boldsymbol{W}^{V\downarrow}_i$
2. 每个注意力头求得 $\overrightarrow{\text{head}}_i=\sum_{k=1}^{n}M_{ki}\boldsymbol{W}_{V\downarrow}\vec{\mathrm{E}_i}$，显然这个 $\vec{\text{head}}_i$ 是 $d_v$ 维的向量。
3. 将 $h$ 个注意力头得到的结果向量拼接得到 $\overrightarrow{\text{MultiHead}}=\mathrm{Concat}(\overrightarrow{\text{head}}_1,\overrightarrow{\text{head}}_2,\dots,\overrightarrow{\text{head}}_N)$，显然这个 $\overrightarrow{\text{MultiHead}}$ 是 $h\times d_v$ 维的向量。
4. 然后将 $\overrightarrow{\text{MultiHead}}$ 用 $\boldsymbol{W}_{V\uparrow}$ 重新映射回 $d_{\text{model}}$ 维的空间，显然此时 $\boldsymbol{W}_{V\uparrow}$ 应当是 $d_{\text{model}}\times hd_v$ 的矩阵。

综上，对于 $h$ 个注意力头的多头注意力，参数量为：

$$ hd_{\text{model}}(d_q+d_k+d_v)+hd_{\text{model}}d_v=hd_{\text{model}}(d_q+d_k+2d_v) $$

对于原论文，$d_q=d_k=d_v$，那么参数量为：$4hd_{\text{model}}d_q$​

上述过程的结构可以以下图方式展现：

需要注意的是，上面原论文的图将 $\boldsymbol{W}_{V\downarrow}$ 表示为 Linear。虽说看着不一样，实际上是完全一样的，只是符号的区别。