深入解析Transformer模型实现：基于nlp-tutorial项目的代码剖析

前言

Transformer模型自2017年由Vaswani等人提出以来，已成为自然语言处理领域的基石架构。本文将通过一个简洁而完整的实现案例，深入解析Transformer的核心组件和工作原理。这个实现来自一个专注于NLP教程的项目，它以教学为目的，清晰地展现了Transformer的关键技术细节。

Transformer模型概述

Transformer是一种完全基于注意力机制的序列转换模型，摒弃了传统的循环和卷积结构。其主要特点包括：

1. 自注意力机制(Self-Attention)
2. 多头注意力(Multi-Head Attention)
3. 位置编码(Positional Encoding)
4. 残差连接和层归一化
5. 前馈神经网络

核心组件实现解析

1. 位置编码

位置编码是Transformer的关键创新之一，它通过正弦和余弦函数为序列中的每个位置生成独特的编码：

def get_sinusoid_encoding_table(n_position, d_model):
    def cal_angle(position, hid_idx):
        return position / np.power(10000, 2 * (hid_idx // 2) / d_model)
    def get_posi_angle_vec(position):
        return [cal_angle(position, hid_j) for hid_j in range(d_model)]

    sinusoid_table = np.array([get_posi_angle_vec(pos_i) for pos_i in range(n_position)])
    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # 偶数位置使用sin
    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # 奇数位置使用cos
    return torch.FloatTensor(sinusoid_table)

这种编码方式使模型能够捕获序列中单词的相对位置信息，同时可以处理比训练时更长的序列。

2. 注意力掩码

Transformer中使用了两种重要的注意力掩码：

def get_attn_pad_mask(seq_q, seq_k):
    # 创建填充(PAD)标记的掩码
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)
    return pad_attn_mask.expand(batch_size, len_q, len_k)

def get_attn_subsequent_mask(seq):
    # 创建上三角矩阵，防止解码器看到未来信息
    subsequent_mask = np.triu(np.ones(attn_shape), k=1)
    return torch.from_numpy(subsequent_mask).byte()

• 填充掩码：防止注意力机制处理填充标记
• 后续掩码：确保解码器只能关注当前位置及之前的输出

3. 缩放点积注意力

这是Transformer中最核心的运算：

class ScaledDotProductAttention(nn.Module):
    def forward(self, Q, K, V, attn_mask):
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)
        scores.masked_fill_(attn_mask, -1e9)  # 应用掩码
        attn = nn.Softmax(dim=-1)(scores)
        context = torch.matmul(attn, V)
        return context, attn

关键点：

• 除以√d_k防止梯度消失
• 使用softmax计算注意力权重
• 最终输出是值的加权和

4. 多头注意力机制

多头注意力并行运行多个注意力头，增强模型捕获不同子空间信息的能力：

class MultiHeadAttention(nn.Module):
    def __init__(self):
        self.W_Q = nn.Linear(d_model, d_k * n_heads)  # 查询变换
        self.W_K = nn.Linear(d_model, d_k * n_heads)  # 键变换
        self.W_V = nn.Linear(d_model, d_v * n_heads)  # 值变换
        self.linear = nn.Linear(n_heads * d_v, d_model)  # 输出投影
        self.layer_norm = nn.LayerNorm(d_model)  # 层归一化

    def forward(self, Q, K, V, attn_mask):
        residual = Q  # 残差连接
        # 分割为多头
        q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)
        # ...类似处理K和V...
        
        # 计算注意力
        context, attn = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)
        # 合并多头输出
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v)
        output = self.linear(context)
        return self.layer_norm(output + residual), attn

5. 前馈网络

每个编码器和解码器层都包含一个位置式前馈网络：

class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        self.conv1 = nn.Conv1d(d_model, d_ff, kernel_size=1)  # 扩展维度
        self.conv2 = nn.Conv1d(d_ff, d_model, kernel_size=1)  # 压缩维度
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, inputs):
        residual = inputs
        output = nn.ReLU()(self.conv1(inputs.transpose(1, 2)))
        output = self.conv2(output).transpose(1, 2)
        return self.layer_norm(output + residual)

完整Transformer架构

编码器实现

class Encoder(nn.Module):
    def __init__(self):
        self.src_emb = nn.Embedding(src_vocab_size, d_model)  # 词嵌入
        self.pos_emb = nn.Embedding.from_pretrained(
            get_sinusoid_encoding_table(src_len+1, d_model), freeze=True)  # 位置编码
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self, enc_inputs):
        enc_outputs = self.src_emb(enc_inputs) + self.pos_emb(...)  # 词嵌入+位置编码
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs)
        for layer in self.layers:
            enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask)
        return enc_outputs, enc_self_attns

解码器实现

class Decoder(nn.Module):
    def __init__(self):
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_emb = nn.Embedding.from_pretrained(
            get_sinusoid_encoding_table(tgt_len+1, d_model), freeze=True)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        dec_outputs = self.tgt_emb(dec_inputs) + self.pos_emb(...)
        # 创建两种掩码
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs)
        dec_self_attn_subsequent_mask = get_attn_subsequent_mask(dec_inputs)
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequent_mask), 0)
        
        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)
        
        for layer in self.layers:
            dec_outputs, dec_self_attn, dec_enc_attn = layer(
                dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask)
        return dec_outputs, dec_self_attns, dec_enc_attns

训练过程

示例中展示了一个简单的训练循环：

model = Transformer()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(20):
    optimizer.zero_grad()
    outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)
    loss = criterion(outputs, target_batch.contiguous().view(-1))
    loss.backward()
    optimizer.step()

可视化注意力

实现中包含了注意力权重的可视化功能：

def showgraph(attn):
    attn = attn[-1].squeeze(0)[0]
    attn = attn.squeeze(0).data.numpy()
    fig = plt.figure(figsize=(n_heads, n_heads))
    ax = fig.add_subplot(1, 1, 1)
    ax.matshow(attn, cmap='viridis')
    ax.set_xticklabels(['']+sentences[0].split(), rotation=90)
    ax.set_yticklabels(['']+sentences[2].split())
    plt.show()

总结

通过这个简洁的实现，我们可以清晰地看到Transformer模型的几个关键特点：

1. 完全基于注意力机制，没有循环或卷积操作
2. 使用多头注意力并行处理信息
3. 位置编码代替传统的位置信息
4. 残差连接和层归一化促进深层网络训练
5. 掩码机制实现序列处理的自回归特性

这个实现虽然精简，但包含了Transformer的所有核心组件，是理解这一重要模型的绝佳起点。通过研究代码，我们可以更深入地理解Transformer的工作原理，为进一步的模型改进和应用开发奠定基础。