Transformer

首先我们看看网络架构：

Transformer改进了RNN被人诟病的训练慢的特点，利用self-attention可以实现快速并行。

直观认识

Transformer主要由encoder和decoder两部分组成。在Transformer的论文中，encoder和decoder均由6个encoder layer和decoder layer组成，通常我们称之为encoder block。

每一个encoder和decoder的内部简版结构如下图

对于Encoder，包含两层，一个self-attention层和一个前馈神经网络，self-attention能帮助当前节点不仅仅只关注当前的词，从而能获取到上下文的语义。
decoder也包含encoder提到的两层网络，但是在这两层中间还有一层attention层，帮助当前节点获取到当前需要关注的重点内容。
首先，模型需要对输入的数据进行一个embedding操作，enmbedding结束之后，输入到encoder层，self-attention处理完数据后把数据送给前馈神经网络，前馈神经网络的计算可以并行，得到的输出会输入到下一个encoder。

Transformer的结构

Transformer的结构解析出来如下图表示，包括Input Embedding, Position Embedding, Encoder, Decoder。

Embedding

字向量与位置编码的公式表示如下：

\[X=EmbeddingLookup(X)+PositionEncoding \]

Input Embedding

可以将Input Embedding看作是一个 lookup table，对于每个 word，进行 word embedding 就相当于一个lookup操作，查出一个对应结果。返回的结果为\(batch\_size \times sr\_len \times embed\_size\)

Position Encoding

Transformer模型中还缺少一种解释输入序列中单词顺序的方法。为了处理这个问题，transformer给encoder层和decoder层的输入添加了一个额外的向量Positional Encoding，维度和embedding的维度一样，这个向量采用了一种很独特的方法来让模型学习到这个值，这个向量能决定当前词的位置，或者说在一个句子中不同的词之间的距离。这个位置向量的具体计算方法有很多种，论文中的计算方法如下

\[p_{i,2j}=sin(\frac{i}{10000^{\frac{2j}{d}}}) \]

\[p_{i,2j+1}=cos(\frac{i}{10000^{\frac{2j}{d}}}) \]

其中pos是指当前词在句子中的位置，i是指向量中每个值的index，可以看出，在偶数位置，使用正弦编码，在奇数位置，使用余弦编码.

Encoder

用公式把一个Transformer Encoder block 的计算过程整理一下:

• 自注意力机制:

\[Q=XW_Q \]

\[K=XW_K \]

\[V=XW_V \]

\[X_{attention}=self-Attention(Q,K,V) \]

• self-attention 残差连接与 Layer Normalization

\[X_{attention}=LayerNorm(X_{attention}) \]

• FeedForward，其实就是两层线性映射并用激活函数激活，比如说RELU

\[X_{hidden}=Linear(RELU(Linear(X_{attention}))) \]

• FeedForward 残差连接与 Layer Normalization

\[X_{hidden}=X_{attention}+X_{hidden} \]

\[X_{hidden}=LayerNorm(X_{hidden}) \]

其中：\(X_{hidden} \in R^{batch\_size∗seq\_len∗embed\_dim}\)

自注意力机制

1. 首先，自注意力机制（self-attention）会计算出三个新的向量，在论文中，向量的维度是512维，我们把这三个向量分别称为Query、Key、Value，这三个向量是用embedding向量与一个矩阵相乘得到的结果，这个矩阵是随机初始化的，维度为（64，512）注意第二个维度需要和embedding的维度一样，其值在反向传播的过程中会一直进行更新，得到的这三个向量的维度是64低于embedding维度的。
   
2. 计算self-attention的分数值，该分数值决定了当我们在某个位置encode一个词时，对输入句子的其他部分的关注程度。这个分数值的计算方法是Query与Key做点乘，以下图为例，首先我们需要针对Thinking这个词，计算出其他词对于该词的一个分数值，首先是针对于自己本身即q1·k1，然后是针对于第二个词即q1·k2。
   
3. 接下来，把点乘的结果除以一个常数，这里我们除以8，这个值一般是采用上文提到的矩阵的第一个维度的开方即64的开方8，当然也可以选择其他的值，然后把得到的结果做一个softmax的计算。得到的结果即是每个词对于当前位置的词的相关性大小，当然，当前位置的词相关性肯定会很大
   
4. 下一步就是把Value和softmax得到的值进行相乘，并相加，得到的结果即是self-attetion在当前节点的值。
   
   在实际的应用场景，为了提高计算速度，我们采用的是矩阵的方式，直接计算出Query, Key, Value的矩阵，然后把embedding的值与三个矩阵直接相乘，把得到的新矩阵Q与K相乘，乘以一个常数，做softmax操作，最后乘上V矩阵
   
   这种通过 query 和 key 的相似性程度来确定 value 的权重分布的方法被称为scaled dot-product attention。

用公式表达如下：

\[Q=XW_Q \]

\[K=XW_K \]

\[V=XW_V \]

\[X_{attention}=selfAttention(Q,K,V) \]

Self-Attention 复杂度

Self-Attention时间复杂度：\(O(n^2⋅d)\)，这里，n是序列的长度，d是embedding的维度。
Self-Attention包括三个步骤：相似度计算，softmax和加权平均，它们分别的时间复杂度是：
相似度计算可以看作大小为\((n,d)\)和\((d,n)\)的两个矩阵相乘： \((n,d)∗(d,n)=(n^2⋅d)\)，得到一个\((n,n)\)的矩阵
softmax就是直接计算了，时间复杂度为: \(O(n^2)\)
加权平均可以看作大小为\((n,n)\)和\((n,d)\)的两个矩阵相乘： \((n,d)∗(d,n)=(n^2⋅d)\)，得到一个\((n,d)\)的矩阵
因此，Self-Attention的时间复杂度是: \(O(n^2⋅d)\)

Multi-head Attention

不仅仅只初始化一组Q、K、V的矩阵，而是初始化多组，tranformer是使用了8组，所以最后得到的结果是8个矩阵。

multi-head注意力的全过程如下，首先输入句子，"Thinking Machines",在embedding模块把句子中的每个单词变成向量X，在encoder层中，除了第0层有embedding操作外，其他的层没有embedding操作；接着把X分成8个head，

Multi-Head Attention复杂度

多头的实现不是循环的计算每个头，而是通过transposes and reshapes，用矩阵乘法来完成的。
Transformer/BERT中把\(d\)，也就是\(\frac{hidden\_size}{embedding\_size}\)这个维度做了reshape拆分。并将\(num\_attention\_heads\)维度\(transpose\)到前面，使得Q和K的维度都是\((m,n,a)\)，这里不考虑\(batch\)维度。
这样点积可以看作大小为\((m,n,a)\)和\((m,a,n)\)的两个张量相乘，得到一个\((m,n,n)\)的矩阵，其实就相当于\((n,a)\)和\((a,n)\)的两个矩阵相乘，做了\(m\)次，时间复杂度是：

\[O(n^2⋅m⋅a)=O(n^2⋅d) \]

因此Multi-Head Attention时间复杂度也是\(O(n^2⋅d)\)，复杂度相较单头并没有变化，主要还是transposes and reshapes 的操作，相当于把一个大矩阵相乘变成了多个小矩阵的相乘。
详细的维度变化可以看下面。

残差连接（避免梯度消失）

其实这个之前再ResNet中说过。
残差网络（避免梯度消失，例如：\(w_3(w_2(w_1 * X+b_1)+b_2)+b_3\)，如果\(w_1，w_2，w_3\)特别小，0.0000000000000000……1，\(x\)就没了\(w_3(w_2(w_1 * X+b_1)+b_2)+b_3+X\)）所以这里加上原来的\(X\)。

\[X_{hidden}=X_{embedding}+self Attention(Q,K,V) \]

除了self-attention这里做残差连接外，feed forward那个地方也需要残差连接，公式类似：

\[X_{hidden}=X_{feed_forward}+X_{hidden} \]

Layer Normalization（避免梯度爆炸）

归一化（LayerNorm Normalization），做标准化（避免梯度爆炸）。
经过 self-attention 加权之后输出，也就是Attention(Q,K,V) ，然后把他们加起来做残差连接。
Layer Normalization 的作用是把神经网络中隐藏层归一为标准正态分布，也就是独立同分布，以起到加快训练速度，加速收敛的作用\(X_{hidden}=LayerNorm(X_{hidden})\)
其中：\(X_{hidden} \in R^{batch\_size*seq\_len*embed\_dim}\)
LayerNorm的详细操作如下：

\[\mu_{L}=\dfrac{1}{m}\sum_{i=1}^{m}x_{i} \]

上式以矩阵为单位求均值；

\[\delta^{2}=\dfrac{1}{m}\sum_{i=1}^{m}(x_{i}-\mu)^2 \]

上式以矩阵为单位求方差

\[LN(x_{i})=\alpha \dfrac{x_{i}-\mu_{L}}{\sqrt{\delta^{2}+\epsilon}}+\beta \]

然后用每一列的每一个元素减去这列的均值，再除以这列的标准差，从而得到归一化后的数值，加\(\epsilon\)是为了防止分母为0.此处一般初始化\(\alpha\)为全1，而\(\beta\)为全0.

Feed Forward

前面每一步都在做线性变换，\(wx+b\)，线性变化的叠加永远都是线性变化（线性变化就是空间中平移和扩大缩小），通过 Feed Forward中的 Relu 做一次非线性变换，这样的空间变换可以无限拟合任何一种状态了）

将Multi-Head Attention得到的向量再投影到一个更大的空间（论文里将空间放大了4倍）在那个大空间里可以更方便地提取需要的信息（使用Relu激活函数），最后再投影回token向量原来的空间

\[FFN(x)=ReLU(W_{1}x+b_{1})W_{2}+b_{2} \]

借鉴SVM来理解：SVM对于比较复杂的问题通过将特征其投影到更高维的空间使得问题简单到一个超平面就能解决。这里token向量里的信息通过Feed Forward Layer被投影到更高维的空间，在高维空间里向量的各类信息彼此之间更容易区别。

Encoder总结

Eecoders 是\(N=6\)层，通过上图我们可以看到每层\(Encoder\)包括两个\(sub-layers\)：
第一个\(sub-layer\)是\(multi-head self-attention\)，用来计算输入的\(self-attention\)；
第二个\(sub-layer\)是简单的前馈神经网络层\(Feed Forward\)；
注意：在每个\(sub-layer\)我们都模拟了残差网络（在下面的数据流示意图中会细讲），每个\(sub-layer\)的输出都是\(LayerNorm(x+Sub\_layer(x))\)，其中\(sub\_layer\)表示的是该层的上一层的输出。
现在我们给出 Encoder 的数据流示意图，一步一步去剖析

1. 深绿色的\(x_1\)表示\(Embedding\)层的输出，加上代表\(Positional Embedding\)的向量之后，得到最后输入\(Encoder\)中的特征向量，也就是浅绿色向量\(x_1\)；
2. 浅绿色向量\(x_1\)表示单词"Thinking"的特征向量，其中\(x_1\)经过\(Self-Attention\)层，变成浅粉色向量\(z_1\)；
3. \(x_1\)作为残差结构的直连向量，直接和\(z_1\)相加，之后进行\(Layer Norm\)操作，得到粉色向量\(z_1\)；
   • 残差结构的作用：避免出现梯度消失的情况
   • \(Layer Norm\)的作用：避免梯度爆炸,保证数据特征分布的稳定性，并且可以加速模型的收敛
4. \(z_1\)经过前馈神经网络（Feed Forward）层，经过残差结构与自身相加，之后经过 LN 层，得到一个输出向量\(r_1\)；
   • 该前馈神经网络包括两个线性变换和一个ReLU激活函数：\(FFN(x) = max(0,xW_1+b_1)W_2+b2\)
5. 由于\(Transformer\)的\(Encoders\)具有 6 个\(Encoder\)，\(r_1\)也将会作为下一层 \(Encoder\)的输入，代替\(x_1\)的角色，如此循环，直至最后一层\(Encoder\)。
   需要注意的是，上述的\(x、z、r\)都具有相同的维数，论文中为 512 维。

Decoder

和 Encoder 一样，上面三个部分的每一个部分，都有一个残差连接，后接一个 Layer Normalization。Decoder 的中间部件并不复杂，大部分在前面 Encoder 里我们已经介绍过了，但是 Decoder 由于其特殊的功能，因此在训练时会涉及到一些细节，下面会介绍Decoder的Masked Self-Attention和Encoder-Decoder Attention两部分，其结构图如下图所示

Masked Self-Attention

传统Seq2Seq中Decoder使用的是RNN模型，因此在训练过程中输入因此在训练过程中输入t时刻的词，模型无论如何也看不到未来时刻的词，因为循环神经网络是时间驱动的，只有当\(t\)时刻运算结束了，才能看到\(t+1\)时刻的词。而 Transformer Decoder 抛弃了 RNN，改为 Self-Attention，由此就产生了一个问题，在训练过程中，整个 ground truth 都暴露在 Decoder 中，这显然是不对的，我们需要对 Decoder 的输入进行一些处理，该处理被称为 Mask。

Mask 非常简单，首先生成一个下三角全 0，上三角全为负无穷的矩阵，然后将其与 Scaled Scores 相加即可，之后再做 softmax，就能将 -inf 变为 0，得到的这个矩阵即为每个字之间的权重。

Mask 的作用：
遮挡未来的信息：在Transformer模型中，输入序列是逐步解码的，因此在每个解码的时间步，我们只能使用当前和过去的信息，未来的信息是不可见的。因此，通过将未来的位置的attention权重掩盖为负无穷，可以确保模型在生成每个词时只能依赖于过去的信息，避免了数据泄露。这种遮挡未来信息的mask通常被称为"sequence mask"。
下面，让我们看看带掩码的多头注意力层是如何工作的。假设传给解码器的输入句是<sos>Je vais bien。我们知道，自注意力机制将一个单词与句子中的所有单词联系起来，从而提取每个词的更多信息。但这里有一个小问题。在测试期间，解码器只将上一步生成的词作为输入。比如，在测试期间，当t = 2时，解码器的输入中只有[<sos>,Je]，并没有任何其他词。因此，我们也需要以同样的方式来训练模型。模型的注意力机制应该只与该词之前的单词有关，而不是其后的单词。要做到这一点，我们可以掩盖后边所有还没有被模型预测的词。比如，我们想预测与<sos>相邻的单词。在这种情况下，模型应该只看到<sos>，所以我们应该掩盖<sos>后边的所有词。再比如，我们想预测Je后边的词。在这种情况下，模型应该只看到Je之前的词，所以我们应该掩盖Je后边的所有词。其他行同理，如下图所示。

Masked Encoder-Decoder Attention

其实这一部分的计算流程和前面 Masked Self-Attention 很相似，结构也一摸一样，唯一不同的是这里的\(K,V\)为Encoder的输出，\(Q\)为 Decoder 中Masked Self-Attention的输出。

Decoder的解码

下图展示了Decoder的解码过程，Decoder中的字符预测完之后，会当成输入预测下一个字符，知道遇见终止符号为止。

Transformer的最后一层和Softmax

线性层是一个简单的全连接的神经网络，它将解码器堆栈生成的向量投影到一个更大的向量，称为logits向量。如图linear的输出

softmax层将这些分数转换为概率（全部为正值，总和为1.0）。选择概率最高的单元，并生成与其关联的单词作为此时间步的输出。如图softmax的输出。

Transformer的权重共享

Transformer在两个地方进行了权重共享：
（1）Encoder和Decoder间的Embedding层权重共享；
《Attention is all you need》中Transformer被应用在机器翻译任务中，源语言和目标语言是不一样的，但它们可以共用一张大词表，对于两种语言中共同出现的词（比如：数字，标点等等）可以得到更好的表示，而且对于Encoder和Decoder，嵌入时都只有对应语言的embedding会被激活，因此是可以共用一张词表做权重共享的。

论文中，Transformer词表用了bpe来处理，所以最小的单元是subword。英语和德语同属日耳曼语族，有很多相同的subword，可以共享类似的语义。而像中英这样相差较大的语系，语义共享作用可能不会很大。

但是，共用词表会使得词表数量增大，增加softmax的计算时间，因此实际使用中是否共享可能要根据情况权衡。

（2）Decoder中Embedding层和FC层权重共享；

Embedding层可以说是通过onehot去取到对应的embedding向量，FC层可以说是相反的，通过向量（定义为 x）去得到它可能是某个词的softmax概率，取概率最大（贪婪情况下）的作为预测值。

那哪一个会是概率最大的呢？在FC层的每一行量级相同的前提下，理论上和 x 相同的那一行对应的点积和softmax概率会是最大的（可类比本文问题1）。

因此，Embedding层和FC层权重共享，Embedding层中和向量 x 最接近的那一行对应的词，会获得更大的预测概率。实际上，Decoder中的Embedding层和FC层有点像互为逆过程。

通过这样的权重共享可以减少参数的数量，加快收敛。

其实这个Transfomer结构就是：
encoder层： {多头自注意力 + 前馈网络}$ \times n$
decoder层： {Masked 多头自注意力 + encoder-decoder多头自注意力 + 前馈网络} $ \times n$

Transfomer矩阵维度分析

对于一个batch的数据，encoder端的输入大小为：(batch_size, sr_len)；decoder端的输入大小为：(batch_size, tar_len)。不妨假设 encoder layer 及 decoder layer 都只有一层，下面是训练阶段的矩阵维度变化：

训练阶段

训练阶段 Encoder

input_size	Layer	output_size	Layer_parameter_size	Note
batch_size × sr_len	Input Embedding	batch_size × sr_len × embed_size	sr_vocab_size × embed_size	Embedding层的参数即可设为可学习的，也可设为固定参数
batch_size × sr_len × embed_size	Postion Embedding	batch_size × sr_len × embed_size	1 × sr_len × embed_size	固定参数
batch_size × sr_len × embed_size	MultiHead Attention	batch_size × sr_len × hidden_size	（embed_size × hidden_size） × 3 +(hidden_size × hidden_size)	可学习参数
batch_size × sr_len × hidden_size	AddNorm1	batch_size × sr_len × hidden_size	None
batch_size × sr_len × hidden_size	Feed Forward	batch_size × sr_len × hidden_size	(hidden_size × filter_size)+ (filter_size×hidden_size)	可学习参数
batch_size × sr_len × hidden_size	AddNorm2	batch_size × sr_len × hidden_size	None

训练阶段 Decoder

input_size	Layer	output_size	Layer_parameter_size	Note
batch_size × tar_len	Output Embedding	batch_size × tar_len × embed_size	tar_vocab_size × embed_size	Embedding层的参数即可设为可学习的，也可设为固定参数
batch_size × tar_len × embed_size	Postion Embedding	batch_size × tar_len × embed_size	1 × tar_len × embed_size	固定参数
batch_size × tar_len × embed_size	Masked MultiHead Attention	batch_size × tar_len × hidden_size	{embed_size × hidden_size} × 3 +	可学习参数
batch_size × tar_len × hidden_size	AddNorm1	batch_size × tar_len × hidden_size	None
batch_size × tar_len × hidden_size	Encoder-Decoder MultiHead Attention	batch_size × tar_len × hidden_size	{hidden_size × hidden_size} × 4	可学习参数
batch_size × tar_len × hidden_size	AddNorm2	batch_size × tar_len × hidden_size	None
batch_size × tar_len × hidden_size	Feed Forward	batch_size × tar_len × hidden_size	{hidden_size × filter_size}+	可学习参数
batch_size × tar_len × hidden_size	AddNorm3	batch_size × tar_len × hidden_size	None

注意到，为了保持encoder及decoder的层可以堆叠，需要保证每个层的输入和输出的维度一致，因此，需要保证 embed_size = hidden_size

预测阶段

预测阶段的encoder与训练阶段是相同的，只是batch_size=1；而decoder部分由于每个step只能看到当前位置之前的信息，因此每次输入的tar_len也等于1。

input_size	Layer	output_size
\(1 \times 1\)	Output Embedding	\(1 \times 1 \times embed\_size\)
\(1 \times 1 \times embed\_size\)	Postion Embedding	\(1 \times 1 \times embed\_size\)
\(1 \times 1 \times embed\_size\)	Masked MultiHead Attention	\(1 \times 1 \times hidden\_size\)
\(1 \times 1 \times hidden\_size\)	AddNorm1	\(1 \times 1 \times hidden\_size\)
\(1 \times 1 \times hidden\_size\)	Encoder-Decoder MultiHead Attention	\(1 \times 1 \times hidden\_size\)
\(1 \times 1 \times hidden\_size\)	AddNorm2	\(1 \times 1 \times hidden\_size\)
\(1 \times 1 \times hidden\_size\)	Feed Forward	\(1 \times 1 \times hidden\_size\)
\(1 \times 1 \times hidden\_size\)	AddNorm3	\(1 \times 1 \times hidden\_size\)

Multihead Attention解析

训练阶段

Encoder Multihead Attention

1.Input： Encoder Multihead Attention 输入的 query, key, value 是相同的，都是经过了word embedding和pos embedding之后的source sentence，其维度为\(batch\_size \times sr\_len \times hidden\_size\)。由于有num_heads个头需要并行计算，首先query, key, value分别经过一个线性变换，再将数据split给num_heads个头分别做注意力查询，即：
query:
\(batch\_size \times sr\_len\_q \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden\_size}{num\_heads}\)
key:
\(batch\_size \times sr\_len\_q \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden\_size}{num\_heads}\)
value:
\(batch\_size \times sr\_len\_q \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden\_size}{num\_heads}\)

由于query, key, value 是相同的，因此有 sr_len_q = sr_len_k = sr_len_v
2.DotProductAttention： num_heads 个头的计算是并行的，即：

Encoder Multihead Attention中在计算softmax之前对 key 进行了 mask，目的是消除 padding 的影响。事实上，padding不仅对key有影响，对query也有影响，但在实际代码中mask仅针对key，而没有针对query。其实最原始代码是既有key mask，也有query mask的，但后来作者将query mask删去了，因为在最后计算 loss 的时候对 padding 位置的 loss 进行mask，也可达到相同的效果。

假设 batch_size = num_heads = 1，sr_len_q = sr_len_k = 6，source sentence 的最后两个位置是padding，那么Encoder Multihead Attention 中的 mask 为：

\[\begin{pmatrix} 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ \end{pmatrix} \]

即只对 key 的 padding 位置进行了 mask

3.Output： 需要将上面输出的num_heads个头的结果堆叠之后，再做一个线性变换：

\[batch\_size \times num\_heads \times sr\_len\_q \times \frac{hidden_size}{num_heads} \]

\[reshape \]

\[\Downarrow \]

\[batch\_size \times sr\_len\_q \times hidden\_size \]

\[线性变换 \]

\[\Downarrow \]

\[batch\_size \times sr\_len\_q \times hidden\_size \]

Masked Multihead Attention

与 Encoder Multihead Attention 类似，Masked Multihead Attention 输入的 query, key, value 也是相同的，都是经过了word embedding和pos embedding之后的 target sentence。包括后面的计算流程也基本一致。

主要的区别在于：由于在 inference 时，每个 step 位置只能看到它之前的 steps 的信息，而看不到它之后的 steps的信息。因此 Masked Multihead Attention 中的 mask 除了要消除 key 信息里 padding 的影响，还需要消除当前 step 后面的所有 step 的信息：

假设 batch_size = num_heads = 1，tar_len_q = tar_len_k = 5，target sentence 的最后两个位置是 padding，那么Masked Multihead Attention 中的 mask 为：

\[\begin{pmatrix} 1 & 0 & 0 & 0 & 0 & 0\\ 1 & 1 & 0 & 0 & 0 & 0\\ 1 & 1 & 1 & 0 & 0 & 0\\ 1 & 1 & 1 & 0 & 0 & 0\\ 1 & 1 & 1 & 0 & 0 & 0\\ \end{pmatrix} \]

注意到上述 mask 并不是一个单纯的下三角矩阵，因为最后两个位置都是padding，因此无论如何都要被 mask 掉

Encoder-Decoder Multihead Attention

1.Input： Encoder-Decoder Multihead Attention 输入的 query 来自于 target sentence，其维度为\(batch\_size \times tar\_len \times hidden\_size\)；而 key 和 value 则来自于 encoder layer 的输出，其维度为\(batch\_size \times sr\_len \times hidden\_size\)。同样是先做线性变换，再 split 成 num_heads 个头：
query:
\(batch\_size \times tar\_len\_q \times hidden\_size \stackrel{线性变化}{\rightarrow} batch\_size \times tar\_len\_q \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times tar\_len\_q \times \frac{hidden\_size}{num\_heads}\)
key:
\(batch\_size \times sr\_len\_k \times hidden\_size \stackrel{线性变化}{\rightarrow} batch\_size \times sr\_len\_k \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_k \times \frac{hidden\_size}{num\_heads}\)
value:
\(batch\_size \times sr\_len\_v \times hidden\_size \stackrel{线性变化}{\rightarrow} batch\_size \times sr\_len\_v \times hidden\_size \stackrel{reshape}{\rightarrow} batch\_size \times num\_heads \times sr\_len\_v \times \frac{hidden\_size}{num\_heads}\)

这里 sr_len_q ≠ sr_len_k = sr_len_v

2.DotProductAttention： num_heads 个头的计算依然可以并行：

\[query∗key^T=batch\_size×num\_heads×tar\_len\_q×sr\_len\_k \]

\[消除key中padding的影响，对其做mask \]

\[\Downarrow \]

\[masked\_softmax(query∗key^T)=batch\_size×num\_heads×tar\_len\_q×sr\_len\_k \]

\[\Downarrow \]

\[masked\_softmax(query∗key^T)∗value=batch\_size×num\_heads×tar\_len\_q×\frac{hidden_size}{num_heads} \]

假设 batch_size = num_heads = 1，这里sr_len_q可以不等于sr_len_k，不妨假设 sr_len_q = 5，sr_len_k = 6，因为 mask 只针对key，因此这里只需要关注 source sentence 中的padding， 假设 source sentence 的最后两个位置是padding，那么Masked Multihead Attention 中的 mask 为：

\[\begin{pmatrix} 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ 1 & 1 & 1 & 1 & 0 & 0\\ \end{pmatrix} \]

Output： 需要将上面输出的 num_heads 个头的结果堆叠之后，再做一个线性变换：

\[batch\_size×num\_heads×tar\_len_q×\frac{hidden\_size}{num\_heads} \]

\[reshape \]

\[\Downarrow \]

\[batch\_size×tar\_len\_q×hidden\_size \]

\[线性变化 \]

\[\Downarrow \]

\[batch\_size×tar\_len\_q×hidden\_size \]

预测阶段

Encoder Multihead Attention
与训练阶段的 Encoder Multihead Attention 完全相同
虽然在训练阶段，Masked Multihead Attention 会将当前 step 之后的 steps 信息都 mask 掉，但是由于训练时整个 target sentence 都是已知的，因此还是可以做并行运算的。

Masked Multihead Attention
但是在预测阶段，初始的 query, key, value 都只是一个<bos>起始符号，之后每预测出一个 token，这个 token 直接作为下一个 step 输入的 query，而将这个 token 拼在现有的 key 和 value 之后，就是下一个 step 输入的 key 和 value。也就是说，预测阶段每个 step 输入的 query 是上一 step 输出的token，而 key, value 是之前所有 steps 输出的token。

至于 mask 的部分，由于输入中不再含有未来 steps 的信息，因此不再需要用 mask 来消除这部分信息。而对于 key mask，由于 Masked Multihead Attention 的 key 是 target sentence，而在预测完成前 target sentence 的长度是未知的，因此针对 key 的 mask 也是不需要的。 也就是说，Masked Multihead Attention 是不需要 mask 的。

下面是预测阶段 Masked Multihead Attention 的流程：
1.Input： key, value 是到当前 step 为止的所有 steps 的信息，大小为\(1 \times cur\_tar\_len \times hidden\_size\)；而 query 是上一 step 的输出 token，大小为\(1 \times 1 \times hidden\_size\)：

query:
\(1 \times 1 \times hidden\_size \stackrel{线性变化}{\rightarrow} 1 \times 1 \times hidden\_size \stackrel{reshape}{\rightarrow} 1 \times num\_heads \times 1 \times \frac{hidden\_size}{num\_heads}\)
key:
\(1 \times cur\_tar\_len\_k \times hidden\_size \stackrel{线性变化}{\rightarrow} 1 \times cur\_tar\_len\_k \times hidden\_size \stackrel{reshape}{\rightarrow} 1 \times num\_heads \times cur\_tar\_len\_k \times \frac{hidden\_size}{num\_heads}\)
value:
\(1 \times cur\_tar\_len\_v \times hidden\_size \stackrel{线性变化}{\rightarrow} 1 \times cur\_tar\_len\_v \times hidden\_size \stackrel{reshape}{\rightarrow} 1 \times num\_heads \times cur\_tar\_len\_v \times \frac{hidden\_size}{num\_heads}\)

2.DotProductAttention： num_heads 个头的计算依然可以并行

\[query∗key^T=1×num\_heads×1×cur\_tar\_len\_k \]

\[\Downarrow \]

\[softmax(query∗key^T)=1×num\_heads×1×cur\_tar\_len\_k \]

\[\Downarrow \]

\[softmax(query∗key^T)∗value=1×num\_heads×1×\frac{hidden\_size}{num\_heads} \]

3.Output： 需要将上面输出的 num_heads 个头的结果堆叠之后，再做一个线性变换：

\[1 × num\_heads × 1 ×\frac{hidden\_size}{num\_heads} \]

\[reshape \]

\[\Downarrow \]

\[1 × 1 × hidden\_size \]

\[线性变换 \]

\[\Downarrow \]

\[1 × 1 × hidden\_size \]

Encoder-Decoder Multihead Attention

预测阶段 Encoder-Decoder Multihead Attention 输入的 query 是上一层 Masked Multihead Attention 的输出，大小为 \(1 \times 1 \times hidden\_size\)。而输入的 key 和 value 则是 encoder layer 的输出，大小为：\(1 \times sr_len \times hidden\_size\)。具体流程为：
1.Input：
query:
\(1 \times 1 \times hidden\_size \stackrel{线性变化}{\rightarrow} 1 \times 1 \times hidden\_size \stackrel{reshape}{\rightarrow} 1 \times num\_heads \times 1 \times \frac{hidden\_size}{num\_heads}\)
key:
\(1 \times sr\_len\_len\_k \times hidden\_size \stackrel{线性变化}{\rightarrow} 1 \times sr\_len\_k \times hidden\_size \stackrel{reshape}{\rightarrow} 1 \times num\_heads \times sr\_len\_k \times \frac{hidden\_size}{num\_heads}\)
value:
\(1 \times sr\_len\_len\_v \times hidden\_size \stackrel{线性变化}{\rightarrow} 1 \times sr\_len\_v \times hidden\_size \stackrel{reshape}{\rightarrow} 1 \times num\_heads \times sr\_len\_v \times \frac{hidden\_size}{num\_heads}\)

2.DotProductAttention： num_heads 个头的计算并行：

\[query∗key^T=1×num\_heads×1×sr\_len\_k \]

\[消除key中padding的影响，对其做mask \]

\[\Downarrow \]

\[masked_softmax(query∗key^T)=1×num\_heads×1×sr\_len\_k \]

\[\Downarrow \]

\[masked_softmax(query∗key^T)∗value=1×num\_heads×1×\frac{hidden\_size}{num\_heads} \]

假设 num_heads = 1，sr_len_k = 6，因为 mask 只针对key，因此这里只需要关注 source sentence 中的padding， 假设 source sentence 的最后两个位置是padding，那么Masked Multihead Attention 中的 mask 为：
( 1 1 1 1 0 0 )

3.Output： 需要将上面输出的 num_heads 个头的结果堆叠之后，再做一个线性变换：

\[1 × num\_heads × 1 ×\frac{hidden\_size}{num\_heads} \]

\[reshape \]

\[\Downarrow \]

\[1 × 1 × hidden\_size \]

\[线性变换 \]

\[\Downarrow \]

\[1 × 1 × hidden\_size \]

由于最后的 Feed Forward 层不改变矩阵大小，至此可以总结一下预测阶段的 Decoder layer，输入是上一 step 输出的 token，大小为 \(1 \times 1 \times hidden\_size\)，经过两种MultiHead + Feed Forward 后，大小依然为 \(1 \times 1 \times hidden\_size\)，再经过 Linear+Softmax，其输出就是预测的当前 step 的token，而这个 token 又会作为下一个 step 的输入 query。直到达到最大长度，或者输出的 token 是 “<eos>”

参考
https://paddlepedia.readthedocs.io/en/latest/tutorials/pretrain_model/transformer.html

https://blog.csdn.net/muyuu/article/details/123616233