我从未见过如此全面的Transformer剖析
极市导读
近4万余字、50张图，带你一网打尽Transformer。 >>加入极市CV技术交流群，走在计算机视觉的最前沿
1. 多头注意力机制原理
1.1 动机
今天要和大家介绍的一篇论文是谷歌2017年所发表的一篇论文，名字叫做”Attention is all you need“[1]。虽然，网上已经有了大量的关于这篇论文的解析，不过好菜不怕晚，笔者在这里也会谈谈自己对于它的理解以及运用。按照我们一贯解读论文的顺序，首先让我们先一起来看看作者当时为什么要提出Transformer这个模型？需要解决什么样的问题？现在的模型有什么样的缺陷？
1.1.1 面临问题
在论文的摘要部分作者提到，现在主流的序列模型都是基于复杂的循环神经网络或者是卷积神经网络构造而来的Encoder-Decoder模型，并且就算是目前性能最好的序列模型也都是基于注意力机制下的Encoder-Decoder架构。为什么作者会不停的提及这些传统的Encoder-Decoder模型呢？接着，作者在介绍部分谈到，由于传统的Encoder-Decoder架构在建模过程中，下一个时刻的计算过程会依赖于上一个时刻的输出，而这种固有的属性就限制了传统的Encoder-Decoder模型就不能以并行的方式进行计算，如图1-1所示。
This inherently sequential nature precludes parallelization within training examples, which becomes critical at longer sequence lengths, as memory constraints limit batching across examples.
随后作者谈到，尽管最新的研究工作已经能够使得传统的循环神经网络在计算效率上有了很大的提升，但是本质的问题依旧没有得到解决。
Recent work has achieved significant improvements in computational efficiency through factorization tricks and conditional computation, while also improving model performance in case of the latter. The fundamental constraint of sequential computation, however, remains.
1.1.2 解决思路
因此，在这篇论文中，作者首次提出了一种全新的Transformer架构来解决这一问题，如图1-2所示。当然，Transformer架构的优点在于它完全摈弃了传统的循环结构，取而代之的是只通过注意力机制来计算模型输入与输出的隐含表示，而这种注意力的名字就是大名鼎鼎的自注意力机制（self-attention），也就是图1-2中的Multi-Head Attention模块。
To the best of our knowledge, however, the Transformer is the first transduction model relying entirely on self-attention to compute representations of its input and output without using sequence- aligned RNNs or convolution.
总体来说，所谓自注意力机制就是通过某种运算来直接计算得到句子在编码过程中每个位置上的注意力权重；然后再以权重和的形式来计算得到整个句子的隐含向量表示。最终，Transformer架构就是基于这种的自注意力机制而构建的Encoder-Decoder模型。
1.2 技术手段
在介绍完整篇论文的提出背景后，下面就让我们一起首先来看一看自注意力机制的庐山真面目，然后再来探究整体的网络架构。
1.2.1 什么是self-Attention
首先需要明白一点的是，所谓的自注意力机制其实就是论文中所指代的“Scaled Dot-Product Attention“。在论文中作者说道，注意力机制可以描述为将query和一系列的key-value对映射到某个输出的过程，而这个输出的向量就是根据query和key计算得到的权重作用于value上的权重和。
An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.
不过想要更加深入的理解query、key和value的含义，得需要结合Transformer的解码过程，这部分内容将会在后续进行介绍。具体的，自注意力机制的结构如图1-3所示。
从图1-3可以看出，自注意力机制的核心过程就是通过Q和K计算得到注意力权重；然后再作用于V得到整个权重和输出。具体的，对于输入Q、K和V来说，其输出向量的计算公式为：
其中Q、K和V分别为3个矩阵，且其（第2个）维度分别为 （从后面的计算过程其实可以发现 (1.1)中除以 的过程就是图1-3中所指的Scale过程。
之所以要进行缩放这一步是因为通过实验作者发现，对于较大的 来说在完成 后将会得到很大的值，而这将导致在经过sofrmax操作后产生非常小的梯度，不利于网络的训练。
We suspect that for large values of dk, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients.
如果仅仅只是看着图1-3中的结构以及公式 中的计算过程显然是不那么容易理解自注意力机制的含义，例如初学者最困惑的一个问题就是图1-3中的Q、K和V分别是怎么来的？下面，我们来看一个实际的计算示例。现在，假设输入序列“我 是 谁”，且已经通过某种方式得到了1个形状为 的矩阵来进行表示，那么通过图1-3所示的过程便能够就算得到Q、K以及V[2]。
从图1-4的计算过程可以看出，Q、K和V其实就是输入X分别乘以3个不同的矩阵计算而来（但这仅仅局限于Encoder和Decoder在各自输入部分利用自注意力机制进行编码的过程，Encoder和Decoder交互部分的Q、K和V另有指代）。此处对于计算得到的Q、K、V，你可以理解为这是对于同一个输入进行3次不同的线性变换来表示其不同的3种状态。在计算得到Q、K、V之后，就可以进一步计算得到权重向量，计算过程如图1-5所示。
如图1-5所示，在经过上述过程计算得到了这个注意力权重矩阵之后我们不禁就会问到，这些权重值到底表示的是什么呢？对于权重矩阵的第1行来说，0.7表示的就是“我”与“我”的注意力值；0.2表示的就是“我”与”是”的注意力值；0.1表示的就是“我”与“谁”的注意力值。换句话说，在对序列中的“我“进行编码时，应该将0.7的注意力放在“我”上，0.2的注意力放在“是”上，将0.1的注意力放在谁上。
同理，对于权重矩阵的第3行来说，其表示的含义就是，在对序列中”谁“进行编码时，应该将0.2的注意力放在“我”上，将0.1的注意力放在“是”上，将0.7的注意力放在“谁”上。从这一过程可以看出，通过这个权重矩阵模型就能轻松的知道在编码对应位置上的向量时，应该以何种方式将注意力集中到不同的位置上。
不过从上面的计算结果还可以看到一点就是，模型在对当前位置的信息进行编码时，会过度的将注意力集中于自身的位置（虽然这符合常识）而可能忽略了其它位置[2]。因此，作者采取的一种解决方案就是采用多头注意力机制（MultiHeadAttention），这部分内容我们将在稍后看到。
It expands the model’s ability to focus on different positions. Yes, in the example above, contains a little bit of every other encoding, but it could be dominated by the the actual word itself.
在通过图1-5示的过程计算得到权重矩阵后，便可以将其作用于V ，进而得到最终的编码输出，计算过程如图1-6所示。
根据如图1-6所示的过程，我们便能够得到最后编码后的输出向量。当然，对于上述过程我们还可以换个角度来进行观察，如图1-7所示。
从图1-7可以看出，对于最终输出“是”的编码向量来说，它其实就是原始“我 是 谁”3个向量的加权和，而这也就体现了在对“是”进行编码时注意力权重分配的全过程。
当然，对于整个图1-5到图1-6的过程，我们还可以通过如图1-8所示的过程来进行表示。
可以看出通过这种自注意力机制的方式确实解决了作者在论文伊始所提出的“传统序列模型在编码过程中都需顺序进行的弊端”的问题，有了自注意力机制后，仅仅只需要对原始输入进行几次矩阵变换便能够得到最终包含有不同位置注意力信息的编码向量。
对于自注意力机制的核心部分到这里就介绍完了，不过里面依旧有很多细节之处没有进行介绍。例如Encoder和Decoder在进行交互时的Q、K、V是如何得到的？在图1-3中所标记的Mask操作是什么意思，什么情况下会用到等等？这些内容将会在后续逐一进行介绍。
下面，让我们继续进入到MultiHeadAttention机制的探索中。
1.2.2 为什么要MultiHeadAttention
经过上面内容的介绍，我们算是在一定程度上对于自注意力机制有了清晰的认识，不过在上面我们也提到了自注意力机制的缺陷就是：模型在对当前位置的信息进行编码时，会过度的将注意力集中于自身的位置， 因此作者提出了通过多头注意力机制来解决这一问题。同时，使用多头注意力机制还能够给予注意力层的输出包含有不同子空间中的编码表示信息，从而增强模型的表达能力。
Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions.
在说完为什么需要多头注意力机制以及使用多头注意力机制的好处之后，下面我们就来看一看到底什么是多头注意力机制。
如图1-9所示，可以看到所谓的多头注意力机制其实就是将原始的输入序列进行多组的自注意力处理过程；然后再将每一组自注意力的结果拼接起来进行一次线性变换得到最终的输出结果。具体的，其计算公式为：
其中
同时，在论文中，作者使用了 个并行的自注意力模块（8个头）来构建一个注意力层，并且对于每个自注意力模块都限定了 。从这里其实可以发现，论文中所使用的多头注意力机制其实就是将一个大的高维单头拆分成了 个多头。因此，整个多头注意力机制的计算过程我们可以通过如图1-10所示的过程来进行表示。
注意：图中的d_m 就是指
如图1-10所示，根据输入序列X和 我们就计算得到了 ，进一步根据公式 就得到了单个自注意力模块的输出 ；同理，根据X和 就得到了另外一个自注意力模块输出 。最后，根据公式 将 水平堆叠形成 ，然后再用 乘以 便得到了整个多头注意力层的输出。同时，根据图1-9中的计算过程，还可以得到 。
到此，对于整个Transformer的核心部分，即多头注意力机制的原理就介绍完了。
1.2.3 同维度中的单头与多头的区别
在多头注意力中，对于初学者来说一个比较经典的问题就是，在相同维度下使用单头和多头的区别是什么？这句话什么意思呢？以图1-10中示例为例，此时的自注意力中使用了两个头，每个头的维度为 ，即采用了多头的方式。另外一种做法就是，只是用一个头，但是其维度为 ，即采用单头的方式。那么在这两种情况下有什么区别呢？
首先，从论文中内容可知，作者在头注意力机制与多头个数之间做了如下的限制
从式 可以看出，单个头注意力机制的维度 乘上多头的个数 就等于模型的维度 。
注意：后续的d_m， 以及 都是指代模型的维度。
同时，从图1-10中可以看出，这里使用的多头数量 ，即 。此时，对于第1个头来说有：
对于第2个头来说有：
最后，可以将 在横向堆叠起来进行一个线性变换得到最终的 。因此，对于图1-10所示的计算过程，我们还可以通过图1-13来进行表示。
从图1-13可知，在一开始初始化 这3个权重矩阵时，可以直接同时初始化 个头的权重，然后再进行后续的计算。而且事实上，在真正的代码实现过程中也是采用的这样的方式，这部分内容将在3.3.2节中进行介绍。因此，对图1-13中的多头计算过程，还可以根据图1-14来进行表示。
说了这么多，终于把铺垫做完了。此时，假如有如图1-15所示的头注意力计算过程：
如图1-15所示，该计算过程采用了头注意力机制来进行计算，且头的计算过程还可通过图1-16来进行表示。
那现在的问题是图1-16中的 能够计算得到吗？答案是不能。为什么？因为我没有告诉你这里的 等于多少。如果我告诉你多头 ，那么毫无疑问图1-16的计算过程就等同于图1-14的计算过程，即
且此时 。但是如果我告诉你多头 ，那么图1-16的计算过程会变成
那么此时 则为 。
现在回到一开始的问题上，根据上面的论述我们可以发现，在 固定的情况下，不管是使用单头还是多头的方式，在实际的处理过程中直到进行注意力权重矩阵计算前，两者之前没有任何区别。当进行进行注意力权重矩阵计算时， 越大那么 就会被切分得越小，进而得到的注意力权重分配方式越多，如图1-19所示。
从图1-19可以看出，如果 ，那么最终可能得到的就是一个各个位置只集中于自身位置的注意力权重矩阵；如果 ，那么就还可能得到另外一个注意力权重稍微分配合理的权重矩阵； 同理如此。因而多头这一做法也恰好是论文作者提出用于克服模型在对当前位置的信息进行编码时，会过度的将注意力集中于自身的位置的问题。这里再插入一张真实场景下同一层的不同注意力权重矩阵可视化结果图：
同时，当 不一样时， 的取值也不一样，进而使得对权重矩阵的scale的程度不一样。例如，如果 ，那么当 时，则 ；当 时，则 。
所以，当模型的维度 确定时，一定程度上 越大整个模型的表达能力越强，越能提高模型对于注意力权重的合理分配。
2. 位置编码与编码解码过程
2.1 Embedding机制
在正式介绍Transformer的网络结构之前，我们先来一起看看Transformer如何对字符进行Embedding处理。
2.1.1 Token Embedding
熟悉文本处理的读者可能都知道，在对文本相关的数据进行建模时首先要做的便是对其进行向量化。例如在机器学习中，常见的文本表示方法有one-hot编码、词袋模型以及TF-IDF等。不过在深度学习中，更常见的做法便是将各个词（或者字）通过一个Embedding层映射到低维稠密的向量空间。因此，在Transformer模型中，首先第一步要做的同样是将文本以这样的方式进行向量化表示，并且将其称之为Token Embedding，也就是深度学习中常说的词嵌入（Word Embedding）如图2-1所示。
如果是换做之前的网络模型，例如CNN或者RNN，那么对于文本向量化的步骤就到此结束了，因为这些网络结构本身已经具备了捕捉时序特征的能力，不管是CNN中的n-gram形式还是RNN中的时序形式。但是这对仅仅只有自注意力机制的网络结构来说却不行。为什么呢？根据自注意力机制原理的介绍我们知道，自注意力机制在实际运算过程中不过就是几个矩阵来回相乘进行线性变换而已。因此，这就导致即使是打乱各个词的顺序，那么最终计算得到的结果本质上却没有发生任何变换，换句话说仅仅只使用自注意力机制会丢失文本原有的序列信息。
如图2-2所示，在经过词嵌入表示后，序列“我 在 看 书”经过了一次线性变换。现在，我们将序列变成“书 在 看 我”，然后同样以中间这个权重矩阵来进行线性变换，过程如图2-3所示。
根据图2-3中的计算结果来看，序列在交换位置前和交换位置后计算得到的结果在本质上并没有任何区别，仅仅只是交换了对应的位置。因此，基于这样的原因，Transformer在原始输入文本进行Token Embedding后，又额外的加入了一个Positional Embedding来刻画数据在时序上的特征。
Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence.
2.1.2 Positional Embedding
说了这么多，那到底什么又是Positional Embedding呢？数无形时少直觉，下面我们先来通过一幅图直观看看经过Positional Embedding处理后到底产生了什么样的变化。
如图2-4所示，横坐标表示输入序列中的每一个Token，每一条曲线或者直线表示对应Token在每个维度上对应的位置信息。在左图中，每个维度所对应的位置信息都是一个不变的常数；而在右图中，每个维度所对应的位置信息都是基于某种公式变换所得到。换句话说就是，左图中任意两个Token上的向量都可以进行位置交换而模型却不能捕捉到这一差异，但是加入右图这样的位置信息模型却能够感知到。例如位置20这一处的向量，在左图中无论你将它换到哪个位置，都和原来一模一样；但在右图中，你却再也找不到与位置20处位置信息相同的位置。
下面，笔者通过两个实际的示例来进行说明。
如图2-5所示，原始输入在经过Token Embedding后，又加入了一个常数位置信息的的Positional Embedding。在经过一次线性变换后便得到了图2-5左右边所示的结果。接下来，我们再交换序列的位置，并同时进行Positional Embedding观察其结果。
如图2-6所示，在交换序列位置后，采用同样的Positional Embedding进行处理，并且进行线性变换。可以发现，其计算结果同图2-5中的计算结果本质上也没有发生变换。因此，这就再次证明，如果Positional Embedding中位置信息是以常数形式进行变换，那么这样的Positional Embedding是无效的。
在Transformer中，作者采用了如公式 所示的规则来生成各个维度的位置信息，其可视化结果如图2-4右所示。
其中 就是这个Positional Embedding矩阵， 表示具体的某一个位置， 表示具体的某一维度。
最终，在融入这种非常数的Positional Embedding位置信息后，便可以得到如图2-7所示的对比结果。
从图2-7可以看出，在交换位置前与交换位置后，与同一个权重矩阵进行线性变换后的结果截然不同。因此，这就证明通过Positional Embedding可以弥补自注意力机制不能捕捉序列时序信息的缺陷。
说完Transformer中的Embedding后，接下来我们再来继续探究Transformer的网络结构。
2.2 Transformer网络结构
如图2-8所示便是一个单层Transformer网络结构图。
如图2-8所示，整个Transformer网络包含左右两个部分，即Encoder和Decoder。下面，我们就分别来对其中的各个部分进行介绍。
2.2.1 Encoder层
首先，对于Encoder来说，其网络结构如图2-8左侧所示（尽管论文中是以6个这样相同的模块堆叠而成，但这里我们先以堆叠一层来进行介绍，多层的Transformer结构将在稍后进行介绍）。
如图2-9所示，对于Encoder部分来说其内部主要由两部分网络所构成：多头注意力机制和两层前馈神经网络。
The encoder is composed of a stack of N = 6 identical layers. Each layer has two sub-layers. The first is a multi-head self-attention mechanism, and the second is a simple, position- wise fully connected feed-forward network.
同时，对于这两部分网络来说，都加入了残差连接，并且在残差连接后还进行了层归一化操作。这样，对于每个部分来说其输出均为，并且在都加入了Dropout操作。
We apply dropout to the output of each sub-layer, before it is added to the sub-layer input and normalized.
进一步，为了便于在这些地方使用残差连接，这两部分网络输出向量的维度均为 。
对于第2部分的两层全连接网络来说，其具体计算过程为
其中输入 的维度为 ，第1层全连接层的输出维度为 ，第2层全连接层的输出为 ，且同时对于第1层网络的输出还运用了Relu激活函数。
到此，对于单层Encoder的网络结构就算是介绍完了，接下来让我们继续探究Decoder部分的网络结构。
2.2.2 Decoder层
同Encoder部分一样，论文中也采用了6个完全相同的网络层堆叠而成，不过这里我们依旧只是先看1层时的情况。对于Decoder部分来说，其整体上与Encoder类似，只是多了一个用于与Encoder输出进行交互的多头注意力机制，如图2-10所示。
不同于Encoder部分，在Decoder中一共包含有3个部分的网络结构。最上面的和最下面的部分（暂时忽略Mask）与Encoder相同，只是多了中间这个与Encoder输出（Memory）进行交互的部分，作者称之为“Encoder-Decoder attention”。对于这部分的输入，Q来自于下面多头注意力机制的输出，K和V均是Encoder部分的输出（Memory）经过线性变换后得到。而作者之所以这样设计也是在模仿传统Encoder-Decoder网络模型的解码过程。
In "encoder-decoder attention" layers, the queries come from the previous decoder layer, and the memory keys and values come from the output of the encoder. This mimics the typical encoder-decoder attention mechanisms in sequence-to-sequence models
为了能够更好的理解这里Q、K、V的含义，我们先来看看传统的基于Encoder-Decoder的Seq2Seq翻译模型是如何进行解码的，如图2-11所示。
如图2-11所示是一个经典的基于Encoder-Decoder的机器翻译模型。左下边部分为编码器，右下边部分为解码器，左上边部分便是注意力机制部分。在图2-11中， 表示的是在编码过程中，各个时刻的隐含状态，称之为每个时刻的Memory； 表示解码当前时刻时的隐含状态。此时注意力机制的思想在于，希望模型在解码的时刻能够参考编码阶段每个时刻的记忆。
因此，在解码第一个时刻"<s>"时，
会首先同每个记忆状态进行相似度比较得到注意力权重。这个注意力权重所蕴含的意思就是，在解码第一个时刻时应该将
的注意力放在编码第一个时刻的记忆上（其它的同理），最终通过加权求和得到4个Memory的权重和，即context vector。同理，在解码第二时刻"我"时，也会遵循上面的这一解码过程。可以看出，此时注意力机制扮演的就是能够使得Encoder与Decoder进行交互的角色。
回到Transformer的Encoder-Decoder attention中，K和V均是编码部分的输出Memory经过线性变换后的结果（此时的Memory中包含了原始输入序列每个位置的编码信息），而Q是解码部分多头注意力机制输出的隐含向量经过线性变换后的结果。在Decoder对每一个时刻进行解码时，首先需要做的便是通过Q与 K进行交互（query查询），并计算得到注意力权重矩阵；然后再通过注意力权重与V进行计算得到一个权重向量，该权重向量所表示的含义就是在解码时如何将注意力分配到Memory的各个位置上。这一过程我们可以通过如图2-12和图2-13所示的过程来进行表示。
如图2-12所示，待解码向量和Memory分别各自乘上一个矩阵后得到Q、K、V。
如图2-13所示，在解码第1个时刻时，首先Q通过与K进行交互得到权重向量，此时可以看做是Q（待解码向量）在K（本质上也就是Memory）中查询Memory中各个位置与Q有关的信息；然后将权重向量与V进行运算得到解码向量，此时这个解码向量可以看作是考虑了Memory中各个位置编码信息的输出向量，也就是说它包含了在解码当前时刻时应该将注意力放在Memory中哪些位置上的信息。
进一步，在得到这个解码向量并经过图2-10中最上面的两层全连接层后，便将其输入到分类层中进行分类得到当前时刻的解码输出值。
2.2.3 Decoder预测解码过程
当第1个时刻的解码过程完成之后，解码器便会将解码第1个时刻时的输入，以及解码第1个时刻后的输出均作为解码器的输入来解码预测第2个时刻的输出。整个过程可以通过如图2-14所示的过程来进行表示。
如图2-14所示，Decoder在对当前时刻进行解码输出时，都会将当前时刻之前所有的预测结果作为输入来对下一个时刻的输出进行预测。假设现在需要将"我 是 谁"翻译成英语"who am i"，且解码预测后前两个时刻的结果为"who am"，接下来需要对下一时刻的输出"i"进行预测，那么整个过程就可以通过图2-15和图2-16来进行表示。
如图2-15所示，左上角的矩阵是解码器对输入"<s> who am"这3个词经过解码器中自注意力机制编码后的结果；左下角是编码器对输入"我 是 谁"这3个词编码后的结果（同图2-12中的一样）；两者分别在经过线性变换后便得到了Q、K和V这3个矩阵。此时值得注意的是，左上角矩阵中的每一个向量在经过自注意力机制编码后，每个向量同样也包含了其它位置上的编码信息。
进一步，Q与K作用和便得到了一个权重矩阵；再将其与V进行线性组合便得到了Encoder-Decoder attention部分的输出，如图2-16所示。
如图2-16所示，左下角便是Q与K作用后的权重矩阵，它的每一行就表示在对Memory（这里指图2-16中的V）中的每一位置进行解码时，应该如何对注意力进行分配。例如第3行
的含义就是在解码当前时刻时应该将
的注意力放在Memory中的"我"上，其它同理。这样，在经过解码器中的两个全连接层后，便得到了解码器最终的输出结果。接着，解码器会循环对下一个时刻的输出进行解码预测，直到预测结果为"<e>"或者达到指定长度后停止。
同时，这里需要注意的是，在通过模型进行实际的预测时，只会取解码器输出的其中一个向量进行分类，然后作为当前时刻的解码输出。例如图2-16中解码器最终会输出一个形状为[3,tgt_vocab_len]的矩阵，那么只会取其最后一个向量喂入到分类器中进行分类得到当前时刻的解码输出。具体细节见后续代码实现。
2.2.4 Decoder训练解码过程
在介绍完预测时Decoder的解码过程后，下面就继续来看在网络在训练过程中是如何进行解码的。从2.2.3小节的内容可以看出，在真实预测时解码器需要将上一个时刻的输出作为下一个时刻解码的输入，然后一个时刻一个时刻的进行解码操作。显然，如果训练时也采用同样的方法那将是十分费时的。因此，在训练过程中，解码器也同编码器一样，一次接收解码时所有时刻的输入进行计算。这样做的好处，一是通过多样本并行计算能够加快网络的训练速度；二是在训练过程中直接喂入解码器正确的结果而不是上一时刻的预测值（因为训练时上一时刻的预测值可能是错误的）能够更好的训练网络。
例如在用平行预料"我 是 谁"<==>"who am i"对网络进行训练时，编码器的输入便是"我 是 谁"，而解码器的输入则是"<s> who am i"，对应的正确标签则是"who am i <e>"。
假设现在解码器的输入"<s> who am i"在分别乘上一个矩阵进行线性变换后得到了Q、K、V，且Q与K作用后得到了注意力权重矩阵（此时还未进行softmax操作），如图2-17所示。
从图2-17可以看出，此时已经计算得到了注意力权重矩阵。由第1行的权重向量可知，在解码第1个时刻时应该将
（严格来说应该是经过softmax后的值）的注意力放到"<s>"上，
的注意力放到"who"上等等。不过此时有一个问题就是，在2.2.3节中笔者介绍到，模型在实际的预测过程中只是将当前时刻之前（包括当前时刻）的所有时刻作为输入来预测下一个时刻，也就是说模型在预测时是看不到当前时刻之后的信息。因此，Transformer中的Decoder通过加入注意力掩码机制来解决了这一问题。
self-attention layers in the decoder allow each position in the decoder to attend to all positions in the decoder up to and including that position. We need to prevent leftward information flow in the decoder to preserve the auto-regressive property. We implement this inside of scaled dot-product attention by masking out (setting to −∞) all values in the input of the softmax which correspond to illegal connections.
如图2-18所示，左边依旧是通过Q和K计算得到了注意力权重矩阵（此时还未进行softmax操作），而中间的就是所谓的注意力掩码矩阵，两者在相加之后再乘上矩阵V便得到了整个自注意力机制的输出，也就是图2-10中的Masked Multi-Head Attention。
那为什么注意力权重矩阵加上这个注意力掩码矩阵就能够达到这样的效果呢？以图2-18中第1行权重为例，当解码器对第1个时刻进行解码时其对应的输入只有"<s>"，因此这就意味着此时应该将所有的注意力放在第1个位置上（尽管在训练时解码器一次喂入了所有的输入），换句话说也就是第1个位置上的权重应该是1，而其它位置则是0。从图2-17可以看出，第1行注意力向量在加上第1行注意力掩码，再经过softmax操作后便得到了一个类似
的向量。那么，通过这个向量就能够保证在解码第1个时刻时只能将注意力放在第1个位置上的特性。同理，在解码后续的时刻也是类似的过程。
到此，对于整个单层Transformer的网络结构以及编码解码过程就介绍完了，更多细节内容见后续代码实现。
2.2.5 位置编码与Attention Mask
在刚接触Transformer的时候，有的人会认为在Decoder中，既然已经有了Attention mask那么为什么还需要Positional Embedding呢？如图2-18所示，持这种观点的朋友认为，Attention mask已经有了使得输入序列依次输入解码器的能力，因此就不再需要Positional Embedding了。这样想对吗？
根据2.2.4节内容的介绍可以知道，Attention mask的作用只有一个，那就是在训练过程中掩盖掉当前时刻之后所有位置上的信息，而这也是在模仿模型在预测时只能看到当前时刻及其之前位置上的信息。因此，持有上述观点的朋友可能是把“能看见”和“能看见且有序”混在一起了。
虽然看似有了Attention mask这个掩码矩阵能够使得Decoder在解码过程中可以有序地看到当前位置之前的所有信息，但是事实上没有Positional Embedding的Attention mask只能做到看到当前位置之前的所有信息，而做不到有序。前者的“有序”指的是喂入解码器中序列的顺序，而后者的“有序”指的是序列本身固有的语序。
如果不加Postional Embedding的话，那么以下序列对于模型来说就是一回事：
<s> → 北 → 京 → 欢 → 迎 → 你 → <e>
<s> → 北 → 京 → 迎 → 欢 → 你 → <e>
<s> → 北 → 京 → 你 → 迎 → 欢 → <e>
虽然此时Attention mask具有能够让上述序列一个时刻一个时刻的按序喂入到解码器中，但是它却无法识别出这句话本身固有的语序。
2.2.6 原始Q、K、V来源
在Transformer中各个部分的Q、K、V到底是怎么来的一直以来都是初学者最大的一个疑问，并且这部分内容在原论文中也没有进行交代，只是交代了如何根据Q、K、V来进行自注意力机制的计算。虽然在第2部分的前面几个小节已经提及过了这部分内容，但是这里再给大家进行一次总结。
根据图2-8（Transformer结构图）可知，在整个Transformer中涉及到自注意力机制的一共有3个部分：Encoder中的Multi-Head Attention；Decoder中的Masked Multi-Head Attention；Encoder和Decoder交互部分的Multi-Head Attention。
① 对于Encoder中的Multi-Head Attention来说，其原始q、k、v均是Encoder的Token输入经过Embedding后的结果。q、k、v分别经过一次线性变换（各自乘以一个权重矩阵）后得到了Q、K、V（也就是图1-4中的示例），然后再进行自注意力运算得到Encoder部分的输出结果Memory。
② 对于Decoder中的Masked Multi-Head Attention来说，其原始q、k、v均是Decoder的Token输入经过Embedding后的结果。q、k、v分别经过一次线性变换后得到了Q、K、V，然后再进行自注意力运算得到Masked Multi-Head Attention部分的输出结果，即待解码向量。
对于Encoder和Decoder交互部分的Multi-Head Attention，其原始q、k、v分别是上面的带解码向量、Memory和Memory。q、k、v分别经过一次线性变换后得到了Q、K、V（也就是图2-12中的示例），然后再进行自注意力运算得到Decoder部分的输出结果。之所以这样设计也是在模仿传统Encoder-Decoder网络模型的解码过程。
3. 网络结构与自注意力实现
在通过前面几部分内容详细介绍完Transformer网络结构的原理后，接下来就让我们来看一看如何借用Pytorch框架来实现MultiHeadAttention这一结构。同时，需要说明的一点是，下面所有的实现代码都是笔者直接从Pytorch 1.4版本中torch.nn.Transformer模块里摘取出来的简略版，目的就是为了让大家对于整个实现过程有一个清晰的认识。并且为了使得大家在阅读完以下内容后也能够对Pytorch中的相关模块有一定的了解，所以下面的代码在变量名方面也与Pytorch保持了一致。
3.1 多层Transformer
在第2部分中，笔者详细介绍了单层Transformer网络结构中的各个组成部分。尽管多层Transformer就是在此基础上堆叠而来，不过笔者认为还是有必要在这里稍微提及一下。
如图3-1所示便是一个单层Transformer网络结构图，左边是编码器右边是解码器。而多层的Transformer网络就是在两边分别堆叠了多个编码器和解码器的网络模型，如图3-2所示。
如图3-2所示便是一个多层的Transformer网络结构图（原论文中采用了6个编码器和6个解码器），其中的每一个Encoder都是图3-1中左边所示的网络结构（Decoder同理）。可以发现，它真的就是图3-1堆叠后的形式。不过需要注意的是其整个解码过程。
在多层Transformer中，多层编码器先对输入序列进行编码，然后得到最后一个Encoder的输出Memory；解码器先通过Masked Multi-Head Attention对输入序列进行编码，然后将输出结果同Memory通过Encoder-Decoder Attention后得到第1层解码器的输出；接着再将第1层Decoder的输出通过Masked Multi-Head Attention进行编码，最后再将编码后的结果同Memory通过Encoder-Decoder Attention后得到第2层解码器的输出，以此类推得到最后一个Decoder的输出。
值得注意的是，在多层Transformer的解码过程中，每一个Decoder在Encoder-Decoder Attention中所使用的Memory均是同一个。
3.2 Transformer中的掩码
由于在实现多头注意力时需要考虑到各种情况下的掩码，因此在这里需要先对这部分内容进行介绍。在Transformer中，主要有两个地方会用到掩码这一机制。第1个地方就是在2.2.4节中介绍到的Attention Mask，用于在训练过程中解码的时候掩盖掉当前时刻之后的信息；第2个地方便是对一个batch中不同长度的序列在Padding到相同长度后，对Padding部分的信息进行掩盖。下面分别就这两种情况进行介绍。
3.2.1 Attention Mask
如图3-3所示，在训练过程中对于每一个样本来说都需要这样一个对称矩阵来掩盖掉当前时刻之后所有位置的信息。
从图3-3可以看出，这个注意力掩码矩阵的形状为[tgt_len,tgt_len]，其具体Mask原理在2.2.4节中笔者已经介绍过l，这里就不再赘述。在后续实现过程中，我们将通过generate_square_subsequent_mask方法来生成这样一个矩阵。同时，在后续多头注意力机制实现中，将通过attn_mask这一变量名来指代这个矩阵。
3.2.2 Padding Mask
在Transformer中，使用到掩码的第2个地方便是Padding Mask。由于在网络的训练过程中同一个batch会包含有多个文本序列，而不同的序列长度并不一致。因此在数据集的生成过程中，就需要将同一个batch中的序列Padding到相同的长度。但是，这样就会导致在注意力的计算过程中会考虑到Padding位置上的信息。
如图3-4所示，P表示Padding的位置，右边的矩阵表示计算得到的注意力权重矩阵。可以看到，此时的注意力权重对于Padding位置山的信息也会加以考虑。因此在Transformer中，作者通过在生成训练集的过程中记录下每个样本Padding的实际位置；然后再将注意力权重矩阵中对应位置的权重替换成负无穷，经softmax操作后对应Padding位置上的权重就变成了0，从而达到了忽略Padding位置信息的目的。这种做法也是Encoder-Decoder网络结构中通用的一种办法。
如图3-5所示，对于"我 是 谁 P P"这个序列来说，前3个字符是正常的，后2个字符是Padding后的结果。因此，其Mask向量便为[True, True, True, False, False]。通过这个Mask向量可知，需要将权重矩阵的最后两列替换成负无穷，在后续我们会通过torch.masked_fill这个方法来完成这一步，并且在实现时将使用key_padding_mask来指代这一向量。
到此，对于Transformer中所要用到Mask的地方就介绍完了，下面正式来看如何实现多头注意力机制。
3.3 实现多头注意力机制
3.3.1 多头注意力机制
根据前面的介绍可以知道，多头注意力机制中最为重要的就是自注意力机制，也就是需要前计算得到Q、K和V，如图3-6所示。
然后再根据Q、K、V来计算得到最终的注意力编码，如图3-7所示：
同时，为了避免单个自注意力机制计算得到的注意力权重过度集中于当前编码位置自己所在的位置（同时更应该关注于其它位置），所以作者在论文中提到通过采用多头注意力机制来解决这一问题，如图3-8所示。
3.3.2 定义类MyMultiHeadAttention
综上所述，我们可以给出类MyMultiHeadAttentiond的定义为
class MyMultiheadAttention(nn.Module):
def __init__(self, embed_dim, num_heads, dropout=0., bias=True):
super(MyMultiheadAttention, self).__init__()
"""
:param embed_dim: 词嵌入的维度，也就是前面的d_model参数，论文中的默认值为512
:param num_heads: 多头注意力机制中多头的数量，也就是前面的nhead参数， 论文默认值为 8
:param bias: 最后对多头的注意力（组合）输出进行线性变换时，是否使用偏置
"""
self.embed_dim = embed_dim # 前面的d_model参数
self.head_dim = embed_dim // num_heads # head_dim 指的就是d_k,d_v
self.kdim = self.head_dim
self.vdim = self.head_dim
self.num_heads = num_heads # 多头个数
self.dropout = dropout
assert self.head_dim * num_heads == self.embed_dim, "embed_dim 除以 num_heads必须为整数"
# 上面的限制条件就是论文中的 d_k = d_v = d_model/n_head 条件
self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
# embed_dim = kdim * num_heads
# 这里第二个维度之所以是embed_dim，实际上这里是同时初始化了num_heads个W_q堆叠起来的, 也就是num_heads个头
self.k_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
# W_k, embed_dim = kdim * num_heads
self.v_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
# W_v, embed_dim = vdim * num_heads
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
# 最后将所有的Z组合起来的时候，也是一次性完成， embed_dim = vdim * num_heads
在上述代码中，embed_dim表示模型的维度（图3-8中的d_m）；num_heads表示多头的个数；bias表示是否在多头线性组合时使用偏置。同时，为了使得实现代码更加高效，所以Pytorch在实现的时候是多个头注意力机制一起进行的计算，也就上面代码的第17-22行，分别用来初始化了多个头的权重值（这一过程从图3-8也可以看出）。当多头注意力机制计算完成后，将会得到一个形状为[src_len,embed_dim]的矩阵，也就是图3-8中多个
水平堆叠后的结果。因此，第24行代码将会初始化一个线性层来对这一结果进行一个线性变换。
3.3.3 定义前向传播过程
在定义完初始化函数后，便可以定义如下所示的多头注意力前向传播的过程
def forward(self, query, key, value, attn_mask=None, key_padding_mask=None):
"""
在论文中，编码时query, key, value 都是同一个输入，
解码时 输入的部分也都是同一个输入，
解码和编码交互时 key,value指的是 memory, query指的是tgt
:param query: # [tgt_len, batch_size, embed_dim], tgt_len 表示目标序列的长度
:param key: # [src_len, batch_size, embed_dim], src_len 表示源序列的长度
:param value: # [src_len, batch_size, embed_dim], src_len 表示源序列的长度
:param attn_mask: # [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
一般只在解码时使用，为了并行一次喂入所有解码部分的输入，所以要用mask来进行掩盖当前时刻之后的位置信息
:param key_padding_mask: [batch_size, src_len], src_len 表示源序列的长度
:return:
attn_output: [tgt_len, batch_size, embed_dim]
attn_output_weights: # [batch_size, tgt_len, src_len]
"""
return multi_head_attention_forward(query, key, value, self.num_heads,
self.dropout, self.out_proj.weight, self.out_proj.bias,
training=self.training,
key_padding_mask=key_padding_mask,
q_proj_weight=self.q_proj_weight,
k_proj_weight=self.k_proj_weight,
v_proj_weight=self.v_proj_weight,
attn_mask=attn_mask)
在上述代码中，query、key、value指的并不是图3-6中的Q、K和V，而是没有经过线性变换前的输入。例如在编码时三者指的均是原始输入序列src；在解码时的Mask Multi-Head Attention中三者指的均是目标输入序列tgt；在解码时的Encoder-Decoder Attention中三者分别指的是Mask Multi-Head Attention的输出、Memory和Memory。key_padding_mask指的是编码或解码部分，输入序列的Padding情况，形状为[batch_size,src_len]或者[batch_size,tgt_len];attn_mask指的就是注意力掩码矩阵，形状为[tgt_len,src_len]，它只会在解码时使用。
注意，在上面的这些维度中，tgt_len本质上指的其实是query_len；src_len本质上指的是key_len。只是在不同情况下两者可能会是一样，也可能会是不一样。
3.3.4 多头注意力计算过程
在定义完类MyMultiHeadAttentiond后，就需要定义出多头注意力的实际计算过程。由于这部分代码较长，所以就分层次进行介绍。
def multi_head_attention_forward(
query, # [tgt_len,batch_size, embed_dim]
key, # [src_len, batch_size, embed_dim]
value, # [src_len, batch_size, embed_dim]
num_heads,
dropout_p,
out_proj_weight, # [embed_dim = vdim * num_heads, embed_dim]
out_proj_bias,
training=True,
key_padding_mask=None, # [batch_size,src_len/tgt_len]
q_proj_weight=None, # [embed_dim,kdim * num_heads]
k_proj_weight=None, # [embed_dim, kdim * num_heads]
v_proj_weight=None, # [embed_dim, vdim * num_heads]
attn_mask=None, # [tgt_len,src_len]
):
# 第一阶段：计算得到Q、K、V
q = F.linear(query, q_proj_weight)
# [tgt_len,batch_size,embed_dim] x [embed_dim,kdim * num_heads]
# = [tgt_len,batch_size,kdim * num_heads]
k = F.linear(key, k_proj_weight)
# [src_len, batch_size,embed_dim] x [embed_dim,kdim * num_heads]
# = [src_len,batch_size,kdim * num_heads]
v = F.linear(value, v_proj_weight)
# [src_len, batch_size,embed_dim] x [embed_dim,vdim * num_heads]
# = [src_len,batch_size,vdim * num_heads]
在上述代码中，第17-23行所做的就是根据输入进行线性变换得到图3-6中的Q、K和V。
# 第二阶段：缩放，以及attn_mask维度判断
tgt_len, bsz, embed_dim = query.size() # [tgt_len,batch_size, embed_dim]
src_len = key.size(0)
head_dim = embed_dim // num_heads # num_heads * head_dim = embed_dim
scaling = float(head_dim) ** -0.5
q = q * scaling # [query_len,batch_size,kdim * num_heads]
if attn_mask is not None:
# [tgt_len,src_len] or [num_heads*batch_size,tgt_len, src_len]
if attn_mask.dim() == 2:
attn_mask = attn_mask.unsqueeze(0) # [1, tgt_len,src_len] 扩充维度
if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
raise RuntimeError('The size of the 2D attn_mask is not correct.')
elif attn_mask.dim() == 3:
if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
raise RuntimeError('The size of the 3D attn_mask is not correct.')
# 现在 atten_mask 的维度就变成了3D
接着，在上述代码中第5-6行所完成的就是图3-7中的缩放过程；第8-16行用来判断或修改attn_mask的维度，当然这几行代码只会在解码器中的Masked Multi-Head Attention中用到。
# 第三阶段：计算得到注意力权重矩阵
q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
# [batch_size * num_heads,tgt_len,kdim]
# 因为前面是num_heads个头一起参与的计算，所以这里要进行一下变形，以便于后面计算。且同时交换了0，1两个维度
k = k.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)
#[batch_size * num_heads,src_len,kdim]
v = v.contiguous().view(-1, bsz*num_heads, head_dim).transpose(0,1)
#[batch_size * num_heads,src_len,vdim]
attn_output_weights = torch.bmm(q, k.transpose(1, 2))
# [batch_size * num_heads,tgt_len,kdim] x [batch_size * num_heads, kdim, src_len]
# = [batch_size * num_heads, tgt_len, src_len] 这就num_heads个QK相乘后的注意力矩阵
继续，在上述代码中第1-7行所做的就是交换Q、K、V中的维度，以便于多个样本同时进行计算；第9行代码便是用来计算注意力权重矩阵；其中上contiguous()方法是将变量放到一块连续的物理内存中；bmm的作用是用来计算两个三维矩阵的乘法操作[4]。
需要提示的是，大家在看代码的时候，最好是仔细观察一下各个变量维度的变化过程，笔者也在每次运算后进行了批注。
# 第四阶段：进行相关掩码操作
if attn_mask is not None:
attn_output_weights += attn_mask # [batch_size * num_heads, tgt_len, src_len]
if key_padding_mask is not None:
attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
# 变成 [batch_size, num_heads, tgt_len, src_len]的形状
attn_output_weights = attn_output_weights.masked_fill(
key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
# 扩展维度，从[batch_size,src_len]变成[batch_size,1,1,src_len]
attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len,src_len)
# [batch_size * num_heads, tgt_len, src_len]
进一步，在上述代码中第2-3行便是用来执行图3-3中的步骤；第4-8行便是用来执行图3-5中的步骤，同时还进行了维度扩充。
attn_output_weights = F.softmax(attn_output_weights, dim=-1)
# [batch_size * num_heads, tgt_len, src_len]
attn_output_weights = F.dropout(attn_output_weights, p=dropout_p, training=training)
attn_output = torch.bmm(attn_output_weights, v)
# Z = [batch_size * num_heads, tgt_len, src_len] x [batch_size * num_heads,src_len,vdim]
# = # [batch_size * num_heads,tgt_len,vdim]
# 这就num_heads个Attention(Q,K,V)结果
attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
# 先transpose成 [tgt_len, batch_size* num_heads ,kdim]
# 再view成 [tgt_len,batch_size,num_heads*kdim]
attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
Z = F.linear(attn_output, out_proj_weight, out_proj_bias)
# 这里就是多个z 线性组合成Z [tgt_len,batch_size,embed_dim]
return Z, attn_output_weights.sum(dim=1) / num_heads # 将num_heads个注意力权重矩阵按对应维度取平均
最后，在上述代码中第1-4行便是用来对权重矩阵进行归一化操作，以及计算得到多头注意力机制的输出；第14行代码便是用来对多个注意力的输出结果进行线性组合；第16行代码用来返回线性组合后的结果，以及多个注意力权重矩阵的平均值。
3.3.5 示例代码
在实现完类MyMultiHeadAttention的全部代码后，便可以通过类似如下的方式进行使用。
if __name__ == '__main__': src_len = 5 batch_size = 2 dmodel = 32 num_head = 1 src = torch.rand((src_len, batch_size, dmodel)) # shape: [src_len, batch_size, embed_dim] src_key_padding_mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]]) # shape: [src_len, src_len] my_mh = MyMultiheadAttention(embed_dim=dmodel, num_heads=num_head) r = my_mh(src, src, src,key_padding_mask = src_key_padding_mask)
在上述代码中，第6-11行其实也就是Encoder中多头注意力机制的实现过程。同时，在计算过程中还可以打印出各个变量的维度变化信息：
进入多头注意力计算: 多头num_heads = 1, d_model=32, d_k = d_v = d_model/num_heads=32 query的shape([tgt_len, batch_size, embed_dim]):torch.Size([5, 2, 32]) W_q 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32]) Q 的shape([tgt_len, batch_size,kdim * num_heads]):torch.Size([5, 2, 32]) ---------------------------------------------------------------------- key 的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32]) W_k 的shape([embed_dim,kdim * num_heads]):torch.Size([32, 32]) K 的shape([src_len,batch_size,kdim * num_heads]):torch.Size([5, 2, 32]) ---------------------------------------------------------------------- value的shape([src_len,batch_size, embed_dim]):torch.Size([5, 2, 32]) W_v 的shape([embed_dim,vdim * num_heads]):torch.Size([32, 32]) V 的shape([src_len,batch_size,vdim * num_heads]):torch.Size([5, 2, 32]) ---------------------------------------------------------------------- ***** 注意，这里的W_q, W_k, W_v是多个head同时进行计算的. 因此，Q,K,V分别也是包含了多个head的q,k,v堆叠起来的结果 ***** 多头注意力中,多头计算结束后的形状（堆叠）为([tgt_len,batch_size,num_heads*kdim])torch.Size([5, 2, 32]) 多头计算结束后，再进行线性变换时的权重W_o的形状为([num_heads*vdim, num_heads*vdim ])torch.Size([32, 32]) 多头线性变化后的形状为([tgt_len,batch_size,embed_dim]) torch.Size([5, 2, 32])
4 Transformer的实现过程
在前面几部分内容中，笔者陆续介绍了多头注意力机制的原理、Transformer中编码器和解码器的工作流程以及多头注意力的实现过程等。接下来，笔者将会一步一步地来详细介绍如何通过Pytorch框架实现Transformer的整体网络结构，包括Token Embedding、Positional Embedding、编码器和解码器等。
下面，首先要介绍的就是对于Embedding部分的编码实现。
4.1 Embedding 实现
4.1.1 Token Embedding
这里首先要实现的便是最基础的Token Enbedding，也是字符转向量的一种常用做法，如下所示：
class TokenEmbedding(nn.Module): def __init__(self, vocab_size: int, emb_size): super(TokenEmbedding, self).__init__() self.embedding = nn.Embedding(vocab_size, emb_size) self.emb_size = emb_size def forward(self, tokens): """ :param tokens: shape : [len, batch_size] :return: shape: [len, batch_size, emb_size] """ return self.embedding(tokens.long()) * math.sqrt(self.emb_size)
如上代码所示便是TokenEmbedding的实现过程，由于这部分代码并不复杂所以就不再逐行进行介绍。注意，第12行代码对原始向量进行缩放是出自论文中3.4部分的描述。
4.1.2 Positional Embedding
在2.1.2节中笔者已经对Positional Embedding的原理做了详细的介绍，其每个位置的变化方式如式 所示。
进一步，我们还可以对式 中括号内的参数进行化简得到如式 中的形式。
由此，根据式 便可以实现Positional Embedding部分的代码，如下所示：
class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout=0.1, max_len=5000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) pe = torch.zeros(max_len, d_model) # [max_len, d_model] position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) # [max_len, 1] div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # [d_model/2] pe[:, 0::2] = torch.sin(position * div_term) # [max_len, d_model/2] pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0).transpose(0, 1) # [max_len, 1, d_model] self.register_buffer('pe', pe) def forward(self, x): """ :param x: [x_len, batch_size, emb_size] :return: [x_len, batch_size, emb_size] """ x = x + self.pe[:x.size(0), :] # [x_len, batch_size, d_model] return self.dropout(x)
如上代码所示便是整个Positional Embedding的实现过程，其中第5行代码是用来初始化一个全0的位置矩阵（也就是图1中从左往右数第2个矩阵），同时还指定了一个序列的最大长度；第6-10行是用来计算每个维度（每一列）的相关位置信息；第19行代码首先是在位置矩阵中取与输入序列长度相等的前x_len行，然后在加上Token Embedding的结果；第20行是用来返回最后得到的结果并进行Dropout操作。同时，这里需要注意的一点便是，在输入x的维度中batch_size并不是第1个维度。
4.1.3 Embedding代码示例
在实现完这部分代码后，便可以通过如下方式进行使用：
if __name__ == '__main__': x = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]], dtype=torch.long) x = x.reshape(5, 2) # [src_len, batch_size] token_embedding = TokenEmbedding(vocab_size=11, emb_size=512) x = token_embedding(tokens=x) pos_embedding = PositionalEncoding(d_model=512) x = pos_embedding(x=x) print(x.shape) # torch.Size([5, 2, 512])
4.2 Transformer实现
在介绍完Embedding部分的编码工作后，下面就开始正式如何来实现Transformer网络结构。如图4-2所示，对于Transformer网络的实现一共会包含4个部分：TransformerEncoderLayer、TransformerEncoder、TransformerDecoderLayer和TransformerDecoder，其分别表示定义一个单独编码层、构造由多个编码层组合得到的编码器、定义一个单独的解码层以及构造由多个解码层得到的解码器。
需要注意的是，图4-2中的一个EncoderLayer指的就是图3-2中的一个对应的Encoder，DecoderLayer同理。
4.2.1 编码层的实现
首先，我们需要实现最基本的编码层单元，也就是图4-2中的TransformerEncoderLayer，其内部结构为图4-3所示的前向传播过程（不包括Embedding部分）。
对于这部分前向传播过程，可以通过如下代码来进行实现：
class MyTransformerEncoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1): super(MyTransformerEncoderLayer, self).__init__() """ :param d_model: d_k = d_v = d_model/nhead = 64, 模型中向量的维度，论文默认值为 512 :param nhead: 多头注意力机制中多头的数量，论文默认为值 8 :param dim_feedforward: 全连接中向量的维度，论文默认值为 2048 :param dropout: 丢弃率，论文中的默认值为 0.1 """ self.self_attn = MyMultiheadAttention(d_model, nhead, dropout=dropout) self.dropout1 = nn.Dropout(dropout) self.norm1 = nn.LayerNorm(d_model) self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.activation = F.relu self.dropout2 = nn.Dropout(dropout) self.norm2 = nn.LayerNorm(d_model)
在上述代码中，第10行用来定义一个多头注意力机制模块，并传入相应的参数；第11-20行代码便是用来定义其它层归一化和线性变换的模块。在完成类MyTransformerEncoderLayer的初始化后，便可以实现整个前向传播的forward方法：
def forward(self, src, src_mask=None, src_key_padding_mask=None): """ :param src: 编码部分的输入，形状为 [src_len,batch_size, embed_dim] :param src_mask: 编码部分输入的padding情况，形状为 [batch_size, src_len] :return: # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim] """ src2 = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask, )[0] # 计算多头注意力 # src2: [src_len,batch_size,num_heads*kdim] num_heads*kdim = embed_dim src = src + self.dropout1(src2) # 残差连接 src = self.norm1(src) # [src_len,batch_size,num_heads*kdim] src2 = self.activation(self.linear1(src)) # [src_len,batch_size,dim_feedforward] src2 = self.linear2(self.dropout(src2)) # [src_len,batch_size,num_heads*kdim] src = src + self.dropout2(src2) src = self.norm2(src) return src # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]
在上述代码中，第7-8行便是用来实现图4-3中Multi-Head Attention部分的前向传播过程；第10-11行用来实现多头注意力后的Add&Norm部分；第13-16行用来实现图4-3中最上面的Feed Forward部分和Add&Norm部分。
这里再次提醒大家，在阅读代码的时候最好是将对应的维度信息带入以便于理解。
4.2.2 编码器实现
在实现完一个标准的编码层之后，便可以基于此来实现堆叠多个编码层，从而得到Transformer中的编码器。对于这部分内容，可以通过如下代码来实现：
def _get_clones(module, N): return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])class MyTransformerEncoder(nn.Module): def __init__(self, encoder_layer, num_layers, norm=None): super(MyTransformerEncoder, self).__init__() """ encoder_layer: 就是包含有多头注意力机制的一个编码层 num_layers: 克隆得到多个encoder layers 论文中默认为6 norm: 归一化层 """ self.layers = _get_clones(encoder_layer, num_layers) # 克隆得到多个encoder layers 论文中默认为6 self.num_layers = num_layers self.norm = norm
在上述代码中，第1-2行是用来定义一个克隆多个编码层或解码层功能函数；第12行中的encoder_layer便是一个实例化的编码层，self.layers中保存的便是一个包含有多个编码层的ModuleList。在完成类MyTransformerEncoder的初始化后，便可以实现整个前向传播的forward方法：
def forward(self, src, mask=None, src_key_padding_mask=None): """ :param src: 编码部分的输入，形状为 [src_len,batch_size, embed_dim] :param mask: 编码部分输入的padding情况，形状为 [batch_size, src_len] :return:# [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim] """ output = src for mod in self.layers: output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask) # 多个encoder layers层堆叠后的前向传播过程 if self.norm is not None: output = self.norm(output) return output # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim]
在上述代码中，第8-10行便是用来实现多个编码层堆叠起来的效果，并完成整个前向传播过程；第11-13行用来对多个编码层的输出结果进行层归一化并返回最终的结果。
4.2.3 编码器使用示例
在完成Transformer中编码器的实现过程后，便可以将其用于对输入序列进行编码。例如可以仅仅通过一个编码器对输入序列进行编码，然后将最后的输出喂入到分类器当中进行分类处理，这部分内容在后续也会进行介绍。下面先看一个使用示例。
if __name__ == '__main__': src_len = 5 batch_size = 2 dmodel = 32 num_head = 3 num_layers = 2 src = torch.rand((src_len, batch_size, dmodel)) # shape: [src_len, batch_size, embed_dim] src_key_padding_mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]]) # shape: [batch_size, src_len] my_transformer_encoder_layer = MyTransformerEncoderLayer(d_model=dmodel, nhead=num_head) my_transformer_encoder = MyTransformerEncoder(encoder_layer=my_transformer_encoder_layer, num_layers=num_layers, norm=nn.LayerNorm(dmodel)) memory = my_transformer_encoder(src=src, mask=None, src_key_padding_mask=src_key_padding_mask) print(memory.shape) # torch.Size([5, 2, 32])
在上述代码中，第2-6行定义了编码器中各个部分的参数值；第11-12行则是首先定义一个编码层，然后再定义由多个编码层组成的编码器；第15-16行便是用来得到整个编码器的前向传播输出结果，并且需要注意的是在编码器中不需要掩盖当前时刻之后的位置信息，所以mask=None。
4.2.4 解码层实现
在介绍完编码器的实现后，下面就开始介绍如何实现Transformer中的解码器部分。同编码器的实现流程一样，首先需要实现的依旧是一个标准的解码层，也就是图4-4所示的前向传播过程（不包括Embedding部分）。
对于这部分前向传播过程，可以通过如下代码来进行实现：
class MyTransformerDecoderLayer(nn.Module): def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1): super(MyTransformerDecoderLayer, self).__init__() """ :param d_model: d_k = d_v = d_model/nhead = 64, 模型中向量的维度，论文默认值为 512 :param nhead: 多头注意力机制中多头的数量，论文默认为值 8 :param dim_feedforward: 全连接中向量的维度，论文默认值为 2048 :param dropout: 丢弃率，论文中的默认值为 0.1 """ self.self_attn = MyMultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout) # 解码部分输入序列之间的多头注意力（也就是论文结构图中的Masked Multi-head attention) self.multihead_attn = MyMultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout) # 编码部分输出（memory）和解码部分之间的多头注意力机制。self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) self.activation = F.relu
在上述代码中，第10行代码用来定义图4-4中Masked Multi-head Attention部分的前向传播过程；第12行则是用来定义图4-4中编码器与解码器交互的多头注意力机制模块；第14-24行是用来定义剩余的全连接层以及层归一化相关操作。在完成类MyTransformerDecoderLayer的初始化后，便可以实现整个前向传播的forward方法：
def forward(self, tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None): """ :param tgt: 解码部分的输入，形状为 [tgt_len,batch_size, embed_dim] :param memory: 编码部分的输出（memory）, [src_len,batch_size,embed_dim] :param tgt_mask: 注意力Mask输入，用于掩盖当前position之后的信息, [tgt_len, tgt_len] :param memory_mask: 编码器-解码器交互时的注意力掩码，一般为None :param tgt_key_padding_mask: 解码部分输入的padding情况，形状为 [batch_size, tgt_len] :param memory_key_padding_mask: 编码部分输入的padding情况，形状为 [batch_size, src_len] :return:# [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim] """ tgt2 = self.self_attn(tgt, tgt, tgt, # [tgt_len,batch_size, embed_dim] attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0] # 解码部分输入序列之间的多头注意力（也就是论文结构图中的Masked Multi-head attention) tgt = tgt + self.dropout1(tgt2) # 接着是残差连接 tgt = self.norm1(tgt) # [tgt_len,batch_size, embed_dim] tgt2 = self.multihead_attn(tgt, memory, memory, # [tgt_len, batch_size, embed_dim] attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0] # 解码部分的输入经过多头注意力后同编码部分的输出（memory）通过多头注意力机制进行交互 tgt = tgt + self.dropout2(tgt2) # 残差连接 tgt = self.norm2(tgt) # [tgt_len, batch_size, embed_dim] tgt2 = self.activation(self.linear1(tgt)) # [tgt_len, batch_size, dim_feedforward] tgt2 = self.linear2(self.dropout(tgt2)) # [tgt_len, batch_size, embed_dim] # 最后的两层全连接 tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]
在上述代码中，第12-14行用来完成图4-4中Masked Multi-head Attention部分的前向传播过程，其中的tgt_mask就是在训练时用来掩盖当前时刻之后位置的注意力掩码；第16-17行用来完成图4-4中Masked Multi-head Attention之后Add&Norm部分的前向传播过程；第19-21行用来实现解码器与编码器之间的交互过程，其中memory_mask为None，memory_key_padding_mask为src_key_padding_mask用来对编码器的输出进行（序列）填充部分的掩盖，这一点同编码器中的key_padding_mask原理一样；第23-31行便是用来实现余下的其它过程。
4.2.5 解码器实现
在实现完一个标准的解码层之后，便可以基于此来实现堆叠多个解码层，从而得到Transformer中的解码器。对于这部分内容，可以通过如下代码来实现：
class MyTransformerDecoder(nn.Module):
def __init__(self, decoder_layer, num_layers, norm=None):
super(MyTransformerDecoder, self).__init__()
self.layers = _get_clones(decoder_layer, num_layers)
self.num_layers = num_layers
self.norm = norm
def forward(self, tgt, memory, tgt_mask=None, memory_mask=None, tgt_key_padding_mask=None,
memory_key_padding_mask=None):
"""
:param tgt: 解码部分的输入，形状为 [tgt_len,batch_size, embed_dim]
:param memory: 编码部分最后一层的输出 [src_len,batch_size, embed_dim]
:param tgt_mask: 注意力Mask输入，用于掩盖当前position之后的信息, [tgt_len, tgt_len]
:param memory_mask: 编码器-解码器交互时的注意力掩码，一般为None
:param tgt_key_padding_mask: 解码部分输入的padding情况，形状为 [batch_size, tgt_len]
:param memory_key_padding_mask: 编码部分输入的padding情况，形状为 [batch_size, src_len]
:return: # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]
"""
output = tgt # [tgt_len,batch_size, embed_dim]
for mod in self.layers: # 这里的layers就是N层解码层堆叠起来的
output = mod(output, memory,
tgt_mask=tgt_mask,
memory_mask=memory_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
memory_key_padding_mask=memory_key_padding_mask)
if self.norm is not None:
output = self.norm(output)
return output # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]
在上述代码中，第4行用来克隆得到多个解码层；第20-25行用来实现多层解码层的前向传播过程；第28行便是用来返回最后的结果。
4.2.6 Transformer网络实现
在实现完Transformer中各个基础模块的话，下面就可以来搭建最后的Transformer模型了。总体来说这部分的代码也相对简单，只需要将上述编码器解码器组合到一起即可，具体代码如下所示：
class MyTransformer(nn.Module):
def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
num_decoder_layers=6, dim_feedforward=2048, dropout=0.1):
super(MyTransformer, self).__init__()
"""
:param d_model: d_k = d_v = d_model/nhead = 64, 模型中向量的维度，论文默认值为 512
:param nhead: 多头注意力机制中多头的数量，论文默认为值 8
:param num_encoder_layers: encoder堆叠的数量，也就是论文中的N，论文默认值为6
:param num_decoder_layers: decoder堆叠的数量，也就是论文中的N，论文默认值为6
:param dim_feedforward: 全连接中向量的维度，论文默认值为 2048
:param dropout: 丢弃率，论文中的默认值为 0.1
"""
# ================ 编码部分 =====================
encoder_layer = MyTransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
encoder_norm = nn.LayerNorm(d_model)
self.encoder = MyTransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
# ================ 解码部分 =====================
decoder_layer = MyTransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
decoder_norm = nn.LayerNorm(d_model)
self.decoder = MyTransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
self._reset_parameters() # 初始化模型参数
self.d_model = d_model
self.nhead = nhead
在上述代码中，第15-17行是用来定义编码器部分；第19-21行是用来定义解码器部分；第22行用来以某种方式初始化Transformer中的权重参数，具体实现在稍后的内容中。在定义完类MyTransformer的初始化函数后，便可以来实现Transformer的整个前向传播过程，代码如下：
def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None): """ :param src: [src_len,batch_size,embed_dim] :param tgt: [tgt_len, batch_size, embed_dim] :param src_mask: None :param tgt_mask: [tgt_len, tgt_len] :param memory_mask: None :param src_key_padding_mask: [batch_size, src_len] :param tgt_key_padding_mask: [batch_size, tgt_len] :param memory_key_padding_mask: [batch_size, src_len] :return: [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim] """ memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask) # [src_len, batch_size, num_heads * kdim] <==> [src_len,batch_size,embed_dim] output = self.decoder(tgt=tgt, memory=memory, tgt_mask=tgt_mask, memory_mask=memory_mask, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask) return output # [tgt_len, batch_size, num_heads * kdim] <==> [tgt_len,batch_size,embed_dim]
在上述代码中，src表示编码器的输入；tgt表示解码器的输入；src_mask为空，因为编码时不需要对当前时刻之后的位置信息进行掩盖；tgt_mask用于掩盖解码输入中当前时刻以后的所有位置信息；memory_mask为空；src_key_padding_mask表示对编码输入序列填充部分的Token进行mask；tgt_key_padding_mask表示对解码输入序列填充部分的Token进行掩盖；memory_key_padding_mask表示对编码器的输出部分进行掩盖，掩盖原因等同于编码输入时的mask操作。
到此，对于整个Transformer的网络结构就算是搭建完毕了，不过这还没有实现论文中基于Transformer结构的翻译模型，而这部分内容笔者也将会在下一节中进行详细的介绍。当然，出了上述模块之外，Transformer中还有两个部分需要实现的就是参数初始化方法和注意力掩码矩阵生成方法，具体代码如下：
def _reset_parameters(self): for p in self.parameters(): if p.dim() > 1: xavier_uniform_(p) def generate_square_subsequent_mask(self, sz): mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) return mask # [sz,sz]
4.2.7 Transfromer使用示例
在实现完Transformer的整个完了结构后，便可以通过如下步骤进行使用：
if __name__ == '__main__': src_len = 5 batch_size = 2 dmodel = 32 tgt_len = 6 num_head = 8 src = torch.rand((src_len, batch_size, dmodel)) # shape: [src_len, batch_size, embed_dim] src_key_padding_mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]]) # shape: [batch_size, src_len] tgt = torch.rand((tgt_len, batch_size, dmodel)) # shape: [tgt_len, batch_size, embed_dim] tgt_key_padding_mask = torch.tensor([[True, True, True, False, False, False], [True, True, True, True, False, False]]) # shape: [batch_size, tgt_len] my_transformer = MyTransformer(d_model=dmodel, nhead=num_head, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=500) tgt_mask = my_transformer.generate_square_subsequent_mask(tgt_len) out = my_transformer(src=src, tgt=tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=src_key_padding_mask) print(out.shape) #torch.Size([6, 2, 32])
在上述代码中，第7-13行用来生成模拟的输入数据；第15-16行用来实例化类MyTransformer；第17行用来生成解码输入时的注意力掩码矩阵；第18-21行用来执行Transformer网络结构的前向传播过程。
5 总结
在本篇文章中，笔者首先详细地介绍了Transformer论文的动机以及自注意力机制的原理与多头注意力机制的作用；然后介绍了Transformer中的位置编码以及整个预测和训练过程中模型的编码解码过程；接着进一步介绍了Transformer的网络结构、自注意力机制的原理实现以及Transformer网络的实现。
引用
本文亮点总结
如果觉得有用，就请分享到朋友圈吧！
公众号后台回复“CVPR21检测”获取CVPR2021目标检测论文下载～
# CV技术社群邀请函 #
备注：姓名-学校/公司-研究方向-城市（如：小极-北大-目标检测-深圳）
即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群
每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~
相关内容
微信扫码咨询专知VIP会员