构建用于语言翻译的 Transformer 模型

Transformer 架构于 2017 年推出，通过消除循环神经网络的需要，彻底改变了序列到序列的任务，例如语言翻译。相反，它依赖于自注意力机制来处理输入序列。在本帖中，您将学习如何从头开始构建 Transformer 模型。特别是，您将理解

• 自注意力机制如何处理输入序列
• Transformer 编码器和解码器如何工作
• 如何实现一个完整的 Transformer 翻译系统

让我们开始吧。

构建用于语言翻译的 Transformer 模型
照片来自 Sorasak。部分权利保留。

概述

这篇文章分为六个部分；它们是：

• 为什么 Transformer 优于 Seq2Seq
• 数据准备与分词
• Transformer 模型的设计
• 构建 Transformer 模型
• 因果掩码和填充掩码
• 训练与评估

为什么 Transformer 优于 Seq2Seq

传统的基于循环神经网络的 seq2seq 模型有两个主要局限性

1. 顺序处理阻碍了并行化
2. 捕捉长期依赖关系的能力有限，因为每当处理一个元素时，隐藏状态都会被覆盖

Transformer 架构在 2017 年的论文“Attention is All You Need”中提出，克服了这些局限性。它可以使用自注意力机制捕捉序列中任何位置之间的依赖关系。它可以并行处理整个序列。Transformer 模型的序列处理能力不依赖于循环连接。

数据准备与分词

在本帖中，您将构建一个用于翻译的 Transformer 模型，因为这是完整 Transformer 的典型用例。

您将使用的数据集是 Anki 的英法翻译数据集，其中包含英法句子对。这与您在之前帖子中使用的相同数据集，并且准备步骤也类似。

法语文本包含重音符号和复杂的动词变位，比简单的单词分割需要更复杂的分词。字节对编码 (BPE) 有效地处理这些子词单元和形态丰富的语言。它也是处理未知词语的好解决方案。

首先，您需要下载数据集并将其读入内存。该数据集是纯文本文件，每行包含一个由制表符分隔的英语和法语句子。以下是如何下载和读取数据集

法语句子使用 Unicode 字符，这些字符可以有多种表示形式。我们在处理之前将文本规范化为“NFKC”形式以确保一致的表示。这是一个确保文本“干净”的好习惯，以便模型可以专注于文本的实际内容。

text_pairs 中的翻译对是完整句子的字符串对。您可以使用它们来训练 BPE 分词器，然后使用该分词器对后续句子进行分词。

上面的代码使用 Hugging Face 的 tokenizers 库来训练分词器。训练好的分词器会保存为 JSON 文件以供重复使用。在训练分词器时，您添加了三个特殊标记：[start]、[end] 和 [pad]。这些标记用于标记句子的开头和结尾，以及将序列填充到相同的长度。分词器设置为 enable_padding()，因此当您使用分词器处理字符串时，会添加填充标记。您将在后面的部分中看到它们是如何使用的。

下面是如何使用分词器的示例

分词器不仅将文本分割成标记，还提供了一种将标记编码为整数 ID 的方法。这对于 Transformer 模型至关重要，因为模型需要将输入序列作为数字序列进行处理。

Transformer 模型的设计

Transformer 结合了编码器和解码器。编码器包含多层自注意力和前馈网络，而解码器还包含交叉注意力。编码器处理输入序列，解码器生成输出序列，这与 seq2seq 模型的情况相同。然而，Transformer 模型中有许多变体。常见的架构变体包括

• 位置编码：提供位置信息，因为 Transformer 会并行处理序列。有多种策略可以将序列中元素的��置传递给模型。
• 注意力机制：虽然缩放点积注意力是标准方法，但在模型级别存在其实现的变体，例如多头注意力 (MHA)、多查询注意力 (MQA)、分组查询注意力 (GQA) 和多头潜在注意力 (MLA)。这是因为 Transformer 模型中的每个注意力层都包含多个并行运行的注意力“头”。这些是将输入应用于不同头的不同方法。
• 前馈网络：这是一个多层感知器网络，但您可以选择不同的激活函数或层数。在需要处理各种输入的模型较大的情况下，可以使用专家混合网络作为前馈网络的替代方案。
• 层归一化：应在注意力和前馈网络之间应用层归一化或 RMS 归一化。您可以使用带跳跃连接的“预归一化”或“后归一化”。
• 超参数：对于相同的设计，您可以调整隐藏维度的大小、头/层的数量、dropout 率以及模型应支持的最大序列长度来扩展模型。

在本帖中，我们将使用以下内容

• 位置编码：旋转位置编码，最大序列长度为 768
• 注意力机制：分组查询注意力，具有 8 个查询头和 4 个键值头
• 前馈网络：两层 SwiGLU，隐藏层维度为 512
• 层归一化：RMS Norm，预归一化
• 隐藏维度：128
• 编码器和解码器层数：4
• Dropout 率：0.1

要构建的 Transformer 模型

我们将构建的模型如下图所示

构建 Transformer 模型

各种位置编码方法及其实现已在之前的帖子中介绍。对于 RoPE，这是 PyTorch 的实现

旋转位置编码通过将向量的每两个元素乘以一个 2x2 的旋转矩阵来改变输入向量。

$$
\mathbf{\hat{x}}_m = \mathbf{R}_m\mathbf{x}_m = \begin{bmatrix}
\cos(m\theta_i) & -\sin(m\theta_i) \\
\sin(m\theta_i) & \cos(m\theta_i)
\end{bmatrix} \mathbf{x}_m
$$

其中 $\mathbf{x}_m$ 代表位置 $m$ 上向量的元素对 $(i, d/2+i)$。实际使用的矩阵取决于向量在序列中的位置 $m$。

RoPE 与原始 Transformer 的正弦位置编码不同之处在于，它是在注意力子层内部应用的，而不是在外部。

您将使用的注意力是分组查询注意力 (GQA)。PyTorch 支持 GQA，但在注意力子层中，您应该实现查询、键和值的投影。GQA 的实现已在之前的帖子中介绍，但以下是一个扩展版本，允许您不仅在自注意力中，还在交叉注意力中使用它。

请注意，在GQA类的forward()方法中，您可以在rope参数中指定位置编码模块。这使得位置编码成为可选的。在PyTorch中，为了优化注意力计算，输入张量应该是内存中的连续块。如果张量不是连续的，则使用q = q.contiguous()这行代码来对其进行重塑。

您将使用的前馈网络是双层SwiGLU。SwiGLU激活函数很独特，因为PyTorch不支持它，但可以通过SiLU激活来实现。下面是来自之前一篇帖子的SwiGLU前馈网络实现。

这样，您就可以构建编码器和解码器层了。编码器层更简单，因为它由一个自注意力层和一个前馈网络组成。但是，您仍然需要实现跳跃连接和使用RMS范数的预归一化。下面是编码器层的实现。

前馈网络已实现为之前定义的SwiGLU模块。您可以看到中间维度被定义为隐藏维度大小的4倍。这是行业中的常见设计，但您可以尝试不同的比例。

解码器层更复杂，因为它由一个自注意力层、一个交叉注意力层和一个前馈网络组成。实现如下：

您可以看到，自注意力和交叉注意力子层都是使用GQA类实现的。区别在于它们在forward()方法中的使用方式。RoPE应用于两者，但掩码仅用于自注意力子层。

Transformer模型构建用于连接编码器和解码器，但在将序列传递给编码器或解码器之前，输入序列的token ID首先被转换为嵌入向量。实现如下：

可以看到，Transformer类在其构造函数中拥有大量的参数。这是因为它作为创建整个模型的入口点，Transformer类将初始化所有子层。这是一个很好的设计，因为您可以将Python字典定义为模型配置。下面是如何使用上面定义的类创建模型的示例。

因果掩码和填充掩码

训练模型的第一个步骤是创建一个数据集对象，该对象可用于以批次和随机顺序迭代数据集。在上一个部分，您已将数据集读入列表text_pairs。您还为英语和法语创建了分词器。现在，您可以使用PyTorch的Dataset类来创建数据集对象。下面是数据集对象的实现。

您可以尝试从数据集中打印一个样本。

TranslationDataset类封装了text_pairs，并向法语句子添加了[start]和[end]标记。dataloader对象在分词后提供批处理的、随机化的样本。collate_fn()函数负责分词和填充，以确保每个批次中的序列长度一致。

对于训练，我们使用交叉熵损失和Adam优化器。模型采用教师强制技术，在训练过程中向解码器提供真实序列，而不是重复使用其自身的输出。请注意，在教师强制中，解码器在生成第 N 个标记时应该只能看到前 N-1 个标记。

Transformer是一种可以并行化的架构。当您向解码器提供长度为N的序列时，它可以并行处理序列的所有元素并输出长度为N的序列。通常，我们将此输出序列的最后一个元素视为输出。或者，为了节省计算量，您可以使用输入序列的最后一个元素作为注意力中的“查询”，同时使用完整的输入序列作为“键”和“值”。

如果您仔细观察，您会发现对于长度为N的序列，您可以训练模型N次。如果模型可以并行化，您就可以为相同的输入序列N并行生成N个输出。但是，有一个问题：当模型生成输出N时，您希望它只使用直到位置N-1的序列，而不是位置N或之后的任何内容。

训练Transformer时的因果预测：迭代地向解码器提供更长的序列（白色方块）。每一步，解码器预测一个额外的输出（蓝色方块）。灰色方块在相应步骤中未提供给模型。

为了实现这一点，可以使用因果掩码。因果掩码是一个形状为（N, N）的方阵，其中N是序列长度。如今，因果掩码被实现为三角矩阵，对角线以上的元素设置为$-\infty$，对角线或对角线下方的元素设置为0，如下所示。

$$
M = \begin{bmatrix}
0 & -\infty & -\infty & \cdots & -\infty \\
0 & 0 & -\infty & \cdots & -\infty \\
0 & 0 & 0 & \cdots & -\infty \\
\vdots & \vdots & \vdots & \ddots & \vdots \\
0 & 0 & 0 & \cdots & 0
\end{bmatrix}
$$

因果掩码通过注意力类GQA在解码器中使用，并由PyTorch中的scaled_dot_product_attention()函数调用。它将“掩盖”不允许注意到的位置（即“未来”位置）的注意力分数，以便softmax操作将那些位置设置为零。上面所示的矩阵M垂直表示“查询”，水平表示“键”。矩阵中的0位置表示查询只能注意到不晚于其自身位置的键。因此得名“因果”。

因果掩码应用于解码器的自注意力，其中查询和键是相同的序列。因此，M是一个方阵。您可以在PyTorch中创建这样的矩阵，如下所示。

除了因果掩码之外，您还想跳过序列中的填充标记。填充标记是在批次中的序列长度不一致时添加的。由于它们不应携带任何信息，因此应将其从注意或输出的损失计算中排除。填充掩码也是每个序列的方阵。以下是根据张量中的序列批次创建填充掩码的Python代码。

此代码首先创建一个与张量batch形状匹配的二维张量padded。张量padded在所有地方都为零，除了原始张量batch等于填充标记ID的位置。然后创建一个三维张量mask，其形状为(batch_size, seq_len, seq_len)。张量mask是方阵的批次。在每个方阵中，行和列都用padded进行设置，使得对应于填充标记的位置被设置为$-\infty$。

上述函数使用了PyTorch中的维度扩展技术。用None索引一个张量会在该位置添加一个新的维度。它还利用PyTorch的广播功能用padded张量填充mask。

创建的填充掩码的形状为（batch_size, 1, seq_len, seq_len）。然而，因果掩码的形状为（seq_len, seq_len）。当应用自注意力时，它们可以被广播并相加。

训练与评估

现在您可以按如下方式实现训练循环。

训练循环实现为嵌套的for循环。每个epoch扫描一次整个数据集。从数据集中提取的每个批次都用于创建掩码。然后，将数据和掩码传递给模型以生成输出。然后通过将输出与真实值进行比较来计算损失。损失反向传播以更新模型参数。

用于编码器的掩码是源（英语）序列的填充掩码。用于解码器的掩码是因果掩码加上目标（法语）序列的填充掩码。反向传播计算的梯度被裁剪以减轻梯度爆炸问题。

虽然训练循环使模型能够学会生成目标序列，但在每个epoch之后运行评估以评估模型的性能并保存最佳模型也很有益。评估实现如下。

此评估复用了与训练相同的评估集，因为你没有单独的测试集。代码与训练循环类似，只是你不需要反向传播，并且在 torch.no_grad() 上下文中运行模型。

损失是针对整个数据集平均计算的。变量 best_loss 跟踪最低损失。每当损失得到改善时，模型的一个副本就会被保存。

一个训练良好的模型应该能达到 0.1 或更低的平均损失。训练完成后，你的模型就可以使用了。下面是一个例子

这比训练循环更复杂，因为你没有使用模型的 forward() 方法，而是分别使用了编码器和解码器。你首先使用编码器获取上下文 enc_out。然后，你以 fr_ids 作为起始 token 开始，并从 transformer 的解码器部分迭代生成输出。每一步都会将 fr_ids 扩展一个 token。生成将在生成结束 token 或达到最大长度时结束。

你也可以使用模型的 forward() 方法，但你会在每一步都用相同的源序列调用编码器。这上面代码中已优化了不必要的计算。实际上，你可能希望在模型类中实现一个仅用于推理的方法。

运行上面的代码，你会看到以下输出

为了完整起见，完整的代码如下