【从 0 开始学习 Transformer】上：搭建与理解

在我的码云上阅读

• 1. 前言
• 2. 参考代码、文章及部分插图来源
• 3. 在开始前的推荐了解
  • 3.1. 循环神经网络（RNN）
  • 3.2. 基于编码-解码（encoder-decoder）的序列到序列（sequence2sequence）模型
  • 3.3. 注意力机制
  • 3.4. 词嵌入（Word Embedding）
• 4. 初探 Transformer
• 5. 基础算法和模块
  • 5.1. 位置编码（Positional encoding）
  • 5.2. 注意力机制
    • 5.2.1. 计算步骤细节
    • 5.2.2. 使用向量化来提升效率
    • 5.2.3. 如何作为自注意力使用
  • 5.3. 遮挡 Mask
    • 5.3.1. 填充遮挡
    • 5.3.2. 前瞻遮挡（look-ahead mask）
  • 5.4. 多头注意力（Multi-head attention）
    • 5.4.1. 代码分析
  • 5.5. 点式前馈网络（Point wise feed forward network）
• 6. 编码器解码器
  • 6.1. 编码器层
  • 6.2. 解码器层
  • 6.3. 编码器
  • 6.4. 解码器
• 7. 创建 Transformer
• 8. 小结

1. 前言

本文章结合代码讲解如何使用 Tensorflow 从零开始学习理解及搭建一个Transformer，本文代码基于Tensorflow 2.0版本，若使用其他版本或Pytorch亦不妨参考。

相较于 Tensorflow 官方指南，本文将对模块细节和使用进行更精细的讲解；将对代码构建过程中代码进行逐行解释以及自注意力机制等细节设计进行讲解。除此之外，也加入了笔者对于 Transformer 细节部分的一些个人疑惑和理解。简而言之，是融合了代码指南和原理精讲的一篇文章，力求将笔者对于 Transformer 的理解精粹于一篇文章，本人才疏学浅，欢迎批评指正。

整个文章结构参考 Tensorflow 官方指南，按照自底向上的顺序来逐渐搭建一个用于将葡萄牙语翻译为英语的Transformer模型：先从最基本的算法模块实现，然后组装。在组装过程中讲解各个子模块产生的作用。

2. 参考代码、文章及部分插图来源

本文大量参考及使用他人文章和官方文档中代码和图片，在此表示感谢：

• The Illustrated Transformer
• Transformer model for language understanding
• TensorFlow Core r2.0 API
• Datasets v1.3.0 API

3. 在开始前的推荐了解

3.1. 循环神经网络（RNN）

RNN是一种可用于抽取序列数据特征的神经网络结构，一个最基本的RNN是非常容易理解的，相当于 NLP 领域的 “Hello world“。很多 NLP 领域的知识点（如本文可能用到的注意力机制与自注意力机制、encoder-decoder架构）的相关文章很多都会使用一个最基本的 RNN 来类比），所以建议阅读本文前对最基本的循环神经网络有一个初步的了解。

这部分内容推荐阅读 Ian Goodfellow 等人编写的 Deep Learning 一书中的 10.1 与 10.2 节。

3.2. 基于编码-解码（encoder-decoder）的序列到序列（sequence2sequence）模型

在⾃然语⾔处理的很多应⽤中，输⼊和输出都可以是不定⻓序列。以机器翻译为例，输⼊可以是⼀段不定⻓的英语⽂本序列，输出可以是⼀段不定⻓的法语⽂本序列，当输⼊和输出都是不定⻓序列时，我们可以使⽤编码器—解码器（encoder-decoder）架构或者sequence2sequence模型。

这部分内容推荐阅读 Ian Goodfellow 等人编写的 Deep Learning 一书中的 10.4 节。

3.3. 注意力机制

Transformer 的核心思想——自注意力机制，实际上是受启发于注意力机制。对比注意力机制和自注意力机制的异同，可以更加的深刻理解自注意力机制的作用机理。

强烈推荐阅读 张俊林 - 深度学习中的注意力模型。

对英语有自信的同学可以阅读 Visualizing A Neural Machine Translation Model (Mechanics of Seq2seq Models With Attention)，其中包含大量的动画实例。

3.4. 词嵌入（Word Embedding）

推荐阅读此文章 张俊林 - 从Word Embedding到Bert模型—自然语言处理中的预训练技术发展史 的前半部分对 Word Embedding 的介绍。

4. 初探 Transformer

Transformer 是基于自注意力机制的全新 encoder-decoder 模型。相较于传统循环神经网络搭建的同类模型，Transformer 具有多重优势：解决长期依赖的问题，可以并行化等等。

现在我们已经知道了一个 encoder-decoder 可以完成多种自然语言处理任务。为了举例方便，假设要完成一个机器翻译任务：从法语翻译到英语。Transformer 做的事情就是这样的：

作为一个 encoder-decoder 模型，Transformer 将会被分为编码器（左侧浅绿色）和解码器（右侧粉色）两部分：

如何设计这两个部分呢？论文中的图示是这样的：

这一张图掩盖了许多细节，大量的箭头也使人不明就里。暂时来看，我们可以得到以下信息：

1. 解码器和编码器的结构类似，都是多头注意力机制（图中 Multi-Head Attention 块，是对于自注意力机制的变化使用）和前馈神经网络（Feed Forward）的堆叠多层。
2. 输入不仅仅是句子的词嵌入表示，还额外增加了称作位置编码（Positional Encoding）的额外信息。

其他姑且按下不表。

5. 基础算法和模块

5.1. 位置编码（Positional encoding）

一个传统的，使用 RNN 构建的 Encoder-Decoder 模型对于输入句子中的单词是逐个读取的：读取完前一个单词，更新模型的状态，然后在此状态下再读取下一个单词。这种方式天然的包含了句子的位置前后关系。

我们后面会发现 Transformer 对于输入却并非逐个读取，而是对整个句子的每个单词进行同时读取。这种方式就显然丢失了句子的前后位置关系。

为了能够保留句子中单词和单词之间的位置关系，需要将位置也融合进入输入的句子中，需要对位置进行编码。

Transformer 使用的位置编码算法如下定义：

$$P{E}_{(pos,2i)}=\sin \left(\text{pos}/10000^{2i/d_{\text{model}}}\right)$$

$$P{E}_{(pos,2i+1)}=\cos \left(\text{pos}/10000^{2i/d_{\text{matel}}}\right)$$

由公式来看，对于一个句子，此编码算法对于偶数位置 2i 和奇数位置 2i + 1 分开进行编码。编码的结果是每个位置最终转化为 d_model 维度的向量。

一个对50个位置编码至 512 维度位置向量的图示例子如下（其作用于的输入就是句子，句子长至 50 单词，且每个词嵌入的维度为 512）：

这个位置向量会直接加在词嵌入上使用，从而使词嵌入在含义远近的表示能力之外增加了句子中的位置关系的表示能力：

位置编码向量被加到嵌入（embedding）向量中。嵌入表示一个 d 维空间的标记，在 d 维空间中有着相似含义的标记会离彼此更近。但是，嵌入并没有对在一句话中的词的相对位置进行编码。因此，当加上位置编码后，词将基于它们含义的相似度以及它们在句子中的位置，在 d 维空间中离彼此更近。

才疏学浅，没法一个直观的角度来解释为何这样设计编码以及为何这样编码可以起到这样的效果。请各位指教。

考虑此函数需要有大量的切片赋值操作（tensorflow api 实现起来较为困难，需要调用assign来实现numpy的赋值运算符=），所以使用numpy的api来实现，最后使用tf.cast将其自动转换为tensor（此处为一个eager tensor）。

def get_angles(pos, i, d_model):
  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
  return pos * angle_rates

def positional_encoding(position, d_model):
  angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)
  
  # 将 sin 应用于数组中的偶数索引（indices）；2i
  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
  
  # 将 cos 应用于数组中的奇数索引；2i+1
  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    
  pos_encoding = angle_rads[np.newaxis, ...] 
  # 在这里增加了一个维度。eg: (50, 512) -> (1, 50, 512)
  # 为什么要增加一个维度呢？因为输入到 Transformer 的输入通常是多个句子层叠成一个批次。
  # 增加一个维度就可以利用广播机制一次加法将一个批次的每个句子添加上位置编码。
    
  return tf.cast(pos_encoding, dtype=tf.float32)

5.2. 注意力机制

准确来说，是按比缩放的点积注意力机制（Scaled dot product attention）。

顾名思义，Transformer 一种点积注意力机制，并且在普通的点积注意力机制上增加了“按比缩放”这个特性。

根据输入的不同，这个模块既可以是自注意力，也可以是非自注意力。

如果已经阅读了本文第三部分推荐的注意力机制相关的文章张俊林 - 深度学习中的注意力模型，应该已经理解了注意力机制的本质：

对于三个输入：

• Q：请求 query
• K：主键 key
• V：数值 value

注意力机制本质就是使用 Q 和 K 来计算出“注意力权重“，然后利用注意力权重对Ｖ进行加权求和。

按比缩放的点积注意力定义如下：

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

使用计算图来表示如下：

可以看出，注意力权重是通过 Q 和 K 直接进行点积MatMul，并按比缩放Scale（除以深度的平方根），然后进行 softmax 得到的。

同样的，使用注意力权重乘上 V 便得到了一次输出。

有两点需要讨论：

1. 为什么要按比缩放Scale：这是原论文中的一个推测——如果输入的Q，K维度过大，则会导致点积后的结果很大softmax函数有一个特点，当输入的 x 越大，其梯度会趋近于 0。这对于基于梯度下降法的优化非常不利。（这个是一个有根据的推测：假设q和k都是独立的随机变量，那么q 乘上 k 是均值的 0 方差为 $d_k$ 的。除以深度的平方根，可以让方差为 1）
2. 在 softmax 前，为什么有一个可选的 Mask 过程： 这是由于在整个模型的运行过程中，可能根据设计，要忽略掉一些输入。关于Mask的生成，将在下一部分详细解释。关于Mask的使用，将在#TODO时提及。

5.2.1. 计算步骤细节

注意：这个例子和实际 Transformer 使用的实际输入不同（并非直接使用词嵌入来作为Query、Key等输入），仅为了能够用自然语言阐述而设计。精确的解释我们放在后文。

假设 Query 序列是一个句子，长度为 seq_len_q。

而 Key 序列也是一些一个句子，长度为 seq_len_k。

Value 序列中的每个值都对应 Key 句子中的一个单词，所以 v 的长度seq_len_v 等于 seq_len_k 。

根据上文描述的注意力机制的计算方法：依次取 Query 中的单词，来和 Key 中的 seq_len_v 个单词一一计算注意力权重，得到 seq_len_v 个注意力权重，使用这个权重对 seq_len_k （再次强调，长度必须相同，这部分描述对应着下面代码注释中的要求2.）个 Value 进行加权求和，得到一个输出。

一个 Query 序列将产生 seq_len_q 个加权求和。

让我们用矩阵点乘的模式来描述这个过程：

假设每个单词的嵌入为depth维度，显然这个步骤是:

Q 矩阵(seq_len_q, depth) 和 K 矩阵(seq_len_k, depth) 的转置的矩阵乘法。

最终得到注意力权重 (seq_len_q, seq_len_k) 矩阵。这个矩阵乘上 Value 矩阵(seq_len_v, depth_v) 便是加权求和过程，最终得到了输出Output (seq_len_q, depth_v)。

depth_v 表示 Value 序列中每个值的维度，显然，这个维度不一定要和上文的depth一致

5.2.2. 使用向量化来提升效率

每个 Query 序列对应着一个 Key 序列，但这 Query-Key 组合彼此之间是独立的。完全可以将 Query、Key、Value 堆叠成批，一次运算搞定。矩阵乘法或是转置是针对最后的两个维度，所以只需要保持前置维度匹配（对应，下方注释的要求1.），计算结果和上面完全等效。

举个例子：若有 batch_size 批次，每批次 N 条的 Query，Key。 其计算完全可以组织成 (batch_size, N, seq_len_q, depth)点乘(batch_size, N, seq_len_q, depth)的转置，最终得到(batch_size, N, seq_len_q, seq_len_k) 形状的张量。

mask 以乘以一个极大的负数-1e9，然后在加上注意力权重，最终达到使一些位置的 Value 失效的效果。只需要保证其可通过广播机制能够自动转换为 (..., seq_len_q, seq_len_k) 形状（代码注释中的要求3.)。

这一要求将会影响下一节的遮挡（Masking）的实现。

def scaled_dot_product_attention(q, k, v, mask):
  """ 这部分是对于输入的张量维度进行描述

  1. q, k, v 必须具有匹配的前置维度。

  2. k, v 必须有匹配的倒数第二个维度，例如：seq_len_k = seq_len_v。
  
  3. 虽然 mask 根据其类型（填充或前瞻）有不同的形状，但是 mask 必须能进行广播转换以便求和。

  
  参数:
    q: 请求的形状 == (..., seq_len_q, depth)
    k: 主键的形状 == (..., seq_len_k, depth)
    v: 数值的形状 == (..., seq_len_v, depth_v)
    mask: Float 张量，其形状能转换成
          (..., seq_len_q, seq_len_k)。默认为None。
    
  返回值:
    输出，注意力权重
  """

  matmul_qk = tf.matmul(q, k, transpose_b=True)  # 得到张量形状为 (..., seq_len_q, seq_len_k)
  
    # 缩放 matmul_qk
  dk = tf.cast(tf.shape(k)[-1], tf.float32) # dk 即 depth
  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

    # 将 mask 加入到缩放的张量上。
  if mask is not None:
    scaled_attention_logits += (mask * -1e9)  

    # softmax 在最后一个轴（seq_len_k）上归一化，因此分数
    # 相加等于1。
  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., seq_len_q, seq_len_k)

  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)

  return output, attention_weights

5.2.3. 如何作为自注意力使用

有人会这样总结：自注意力就是将同一个句子同时作为 Query 和 Key 来使用。我认为还不够精确。

事实上，输入自注意力的 Query、Key 和 Value 都来自原始嵌入的线性变换：创建三个会随着模型训练过程不断优化的矩阵$W^Q$，$W^K$ 和 $W^V$，原始词嵌入乘上三个矩阵从而得到真正的 Query、Key 和 Value（也可以理解为三个独立的全连接神经网络，输入节点数量为词嵌入的维度，输出节点数量分别为Query，Key 和 Value 的维度）。

其过程如下：

其余过程则完全和上述过程一致：

而自注意力的注意效果也是逐层变得集中。使用图形化来对自注意力的效果进行初步理解。

当只经过线性变换的注意力效果：

注意力显然非常分散，有种抓不到要领的感觉。

而经过了5层自注意力机制，单个单词的注意力开始集中于少数部分。

5.3. 遮挡 Mask

5.3.1. 填充遮挡

如果一个输入句子由于长短不一不方便计算或是其他原因需要补充一些填充标记（pad tokens），显然在输出结果的时候应该把这些无意义的填充标记排除，因此需要一个函数产生此用途的 mask。

def create_padding_mask(seq):

  # 对于一个序列，如果某位置其值为0，则认定为填充标记，在此位置标 1
  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)
  
  # 为了能够利用广播机制匹配注意力权重张量，需要增加必要的维度
  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)

测试效果：

x = tf.constant([[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]])
create_padding_mask(x)
"""

"""

5.3.2. 前瞻遮挡（look-ahead mask）

前瞻遮挡通常用于需要只考虑序列中的前一部分的时候，这个遮挡将会用在 Transform 的解码器部分，其设计原理是预测一个单词只考虑此单词前的单词，而不考虑此单词后的部分。

def create_look_ahead_mask(size):
  # tf.linalg.band_path(Tensor, -1, 0) 是取Tensor的左下半矩阵
  mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
  return mask  # (seq_len, seq_len)

最终效果得到：

x = tf.random.uniform((1, 3))
temp = create_look_ahead_mask(x.shape[1])
temp
"""

"""

5.4. 多头注意力（Multi-head attention）

前文说到，使用一组会随着模型训练过程不断优化的矩阵$W^Q$，$W^K$ 和 $W^V$，可以通过对词嵌入 X 或者前一层的输出 R 相乘而得到一套可以输入进注意力机制的 Query, Key 和 Value。

多头注意力机制是将初始的词向量（第一层）或前一层的输入（第二层开始）通过线性变换转换为多组 Query, Key 和 Value，从而得到不同的输出 Z。最后将所有的输出拼合起来，通过可训练的线性变换 $W^O$ 融合为一个输出：

从注意力角度看，由于矩阵是随机初始化的，随着训练的过程，最终不同的Query, Key可能得到不同的注意力结果。

论文认为：

• 这种方式拓展了模型专注于不同位置的能力。
• 模型最终的“注意力”实际上是来自不同“表示子空间”的注意力的综合。

做个比喻来说，这就好像是八个有不同阅读习惯的翻译家一同翻译同一个句子，他们每个人可能翻译时阅读顺序和关注点都有所不同，综合他们八个人的意见，最终得出来的翻译结果可能会更加准确。

通过继承 tf.keras.layers.Layer 可以对多头注意力机制进行结构清晰的实现。

class MultiHeadAttention(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads):
  """
    参数: d_model 必须能被 num_heads 整除
    d_model: 由于要映射 num_heads 组 Q,K,V. d_model 的值需要为 num_heads * depth
    num_heads: 代表注意力机制的头数
  """
    super(MultiHeadAttention, self).__init__()
    self.num_heads = num_heads
    self.d_model = d_model
    
    assert d_model % self.num_heads == 0
    
    self.depth = d_model // self.num_heads
    
    self.wq = tf.keras.layers.Dense(d_model)
    self.wk = tf.keras.layers.Dense(d_model)
    self.wv = tf.keras.layers.Dense(d_model)
    
    self.dense = tf.keras.layers.Dense(d_model)
        
  def split_heads(self, x, batch_size):
    """分拆最后一个维度到 (num_heads, depth).
    转置结果使得形状为 (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])
    
  def call(self, v, k, q, mask):
    batch_size = tf.shape(q)[0]
    
    # 这里采取将 q,k,v 先线性变换到 (..., seq_len, d_model)
    # 再拆分成 num_heads 份 (..., nums_heads, seq_ken, d_model / nums_heads)
    # 这和直接将原始 q,k,v 线性变换成 nums_heads 组是等效的！这样写效率更高！

    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)
    
    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)
    
    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
    scaled_attention, attention_weights = scaled_dot_product_attention(
        q, k, v, mask)
    
    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention, 
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)

    output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)
        
    return output, attention_weights

5.4.1. 代码分析

由于维度多，两处更改reshape 和转置操作transpose令人头秃。所以来详细讲解一下Tensor是如何在多头注意力机制力流动的。如果对模型实现手到擒来的的同学可以直接跳过：

首先，我们确定要输送给注意力机制 scaled_dot_product_attention 的形状是 (batch_size, nums_heads, seq_len_q, depth)，通俗来讲，注意力机制函数将会并行处理 batch_size 批句子，每批句子 nums_heads 句。为了得到合适的输出，要通过以下几步：

1. 将输入映射至足够维度(batch_size, seq_len, input_depth) -> (batch_size, seq_len, d_model)

...
    q = self.wq(q)  # (batch_size, seq_len, d_model)
    k = self.wk(k)  # (batch_size, seq_len, d_model)
    v = self.wv(v)  # (batch_size, seq_len, d_model)
...

这个 d_model 的深度显然大于我们需要的 depth，这是为了下一步拆成 num_heads，将多次线性变换，转换为一次。

2. 将输入拆分为多头

  def split_heads(self, x, batch_size):
    """分拆最后一个维度到 (num_heads, depth).
    转置结果使得形状为 (batch_size, num_heads, seq_len, depth)
    """
    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
    return tf.transpose(x, perm=[0, 2, 1, 3])
...
    # 这里记得要把 batch_size 传进去
    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)
    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)
    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)
...

分拆最后一个维度到 (num_heads, depth)

向量角度而言： reshape 操作将张量中每一行 d_model 都拆成了 num_heads 个 depth 长度的行向量。即：(batch_size, seq_len, d_model) -> (batch_size, seq_len, nums_heads, depth)。

从神经网络角度而言：由于对于单层全连接网络，输入层与隐层节点的任何一个子集结合，都是一个完整的单隐层全连接网络。也就是说，这种拆分完全可以看做将前一步input_depth 个节点到 d_model 个节点的全连接网络，拆分成了 nums_heads 个小的 input_depth 个节点到 depth 个节点的全连接网络。

然后，我们处理的仍然是序列本身，因此，通过转置 transpose 将 seq_len 放回它原来的位置，让 nums_heads 成为一个前置维度：(batch_size, seq_len, nums_heads, depth) -> (batch_size, num_heads, seq_len, depth)

自注意力机制的详细运算过程已经在前文说的很清楚了，接下来是对输出的处理。

3. 将多头输出通过全连接映射为一个输出

显然自注意力机制函数的输出形状将是(batch_size, num_heads, seq_len_q, depth)。

为了能够方便地将多头结果拼合起来，首先我们将其转置到倒数第二个维度。

然后，2.中怎么拆开的，就怎么拼回去。

...
    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])  # (batch_size, seq_len_q, num_heads, depth)

    concat_attention = tf.reshape(scaled_attention, 
                                  (batch_size, -1, self.d_model))  # (batch_size, seq_len_q, d_model)
...

最后，通过全连接层将拼合好的输出融合起来：

...
output = self.dense(concat_attention)  # (batch_size, seq_len_q, d_model)
...

这里，经过全连接层融合后的最后一维仍然是 d_model。

5.5. 点式前馈网络（Point wise feed forward network）

点式前馈网络由两层全联接层组成，两层之间有一个 ReLU 激活函数。

def point_wise_feed_forward_network(d_model, dff):
  return tf.keras.Sequential([
      tf.keras.layers.Dense(dff, activation='relu'),  # 第一层输出尺寸 (batch_size, seq_len, dff)
      tf.keras.layers.Dense(d_model)  # 第二层输出尺寸 (batch_size, seq_len, d_model)
  ])

dff: 规定了点式前馈神经网络的内部第一层输出节点。

在论文中，这个神经网络被称作位置式前馈神经网络（Position-wise Feed-Forward Networks
In，不知道为什么Tensorflow的文档要改名），其定义如下：

$$\mathrm{F}\mathrm{F}\mathrm{N}(x)=\max \left(0,x{W}_1+b_1\right)W_2+b_2$$

可以看出其本质上是两次线性变换的串联（不考虑激活函数的情况下）。

我们考虑一个(..., seq_len, depth) 的输入，第一次线性变换相当于针对每一行（换句话说，句子的每个位置），做了一个相同（但隐层参数不同）的全连接网络，输入节点数 depth 和 输出节点数 dff。得到 （..., seq_len, dff)。同理第二次线性变换类似，得到(..., seq_len, d_model)。

事实上，它只是一个多层神经网络，但考虑它等同地处理句子的每个位置，起名如此也算合理。

6. 编码器解码器

有了上面的诸多模块，终于可以召唤此图：

左侧编码器，右侧解码器。

并开始组装编码器和解码器：

Transformer 模型与标准的具有注意力机制的序列到序列模型（sequence to sequence with attention model），遵循相同的一般模式。

输入语句经过 N 个编码器层，为序列中的每个词/标记生成一个输出。

解码器关注编码器的输出以及它自身的输入（自注意力）来预测下一个词。

从这里一直到到 Transformer 的完成，我们始终忽略 mask 的实际细节。由于这一部分涉及到模型的训练优化所以我们放在下一篇文章展开来讲。在此我们只把它当做一个参数暴露给外部模块。并传给内部模块。

6.1. 编码器层

对于编码器，一个编码器层是其核心的最小单位。

每个编码器层包括以下子层：

1. 多头注意力（有填充遮挡）
2. 点式前馈网络（Point wise feed forward networks）。

每个子层在其周围有一个残差连接（图中最左侧的上下两个黑箭头），然后进行层归一化。残差连接有助于避免深度网络中的梯度消失问题。

每个子层的输出是 LayerNorm(x + Sublayer(x))。归一化是在 d_model（最后一个）维度完成的。Transformer 中有 N 个编码器层。

此外对于每个子层的输出，都使用了0.1概率的的dropout。

class EncoderLayer(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(EncoderLayer, self).__init__()

    self.mha = MultiHeadAttention(d_model, num_heads)
    self.ffn = point_wise_feed_forward_network(d_model, dff)

    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    
    self.dropout1 = tf.keras.layers.Dropout(rate)
    self.dropout2 = tf.keras.layers.Dropout(rate)
    
  def call(self, x, training, mask): # 填充遮挡将在调用时传入

    attn_output, _ = self.mha(x, x, x, mask)  # (batch_size, input_seq_len, d_model)
    attn_output = self.dropout1(attn_output, training=training)
    out1 = self.layernorm1(x + attn_output)  # (batch_size, input_seq_len, d_model)
    
    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)
    ffn_output = self.dropout2(ffn_output, training=training)
    out2 = self.layernorm2(out1 + ffn_output)  # (batch_size, input_seq_len, d_model)
    
    return out2

6.2. 解码器层

每个解码器层包括以下子层：

1. 遮挡的多头注意力（前瞻遮挡和填充遮挡）
2. 多头注意力（用填充遮挡）。V（数值）和 K（主键）接收编码器输出作为输入。Q（请求）接收遮挡的多头注意力子层的输出。
3. 点式前馈网络
   每个子层在其周围有一个残差连接，然后进行层归一化。每个子层的输出是 LayerNorm(x + Sublayer(x))。归一化是在 d_model（最后一个）维度完成的。

Transformer 中共有 N 个解码器层。

当 Q 接收到解码器的第一个注意力块的输出，并且 K 接收到编码器的输出时，注意力权重表示根据编码器的输出赋予解码器输入的重要性。换一种说法，解码器通过查看编码器输出和对其自身输出的自注意力，预测下一个词。

class DecoderLayer(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(DecoderLayer, self).__init__()

    self.mha1 = MultiHeadAttention(d_model, num_heads)
    self.mha2 = MultiHeadAttention(d_model, num_heads)

    self.ffn = point_wise_feed_forward_network(d_model, dff)
 
    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    
    self.dropout1 = tf.keras.layers.Dropout(rate)
    self.dropout2 = tf.keras.layers.Dropout(rate)
    self.dropout3 = tf.keras.layers.Dropout(rate)
    
    
  def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):
    # enc_output.shape == (batch_size, input_seq_len, d_model)

    attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)  # (batch_size, target_seq_len, d_model)
    attn1 = self.dropout1(attn1, training=training)
    out1 = self.layernorm1(attn1 + x)
    
    attn2, attn_weights_block2 = self.mha2(
        enc_output, enc_output, out1, padding_mask)  # (batch_size, target_seq_len, d_model)
    attn2 = self.dropout2(attn2, training=training)
    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)
    
    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)
    ffn_output = self.dropout3(ffn_output, training=training)
    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)
    
    return out3, attn_weights_block1, attn_weights_block2

显然，相较于编码器，除了用于自注意力的多头注意力层。解码器增加了一层注意力层用于实现编码器和解码器之间的注意力。其他与编码器完全一致。

注意：两层注意力使用的遮挡类型略有不同！

6.3. 编码器

由于封装了上文实现的所有模块。所以这个模块的参数显得有些多，我们只关注此层特有的参数：

• num_layers：规定编码器使用多少个编码器层。
• input_vocab_size: 原语料的词汇量
• maximum_position_encoding：最大的位置编码，代表位置编码最长匹配的位置长度

其总体步骤包括三步：

1. 将原始句子单词编码更换为词嵌入
2. 将词嵌入加上位置编码
3. 将处理过后的输出输入规定好的多个编码器层

其中：在原始句子的词嵌入上需要乘上 $\sqrt{d\_model}$，这是原论文中规定的。文章中没有对这个变量的解释。

论文原文： In the embedding layers, we multiply those weights by $\sqrt{d_{model}}$

class Encoder(tf.keras.layers.Layer):
  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,
               maximum_position_encoding, rate=0.1):
    super(Encoder, self).__init__()

    self.d_model = d_model
    self.num_layers = num_layers
    
    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
    self.pos_encoding = positional_encoding(maximum_position_encoding, 
                                            self.d_model)
    
    
    self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) 
                       for _ in range(num_layers)]
  
    self.dropout = tf.keras.layers.Dropout(rate)
        
  def call(self, x, training, mask):

    seq_len = tf.shape(x)[1]
    
    # 将嵌入和位置编码相加。
    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]

    x = self.dropout(x, training=training)
    
    for i in range(self.num_layers):
      x = self.enc_layers[i](x, training, mask)
    
    return x  # (batch_size, input_seq_len, d_model)

6.4. 解码器

解码器和编码器参数类似。

1. 将输出序列同样转换为同维度词嵌入
2. 加上位置编码
3. 和编码器的输出一同输入给多个解码器层

在调用时，enc_output 是来自编码器层的输出。

class Decoder(tf.keras.layers.Layer):
  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,
               maximum_position_encoding, rate=0.1):
    super(Decoder, self).__init__()

    self.d_model = d_model
    self.num_layers = num_layers
    
    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)
    self.pos_encoding = positional_encoding(maximum_position_encoding, d_model)
    
    self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) 
                       for _ in range(num_layers)]
    self.dropout = tf.keras.layers.Dropout(rate)
    
  def call(self, x, enc_output, training, 
           look_ahead_mask, padding_mask):

    seq_len = tf.shape(x)[1]
    attention_weights = {}
    
    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :seq_len, :]
    
    x = self.dropout(x, training=training)

    for i in range(self.num_layers):
      x, block1, block2 = self.dec_layers[i](x, enc_output, training,
                                             look_ahead_mask, padding_mask)
      
      attention_weights['decoder_layer{}_block1'.format(i+1)] = block1
      attention_weights['decoder_layer{}_block2'.format(i+1)] = block2
    
    # x.shape == (batch_size, target_seq_len, d_model)
    return x, attention_weights

7. 创建 Transformer

将编码器和解码器组合起来，连接最后的线性输出层，就得到了整体的 Transformer：

class Transformer(tf.keras.Model):
  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, 
               target_vocab_size, pe_input, pe_target, rate=0.1):
    super(Transformer, self).__init__()

    self.encoder = Encoder(num_layers, d_model, num_heads, dff, 
                           input_vocab_size, pe_input, rate)

    self.decoder = Decoder(num_layers, d_model, num_heads, dff, 
                           target_vocab_size, pe_target, rate)

    self.final_layer = tf.keras.layers.Dense(target_vocab_size)
    
  def call(self, inp, tar, training, enc_padding_mask, 
           look_ahead_mask, dec_padding_mask):

    enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)
    
    # dec_output.shape == (batch_size, tar_seq_len, d_model)
    dec_output, attention_weights = self.decoder(
        tar, enc_output, training, look_ahead_mask, dec_padding_mask)
    
    final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)
    
    return final_output, attention_weights

8. 小结

到此为止，我们已经描述了 Transformer 的整个模型搭建过程，并逐层逐行地解释了其正向传播的原理和细节。本文仍未讲到的是：

• 如何训练一个Transformer：
  • 前瞻遮挡和填充遮挡如何使用。
  • 超参数如何配置。
  • 如何设计损失函数。
  • 如何优化和评估。

后篇将结合实例详细描述这些内容。