数学模型#

注意力机制的核心思想是，句子中的每个词，都可以直接与句子中的所有其他词进行交互，“注意到”那些与自己相关的部分。

Q、K、V 模型

注意力机制借用了数据库检索的概念：

• 查询张量(Query, $q \in \mathbb{R}^{N \times d_k}$ 或 $\mathbb{R}^{M \times d_k}$)：当前正在寻找相关信息的词(我该关注谁？)。
• 键张量(Key, $k \in \mathbb{R}^{N \times d_k}$)：被考察的词(我包含什么信息？)。
• 值张量(Value, $v \in \mathbb{R}^{N \times d_v}$)：被考察词的实际内容(如果你关注我，我就把这段信息给你)。

从张量形状的角度看，$N$ 是输入序列的长度，$d_k$ 和 $d_v$ 是键和值的维度。这表明，对于每一个 token，我们都为它生成了一个 $d_k$ 维度的查询向量、一个 $d_k$ 维度的键向量和一个 $d_v$ 维度的值向量。

根据网络构建目的的不同，注意力机制主要有两种模式，分别是自注意力(Self-Attention) 和 交叉注意力(Cross-Attention)。这两种注意力机制的 $q$、$k$、$v$ 特征向量的来源不同：

1. 对于自注意力(Self-Attention)，$q$、$k$、$v$ 都来自同一个输入序列 $x \in \mathbb{R}^{N \times d}$：

   $$q = xW_Q, \quad k = xW_K, \quad v = xW_V$$

   其中 $W_Q \in \mathbb{R}^{d \times d_k}, W_K \in \mathbb{R}^{d \times d_k}, W_V \in \mathbb{R}^{d \times d_v}$ 是模型自动学习的权重矩阵。

2. 对于交叉注意力(Cross-Attention)，$q$ 来自一个序列(例如目标语言的输入)，而 $k$ 和 $v$ 来自另一个序列(例如源语言的输入)：

   $$q = y_{\text{target}}W_Q, \quad k = x_{\text{source}}W_K, \quad v = x_{\text{source}}W_V$$

   其中 $y_{\text{target}} \in \mathbb{R}^{M \times d}$ 是目标输入序列，$x_{\text{source}} \in \mathbb{R}^{N \times d}$ 是源输入序列，$W_Q \in \mathbb{R}^{d \times d_k}, W_K \in \mathbb{R}^{d \times d_k}, W_V \in \mathbb{R}^{d \times d_v}$ 是模型自动学习的权重矩阵。

缩放点积注意力(Scaled Dot-Product Attention)

注意力机制的目标是为了提取出一整个序列关于某一个 token 的信息。为了衡量信息的相关性，这里采用了点积的方式。有了特征张量 $q$ 和 $k$ 后，依次计算它们各个 tokens 之间的点积。点积结果越大，说明两个 tokens 越相关，也称作注意力得分越高。

单个点积的计算如下（这里把 $\mathbb{R}^{d_k}$ 的向量看作是行向量）

其中 $q_i \in \mathbb{R}^{d_k}$ 是查询张量 $q$ 中的第 $i$ 个 token 的查询向量，$k_j \in \mathbb{R}^{d_k}$ 是键张量 $k$ 中的第 $j$ 个 token 的键向量。通过计算所有 $q_i$ 与所有 $k_j$ 的点积组合，我们可以得到一个注意力得分矩阵

为了方便统一相似性度量的尺度，需要沿着键张量的序列长度维度 $N$ 做 Softmax 归一化，将得分转换为概率分布（每个查询 token 都会有一个概率分布），表示键张量对应序列中的每个 token 对当前查询 token 的关注程度：

然而，随着 $d_k$ 的增加，点积的结果可能会变得非常大，这会导致 Softmax 函数的容易出现梯度消失。为了解决这个问题，我们对点积结果进行缩放，除以 $\sqrt{d_k}$，从而得到缩放点积注意力的计算公式

缩放点积注意力的本质是关于值张量 $v$ 的权重，加权求和，就是最终的注意力结果

缩放点积注意力的 PyTorch 实现#

PyTorch 实现核心缩放点积注意力

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    # 1. 计算未缩放的分数 q * k^\top 
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 2. 如果存在掩码，将特定位置屏蔽，从而不参与注意力计算（后续会解释这步操作）
    if mask is not None:
        scores = scores.masked_fill(mask == False, -1e9)
        
    # 3. Softmax 归一化
    attention_weights = F.softmax(scores, dim=-1)
    
    # 4. 加权求和 Value
    output = torch.matmul(attention_weights, v)
    return output, attention_weights

多头注意力(Multi-Head Attention)#

为了让模型能从多个不同的“角度”去观察词与词间的关系，通常会将输入拆分成多个并行的“头”(Heads)来计算上述 Attention，然后将结果拼接在一起。

这里的头可以理解为模型关注的不同的特征领域，每个头都有自己独立的查询、键、值的权重矩阵。通过多头注意力，模型能够捕捉到输入序列中不同层次、不同类型的依赖关系，从而提升模型的表达能力和性能。

打个比方，使用多头注意力后，有的头可能专注于捕捉语法关系(如主谓宾)，有的头可能专注于捕捉语义关系(如同义词)，还有的头可能专注于捕捉长距离依赖(如跨句子关系)。通过这种方式，模型能够更全面地理解输入文本的结构和意义，从而生成更准确、更自然的翻译结果。

不论是自注意力机制还是交叉注意力机制，都可以封装成下面的 MultiHeadAttention 层

class MultiHeadAttention(nn.Module):
    """多头注意力模块"""
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def _split_heads(self, x):
        """将最后一维拆分为 (n_heads, d_k)，并转置为 (batch, n_heads, seq_len, d_k)"""
        batch, seq_len, _ = x.shape
        x = x.view(batch, seq_len, self.n_heads, self.d_k)
        return x.transpose(1, 2)  # (batch, n_heads, seq_len, d_k)

    def _combine_heads(self, x):
        """合并头，形状 (batch, seq_len, d_model)"""
        batch, n_heads, seq_len, d_k = x.shape
        x = x.transpose(1, 2).contiguous()
        return x.view(batch, seq_len, self.d_model)

    def forward(self, q, k, v, mask=None):
        """
        q, k, v: (batch, seq_len_q/k/v, d_model)
        mask: 可选，形状 (batch, 1, seq_len_q, seq_len_k) 或可广播的形状
        """
        # 线性变换并分头
        Q = self._split_heads(self.W_q(q))
        K = self._split_heads(self.W_k(k))
        V = self._split_heads(self.W_v(v))

        # 计算注意力
        attn_output, _ = scaled_dot_product_attention(Q, K, V, mask=mask)
        attn_output = self.dropout(attn_output)

        # 合并头并输出
        output = self._combine_heads(attn_output)
        output = self.W_o(output)
        return output

输入序列的文本嵌入 (Text Embeddings) 和位置编码 (Positional Encoding)#

Transformer 模型本身不包含任何循环或卷积结构，因此它无法天然感知序列中 token 的顺序信息，这一点从注意力结果的序列次序可交换性上可以看出来。但是，位置关系在语言建模中是相当重要的，一个简单的例子就足以体现

• 文本 1： “我爱自然语言处理”
• 文本 2： “自然语言处理爱我” 这两句话的 tokens 是完全一样的，但是它们的语义是完全不同的，这就是位置关系导致的差异。

为了让模型能够利用 token 在序列中的位置关系，我们必须在输入端显式地注入位置信息。具体做法是，将 文本嵌入 与 位置编码 相加，作为编码器和解码器真正的输入。其中位置编码必须要能体现出不同位置的区别。

数学描述上，位置编码的目标是为每个绝对位置 $pos$（$0 \le pos < L$，$L$ 为序列长度）生成一个 $d$ 维向量，使得模型能够区分不同位置的 token。有两种常见的设计方案：

1. 正余弦型位置编码 (Sinusoidal Positional Encoding)

其中 $\text{PE}(pos, k)$ 表示位置在 $pos$ 处的 token 的，在文本嵌入向量的第 $k$ 个维度上的位置编码。这种设计的优点包括：

1. 每个位置的编码是唯一的；
2. 对于任意固定的偏移量 $k$，$\text{PE}(pos+k)$ 可以表示为 $\text{PE}(pos)$ 的线性函数，这有助于模型学习相对位置信息；
3. 值域在 $[-1, 1]$ 之间，与文本嵌入相加时不会造成数值不稳定。

最后，将文本嵌入与位置编码按元素相加，得到编码器和解码器的真正输入：

PyTorch 实现

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        # 预计算所有位置的位置编码
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # (max_len, 1)
        
        # 计算分母项 10000^(2i/d_model)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        # 偶数索引使用 sin，奇数索引使用 cos
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # 增加 batch 维度方便广播，形状 (1, max_len, d_model)
        self.register_buffer('pe', pe.unsqueeze(0))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x 形状: (batch_size, seq_len, d_model)
        return x + self.pe[:, :x.size(1), :]

# 示例：生成一个 batch 的输入嵌入并加上位置编码
batch_size, seq_len, d_model = 2, 10, 512
embedding = nn.Embedding(num_embeddings=10000, embedding_dim=d_model)
x_ids = torch.randint(0, 10000, (batch_size, seq_len))
x_emb = embedding(x_ids)                     # 文本嵌入
pos_enc = PositionalEncoding(d_model)
x_input = pos_enc(x_emb)                     # 最终输入

2. 旋转位置编码 (Rotary Position Embedding, RoPE)

正弦余弦位置编码直接将位置向量加到嵌入上，而 RoPE 采用了一种不同的思路——将位置信息通过旋转矩阵融入 query 和 key 向量中，使得注意力分数天然依赖于 token 之间的相对位置。

对于序列中第 $m$, $n$ 个位置的 token，其对应的 query 向量 $q_m$ 和 key 向量 $k_n$ 不直接加上位置编码，而是乘以一个与位置 $m$ 相关的旋转矩阵

其中 $q_m, k_n$ 是未经过位置编码的原始 query/key 向量（来自文本嵌入的线性变换），$R_{\Theta,m}$ 是旋转矩阵。旋转矩阵具有性质

因此，两个旋转后的向量的内积为

这表明 注意力分数仅依赖于相对位置 $n-m$，而不是绝对位置 $m$ 和 $n$，这是 RoPE 能够很好地外推到更长序列的关键。

为了高效实现，RoPE 通常将 $d$ 维向量拆分成 $d/2$ 个二维子空间，每个子空间独立旋转。对于第 $i$ 个子空间（$0 \le i < d/2$），旋转角度为

于是位置 $m$ 的旋转矩阵为

实际计算时不需要构造完整的稀疏矩阵，而是利用复数乘法高效实现。将二维向量视为复数 $q_{2k} + j q_{2k+1}$，乘以旋转因子 $e^{j m \theta_k}$ 即可完成旋转。

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(q, k, cos, sin):
    """
    通用 RoPE 应用函数
    q, k: (batch, n_heads, seq_len, d_k)
    cos, sin: (1, 1, seq_len, d_k)
    """
    # 确保 cos 和 sin 的维度与 q 匹配（截断到当前序列长度）
    seq_len = q.shape[2]
    cos = cos[:, :, :seq_len, :]
    sin = sin[:, :, :seq_len, :]

    # 旋转公式: q_rotated = (q * cos) + (rotate_half(q) * sin)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

class Model(nn.Module):
    def __init__(self, ...):
        super().__init__()
        pass

    def get_rotary_emb(self, seq_len, device):
        """
        计算旋转位置编码的 cos 和 sin 缓存向量
        Args:
            seq_len: 当前输入的序列长度
            device: 张量所在设备
        Returns:
            cos, sin: 形状为 (1, 1, seq_len, d_k) 的旋转向量
        """
        t = torch.arange(seq_len, device=device).float()
        freqs = torch.outer(t, self.inv_freq)  # (seq_len, d_k/2)
        
        # 核心逻辑：这里为了配合 rotate_half 中的 chunk(2, dim=-1) 操作，
        # 我们需要构造一个对称的角度序列 [theta_0, ..., theta_{d_k/2-1}, theta_0, ..., theta_{d_k/2-1}]
        # 这样在 apply_rotary_pos_emb 中，每一对 (x1, x2) 都会乘以相同的旋转频率
        emb = torch.cat([freqs, freqs], dim=-1)         # (seq_len, d_k)
        
        # 整理形状以适配多头注意力：(batch, n_heads, seq_len, d_k)
        cos = emb.cos().view(1, 1, seq_len, self.d_k) 
        sin = emb.sin().view(1, 1, seq_len, self.d_k) 
        return cos, sin

还需注意的一点是，当引入多头注意力机制后，RoPE 位置编码的维度就不再是 $d$，而是每个头的维度 $d_{\text{head}} = d / \text{num\_heads}$ 了，需要注意调整。

在实际应用中，如果任务序列长度固定且较短（如机器翻译），正弦余弦编码足够；如果需要处理超长上下文或要求更好的外推性，RoPE 是更优选择。

注意力掩码 (Attention Mask)#

在实际训练和推理中，我们需要控制注意力机制能看到哪些位置。掩码（Mask）是一个与注意力得分矩阵形状相同（或可广播）的布尔矩阵，用于屏蔽某些位置的注意力分数（通常将屏蔽位置设为 $-\infty$，使其 Softmax 权重为 0）。

主要有两类掩码：

1. 填充掩码 (Padding Mask)

由于批次内序列长度不一，较短的序列会用 <PAD> token 填充至统一长度。这些填充位置不应参与注意力计算，否则模型会学到无意义的模式。

填充掩码通常是一个形状为 (batch_size, seq_len) 的布尔矩阵，其中 True 表示有效 token，False 表示填充 token。在计算注意力时，将填充 token 对应的键/查询位置屏蔽。

2. 因果掩码 (Causal Mask / Look-ahead Mask)

在解码器的自注意力中，为了保持自回归性质（预测第 $t$ 个 token 时不能看到 $t$ 及之后的 token），需要使用因果掩码。它确保位置 $i$ 只能关注到位置 $j \le i$。

因果掩码是一个下三角矩阵：

当与填充掩码结合时，通常将两者取逻辑与，然后在计算注意力得分前将无效位置设为 $-\infty$。

PyTorch 实现

def create_padding_mask(seq, pad_token_id=0):
    """seq: (batch_size, seq_len)  token id 序列"""
    return (seq != pad_token_id).unsqueeze(1).unsqueeze(2)  # (batch, 1, 1, seq_len)

def create_causal_mask(size):
    """生成下三角掩码，shape: (size, size)"""
    mask = torch.triu(torch.ones(size, size), diagonal=1).bool()  # 上三角为 True
    return ~mask   # True 表示允许关注

# 在注意力计算中合并掩码
def combine_masks(padding_mask, causal_mask):
    # padding_mask: (batch, 1, 1, seq_len_k)
    # causal_mask:   (1, seq_len_q, seq_len_k) 或 (seq_len_q, seq_len_k)
    # 输出形状 (batch, 1, seq_len_q, seq_len_k)
    if padding_mask is not None:
        mask = padding_mask & causal_mask.unsqueeze(0)
    else:
        mask = causal_mask
    return mask

# 在 scaled_dot_product_attention 中使用
def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        # mask 中 True 表示允许关注，False 表示屏蔽
        scores = scores.masked_fill(~mask, -1e9)
    attention_weights = F.softmax(scores, dim=-1)
    return torch.matmul(attention_weights, v), attention_weights

注意：在实际实现中，为了提高效率，常常将填充掩码和因果掩码预先计算好并合并为一个张量，传入注意力函数。

前馈层 (Feed-Forward Network)#

在多头注意力机制之后，Transformer 的每个编码器和解码器层都包含一个全连接的前馈网络（FFN）。该网络对每个 token 的表示独立地进行相同的非线性变换，从而增强模型的表达能力。

FFN 由两个线性变换和一个激活函数组成（通常使用 ReLU 或其变体 GELU）。设输入为 $x \in \mathbb{R}^{N \times d}$，则：

其中：

• $W_1 \in \mathbb{R}^{d \times d_{\text{ff}}}$，$b_1 \in \mathbb{R}^{d_{\text{ff}}}$ 是第一个线性层的参数；
• $W_2 \in \mathbb{R}^{d_{\text{ff}} \times d}$，$b_2 \in \mathbb{R}^{d}$ 是第二个线性层的参数；
• $d_{\text{ff}}$ 是前馈层的隐藏维度，通常设为 $4d$（即 2048 当 $d=512$ 时）；
• GELU 是高斯误差线性单元激活函数：$\text{GELU}(x) = x \cdot \Phi(x)$，其中 $\Phi$ 是标准正态分布的累积分布函数。相比 ReLU，GELU 更平滑，常用于现代 Transformer 模型。

值得注意的是，FFN 对序列中的每个位置独立作用（参数共享），即输出矩阵的第 $i$ 行仅依赖于输入矩阵的第 $i$ 行。这使得 FFN 能够在不破坏位置独立性的前提下，为每个 token 引入非线性变换。

PyTorch 实现

class FeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        if d_ff is None:
            d_ff = 4 * d_model
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        x = self.linear1(x)
        x = F.gelu(x)               # GELU 激活
        x = self.dropout(x)
        x = self.linear2(x)
        return x

残差连接与层归一化 (Add & Norm)#

随着神经网络层数的加深，梯度消失和表示退化问题会变得严重。Transformer 采用了 残差连接 (Residual Connection) 与 层归一化 (Layer Normalization) 相结合的策略，使得梯度可以顺畅地流过深层网络，同时稳定训练过程。

每个子层（多头注意力或前馈网络）的输出都会与该子层的输入相加，然后再进行层归一化。形式化地，对于子层 $\text{Sublayer}(\cdot)$，其输出为：

其中 $x$ 是子层的输入。在原始 Transformer 论文中，层归一化位于残差连接之后（即 Post-LN），但后续实践（如 GPT、BERT）常采用 Pre-LN 变体，即先归一化再进入子层，然后残差连接：$x + \text{Sublayer}(\text{LayerNorm}(x))$。Pre-LN 更有利于训练稳定性，尤其当层数很深时。

层归一化 与批归一化 (Batch Normalization) 不同，它对每个样本的每个特征维度独立计算均值和方差，并应用可学习的缩放和平移参数。对于输入 $x \in \mathbb{R}^{N \times d}$，层归一化沿最后一维（特征维度）进行：

其中：

• $\mu = \frac{1}{d} \sum_{j=1}^{d} x_{:,j}$（每个 token 的均值，形状为 $N$）；
• $\sigma^2 = \frac{1}{d} \sum_{j=1}^{d} (x_{:,j} - \mu)^2$（每个 token 的方差）；
• $\gamma, \beta \in \mathbb{R}^d$ 是可学习的缩放和偏移参数；
• $\epsilon$ 是一个小常数（如 $10^{-5}$）防止除零。

PyTorch 实现

class AddNorm(nn.Module):
    def __init__(self, d_model: int, dropout: float = 0.1, pre_ln: bool = True):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        self.pre_ln = pre_ln   # True: Pre-LN, False: Post-LN

    def forward(self, x, sublayer):
        # x: 子层输入, sublayer: 子层（尚未加残差）
        if self.pre_ln:
            # Pre-LN: 先归一化，再经过子层，然后残差连接
            return x + self.dropout(sublayer(self.norm(x)))
        else:
            # Post-LN: 先子层再加残差，最后归一化
            return self.norm(x + self.dropout(sublayer(x)))

实际使用时，Add & Norm 不单独创建一个模块，通常直接封装在 编码器/解码器 层内部：

class EncoderLayer(nn.Module):
    """单个编码器层（Pre-LN 风格）"""
    def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-LN: 先归一化，再子层，然后残差
        residual = x
        x = self.norm1(x)
        x = self.self_attn(x, x, x, mask)
        x = self.dropout1(x)
        x = residual + x

        residual = x
        x = self.norm2(x)
        x = self.ffn(x)
        x = self.dropout2(x)
        x = residual + x
        return x

class DecoderLayer(nn.Module):
    """单个解码器层（Pre-LN 风格）"""
    def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self, x, encoder_output, src_mask=None, tgt_mask=None):
        # 自注意力（带因果掩码）
        residual = x
        x = self.norm1(x)
        x = self.self_attn(x, x, x, tgt_mask)
        x = self.dropout1(x)
        x = residual + x

        # 交叉注意力（q 来自解码器，k/v 来自编码器）
        residual = x
        x = self.norm2(x)
        x = self.cross_attn(x, encoder_output, encoder_output, src_mask)
        x = self.dropout2(x)
        x = residual + x

        # 前馈网络
        residual = x
        x = self.norm3(x)
        x = self.ffn(x)
        x = self.dropout3(x)
        x = residual + x
        return x

动手实践：构建微型翻译机#

1. 首先完成对 Transformer 的搭建：

class Transformer(nn.Module):
    """完整的 Transformer 模型（用于序列到序列任务）"""
    def __init__(self,
                 src_vocab_size: int,
                 tgt_vocab_size: int,
                 d_model: int = 512,
                 n_heads: int = 8,
                 num_encoder_layers: int = 6,
                 num_decoder_layers: int = 6,
                 d_ff: int = 2048,
                 max_seq_len: int = 5000,
                 dropout: float = 0.1,
                 pad_token_id: int = 0):
        super().__init__()
        self.d_model = d_model
        self.pad_token_id = pad_token_id

        # 嵌入层（共享或独立，考虑到做文本翻译任务，两种语言有很大的区别，这里分开）
        self.src_embedding = Embeddings(src_vocab_size, d_model)
        self.tgt_embedding = Embeddings(tgt_vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_seq_len)    # 这里的代码设计是 PE 模块自己内部把加法做好了，其他地方可能是需要自己手动在此处完成与 PE 的相加

        self.encoder = Encoder(num_encoder_layers, d_model, n_heads, d_ff, dropout)     # Encoder 是由若干个 EncoderLayer 堆叠而成的
        self.decoder = Decoder(num_decoder_layers, d_model, n_heads, d_ff, dropout)     # Decoder 是由若干个 DecoderLayer 堆叠而成的

        self.output_proj = nn.Linear(d_model, tgt_vocab_size)
        self.dropout = nn.Dropout(dropout)

        self._init_parameters()

    def _init_parameters(self):
        """参数初始化（Xavier 均匀分布）"""
        for p in self.parameters():
            if p.dim() > 1:
                # xavier 均匀分布是从 [-sqrt(6 / (fan_in + fan_out)), sqrt(6 / (fan_in + fan_out))] 之间均匀采样，其中 fan_in = 输入维度，fan_out = 输出维度
                # 对于嵌入层，它各个维度的取值范围是 [-sqrt(3/d_model), sqrt(3/d_model)]，明显小于位置编码的取值范围 [-1, 1]，在向量范数上看来，位置编码占了主导作用
                # 因此在早期的嵌入层设计中，会给嵌入向量乘上 sqrt(d_model) 来放大嵌入向量的范数，使其与位置编码的范数保持一致
                nn.init.xavier_uniform_(p)

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        """
        src: (batch, src_seq_len) 源语言 token ID
        tgt: (batch, tgt_seq_len) 目标语言 token ID（训练时使用，需要偏移）
        src_mask: 可选，源序列的掩码（通常为填充掩码）
        tgt_mask: 可选，目标序列的掩码（通常为因果掩码 + 填充掩码）
        """
        # 嵌入并添加位置编码
        src_emb = self.src_embedding(src)
        src_emb = self.pos_encoding(src_emb)
        src_emb = self.dropout(src_emb)

        tgt_emb = self.tgt_embedding(tgt)
        tgt_emb = self.pos_encoding(tgt_emb)
        tgt_emb = self.dropout(tgt_emb)

        # 编码器
        encoder_output = self.encoder(src_emb, src_mask)

        # 解码器
        decoder_output = self.decoder(tgt_emb, encoder_output, src_mask, tgt_mask)

        # 输出投影到目标词表
        logits = self.output_proj(decoder_output)  # (batch, tgt_seq_len, tgt_vocab_size)
        return logits

    def encode(self, src, src_mask=None):
        """仅编码（用于推理时复用编码器输出）"""
        src_emb = self.src_embedding(src)
        src_emb = self.pos_encoding(src_emb)
        src_emb = self.dropout(src_emb)
        return self.encoder(src_emb, src_mask)

    def decode(self, tgt, encoder_output, src_mask=None, tgt_mask=None):
        """单步解码（自回归生成时调用）"""
        tgt_emb = self.tgt_embedding(tgt)
        tgt_emb = self.pos_encoding(tgt_emb)
        tgt_emb = self.dropout(tgt_emb)
        return self.decoder(tgt_emb, encoder_output, src_mask, tgt_mask)

    def generate(self, src, start_token_id, end_token_id, max_len=100, device='cpu'):
        """
        自回归生成（简单示例，使用贪婪搜索）
        src: (batch, src_seq_len)
        返回: (batch, generated_seq_len)
        """
        batch_size = src.size(0)
        src_mask = create_padding_mask(src, self.pad_token_id).to(device)
        encoder_output = self.encode(src, src_mask)

        # 初始化解码器输入：只包含起始标记
        tgt = torch.full((batch_size, 1), start_token_id, dtype=torch.long, device=device)

        for _ in range(max_len):
            # 生成目标掩码（因果掩码 + 填充掩码）
            tgt_seq_len = tgt.size(1)
            causal_mask = create_causal_mask(tgt_seq_len, device)  # (1, tgt_len, tgt_len)
            # 填充掩码：tgt 中没有填充（假设起始标记和生成 token 都不含 pad）
            # 但为了统一接口，可以生成一个全 True 的掩码
            tgt_padding_mask = torch.ones(batch_size, tgt_seq_len, dtype=torch.bool, device=device)
            tgt_padding_mask = tgt_padding_mask.unsqueeze(1).unsqueeze(2)  # (batch,1,1,seq_len)
            combined_mask = combine_masks(tgt_padding_mask, causal_mask)

            logits = self.decode(tgt, encoder_output, src_mask, combined_mask)
            next_token_logits = logits[:, -1, :]  # 取最后一个位置的输出
            next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

            # 拼接
            tgt = torch.cat([tgt, next_token], dim=1)

            # 如果所有序列都生成了结束标记，则停止
            if (next_token == end_token_id).all():
                break

        return tgt[:, 1:]  # 去掉起始标记

模型代码详见 codes/models/transformer.py

2. 数据集准备

这里选用 translation2019zh 语料作为训练和测试的数据集。这个数据集包含了大量的中英文平行句对，非常适合用来训练机器翻译模型。

Transformer 模型的输入需要有源语言序列 src 和不完整的目标语言序列 tgt 两个部分，输出是目标语言序列的下一个 token 的预测分布。但是数据集加载器在构建时一般是直接给出完整的源语言序列 src 和完整的目标语言序列 tgt 的，因此在训练前需要对目标语言序列 tgt 进行偏移处理（shifted right），分别构建 tgt_in 和 tgt_out

tgt_in = tgt[:, :-1]   # 目标序列的输入，去掉最后一个 token
tgt_out = tgt[:, 1:]   # 目标序列的输出，去掉第一个 token

训练使用的损失函数使用预测当前 token 的交叉熵损失函数即可，优化器可以选择 AdamW。

optimizer = optim.AdamW(model.parameters(), lr=args.lr)
criterion = nn.CrossEntropyLoss(ignore_index=pad_id)    # ignore_index 用于忽略 <PAD> token 的损失计算

具体的训练循环如下：

for src, tgt in dataloader:
    src, tgt = src.to(device), tgt.to(device)
    tgt_in = tgt[:, :-1]    # (batch, tgt_seq_len-1)
    tgt_out = tgt[:, 1:]    # (batch, tgt_seq_len-1)

    src_mask = create_padding_mask(src, pad_id).to(device)
    tgt_mask = create_causal_mask(tgt_in.size(1), device)  # 因果掩码
    tgt_padding_mask = create_padding_mask(tgt_in, pad_id).to(device)
    combined_mask = combine_masks(tgt_padding_mask, tgt_mask)

    logits = model(src, tgt_in, src_mask, combined_mask)    # (batch, tgt_seq_len-1, tgt_vocab_size)
    loss = criterion(logits.view(-1, logits.size(-1)), tgt_out.view(-1))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

训练部分代码详见 codes/train_transformer.py，可以使用下列命令来运行训练：

python train_transformer.py \
    --data_root "data/translation" \
    --train_file "translation2019zh_train.json" \
    --valid_file "translation2019zh_valid.json" \
    --tokenizer_name "bert-base-multilingual-cased" \
    --batch_size 64 \
    --epochs 20 \
    --lr 0.0001 \
    --max_len 128 \
    --d_model 512 \
    --n_heads 8 \
    --num_layers 6 \
    --d_ff 2048 \
    --dropout 0.1

3. 测试模型推理阶段的能力

在测试阶段，我们可以使用训练好的模型进行推理，生成新的翻译结果。这通常涉及到自回归生成过程，即每次生成一个 token，然后将其添加到输入中继续生成下一个 token，直到生成结束标记 <EOS> 或达到最大长度。

test_sentence = "Hello, how are you?"
src_tokens = [bos_id] + tokenizer.encode(test_sentence, add_special_tokens=False) + [eos_id]
src_tensor = torch.tensor(src_tokens, dtype=torch.long).unsqueeze(0).to(device)

with torch.no_grad():
    generated = model.generate(src_tensor, bos_id, eos_id, max_len=50, device=device)
    generated_list = generated.squeeze(0).tolist()
    # 清除特殊字符并解码
    generated_list = [idx for idx in generated_list if idx not in [pad_id, bos_id, eos_id]]
    translated = tokenizer.decode(generated_list)
    print(f"模型输出预测: {translated}")

为了衡量模型的翻译质量，我们可以使用 BLEU 分数等指标来评估生成的翻译与参考翻译之间的相似度。它通过对比”机器生成的译文”与”人工参考译文”的词序重合程度来打分。具体来说：先统计机器译文中有多少词和短语出现在了参考答案中（比如单个词、连续两个词、三个词、四个词的搭配分别检查），然后计算这些匹配的精确率并取平均；为了防止机器”只答一半”也能拿高分，还加入了一个长度惩罚项——如果生成的句子比参考答案短太多，就会扣分。最终得到一个0到1之间的数值，越接近1表示翻译质量越好。举个例子：如果机器译文是”我吃饭了”而参考答案是”我吃过了”，这两个句子虽然意思相近但用词略有不同，BLEU可能会给出一个中等分数而非满分。因此BLEU不是完美指标，但它能快速给大量翻译结果量化排名，是自然语言处理领域最常用的评测标准之一。

其中 $p_n$ 是 n-gram 的精确率，$w_n$ 是 n-gram 的权重，$\text{BP}$ 是 brevity penalty，用于惩罚过短的生成结果。

下面是手动实现的 BLEU 分数计算函数，供参考：

from collections import Counter
import math

def bleu_score(candidate, references, max_n=4):
    """计算BLEU分数"""
    # 生成n-gram函数
    def ngrams(seq, n):
        return list(zip(*[seq[i:] for i in range(n)])) if len(seq) >= n else []
    
    # 1. 统计各级修正精确率
    precisions = []
    for n in range(1, max_n + 1):
        cand_ng = Counter(ngrams(candidate, n))
        
        # 参考取最大计数
        ref_ng = Counter()
        for ref in references:
            ref_ng.update(ngrams(ref, n))
        
        # 裁剪并求和
        match = sum(min(cand_ng[k], ref_ng.get(k, 0)) for k in cand_ng)
        precision = match / len(candidate) if candidate else 0
        precisions.append(precision)
    
    # 2. 几何平均（加平滑防止log(0)）
    log_sum = sum(math.log(p + 1e-10) for p in precisions)
    geo_mean = math.exp(log_sum / max_n)
    
    # 3. 长度惩罚BP
    ref_len = min(references, key=lambda x: abs(len(x)-len(candidate)))
    if len(candidate) >= len(ref_len):
        bp = 1.0
    elif len(candidate) == 0:
        bp = 0.0
    else:
        bp = math.exp(1 - len(ref_len) / len(candidate))
    
    return bp * geo_mean

# ===== 使用示例 =====
if __name__ == "__main__":
    candidate = "the cat sat on the mat".split()
    references = ["the cat sat on the mat".split(), "a cat sat on the rug".split()]
    
    print(f"BLEU 分数: {bleu_score(candidate, references):.4f}")

也可以用 nltk 库中的 sentence_bleu 函数来计算 BLEU 分数，使用起来更方便：

from nltk.translate.bleu_score import sentence_bleu

score = sentence_bleu(
    ["the cat sat on the mat".split(), "a cat sat on the rug".split()],
    "the cat sat".split()
)
print(f"BLEU: {score:.4f}")

4. 常见疑问

为什么训练时的输入是偏移的 tgt_in (正确答案，其中不含有模型预测的 token)，而不是直接把模型预测的 token 添加到 tgt_in 中，作为下一次训练时的输入？

这是因为：

1. 如果像模型推理时使用自己生成的 token 作为下一步的输入，如果自己预测的 token 是错误的，那么后续的输入就会被污染，导致错误不断累加，训练过程不稳定，模型难以收敛；使用真实的目标序列作为输入可以保证训练过程中的输入是正确的，模型能够更快地学习到正确的映射关系。
2. 观察输入的 tgt_in: (batch, tgt_seq_len-1) 和输出的 logits: (batch, tgt_seq_len-1, tgt_vocab_size) 的形状，发现实际上在训练中，模型是并行地对整个序列的各种位置进行预测的，而不是自回归式的逐步生成。自回归需要迭代完成，而并行训练则可以一次性计算整个序列的损失，效率更高。

BERT 的预训练#

在实际训练中，为了让 BERT 有强大的文本理解能力，需要让 BERT 模型预先完成预训练任务。BERT 的预训练包含两大关键任务：

1. 双向掩码语言建模 (Masked Language Modeling, MLM)：

   本质上就是前面所讲述的选词填空。通过随机掩蔽输入文本中的部分单词（用 [MASK] 替换）的方式，让模型需基于上下文去预测被掩蔽的词语。这一任务迫使模型同时学习词语左侧和右侧的上下文信息，从而获得更全面的语义表示。例如

   输入："苹果公司今天发布了最新的 <MASK> 手机。"

   模型可以同时看到 <MASK> 左侧和右侧的上下文，它需要利用这些上下文信息来预测 <MASK> 位置应该填入的词语（例如 “iPhone”）。

   经过一系列编码器层的处理后，模型会输出一个与输入序列长度相同的隐藏状态序列 sequence_output: (batch_size, seq_len, d_model)。我们希望 sequence_output 中的每个位置都和输入序列中对应位置的 token 相对应，像图像分割模型那样，对每个 token 位置的隐藏层状态都做一系列的线性层处理后，可以得到形状为 (batch_size, seq_len, vocab_size) 的预测分布，其中 vocab_size 是词表大小。对于 <MASK> 位置的 token，我们可以直接从对应位置的预测分布中取出概率最高的词语作为模型的预测结果。

   训练的损失函数直接采用交叉熵损失函数即可，计算模型在 <MASK> 位置的预测分布与真实标签之间的损失，并通过反向传播来优化模型参数。

2. 下一句预测 (Next Sentence Prediction, NSP)：

   本质上属于一个二分类任务。输入连续的两个句子对，模型需判断第二个句子是否在原始文本中紧邻第一个句子（yes 或 no）。该任务帮助模型学习句子间的逻辑连贯性与语义关系，为后续处理涉及多句理解的复杂任务（如问答、文本分类中的长文本分析）奠定基础。

   在文本分类任务中，BERT 需要在输入序列的开头添加一个特殊的分类标记 <CLS>。

   <CLS> 这部电影非常精彩，演员的表演也很到位。

   具体到下一句预测的任务中，还需要在两个句子之间添加一个分隔标记 <SEP>，例如

   <CLS> 这部电影非常精彩，演员的表演也很到位。 <SEP> 你觉得这部电影怎么样？ <SEP>

   这个时候，得到了 sequence_output: (batch_size, seq_len, d_model) 后，模型通过取出 <CLS> 位置的隐藏层状态 sequence_output[:, 0, :] 并进行线性变换和 Softmax 操作，就可以得到全局类别信息的预测分布了。在 NSP 任务中，这个预测分布是一个二分类分布，表示输入的第二个句子是否紧邻第一个句子。训练的损失函数亦可以选用交叉熵损失函数。尽管后续一些工作发现去除 NSP 或改用其他句子级任务效果更好，但理解 NSP 有助于掌握 BERT 的设计思想。

多任务训练

通过上述任务的联合训练，BERT 可以获得强大的双向语义编码能力。更多有关预训练的内容，将会在后续章节中展开。

BERT 的微调#

BERT 模型可以不只用于做选词填空和下一句预测，也可以做诸如情感分析、实体识别等其他的文本理解任务。然而，直接对一个没有经过预训练的 BERT 模型做这些更为高级的文本理解任务的训练，通常很难能得到一个表现良好的模型。这是因为这些任务通常需要模型具备对文本的深层次理解能力，而没有预训练的 BERT 模型在初始状态下并不具备这样的能力。

因此，在实际应用中，往往需要预训练的、具有一定文本理解能力的 BERT 作为初始化的模型，然后再在这些更为高级的文本理解任务上训练，这个过程被称作微调 (Fine-tuning)。通过微调，模型可以在预训练阶段学到的通用语言表示的基础上，进一步适应特定任务的需求，从而获得更好的性能。

微调也将会在后续章节中展开介绍。

BERT 的特殊输入设计#

与简单的 Transformer 编码器不同的是，BERT 的输入由三部分嵌入组成：

1. 语言嵌入 (Word Embeddings)：将输入文本中的每个 token 转换为一个固定维度的向量表示，捕捉词语的语义信息。
2. 位置嵌入 (Position Embeddings)：由于 Transformer 模型本身没有序列信息，位置嵌入用于为每个 token 提供位置信息，使模型能够区分不同位置的 token。
3. 类型嵌入 (Token Type Embeddings)：在 BERT 的预训练任务中，输入通常是由两段文本拼接而成的，例如 [CLS] Sentence A [SEP] Sentence B [SEP]，类型嵌入用于区分这两段文本中的 token，帮助模型理解它们之间的关系。

动手实践：BERT 实现#

构建一个 BERT 模型，这里依然选择 Pre-LN 的设计风格 (实际最早的 BERT 是 Post-LN 架构的)，编码器层的结构与之前介绍的 EncoderLayer 类似：

class BERTEncoderLayer(nn.Module):
    """单个编码器层（Pre-LN 风格）"""
    def __init__(self, d_model: int, n_heads: int, d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-LN: 先归一化，再子层，然后残差
        residual = x
        x = self.norm1(x)
        x = self.self_attn(x, x, x, mask)
        x = self.dropout1(x)
        x = residual + x

        residual = x
        x = self.norm2(x)
        x = self.ffn(x)
        x = self.dropout2(x)
        x = residual + x
        return x

class BERTEncoder(nn.Module):
    """编码器：由 N 个 EncoderLayer 堆叠"""
    def __init__(self, num_layers: int, d_model: int, n_heads: int,
                 d_ff: int = None, dropout: float = 0.1):
        super().__init__()
        self.layers = nn.ModuleList([
            BERTEncoderLayer(d_model, n_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x, mask=None):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class BERT(nn.Module):
    def __init__(self, vocab_size, num_classes=None, d_model=256, n_heads=8, num_layers=4, type_vocab_size=2, max_len=512):
        super().__init__()
        self.d_model = d_model
        self.word_embeddings = nn.Embedding(vocab_size, d_model)
        self.position_embeddings = nn.Embedding(max_len, d_model)
        self.token_type_embeddings = nn.Embedding(type_vocab_size, d_model)
        self.embedding_norm = nn.LayerNorm(d_model)
        self.embedding_dropout = nn.Dropout(0.1)
        
        self.encoder = BERTEncoder(num_layers, d_model, n_heads)
        
        # 1. 增强型 MLM Head
        self.mlm_head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.LayerNorm(d_model),
            nn.Dropout(0.1),
            nn.Linear(d_model, vocab_size)
        )
        
        # 2. CLS Pooler (用于分类/NSP任务)
        self.pooler = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.Tanh()
        )
        self.classifier = nn.Linear(d_model, num_classes) if num_classes else None

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, task='mlm'):
        # input_ids: (batch, seq_len)
        # token_type_ids: (batch, seq_len)
        seq_len = input_ids.size(1)
        
        # 词嵌入 + 类型嵌入 + 位置嵌入
        x = self.word_embeddings(input_ids)
        
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)
        x = x + self.token_type_embeddings(token_type_ids)
        
        position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        x = x + self.position_embeddings(position_ids)
        
        x = self.embedding_norm(x)
        x = self.embedding_dropout(x)
        
        # attention_mask: (batch, seq_len) -> (batch,1,1,seq_len)
        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
        
        sequence_output = self.encoder(x, attention_mask)  # (batch, seq_len, d_model)
        
        if task == 'mlm':
            return self.mlm_head(sequence_output)          # (batch, seq_len, vocab_size)
        else:  # classification / nsp
            cls_hidden = sequence_output[:, 0, :]           # 取 <CLS> 向量
            pooled_output = self.pooler(cls_hidden)         # 通过 Pooler
            return self.classifier(pooled_output)           # (batch, num_classes)