【自然语言处理】从0开始，构建用于语言翻译的 Transformer 模型（一）

秋说

2913人浏览 · 2025-09-10 20:15:26

秋说 · 2025-09-10 20:15:26 发布

【精选优质专栏推荐】

《AI 技术前沿》 —— 紧跟 AI 最新趋势与应用
《网络安全新手快速入门(附漏洞挖掘案例)》 —— 零基础安全入门必看
《BurpSuite 入门教程(附实战图文)》 —— 渗透测试必备工具详解
《网安渗透工具使用教程(全)》 —— 一站式工具手册
《CTF 新手入门实战教程》 —— 从题目讲解到实战技巧
《前后端项目开发(新手必知必会)》 —— 实战驱动快速上手

每个专栏均配有案例与图文讲解，循序渐进，适合新手与进阶学习者，欢迎订阅。

文章目录

Transformer 架构于 2017 年提出，通过消除对循环神经网络的依赖，彻底改变了语言翻译等序列到序列任务。相反，它依赖自注意力机制来处理输入序列。

在本文中，我们将从零构建一个 Transformer 模型。你将理解：

自注意力如何处理输入序列
Transformer 编码器和解码器的工作原理
如何使用 Transformer 实现完整的翻译系统

1. 为什么 Transformer 优于 Seq2Seq

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

顺序处理阻碍了并行化
隐状态在处理每个元素时会被覆盖，因此捕捉长期依赖的能力有限

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

2. 数据准备与分词

在本文中，我们将构建一个用于翻译的 Transformer 模型，因为这是完整 Transformer 的典型使用场景。

我们将使用的数据集是来自 Anki 的英法翻译数据集，其中包含英语和法语句子对。这与之前文章中使用的数据集相同，准备步骤也类似。

法语文本包含重音符号和复杂的动词变位，需要比简单的单词分割更复杂的分词方法。字节对编码（Byte-Pair Encoding, BPE）可以有效处理这些子词单元和形态丰富的语言，同时也是处理未知单词的好方法。

首先，我们需要下载数据集并将其读取到内存中。数据集是一个纯文本文件，每行由一个英语句子和一个法语句子组成，二者之间用制表符分隔。下面是下载并读取数据集的方法：

import os
import unicodedata
import zipfile
import requests

# 使用 requests 下载 Anki 提供的数据集: https://www.manythings.org/anki/
if not os.path.exists("fra-eng.zip"):
    url = "http://storage.googleapis.com/download.tensorflow.org/data/fra-eng.zip"
    response = requests.get(url)
    with open("fra-eng.zip", "wb") as f:
        f.write(response.content)

# 文本归一化
# 文件的每一行格式为 "<english>\t<french>"
# 将文本转换为小写并进行 Unicode 归一化（NFKC）
def normalize(line):
    """归一化一行文本并按制表符拆分"""
    line = unicodedata.normalize("NFKC", line.strip().lower())
    eng, fra = line.split("\t")
    return eng.lower().strip(), fra.lower().strip()

text_pairs = []
with zipfile.ZipFile("fra-eng.zip", "r") as zip_ref:
    for line in zip_ref.read("fra.txt").decode("utf-8").splitlines():
        eng, fra = normalize(line)
        text_pairs.append((eng, fra))

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

text_pairs 中的翻译对是完整句子的字符串对。我们可以使用它们训练 BPE 分词器，以便对未来句子进行分词：

import tokenizers

if os.path.exists("en_tokenizer.json") and os.path.exists("fr_tokenizer.json"):
    en_tokenizer = tokenizers.Tokenizer.from_file("en_tokenizer.json")
    fr_tokenizer = tokenizers.Tokenizer.from_file("fr_tokenizer.json")
else:
    en_tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE())
    fr_tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE())

    # 配置预分词器：按空格和标点分割，并在句首添加空格
    en_tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.ByteLevel(add_prefix_space=True)
    fr_tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.ByteLevel(add_prefix_space=True)

    # 配置解码器：移除单词边界符号 "Ġ"
    en_tokenizer.decoder = tokenizers.decoders.ByteLevel()
    fr_tokenizer.decoder = tokenizers.decoders.ByteLevel()

    # 使用相同的训练器训练英文和法文 BPE
    VOCAB_SIZE = 8000
    trainer = tokenizers.trainers.BpeTrainer(
        vocab_size=VOCAB_SIZE,
        special_tokens=["[start]", "[end]", "[pad]"],
        show_progress=True
    )
    en_tokenizer.train_from_iterator([x[0] for x in text_pairs], trainer=trainer)
    fr_tokenizer.train_from_iterator([x[1] for x in text_pairs], trainer=trainer)

    en_tokenizer.enable_padding(pad_id=en_tokenizer.token_to_id("[pad]"), pad_token="[pad]")
    fr_tokenizer.enable_padding(pad_id=fr_tokenizer.token_to_id("[pad]"), pad_token="[pad]")

    # 保存训练好的分词器
    en_tokenizer.save("en_tokenizer.json", pretty=True)
    fr_tokenizer.save("fr_tokenizer.json", pretty=True)

上面的代码使用 Hugging Face 的 tokenizers 库训练分词器。训练好的分词器以 JSON 文件形式保存以便重用。训练时添加了三个特殊符号：[start]、[end] 和 [pad]，分别用于标记句子开始、结束以及对序列进行填充。通过 enable_padding() 配置，使用分词器处理字符串时会自动添加填充符号。在后续部分，我们将看到它们的使用方法。

下面是使用分词器的示例：

encoded = fr_tokenizer.encode("[start] " + fr_sample + " [end]")
print(f"原句: {fr_sample}")
print(f"分词结果: {encoded.tokens}")
print(f"ID 序列: {encoded.ids}")
print(f"解码结果: {fr_tokenizer.decode(encoded.ids)}")

分词器不仅将文本拆分为 token，还提供将 token 编码为整数 ID 的方法。这对于 Transformer 模型是必要的，因为模型需要将输入序列作为数字序列进行处理。

3. Transformer 模型设计

Transformer 由编码器（Encoder）和解码器（Decoder）组成。编码器包含多层自注意力和前馈网络，而解码器还引入了交叉注意力。编码器负责处理输入序列，解码器生成输出序列，这与传统的 seq2seq 模型类似。但 Transformer 模型中存在许多变体。常见的架构变体包括：

位置编码（Positional Encoding）：提供位置信息，因为 Transformer 并行处理序列。将序列中元素的位置传递给模型有多种策略。
注意力机制（Attention Mechanism）：虽然标准是缩放点积注意力（scaled dot-product attention），但其实现存在多种变体，例如多头注意力（Multi-Head Attention, MHA）、多查询注意力（Multi-Query Attention, MQA）、分组查询注意力（Grouped Query Attention, GQA）以及多头潜在注意力（Multi-Head Latent Attention, MLA）。这是因为 Transformer 中的每一层注意力包含多个并行的注意力“头”。不同的实现方式对应将输入应用到不同头上的方式。
前馈网络（Feed-forward Network）：这是一个多层感知机（MLP）网络，但可以选择不同的激活函数或层数。在需要处理多种输入的大模型中，也可以用专家混合网络（Mixture-of-Experts Network）替代前馈网络。
层归一化（Layer Normalization）：在注意力和前馈网络之间应用 Layer Norm 或 RMS Norm，可以选择使用“预归一化（pre-norm）”或“后归一化（post-norm）”并配合跳跃连接。
超参数（Hyperparameters）：同一设计下，可通过调整隐藏维度大小、注意力头/层数、Dropout 比例以及模型支持的最大序列长度来扩展模型。

在本文中，我们采用以下配置：

位置编码：旋转位置编码（Rotary Positional Encoding），最大序列长度 768
注意力机制：分组查询注意力（Grouped-Query Attention），8 个查询头，4 个键值头
前馈网络：两层 SwiGLU，隐藏层维度 512
层归一化：RMS Norm，预归一化（pre-norm）
隐藏维度：128
编码器和解码器层数：4
Dropout 比例：0.1

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

在这里插入图片描述

4. 构建 Transformer 模型

对于 RoPE（旋转位置编码），其 PyTorch 实现如下：

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(x, cos, sin):
    return (x * cos) + (rotate_half(x) * sin)

class RotaryPositionalEncoding(nn.Module):
    def __init__(self, dim, max_seq_len=1024):
        super().__init__()
        N = 10000
        inv_freq = 1. / (N ** (torch.arange(0, dim, 2).float() / dim))
        position = torch.arange(max_seq_len).float()
        inv_freq = torch.cat((inv_freq, inv_freq), dim=-1)
        sinusoid_inp = torch.outer(position, inv_freq)
        self.register_buffer("cos", sinusoid_inp.cos())
        self.register_buffer("sin", sinusoid_inp.sin())

    def forward(self, x, seq_len=None):
        if seq_len is None:
            seq_len = x.size(1)
        cos = self.cos[:seq_len].view(1, seq_len, 1, -1)
        sin = self.sin[:seq_len].view(1, seq_len, 1, -1)
        return apply_rotary_pos_emb(x, cos, sin)

旋转位置编码通过对向量中的每两个元素乘以 2×2 的旋转矩阵来改变输入向量：

$\hat{\mathbf{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 的正弦位置编码不同，它是在注意力子层内部应用，而不是外部。

我们将使用的注意力机制是分组查询注意力（Grouped-Query Attention, GQA）。PyTorch 支持 GQA，但在注意力子层中，需要实现查询（Q）、键（K）和值（V）的投影。

下面是一个扩展版的 GQA 实现，可以同时用于自注意力和交叉注意力：

class GQA(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_kv_heads=None, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads or num_heads
        self.head_dim = hidden_dim // num_heads
        self.num_groups = num_heads // num_kv_heads
        self.dropout = dropout
        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, q, k, v, mask=None, rope=None):
        q_batch_size, q_seq_len, hidden_dim = q.shape
        k_batch_size, k_seq_len, hidden_dim = k.shape
        v_batch_size, v_seq_len, hidden_dim = v.shape

        # 投影
        q = self.q_proj(q).view(q_batch_size, q_seq_len, -1, self.head_dim).transpose(1, 2)
        k = self.k_proj(k).view(k_batch_size, k_seq_len, -1, self.head_dim).transpose(1, 2)
        v = self.v_proj(v).view(v_batch_size, v_seq_len, -1, self.head_dim).transpose(1, 2)

        # 应用旋转位置编码
        if rope:
            q = rope(q)
            k = rope(k)

        # 计算分组查询注意力
        q = q.contiguous()
        k = k.contiguous()
        v = v.contiguous()
        output = F.scaled_dot_product_attention(q, k, v,
                                                attn_mask=mask,
                                                dropout_p=self.dropout,
                                                enable_gqa=True)
        output = output.transpose(1, 2).reshape(q_batch_size, q_seq_len, hidden_dim).contiguous()
        output = self.out_proj(output)
        return output

注意，在 GQA 类的 forward() 方法中，我们可以通过 rope 参数指定位置编码模块，从而使位置编码可选。在 PyTorch 中，为了优化注意力计算，输入张量应在内存中保持连续块。q = q.contiguous() 用于在张量非连续时重排内存布局。

我们将使用的前馈网络是两层 SwiGLU。SwiGLU 激活函数的特点是 PyTorch 原生不支持，但可以使用 SiLU 激活来实现。下面是使用 SwiGLU 构建前馈网络的实现：

class SwiGLU(nn.Module):
    def __init__(self, hidden_dim, intermediate_dim):
        super().__init__()
        self.gate = nn.Linear(hidden_dim, intermediate_dim)
        self.up = nn.Linear(hidden_dim, intermediate_dim)
        self.down = nn.Linear(intermediate_dim, hidden_dim)
        self.act = nn.SiLU()

    def forward(self, x):
        x = self.act(self.gate(x)) * self.up(x)
        x = self.down(x)
        return x

有了它，我们可以构建编码器和解码器层。编码器层较为简单，由自注意力层和前馈网络组成，但仍需实现跳跃连接和使用 RMSNorm 的预归一化。

编码器层实现如下：

class EncoderLayer(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_kv_heads=None, dropout=0.1):
        super().__init__()
        self.self_attn = GQA(hidden_dim, num_heads, num_kv_heads, dropout)
        self.mlp = SwiGLU(hidden_dim, 4 * hidden_dim)
        self.norm1 = nn.RMSNorm(hidden_dim)
        self.norm2 = nn.RMSNorm(hidden_dim)

    def forward(self, x, mask=None, rope=None):
        # 自注意力子层
        out = self.norm1(x)
        out = self.self_attn(out, out, out, mask, rope)
        x = out + x
        # 前馈子层
        out = self.norm2(x)
        out = self.mlp(out)
        return out + x

前馈网络的中间维度定义为隐藏维度的 4 倍，这是业界常用设计，但我们可以尝试不同的比例。

解码器层较为复杂，由自注意力层、交叉注意力层以及前馈网络组成，代码如下：

class DecoderLayer(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_kv_heads=None, dropout=0.1):
        super().__init__()
        self.self_attn = GQA(hidden_dim, num_heads, num_kv_heads, dropout)
        self.cross_attn = GQA(hidden_dim, num_heads, num_kv_heads, dropout)
        self.mlp = SwiGLU(hidden_dim, 4 * hidden_dim)
        self.norm1 = nn.RMSNorm(hidden_dim)
        self.norm2 = nn.RMSNorm(hidden_dim)
        self.norm3 = nn.RMSNorm(hidden_dim)

    def forward(self, x, enc_out, mask=None, rope=None):
        # 自注意力子层
        out = self.norm1(x)
        out = self.self_attn(out, out, out, mask, rope)
        x = out + x
        # 交叉注意力子层
        out = self.norm2(x)
        out = self.cross_attn(out, enc_out, enc_out, None, rope)
        x = out + x
        # 前馈子层
        out = self.norm3(x)
        out = self.mlp(out)
        return out + x

可以看到，自注意力和交叉注意力子层都使用 GQA 类实现。区别在于 forward() 中的使用方式：RoPE 作用于两者，但 mask 仅在自注意力子层使用。

Transformer 模型通过连接编码器和解码器构建，在将 token ID 序列输入编码器或解码器前，会先将其转换为嵌入向量。实现如下：

class Transformer(nn.Module):
    def __init__(self, num_layers, num_heads, num_kv_heads, hidden_dim,
                 max_seq_len, vocab_size_src, vocab_size_tgt, dropout=0.1):
        super().__init__()
        self.rope = RotaryPositionalEncoding(hidden_dim // num_heads, max_seq_len)
        self.src_embedding = nn.Embedding(vocab_size_src, hidden_dim)
        self.tgt_embedding = nn.Embedding(vocab_size_tgt, hidden_dim)
        self.encoders = nn.ModuleList([
            EncoderLayer(hidden_dim, num_heads, num_kv_heads, dropout) for _ in range(num_layers)
        ])
        self.decoders = nn.ModuleList([
            DecoderLayer(hidden_dim, num_heads, num_kv_heads, dropout) for _ in range(num_layers)
        ])
        self.out = nn.Linear(hidden_dim, vocab_size_tgt)

    def forward(self, src_ids, tgt_ids, src_mask=None, tgt_mask=None):
        # 编码器
        x = self.src_embedding(src_ids)
        for encoder in self.encoders:
            x = encoder(x, src_mask, self.rope)
        enc_out = x
        # 解码器
        x = self.tgt_embedding(tgt_ids)
        for decoder in self.decoders:
            x = decoder(x, enc_out, tgt_mask, self.rope)
        return self.out(x)

Transformer 类的构造函数有很多参数，因为它是整个模型的入口，会初始化所有子层。这种设计便于使用 Python 字典定义模型配置，例如：

model_config = {
    "num_layers": 4,
    "num_heads": 8,
    "num_kv_heads": 4,
    "hidden_dim": 128,
    "max_seq_len": 768,
    "vocab_size_src": len(en_tokenizer.get_vocab()),
    "vocab_size_tgt": len(fr_tokenizer.get_vocab()),
    "dropout": 0.1,
}
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Transformer(**model_config).to(device)