【机器学习】详解 Transformer

何处闻韶

已于 2023-02-15 15:12:32 修改

阅读量7.3k

点赞数 60

分类专栏：【机器学习与深度学习】文章标签： transformer 自然语言处理深度学习人工智能

于 2021-07-12 22:00:45 首次发布

原文链接：https://bbs.cvmart.net/articles/4032

版权

【机器学习与深度学习】专栏收录该内容

39 篇文章

订阅专栏

本文深入探讨Transformer模型，重点解析Self-attention和Multi-headSelf-attention的工作原理。Self-attention机制解决了传统RNN和CNN在处理序列数据时的局限性，实现了并行化训练和全局信息获取。Multi-headSelf-attention通过多个独立的Self-attention层，从不同角度捕捉信息，增强了模型的表达能力。Transformer的Encoder和Decoder结构中，Self-attention层与位置编码结合，形成强大的序列处理能力。此外，文章还提供了Transformer的PyTorch实现代码，帮助读者理解其工作流程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言

一、一切从 Self-attention 开始

1.1 处理 Sequence 数据的模型

1.2 Self-attention

1.3 Multi-head Self-attention

1.4 Positional Encoding ☆

二、Transformer 的实现和代码解读

2.1 Transformer 原理分析

2.2 Transformer 代码解读

2.2.1. Scaled Dot Product Attention

2.2.2. Positional Encoding

2.2.3. Multi Head Attention

2.2.4. Feed Forward Network

前言

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在较火热的 BERT 也是基于 Transformer Encoder 的。Transformer 模型使用了 self-attention 机制，不采用 RNN 顺序结构，使得模型可以 并行化训练，而且能够 拥有全局信息。本文是对 Vision Transformer 的原理和代码全面解读的第一部分。

一、一切从 Self-attention 开始

1.1 处理 Sequence 数据的模型

Transformer 是一个 Sequence to Sequence (Seq2Seq) 模型，特别之处在于它大量用到了 Self-attention。

其实，要处理一个 Sequence，最常想到的是用 RNN，它输入一串 Vector Sequence，输出另一串 Vector Sequence，如下图 1 左所示。若为 Single-Directional RNN，那么输出 $b^4$ 时，默认 $a^1,a^2,a^2,a^4$ 都已 “看过”；若为 Bi-directional RNN，那么输出任意 $b^i , i\in \{1,2,3,4\}$ 时，默认 $a^1,a^2,a^2,a^4$ 都已 “看过”。RNN 非常擅长于处理 Sequence 输入。

然而，问题在于 RNN 很不易并行化 (hard to parallel)。例如，对于 Single-Directional RNN，若要算出 $b^4$ ，就必须依次看 $a^1,a^2,a^2,a^4$ ，故这样先后顺序的串行步骤很难平行化。从而，诸如 RNN 和 LSTM 等网络 难以充分发挥 GPU 的加速优势。

所以，研究者提出 用 CNN 取代 RNN，如下图 1 右所示。其中，橘色三角形表示一个 kernel_size = 3 的 CNN filter，每次卷/滑/扫过 3 个向量。卷/滑/扫过一轮以后，就输出一排结果，用橘色小圆点表示。

以上是第一轮或者说第一层橘色的 filter 的卷积过程，还有更多轮/层的 filter，如图 2 中的黄色的 filter，重复相同的卷积过程并输出一排结果，用黄色小圆点表示。

所以，用 CNN 的确也能做到类似 RNN 输入输出的关系 —— 输入一个 Sequence，输出另一个 Sequence。

然而，CNN 和 RNN 表面上可以做到相同的输入和输出，但是 CNN 因感受野有限，只能考虑非常有限的内容。比如，在图 2 右侧中，CNN filter 每次只考虑 3 个 vector，不像 RNN 可以考虑之前的所有 vector。然而，CNN 也并非无法考虑 长时期/长距离依赖 (Long Term Dependency) —— 通过堆叠更多层 filter (扩大感受野)，上层 (深层) 的 CNN filter 即可考虑更多信息。比如，第二层的 filter (蓝色的三角形) 考虑 6 个 vector，比第一层/最浅层的 filter 考虑更多的信息。所以，堆叠足够多层 CNN filter，就能看到相当长时期/长距离的信息。当然，使用 膨胀卷积/空洞卷积 等扩大 kernel size 也是常见的方式。

此外，CNN 的一个优势在于 可并行化 (parallel) 并使用 GPU 加速，即无需依次等待红色的 filter 算完，再算黄色的 filter。但是，必须要堆叠很多层 CNN filter，才可看到足够长时期/长距离的信息。所以，Self-attention 的产生旨在 使用 Self-attention Layer 取代 RNN 所做的事，如下图 3 所示：

图 3：You can try to replace any thing that has been done by RNNwith self attention

重点在于，Self-attention 这种新 Layer 的 输入和输出和 RNN 一模一样，即 输入输出均为 Sequence，每一个输出 $b^1, b^2, b^3, b^4$ 都看过了、考虑到整个输入 Sequence，这点与 Bi-Directional RNN 相同。但 优于 RNN 之处 在于，其每一个输出 $b^1, b^2, b^3, b^4$ 都可以并行化计算得到。二者的对比如上图 3 所示。

1.2 Self-attention

那么 Self-attention 具体如何呢？

首先，设 Input Vector 是图 4 的一串蓝色 Sequence $x^1, x^2, x^3, x^4$ ，每个 Input Vector 先各乘一个权重参数矩阵 $W$ 得到一串绿色的 Embedding Vector $a^1, a^2, a^3, a^4$ ，然后各馈入 Self-attention 层，即令每个 Embedding Vector $a^1, a^2, a^3, a^4$ 分别乘上 3 个不同的 Transformation Matrix $W_q, W_k, W_v$ (Query, Key, Value)。例如 $a^1$ 将分别得到 3 个不同的 Vector $q^1, k^1, v^1$ ，其余同理。

接下来 用每个 Query $q$ 去对每个 Key $k$ 做 Attention，以衡量任意 2 个 Vector 的相似程度。比如，现在要对 $q^1$ 和 $k^1$ 做 Attention，先对这 2 个 Vector 做 Scaled Inner-Product (Dot-Product) 得到 $\alpha_{1,1}$ 。

同理，对 $q^1$ 和 $k^2$ 做 Attention 得到 $\alpha_{1,2}$ ，对 $q^1$ 和 $k^3$ 做 Attention 得到 $\alpha_{1,3}$ ，对 $q^1$ 和 $k^4$ 做 Attention 得到 $\alpha_{1,4}$ 。其中，Scaled Inner Product 的计算公式如下：

其中， $d$ 是 $q^i$ 和 $k^i$ 的维度 dimension 大小。因为 $q^i \cdot k^i$ 的数值会随 dimension 的增大而增大，所以要除以 $\sqrt {dimension}$ 相当于归一化的效果。

接下来，对所有 Attention 结果 $a_i, i \in \{1,2,3,4\}$ 执行 $Softmax$ 操作，如下所示：

执行 $Softmax$ 操作后得到了 $\hat{ \alpha } _{1,i}$ $, i \in \{1, 2, 3, 4 \}$ ，令 $\hat{ \alpha } _{1,i}$ 与各 $v^i$ 相乘并求和，得到 $b^1$ ，如下所示。因而，产生 $b^1$ 的过程中用到了 整个 Input Sequence 的信息 (considering the whole sequence)。

若要 考虑局部 (local) 信息，则只需学习出相应的 $\hat{ \alpha } _{1,i} = 0$ ， $b^1$ 就不再带有那个对应分支的信息了；

若要 考虑全局 (global) 信息，则只需学习所有的 $\hat{ \alpha } _{1,i} \neq 0$ ， $b^1$ 就带有全部的对应分支的信息了。

同理，可计算出 $b^2, b^3, b^4$ ，如下图 8 所示。

经过了以上一连串计算，Self-attention layer 做的事不仅 和 RNN 一样，还可以 并行地得到 layer 输出的结果，如图 9 所示。

接下来，用矩阵表示上述计算过程。首先输入 Embedding $I = [a^1, a^2, a^3, a^4]$ 。然后，

用 $I$ 乘 Transformation Matrix $W^q$ 得到 $Q = [q^1, q^2, q^3, q^4]$ ，其每一列代表一个 Vector $q$ ；
用 $I$ 乘 Transformation Matrix $W^k$ 得到 $K = [k^1, k^2, k^3, k^4]$ ，其每一列代表一个 Vector $k$ 。
用 $I$ 乘 Transformation Matrix $W^v$ 得到 $V = [v^1, v^2, v^3, v^4]$ ，其每一列代表一个 Vector $v$ 。

计算过程如图 10 所示：

接下来是 $k$ 与 $q$ 的 Attention 过程，可以把 Vector $k$ 转置为行向量与列向量 $q$ 做内积得到标量 $\alpha$ (此处省略 $\sqrt d$ )。整体上看，由 4 个行向量 $k^T$ 拼成的矩阵 $K^T$ 和 4 个列向量 $q$ 拼成的矩阵 $Q$ 做内积将得到由标量 $\alpha$ 构成的 $4 \times 4$ 矩阵 $A$ ，并对其取 $Softmax$ 得到 $\hat{A}$ 。要得到 $b^1$ ，就要用 $\hat{\alpha}_{1,i}$ 分别与 $v^i$ 相乘并求和，故整体上 $\hat{A}$ 要再左乘 $V$ 矩阵。上述过程如图 11 所示：

更进一步地，上述过程自上而下可表示为图 12 的形式：

令输入矩阵 $I \in R ^ {d \times N}$ 分别乘 3 个权重矩阵 $W^q$ ， $W^k$ ， $W^v$ 得到 3 个等尺寸矩阵 $Q, K, V \in R ^{d \times N}$
令 $K$ 转置后右乘 $Q$ 得到 Attention 方阵 $A \in R^{N \times N}$ ，代表每一个位置两两间的 Attention
对方阵 $A$ 取 $Softmax$ 得到方阵 $\hat{A} \in R^{N \times N}$
令方阵 $\tilde{A}$ 乘 $V$ 矩阵得到输出 Vector $O \in R^{d \times N}$ ，其与输入矩阵形状一致

图 12：Self-attention 通过一连串的矩阵乘法，可实现 GPU 加速

总之，以上过程可表示为：

即 Scaled Dot-Product Attention：

其作为单头注意力，与多头注意力 MHA 的关系为：

可见，MHA 通过 Linear 线性投影来初始化多组不同的 (Q, K, V)，并将多个 (图中表示为 h 个) 单头的自注意力结果 Concat 后，再经一个全连接层降维输出。可以将初始化不同的 (Q, K, V) 理解为单头从不同的方向去观察文本，从而使自注意力更具 “大局观”。

放缩因子 $\sqrt{d_k}$ 的含义：

方差的基本性质：

1.3 Multi-head Self-attention

还有一种 Multi-head Self-attention，以 2 个 head 的情况为例：

由 $a^i$ (图 10) 生成的 $q^i$ 进一步乘上 2 个转移矩阵 $W^{q, 1}$ 和 $W^{q, 2}$ 变为 $q^{i, 1}$ 和 $q^{i, 2}$
由 $a^i$ (图 10) 生成的 $k^i$ 进一步乘上 2 个转移矩阵 $W^{k, 1}$ 和 $W^{k, 2}$ 变为 $k^{i, 1}$ 和 $k^{i, 2}$
由 $a^i$ (图 10) 生成的 $v^i$ 进一步乘上 2 个转移矩阵 $W^{v, 1}$ 和 $W^{v, 2}$ 变为 $v^{i, 1}$ 和 $v^{i, 2}$

接下来，令 $q^{i, 1}$ 与 $k^{i, 1}$ 做 Attention 再与 $v^{i, 1}$ 相乘、 $q^{i, 1}$ 与 $k^{j, 1}$ 做 Attention 再与 $v^{j, 1}$ 相乘，二者做 Weighted-sum 得到最终的 $b^{i, 1}, i \in {1, 2, ..., N}$ $\in R^{d,1}$ 。同理可得 $b ^ {i, 2}$ $\in R^{d,1}$ 。现在有了 $b^{i, 1}$ 和 $b ^ {i, 2}$ ，可以 把二者 Concat 起来，再通过一个 Transformation Matrix $W^O$ 调整维度，使之与原来 1.2 节中的 $b^i, i \in \{1,2,...,N\}$ $\in R^{d,1}$ 维度一致，如图 13 所示：

图 13：通过一个 Transformation Matrix 调整 Concat 后的 b 的维度

从下图 14 可见，Multi-Head Attention 包含 多个 Self-attention 层。首先，将输入矩阵 $X$ (蓝色) 分别传递到 2 个不同的 Self-attention 中，计算得到 2 个输出结果 (红色)。得到 2 个输出矩阵后，Multi-Head Attention 将它们 Concat 在一起，然后传入一个 Linear 层 (左乘一个参数矩阵 (黄色))，得到 Multi-Head Attention 最终的输出矩阵 $Z$ (粉色)。可见，Multi-Head Attention 的输出矩阵 $Z$ 与其输入矩阵 $X$ 维度一致。

另一组 Multi-head Self-attention 的结果示意，如下图 15 所示。其中，绿色、红色部分表示两组 Query 和 Key。可见，绿色部分覆盖范围更大，更关注 Global 信息；而红色部分覆盖范围更小，更关注 Local 信息。二者测重各不相同。

图 15：Multi-head Self-attention 的不同 Head 分别关注了 Global 和 Local 的信息

1.4 Positional Encoding ☆

以上是 Multi-head Self-attention 的原理，然而问题是：当前的 Self-attention 中没有位置的信息，例如，一个单词向量的 “近在咫尺” 位置的单词向量，和 “远在天涯” 位置的单词向量效果是无差别的；再如，输入 “A 打了 B” 或 “B 打了 A” 的效果其实是一样的，都是因为 没有表示位置的信息 (No position information in self-attention)。于是，Self-attention 原作者为解决该问题，作出如下图 16 所示的 位置编码 操作：

具体做法是：给每个位置人工设定一个表示位置信息的位置向量 $e^i$ (不是神经网络学出来的)，每个位置 (如第 $i$ 个) 都有一个不同的位置向量 $e^i$ ，令其与输入 Embedding $a^i$ 相加 作为新 $a^i$ 参与后续运算过程。

为什么 $e^i$ 与 $a^i$ 是相加而非拼接？相加后，原来表示位置信息的 $e^i$ 不就混到 $a^i$ 中且难以提取了吗？

以下提供一种解答该问题的思路：

如下图 15 所示，先给每个 输入向量 $x^i \in R^{d, 1}$ 加上一个 独热编码的位置向量 $p^i \in R^{d, 1}$ 得到新向量 $x^i_{p}$ 作为输入，乘上一个 Transformation Matrix $W = [W^I, W^P]$ $\in R^ {d, d+N}$ 。此时有：

$W \cdot x^i_p = [W^I, W^P] \cdot \begin{bmatrix} x^i \\ p^i \end{bmatrix} = W^I \cdot x^i + W^P \cdot p^i = a^i + e^i$

可见，位置向量 $e^i$ 与输入 Embedding $a^i$ 直接相加等同于先原输入向量 $x^i$ 拼接一个表示位置的独热编码 $p^i$ 再做 Transformation 得到 Embedding。

其中，与位置编码相乘的矩阵 $W^P$ 是手工设计的，如下图 17 所示：

Transformer 中除了需要 单词 Embedding 表示输入的内容主题，还需要 位置 Embedding 表示单词出现在句子中的位置。因为 Transformer 不采用 RNN 结构，而是使用全局信息，无法捕获或利用到单词的位置顺序信息，而这部分信息对于 NLP 而言非常重要 (事实上对 CV 也很重要)。所以 Transformer 中 使用位置 Embedding 保存单词在序列中的相对或绝对位置。

位置 Embedding 用 PE 表示，其维度与单词 Embedding 一致。PE 可通过训练得到，也可使用某种公式计算得到。在 Transformer 中采用了后者，公式如下：

换言之，第 $pos$ 个位置的位置编码 $PE_{(pos)}$ 为：

其中， $pos$ 表示 token 在 Sequence 中的位置，例如第一个 token "我" 的位置 $pos=0$ 。

而 $i \in [0, ..., d_{model} / 2)$ ，或者准确意义上是 $2i$ 和 $2i + 1$ 表示了 位置编码的维度。展开则更清晰地表示为：

注意， $w_i$ 是位置编码向量的第 $i$ 组分量的频率，其表达式为：

例如，当 $pos=1$ 时，对应的位置编码可表示为：

式中，模型维度 $d_{model}=512$ (NLP 中表示最大文本长度 max_len=512，CV 中表示 Patch Embedding 维度 D=512)。然而，频率 $w_i$ 底数为什么用 $10000$ 呢？原论文中完全没有提到 (玄学)，这里不得不说说论文的 readability 问题，即便是很多高引的文章，最基本的内容都讨论不清楚，所以才出现像上面提问里的讨论，说实话这些论文还远未做到 easy to follow。这里给出一个假想： $10000^ {1 / 512}$ 是一个比较接近 1 的数 (1.018)，如果用 $100000$ ，则是1.023。这只是猜想，其实可能完全可以使用别的底数。

PE 公式的好处：

每个位置有唯一的位置编码，使之能够适应比训练集中所有句子更长的句子。假设训练集中的最长句子长度为 20，若突然出现一个长度为 21 的句子，则使用 PE 公式 仍可计算出第 21 位的位置编码。
可让模型容易地计算出相对位置。对于固定长度的间距 $k$ ，任意位置的 $PE_{pos+k}$ 都可被 $PE_{pos}$ 的线性函数表示 —— 三角函数特性：

证明：

为什么 BERT 和 Transformer 等许多模型直接将 Embedding 直接 Sum Pooling 加在一起？

最后，可以看到 Self-attention 在 Seq2Seq 模型如何得到应用，可以把 Encoder-Decoder 中的 RNN 用 Self-attention 替换。

二、Transformer 的实现和代码解读

2.1 Transformer 原理分析

图 19 展示了一个 Seq2Seq 的模型，左侧为 Encoder Block，右侧为 Decoder Block。Multi-Head Attention (简称 MHA) 由多个 Self-Attention 组成，其中 Encoder block 包含一个 MHA，Decoder block 包含两个 MHA (其中有一个还用到了 Mask)。MHA 上方还包括一个 Add & Norm 层，Add 表示 Residual Connection 用于防止网络退化，Norm 表示 Layer Normalization，用于对每一层的激活值归一化。

例如，对于翻译任务，在 Encoder Input 处输入“机器学习”，在 Decoder Input 处的输入 “\”，输出是 “machine”。下一个时刻在 Decoder Input 处输入 “machine”，输出是 “learning”。不断重复直到输出句点 (.) 代表翻译结束。

接下来简要分析一下各结构的处理过程：

左半部分的 Encoder Block

首先输入向量 $X \in R(n_x,N)$ 通过一个 Input Embedding 转移矩阵 $W^X \in R(d,n_x)$ 得到一个张量 $I \in R(d,N)$ ，再加上一个表示位置的 Positional Encoding $E \in R(d,N)$ 得到新的张量 $I \in R(d,N)$ ，然后进入重复 $N$ 次的绿色 Block。绿色 Block 中，张量 $I \in R(d,N)$ 先经过一个 MHA 输出 $O \in R(d,N)$ 。然后，通过 Add & Norm 层将 MHA 的输入 $I \in R(d,N)$ 和输出 $O \in R(d,N)$ 按元素相加，并进行 Layer Normalization。

下图 20 展示了常见的几种 Normalization 的对比示意图。

图 20：不同 Normalization 对比

图 21：Batch Normalization 和 Layer Normalization 的对比

Batch Normalization 令 batch 所有 samples / features 的某 channel 的 $\mu = 0, \sigma=1$ (对 batch 内所有数据沿/按 channel 归一化)
Layer Normalization 令 batch 内各 samples / features 的所有 channel 的 $\mu = 0, \sigma=1$ (对 batch 内所有数据沿/按 sample 归一化)

为什么用 LN 而不用 BN？ (max_len 对于 c, emb_dim 对应 h 或 w)

【机器学习】详解 Normalization_何处闻韶的博客-优快云博客_机器学习normalization【机器学习】详解 Normalizationhttps://blog.youkuaiyun.com/qq_39478403/article/details/120547567

接着，是一个 Feed Forward 前馈网络和一个 Add & Norm 层。综上，

一个绿色 Block 的前 2 个 Layer 操作表达式为：

一个绿色 Block 的后 2 个 Layer 操作表达式为：

一个绿色 Block 的所有 Layer 简写为：

Transformer 的 Encoder Block 由 N 个绿色 Block 构成，表达式为：

右半部分的 Decoder Block

Decoder Block 的输入包括 2 个来源，来自 Decoder Block 下方的输入是 前一个 Time Step 的 Input Embedding，即前一个 Time Step 的 $I \in R(d,N)$ 加上一个表示位置的 Positional Encoding $E \in R(d,N)$ 所得的张量。然后，该张量进入了重复 N 次的绿色 Block。

对于 Decoder 的绿色 Block，首先是 Masked Multi-Head Self-attention，Mask 旨在使注意力只关注已产生的 Sequence 而不含未产生的部分。这很合理，因为还未产生的东西不存在，就无法做 Attention。

更具体地：

Decoder 输出： (当前 Time Step) 对应位置 $i$ 的输出词的概率分布。
Decoder 输入： (当前 Time Step) 对应位置 $i$ 的 Encoder 输出 + (前一 Time Step) 对应位置 $i-1$ 的 Decoder 输出。所以中间的 Attention 不是 Self-attention，其 Key 和 Value 来自 Encoder 的输出，Query 来自上一位置 Decoder 的输出。
Decoding 串行：注意，编码可并行计算，一次性全部 Encoding 出来；但解码不同，它是像 RNN 一样一个一个 Decoding的，因为要用上一位置 Decoder 的输出当作当前位置 Attention 的 Query。

Decoder 在组成上的主要的区别是：新增了 Mask 的 Attention —— Masked Multi-Head Self-attention。因为训练时的 Output 都是 Ground Truth，这样可以确保预测第 $i$ 个位置时不会接触到未来 $i+1, i+2, ...$ 个位置的信息。

总之，Decoder 包含的两个 Multi-Head Attention 层，第一个 Multi-Head Attention 层采用了 Masked 操作，第二个 Multi-Head Attention 层的 Key，Value 矩阵使用 Encoder 的编码信息矩阵 C 进行计算，而 Query 使用上一个 Decoder block 的输出计算。最后有一个 Softmax 层计算下一个翻译单词的概率。

下面详细介绍 Masked Multi-Head Self-attention 的具体操作，Masked 在 Scale 操作后、Softmax 操作前。

因为在 翻译任务 中，翻译是按顺序的 —— 翻译完第 $i$ 个单词，才可翻译第 $i+1$ 个单词。通过 Masked 操作可防止第 $i$ 个单词不切实际地了解/接触到第 $i+1$ 个单词及之后的信息。下面以将 "我有一只猫" 翻译成 "I have a cat" 为例，说明 Masked 操作。Decoder 时，需根据之前的翻译，求解当前最可能的翻译 (当前位置最大概率输出)，如下图 23 所示。首先根据输入 <Begin> 预测出第一个单词为 "I"，然后根据输入 "<Begin> I" 预测下一个单词 "have"。

Decoder 可在训练过程中使用 Teacher Forcing 且 并行化训练，即将正确的单词序列 (<Begin> I have a cat) 和对应输出 (I have a cat <End>) 传递到 Decoder。那么在预测第 $i$ 个输出时，就要将第 $i+1$ 之后的单词遮盖住 (Masked)。注意 Mask 操作是在 Self-Attention 的 Softmax 之前使用的，下面用 0 1 2 3 4 5 分别表示 "<Begin> I have a cat <End>"。

注意这里 Transformer 模型 训练和测试的解码方法不同：

Transformer 测试时的解码过程：

输入解码开始标志位 <Begin>，Decoder 输出 I
输入已解码的 <Begin>, I，Decoder 输出 have
以此类推 ...
输入已解码的 <Begin>, I, have, a, cat，Decoder 输出解码结束标志位 <end>
总之，每次解码都会到利用先前已解码的所有单词嵌入信息

Transformer 训练时的解码过程：

不采用上述类似 RNN 的方法 一个一个目标单词嵌入向量顺序输入训练，想采用 类似编码器中的矩阵并行算法，一步就把所有目标单词预测出来。要实现这个功能可以参考 Encoder 的操作，把目标单词嵌入向量组成矩阵一次输入即可 —— 并行化训练。

但在解码 have 时，不能利用到后面单词 a 和 cat 的目标单词嵌入向量信息，否则就是作弊 (测试时候不可能能未卜先知)。为此引入 mask。具体是：在解码器中，Self-attention 层只被允许处理输出序列中更靠前的那些位置，在 Softmax 步骤前，它会把后面的位置给隐去。

Masked Multi-Head Self-attention 的计算：

Step1：Input Matrix $X\in R_{N,d_x}$ 包含 "<Begin> I have a cat" (0, 1, 2, 3, 4) 五个单词的表示向量，Mask Matrix 是一个 5×5 矩阵。在其中可见解码单词 0 时只能使用单词 0 的信息，而解码单词 1 时可使用单词 0, 1 的信息 —— 只能使用先前的信息。Input Matrix $X\in R_{N,d_x}$ 经过 3 个 Transformation Matrix 得到 3 个 Matrix：Query $Q \in R_{N,d}$ ，Key $K \in R_{N,d}$ 和 Value $V \in R_{N,d}$ 。
Step2： $Q^T \cdot K$ 得到 Attention Matrix $A\in R_{N,N}$ ，此时先不进行 Softmax 操作，而是与一个 $Mask \in R_{N, N}$ 矩阵相乘，使 Attention Matrix 的部分位置 (即相对当前位置的未来位置) 为 0，得到 Masked Attention Matrix $Mask \; Attention \in R_{N, N}$ 。Masked Attention Matrix 是个下三角矩阵，使得计算 $Z$ 矩阵的某一行时，只考虑其前面 token 的作用 (即相对当前位置的先前位置) 。例如，在计算 $Z$ 的第一行时，刻意地把 Attention Matrix 第一行的后面所有元素屏蔽掉，只考虑 $A_{0, 0}$ 。在产生单词 have 时，则只考虑之前的 I，不考虑之后的 have、a、cat，即只 attend on 已产生的 Sequence。这很合理，因为还没有产生出来的东西不存在，就无法做 Attention。
Step3： Masked Attention Matrix 进行 Softmax，所得矩阵的每一行之和都为 1 (沿列方向按行归一化)。注意，单词 0 在单词 1, 2, 3, 4 上的 Attention Score 都为 0。所得矩阵再与矩阵 $V$ 相乘得到最终的 Self-attention 层的输出结果 $Z_1 \in R_{N, d}$ 。
Step4： $Z_1 \in R_{N, d}$ 只是第 1 个 Head 的结果，将多个 Head 的结果 Concat 一起后，再进行 Linear Transformation 得到最终的 Masked Multi-Head Self-attention 结果 $Linear \; Transformaion (Concat (Z_1, Z_2, ... , Z_n)) = Z \in R_{N, d}$ 。

图 24：Masked Multi-Head Self-attention 的具体操作

此外，需注意的是：第 1 个 Masked Multi-Head Self-attention 的 $Query$ ， $Key$ ， $Value$ 均来自 Output Embedding；

而第 2 个 Multi-Head Self-attention 的 $Query$ 来自第 1 个 Self-attention 层的输出， $Key$ 和 $Value$ 来自 Encoder 的输出。

关于这种设计的个人理解：

$Key$ 和 $Value$ 来自 Encoder 的输出，所以可看做 句子(Sequence) / 图片 (Image) 等的内容信息 (Content，比如句子含义是："我有一只猫" / 图片内容是："有几辆车，几个人等等")。

$Query$ 表达了一种诉求：希望得到 / 了解 / 寻求什么，可看做 引导信息 (Guide)。

通过 Multi-Head Self-attention 结合在一起的过程，就相当于是 把需要的内容信息指导表达出来。

Self-attention 和 Multi-Head Self-attention 的时间复杂度计算：

Decoder 的最后是 Softmax 预测输出单词。因为 Mask 的存在，使得单词 0 的输出 $Z(0, )$ 只包含单词 0 的信息。Softmax 根据输出矩阵的每一行预测下一个单词，如下图 25 所示。

下图 26 则展示了Transformer 的整体结构：

2.2 Transformer 代码解读

Github：GitHub - jadore801120/attention-is-all-you-need-pytorch: A PyTorch implementation of the Transformer model in "Attention is All You Need".

2.2.1. Scaled Dot Product Attention

实现的是图 22 的操作，先令 $Q \cdot K^T$ ，再对结果按位乘以 $Mask$ 矩阵，再做 $Softmax$ 操作，最后的结果与 $V$ 相乘，得到 Self-attention 的输出。

class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):

        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))  # Q · K

        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)  # Mask

        attn = self.dropout(F.softmax(attn, dim=-1))  # Softmax + Dropout

        output = torch.matmul(attn, v)  # Att · V

        return output, attn

2.2.2. Positional Encoding

实现位置编码，如下所示：

class PositionalEncoding(nn.Module):

    def __init__(self, d_hid, n_position=200):
        super(PositionalEncoding, self).__init__()

        # Not a parameter
        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))

    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        ''' Sinusoid position encoding table '''
        # TODO: make it with torch instead of numpy

        def get_position_angle_vec(position):
            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

        return torch.FloatTensor(sinusoid_table).unsqueeze(0)  # (1,N,d)

    def forward(self, x):
        # x(B, N, d)
        return x + self.pos_table[:, :x.size(1)].clone().detach()

2.2.3. Multi Head Attention

实现图 13、14 的多头 Self-attention。

class MultiHeadAttention(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super().__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)

        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, q, k, v, mask=None):

        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        residual = q

        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)   # For head axis broadcasting.

        q, attn = self.attention(q, k, v, mask=mask)

        #q (sz_b,n_head,N=len_q,d_k)
        #k (sz_b,n_head,N=len_k,d_k)
        #v (sz_b,n_head,N=len_v,d_v)

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)

        #q (sz_b, len_q, n_head, N * d_k)
        # 最终的输出矩阵 Z
        q = self.dropout(self.fc(q))  

        # Add & Norm 层
        q += residual
        q = self.layer_norm(q) 

        return q, attn

2.2.4. Feed Forward Network

实现前向传播网络。


class PositionwiseFeedForward(nn.Module):
    ''' A two-feed-forward-layer module '''

    def __init__(self, d_in, d_hid, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_in, d_hid)  # position-wise
        self.w_2 = nn.Linear(d_hid, d_in)  # position-wise
        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):

        residual = x

        # 两层 FCs + Dropout
        x = self.w_2(F.relu(self.w_1(x)))
        x = self.dropout(x)

        # Add & Norm 层
        x += residual
        x = self.layer_norm(x)

        return x

2.2.5. Encoder Layer

实现图 26 中的一个Encoder Layer，具体的结构见图 19。

class EncoderLayer(nn.Module):
    ''' Compose with two layers '''

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
        super(EncoderLayer, self).__init__()

        # MHA + Add & Norm
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        # FFN + Add & Norm
        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(self, enc_input, slf_attn_mask=None):

        enc_output, enc_slf_attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask)
        enc_output = self.pos_ffn(enc_output)

        return enc_output, enc_slf_attn

2.2.6. Decoder Layer

实现图 26 中的一个 Decoder Layer，具体的结构见图19。

class DecoderLayer(nn.Module):
    ''' Compose with three layers '''

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
        super(DecoderLayer, self).__init__()

        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.enc_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(self, dec_input, enc_output, slf_attn_mask=None, dec_enc_attn_mask=None):

        # MMHA + Add & Norm
        dec_output, dec_slf_attn = self.slf_attn(dec_input, dec_input, dec_input, mask=slf_attn_mask)
        # MHA + Add & Norm
        dec_output, dec_enc_attn = self.enc_attn(dec_output, enc_output, enc_output, mask=dec_enc_attn_mask)
        # FFN + Add & Norm        
        dec_output = self.pos_ffn(dec_output)

        return dec_output, dec_slf_attn, dec_enc_attn

2.2.7. Encoder

实现图 26、19 左侧的 Encoder。

class Encoder(nn.Module):
    ''' A encoder model with self attention mechanism. '''

    def __init__(
            self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, dropout=0.1, n_position=200):

        super().__init__()

        self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx)
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(p=dropout)
        self.layer_stack = nn.ModuleList([
            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, src_seq, src_mask, return_attns=False):

        enc_slf_attn_list = []

        # -- Forward --

        # Input Embedding + Position Embedding + Dropout + Norm
        enc_output = self.dropout(self.position_enc(self.src_word_emb(src_seq)))
        enc_output = self.layer_norm(enc_output)

        # N × Encoder Block
        for enc_layer in self.layer_stack:
            enc_output, enc_slf_attn = enc_layer(enc_output, slf_attn_mask=src_mask)
            enc_slf_attn_list += [enc_slf_attn] if return_attns else []

        if return_attns:
            return enc_output, enc_slf_attn_list
        return enc_output,

2.2.8. Decoder

实现图 26，19 右侧的 Decoder。

class Decoder(nn.Module):
    ''' A decoder model with self attention mechanism. '''

    def forward(self, trg_seq, trg_mask, enc_output, src_mask, return_attns=False):

        dec_slf_attn_list, dec_enc_attn_list = [], []

        # -- Forward --
        dec_output = self.dropout(self.position_enc(self.trg_word_emb(trg_seq)))
        dec_output = self.layer_norm(dec_output)

        for dec_layer in self.layer_stack:
            dec_output, dec_slf_attn, dec_enc_attn = dec_layer(
                dec_output, enc_output, slf_attn_mask=trg_mask, dec_enc_attn_mask=src_mask)
            dec_slf_attn_list += [dec_slf_attn] if return_attns else []
            dec_enc_attn_list += [dec_enc_attn] if return_attns else []

        if return_attns:
            return dec_output, dec_slf_attn_list, dec_enc_attn_list
        return dec_output,

2.2.9. 整体结构

实现图 26，19 整体的 Transformer。

class Transformer(nn.Module):
    ''' A sequence to sequence model with attention mechanism. '''

    def __init__(
            self, n_src_vocab, n_trg_vocab, src_pad_idx, trg_pad_idx,
            d_word_vec=512, d_model=512, d_inner=2048,
            n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1, n_position=200,
            trg_emb_prj_weight_sharing=True, emb_src_trg_weight_sharing=True):

        super().__init__()

        self.src_pad_idx, self.trg_pad_idx = src_pad_idx, trg_pad_idx

        self.encoder = Encoder(
            n_src_vocab=n_src_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            pad_idx=src_pad_idx, dropout=dropout)

        self.decoder = Decoder(
            n_trg_vocab=n_trg_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            pad_idx=trg_pad_idx, dropout=dropout)

        self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p) 

        assert d_model == d_word_vec, \
        'To facilitate the residual connections, \
         the dimensions of all module outputs shall be the same.'

        self.x_logit_scale = 1.
        if trg_emb_prj_weight_sharing:
            # Share the weight between target word embedding & last dense layer
            self.trg_word_prj.weight = self.decoder.trg_word_emb.weight
            self.x_logit_scale = (d_model ** -0.5)

        if emb_src_trg_weight_sharing:
            self.encoder.src_word_emb.weight = self.decoder.trg_word_emb.weight

    def forward(self, src_seq, trg_seq):

        # source mask：用于产生 Encoder 的 mask，它是一列 Bool 值，负责把标点 mask 掉
        src_mask = get_pad_mask(src_seq, self.src_pad_idx)
        # target mask：用于产生 Decoder 的 mask。它是一个矩阵，如图 24 中的 mask 所示，功能已在上文介绍
        trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)

        enc_output, *_ = self.encoder(src_seq, src_mask)
        dec_output, *_ = self.decoder(trg_seq, trg_mask, enc_output, src_mask)
        seq_logit = self.trg_word_prj(dec_output) * self.x_logit_scale

        return seq_logit.view(-1, seq_logit.size(2))

2.2.10. 产生 Mask


def get_pad_mask(seq, pad_idx):
    return (seq != pad_idx).unsqueeze(-2)

def get_subsequent_mask(seq):
    ''' For masking out the subsequent info. '''
    sz_b, len_s = seq.size()
    subsequent_mask = (1 - torch.triu(
        torch.ones((1, len_s, len_s), device=seq.device), diagonal=1)).bool()

    return subsequent_mask

参考资料

搞懂视觉 Transformer 原理和代码，看这篇技术综述就够了 - 极市社区

极市社区

面经：什么是Transformer位置编码？

超细节的BERT/Transformer知识点 - 知乎