千行代码搞定Transformer？这份高效PaddlePaddle官方实现请收下

不同网络的主要性质，其中 n 表示序列长度、d 为隐向量维度、k 为卷积核大小。例如单层计算复杂度，一般句子长度 n 都小于隐向量维度 d，那么自注意力层级的计算复杂度最小。

如上所示，Transformer 使用的自注意力模型主要拥有以下优点，1）网络结构的计算复杂度最低；2）由于序列操作数复杂度低，模型的并行度很高；3）最大路径长度小，能够更好地表示长距离依赖关系；4）模型更容易训练。

现在，如果我们需要训练一个 Transformer，那么最好的方法是什么？当然是直接跑已复现的模型了，下面我们将跑一跑 PaddlePaddle 实现的 Transformer。

1.1 处理数据

在 PaddlePaddle 的复现中，百度采用原论文测试的 WMT'16 EN-DE 数据集，它是一个中等规模的数据集。这里比较方便的是，百度将数据下载和预处理等过程都放到了 gen_data.sh 脚本中，包括 Tokenize 和 BPE 编码。

在这个项目中，我们既可以通过脚本预处理数据，也可以使用百度预处理好的数据集。首先最简单的方式是直接运行 gen_data.sh 脚本，运行后可以生成 gen_data 文件夹，该文件夹主要包含以下文件：

其中 wmt16_ende_data_bpe 文件夹包含最终使用的英德翻译数据。

如果我们从头下载并预处理数据，那么大概需要花 1 到 2 个小时完成预处理。为此，百度也提供了预处理好的 WMT'16 EN-DE 数据集，它包含训练、验证和测试所需要的 BPE 数据和字典。

其中，BPE 策略会把稀疏词拆分为高频的子词，这样既能解决低频词无法训练的问题，也能合理降低词表规模。

如果不采用 BPE 的策略，要么词表的规模变得很大，从而使训练速度变慢或者显存太小而无法训练；要么一些低频词会当作未登录词处理，从而得不到训练。

预处理数据地址：https://transformer-res.bj.bcebos.com/wmt16_ende_data_bpe_clean.tar.gz

如果我们有其它数据集，例如中英翻译数据，也可以根据特定的格式进行定义。例如用空格分隔不同的 token（对于中文而言需要提前用分词工具进行分词），用\t 分隔源语言与目标语句对。

1.2 训练模型

如果需要执行模型训练，我们也可以直接运行训练主函数 train.py。如下简要配置了数据路径以及各种模型参数：

# 显存使用的比例，显存不足可适当增大，最大为1 export FLAGS_fraction_of_gpu_memory_to_use=0.8 # 显存清理的阈值，显存不足可适当减小，最小为0，为负数时不启用 export FLAGS_eager_delete_tensor_gb=0.7 python -u train.py \ --src_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \ --trg_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \ --special_token '<s>' '<e>' '<unk>' \ --train_file_pattern gen_data/wmt16_ende_data_bpe/train.tok.clean.bpe.32000.en-de \ --token_delimiter ' ' \ --use_token_batch True \ --batch_size 1600 \ --sort_type pool \ --pool_size 200000 \ n_head 8 \ n_layer 4 \ d_model 512 \ d_inner_hid 1024 \ prepostprocess_dropout 0.3

此外，如果显存不够大，那么我们可以将 Batch Size 减小一点。为了快速测试训练效果，我们将模型调得比 Base Transformer 还小（降低网络的层数、head 的数量、以及隐层的大小）。

上面仅展示了小部分的超参设置，更多的配置可以在 GitHub 项目 config.py 文件中找到。默认情况下，模型每迭代一万次保存一次模型，每个 epoch 结束后也会保存一次 cheekpoint。此外，在我们训练的过程中，默认每一百次迭代会打印一次模型信息，其中 ppl 表示的是困惑度，困惑度越小模型效果越好。

在单机训练中，默认使用所有 GPU，可以通过 CUDA_VISIBLE_DEVICES 环境变量来设置使用的 GPU，例如 CUDA_VISIBLE_DEVICES='0,1'，表示使用 0 号和 1 号卡进行训练。

1.3 预测推断

训练完 Transformer 后就可以执行推断了，我们需要运行对应的推断文件 infer.py。我们也可以在推断过程中配置超参数，但注意超参需要和前面训练时保持一致。

python -u infer.py \ --src_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \ --trg_vocab_fpath gen_data/wmt16_ende_data_bpe/vocab_all.bpe.32000 \ --special_token '<s>' '<e>' '<unk>' \ --test_file_pattern gen_data/wmt16_ende_data_bpe/newstest2016.tok.bpe.32000.en-de \ --token_delimiter ' ' \ --batch_size 32 \ model_path trained_models/iter_100000.infer.model \ n_head 8 \ n_layer 4 \ d_model 512 \ d_inner_hid 1024 \ prepostprocess_dropout 0.3 beam_size 5 \ max_out_len 255

相比模型的训练，推断过程需要一些额外的超参数，例如配置 model_path 指定模型所在目录、设置 beam_size 和 max_out_len 来指定 Beam Search 每一步候选词的个数和最大翻译长度。这些超参数也可以在 config.py 中找到，该文件对这些超参都有注释说明。

执行以上预测命令会将翻译结果直接打出来，每行输出是对应行输入得分最高的翻译。对于使用 BPE 的英德数据，预测出的翻译结果也将是 BPE 表示的数据，所以需要还原成原始数据才能进行正确评估。如下命令可以将 predict.txt 内的翻译结果（BPE 表示）恢复到 predict.tok.txt 文件中（tokenize 后的数据）：

sed -r 's/(@@ )|(@@ ?$)//g' predict.txt > predict.tok.txt

在未使用集成方法的情况下，百度表示 base model 和 big model 在收敛后，测试集的 BLEU 值参考如下：

这两个预训练模型也提供了下载地址：

• Base：https://transformer-res.bj.bcebos.com/base_model.tar.gz
• Big：https://transformer-res.bj.bcebos.com/big_model.tar.gz

2. Transformer 怎么改

如果我们想要训练自己的 Transformer，那么又该怎样理解并修改 PaddlePaddle 代码呢？如果我们需要根据自己的数据集和任务改代码，除了前面数据预处理过程，模型结构等模块有时也需要修改。这就需要我们先理解源代码了，PaddlePaddle 的源代码基本都是基础的函数或运算，我们很容易理解并使用。

对于 PaddlePaddle 不熟悉的读者可查阅文档，也可以看看入门教程，了解基本编写模式后就可以看懂整个实现了。

PaddlePaddle 官网地址：http://paddlepaddle.org/paddle

如 Seq2Seq 一样，原版 Transformer 也采用了编码器-解码器框架，但它们会使用多个 Multi-Head 注意力、前馈网络、层级归一化和残差连接等。下图从左到右展示了原论文所提出的 Transformer 架构、Multi-Head 注意力和标量点乘注意力。

上图右边的点乘注意力就是标准 Seq2Seq 模型中的注意力机制，中间的 Multi-head 注意力其实就是将一个隐层信息切分为多份，并单独计算注意力信息，使得一个词与其它多个目标词的注意力信息计算更精确。最左边为 Transformer 的整体架构，编码器与解码器由多个类似的模块组成，后面将简要介绍这些模块与对应的 PaddlePaddle 代码。

2.1 点乘注意力

注意力机制目前在机器翻译中已经极其流行了，我们可以认为 Transformer 是一种堆叠多层注意力网络的模型，它采用的是一种名为经缩放的点乘注意力机制。这种注意力机制使用经缩放的点乘作为作为评分函数，从而评估各隐藏状态对当前预测的重要性，如下是该注意力的表达式：

其中 Query 向量与 (Key, Value ) 向量在 NMT 中相当于目标语输入序列与源语输入序列，Query 与 Key 向量的点乘，经过 SoftMax 函数后可得出一组归一化的概率。这些概率相当于给源语输入序列做加权平均，即表示在生成新的隐层信息的时候需要关注哪些词。

在 Transformer 的 PaddlePaddle 实现中，经缩放的点乘注意力是在 Multi-head 注意力函数下实现的，如下所示为上述表达式的实现代码：

在这个函数中，q、k、v 和公式中的一样，attn_bias 用于 Mask 掉选定的特定位置（encode 的 self attention 和 decoder 端的 encode attention 都是屏蔽掉 padding 的词；decoder 的 self attention 屏蔽掉当前词后面的词，目的是为了和解码的过程保持一致），因此在给不同输入加权时忽略该位置的输入。

如上 product 计算的是 q 和 k 之间的点乘，且经过根号下 d_key（key 的维度）的缩放。这里我们可以发现参数 alpha 可以直接对矩阵乘法的结果进行缩放，默认情况下它为 1.0，即不进行缩放。在 Transformer 原论文中，作者表示如果 d_key 比较小，那么直接点乘和带缩放的点乘差别不大，所以他们认为高维情况下可能不带缩放的乘积太大而令 Softmax 函数饱和。

weights 表示对输入的不同元素加权，即不同输入对当前预测的重要性，训练中也可以对该权重进行 Dropout。最后 out 表示按照 weights 对输入 V 进行加权和，得出来就是当前注意力的运算结果。

2.2 Muti-head 注意力

Multi-head 注意力其实就是多个点乘注意力并行地处理并最后将结果拼接在一起。一般而言，我们可以对三个输入矩阵 Q、V、K 分别进行线性变换，然后分别将它们投入 h 个点乘注意力函数并拼接所有的输出结果。

这种注意力允许模型联合关注不同位置的不同表征子空间信息，我们可以理解为在参数不共享的情况下，多次执行点乘注意力。如下所示为 Muti-head 注意力的表达式：

其中每一个 head 都为一个点乘注意力，不同 head 的输入是相同 Q、K、V 的不同线性变换。

总体而言，PaddlePaddle 的 Multi-head 注意力实现分为几个步骤：先为 Q、K、V 执行线性变换；再变换维度以计算点乘注意力；最后计算各 head 的注意力输出并合并在一起。

2.2.1 线性变换

如前公式所示，Muti-head 首先要执行线性变换，从而令不同的 head 关注不同表征空间的信息。这种线性变换即乘上不同的权重矩阵，且模型在训练过程中可以学习和更新这些权重矩阵。在如下的 PaddlePaddle 代码中，我们可以直接调用全连接层 layers.fc() 完成线性变换。

直接调用全连接层会自动为输入创建权重，且我们要求不使用偏置项和激活函数。这里比较方便的是，PaddlePaddle 的 layers.fc() 函数可以接受高维输入，省略了手动展平输入向量的操作。因此这里有 num_flatten_dims=2，即将前两个维度展平为一个维度，第三个维度保持不变。

例如对于输入张量 q 而言，线性变换的输出维度应该是 [batch_size，max_sequence_length，d_key * n_head]，最后一个维度即 n_head 个 d_key 维的 Query 向量。每一个 d_key 维的向量都会馈送到不同的 head，并最后拼接起来。

2.2.2 维度变换

为了进行 Multi-Head 的运算，我们需要将线性变换的结果进行 reshape 和转置操作。现在我们将这几个张量的最后一个维度分割成不同的 head，并做转置以便于后续运算。

具体而言，输入张量 q、k 和 v 的维度信息为 [bs, max_sequence_length, n_head * hidden_dim]，我们希望把它们转换为 [bs, n_head, max_sequence_length, hidden_dim]。

如上使用 layers.reshape() 和 layers.transpose() 函数完成分割与转置。其中 layers.reshape() 在接收输入张量后会按照形状 [0, 0, n_head, d_key] 进行转换，其中 0 表示从输入张量对应维数复制出来。此外，因为 inplace 设置为 True，那么 reshape 操作就不会进行数据的复制，从而提升运算效率。

后面的转置就比较简单了，只需要按照维度索引将第「1」个维度和第「2」个维度交换就行了。此外为了更快地执行推断，PaddlePaddle 实现代码还做了非常多的优化，例如这部分后续会对推断过程的缓存和处理流程进行优化。

2.2.3 合并

前面已经介绍过点乘注意力了，那么上面对 q、k、v 执行维度变换后就可直接传入点乘注意力函数，并计算出 head_1、head_2 等注意力结果。现在最后一步只需要将这些 head 拼接起来就完成了整个过程，也就完成了上面 Multi-head 注意力的计算式。

因为每一个批量、head 和时间步都会计算得出一个注意力向量，因此总体上注意力计算结果的维度信息为 [bs, n_head, max_sequence_length, hidden_dim]。如果要将不同的 head 拼接在一起，即将 head 这个维度合并到 hidden_dim 中去，因此合并的过程和前面维度变换的过程正好相反。

如上合并过程会先检验维度信息，然后先转置再 reshape 合并不同的 head。注意在原论文中，合并不同的 head 后，还需要再做一个线性变换，这个线性变换的结果就是 Muti-head 注意力的输出了。

最后，我们再将上面的四部分串起来就是 Transformer 最核心的 Multi-head 注意力。理解了各个模块后，下面串起来就能愉快地看懂整个过程了：

当然，如果编码器和解码器输入到 Multi-head 注意力的 q 与 (k、v) 是相同的，那么它又可称为自注意力网络。

2.3 前馈网络

对于每一个编码器和解码器模块，除了残差连接与层级归一化外，重要的就是堆叠 Muti-head 注意力和前馈网络（FFN）。前面我们已经解决了 Multi-head 注意力，现在需要理解主位置的前馈网络了。直观而言，FFN 的作用是整合 Multi-head 注意力生成的上下文向量，因此能更好地利用从源语句子和目标语句子抽取的深度信息。

如下所示在原论文中，前馈网络的计算过程可以表达为以下方程：

前馈网络的结构很简单，一个 ReLU 激活函数加两次线性变换就完成了。如下基本上只需要调用PaddlePaddle 的 layers.fc() 就可以了：

现在基本上核心操作就定义完了，后面还有更多模块与架构，例如怎样利用核心操作搭建编码器模块与解码器模块、如何搭建整体 Transformer 模型等，读者可继续阅读原项目中的简洁代码。整体而言，包括上面代码在内，千行代码就可以完全弄懂 Transformer，PaddlePaddle 的 Transformer 复现值得我们仔细读一读。

此外，在这千行模型代码中，为了给训练和推断加速，还有很多特殊技巧。例如在 Decoder 中加入对 Encoder 计算结果的缓存等。加上这些技巧，PaddlePaddle 的实现才能在大 Batch Size 下实现 4 倍推断加速。

因为本身 PaddlePaddle 代码就已经非常精炼，通过它们也很容易理解这些技巧。基本上看函数名称就能知道大致的作用，再结合文档使用就能完全读懂了。

最后，除了模型架构，整个项目还会有其它组成部分，例如训练、推断、数据预处理等等。这些代码同样非常简洁，我们可以根据实际需求阅读并修改它们。总体而言，PaddlePaddle 的 Transformer 实现确实非常适合理解与修改。想要跑一跑神经机器翻译的同学，PaddlePaddle 的 Transformer 实现确实值得推荐。