Transformer架构核心原理与PyTorch实现全解析

Transformer是当今自然语言处理乃至整个AI领域的基石架构，自2017年由Google提出以来，彻底改变了序列建模的范式。它摒弃了传统的循环神经网络（RNN）和卷积神经网络（CNN），完全基于自注意力机制，实现了高效的并行计算和对长距离依赖的出色捕捉。本文将从架构总览、核心组件、代码实现三个层面，为你深入解析Transformer。

一、Transformer架构总览

Transformer是一个标准的编码器-解码器（Encoder-Decoder）架构，专为序列到序列（Seq2Seq）任务设计，如机器翻译。其整体架构如下图所示（以经典的N=6层为例）：

输入序列 -> [编码器 × N] -> [解码器 × N] -> 输出序列

编码器（Encoder）：由N个完全相同的层堆叠而成，负责将输入序列（如源语言句子）编码为一系列富含上下文信息的隐藏表示。
解码器（Decoder）：同样由N个完全相同的层堆叠而成，它接收编码器的输出以及自身已生成的部分输出序列，自回归地（auto-regressive）生成目标序列（如目标语言句子）。

二、核心组件深度剖析

1. 输入编码：词嵌入与位置编码

Transformer的输入由两部分相加而成：词嵌入（Word Embedding）和位置编码（Positional Encoding）。

• 词嵌入层：将输入序列中的每个词（Token）映射为一个高维向量（如512维）。这通常通过一个可学习的查找表（Look-up Table）实现。
• 位置编码层：由于自注意力机制本身不具备感知序列顺序的能力，因此必须显式地注入位置信息。Transformer使用正弦和余弦函数来生成固定的位置编码。

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    """位置编码层实现"""
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        # 创建一个足够长的位置编码矩阵 pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # 使用正弦函数对偶数位置进行编码
        pe[:, 0::2] = torch.sin(position * div_term)
        # 使用余弦函数对奇数位置进行编码
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # 形状: (1, max_len, d_model)
        self.register_buffer('pe', pe)  # 注册为缓冲区，不参与训练 def forward(self, x):
        # x的形状: (batch_size, seq_len, d_model)
        # 将位置编码加到输入x上
        x = x + self.pe[:, :x.size(1)]
        return x

# 测试位置编码
d_model = 512
seq_len = 10
batch_size = 2
pos_encoder = PositionalEncoding(d_model)
x = torch.randn(batch_size, seq_len, d_model)
output = pos_encoder(x)
print(f"输入形状: {x.shape}, 输出形状: {output.shape}")

2. 多头自注意力机制（Multi-Head Self-Attention）

这是Transformer最核心的组件，其作用是让序列中的每个位置都能“关注”到序列中所有其他位置的信息，从而动态地计算每个位置的表示。

单头注意力计算过程：

1. 线性投影：对输入序列X（形状：[batch_size, seq_len, d_model]）分别进行三次线性变换，得到查询（Q）、键（K）、值（V）矩阵。
2. 缩放点积注意力：计算 Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V。其中 sqrt(d_k) 是缩放因子，用于防止点积结果过大导致softmax梯度消失。
3. 多头机制：将 d_model 维度的Q、K、V投影到 h（头数）个不同的子空间（每个子空间维度为 d_k = d_v = d_model / h），在每个头上独立计算注意力，最后将结果拼接并投影回 d_model 维度。

import torch.nn.functional as F

def scaled_dot_product_attention(q, k, v, mask=None):
    """缩放点积注意力计算函数"""
    d_k = q.size(-1)
    # 计算 Q * K^T
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
    # 如果提供了掩码，将需要屏蔽的位置设置为一个极小的负数 if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    # 计算softmax得到注意力权重 attention_weights = F.softmax(scores, dim=-1)
    # 注意力权重乘以V，得到最终的输出 output = torch.matmul(attention_weights, v)
    return output, attention_weights

class MultiHeadAttention(nn.Module):
    """多头注意力层实现"""
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads # 定义线性投影层，用于生成Q, K, V以及最后的输出投影 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)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        # 1. 线性投影并分头
        q = self.w_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 2. 在每个头上计算缩放点积注意力 attention_output, attention_weights = scaled_dot_product_attention(q, k, v, mask)

        # 3. 拼接多头输出并投影 attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        output = self.w_o(attention_output)
        return output, attention_weights

3. 编码器层（Encoder Layer）

每个编码器层包含两个子层：

1. 多头自注意力子层：让输入序列的每个位置关注整个序列。
2. 前馈神经网络子层：一个简单的两层全连接网络（通常中间维度扩大4倍），对每个位置的特征进行独立变换。
   每个子层都遵循“子层输出 = LayerNorm(x + Sublayer(x))”的结构，即残差连接（Residual Connection） 后接层归一化（Layer Normalization）。这有助于缓解梯度消失和加速训练。

class EncoderLayer(nn.Module):
    """单个编码器层实现"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 子层1: 多头自注意力 + 残差 & 层归一化 attn_output, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))

        # 子层2: 前馈网络 + 残差 & 层归一化
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))
        return x

4. 解码器层（Decoder Layer）

每个解码器层包含三个子层：

1. 掩码多头自注意力子层：确保在生成第t个词时，只能“看到”前t-1个词，这是实现自回归生成的关键。通过一个上三角掩码矩阵实现。
2. 编码器-解码器注意力子层：其查询（Q）来自解码器的上一子层输出，而键（K）和值（V）来自编码器的最终输出。这使得解码器可以聚焦于输入序列的相关部分。
3. 前馈神经网络子层：与编码器中的相同。
   每个子层同样采用残差连接和层归一化。

class DecoderLayer(nn.Module):
    """单个解码器层实现"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)  # 编码器-解码器注意力 self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        # 子层1: 掩码自注意力 + 残差 & 层归一化
        attn_output, _ = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))

        # 子层2: 编码器-解码器注意力 + 残差 & 层归一化
        attn_output, _ = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + self.dropout(attn_output))

        # 子层3: 前馈网络 + 残差 & 层归一化
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout(ff_output))
        return x

5. 输出层

解码器的最终输出经过一个线性层（将 d_model 维投影到目标词汇表大小 vocab_size 维）和一个softmax函数，得到每个位置上词汇表中每个词的概率分布，用于预测下一个词。

三、完整Transformer模型搭建

将上述组件组合起来，并添加编码器堆叠、解码器堆叠以及最终的线性投影和softmax，即可构建完整的Transformer模型。

class Transformer(nn.Module):
    """完整的Transformer模型"""
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len, dropout=0.1):
        super(Transformer, self).__init__()
        self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)
        self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_seq_len)

        self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)])

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

    def generate_mask(self, src, tgt):
        # 生成源序列填充掩码 (src_mask) 和目标序列的因果掩码 (tgt_mask)
        src_mask = (src != 0).unsqueeze(1).unsqueeze(2)  # (batch_size, 1, 1, src_len)
        tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(2)  # 填充掩码
        seq_len = tgt.size(1)
        nopeak_mask = (1 - torch.triu(torch.ones(1, seq_len, seq_len), diagonal=1)).bool()  # 因果掩码
        tgt_mask = tgt_mask & nopeak_mask  # 结合填充掩码和因果掩码 return src_mask, tgt_mask

    def forward(self, src, tgt):
        src_mask, tgt_mask = self.generate_mask(src, tgt)

        # 编码器前向传播 src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
        enc_output = src_embedded
        for enc_layer in self.encoder_layers:
            enc_output = enc_layer(enc_output, src_mask)

        # 解码器前向传播 tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
        dec_output = tgt_embedded for dec_layer in self.decoder_layers:
            dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)

        # 最终输出投影 output = self.final_linear(dec_output)
        return output

# 模型参数配置示例
src_vocab_size = 10000
tgt_vocab_size = 10000
d_model = 512
num_heads = 8
num_layers = 6
d_ff = 2048
max_seq_len = 100
dropout = 0.1

model = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len, dropout)
print(f"模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

四、Transformer的优势与影响

1. 并行计算：自注意力机制允许序列中所有位置同时计算，极大提升了训练效率，尤其是在GPU上。
2. 长距离依赖：无论序列中两个位置相距多远，自注意力都能直接建立连接，有效解决了RNN中的梯度消失/爆炸问题。
3. 可解释性：注意力权重矩阵可以可视化，展示模型在做出决策时“关注”了输入序列的哪些部分。

正是这些优势，使得Transformer成为BERT、GPT、T5等几乎所有现代大语言模型（LLM）的基础架构，并成功扩展到计算机视觉（ViT）、语音处理等多个领域，奠定了当前AI发展的技术基石。