Transformer架构每层功能与实现深度解析

2026年1月8日 互联网

Transformer架构每层功能与实现深度解析

Transformer架构自2017年提出以来，已成为自然语言处理（NLP）领域的基石模型。其核心优势在于通过自注意力机制（Self-Attention）实现并行化计算，同时通过多层堆叠捕捉文本的深层语义关系。本文将从输入层到输出层，逐层解析Transformer的内部结构，结合数学公式与代码示例，揭示其设计原理与工程实现细节。

一、输入层：嵌入与位置编码

1.1 输入嵌入（Input Embedding）

输入层的核心任务是将离散的词符号映射为连续的向量表示。对于输入序列中的每个词，通过查表操作（Lookup Table）获取其对应的词向量：

  1. # 伪代码示例：词嵌入查表
  2. word_embeddings = nn.Embedding(vocab_size, embedding_dim)
  3. input_vectors = word_embeddings(input_ids)  # shape: [seq_len, embedding_dim]

其中，vocab_size为词汇表大小，embedding_dim为词向量维度（通常为512或768）。词嵌入层本质是一个参数矩阵，每一行对应一个词的向量表示。

1.2 位置编码（Positional Encoding）

由于自注意力机制本身不包含位置信息，Transformer通过正弦/余弦函数生成位置编码，将其与词向量相加：

  1. import torch
  2. import math
  4. def positional_encoding(max_len, d_model):
  5.     position = torch.arange(max_len).unsqueeze(1)
  6.     div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
  7.     pe = torch.zeros(max_len, d_model)
  8.     pe[:, 0::2] = torch.sin(position * div_term)  # 偶数位置
  9.     pe[:, 1::2] = torch.cos(position * div_term)  # 奇数位置
  10.     return pe

位置编码的公式为：
[
PE(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d{model}}}\right), \quad PE(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d{model}}}\right)
]
其中，pos为词在序列中的位置，i为维度索引。这种设计使得模型能够通过相对位置推断词序关系。

二、编码器层：自注意力与前馈网络

Transformer编码器由N个相同层堆叠而成，每层包含两个子层：多头自注意力机制（Multi-Head Self-Attention）和前馈神经网络（Feed-Forward Network），并通过残差连接（Residual Connection）与层归一化（Layer Normalization）增强训练稳定性。

2.1 多头自注意力机制

自注意力机制的核心是计算词与词之间的相关性权重。对于输入序列 ( X \in \mathbb{R}^{n \times d} )（n为序列长度，d为词向量维度），通过线性变换生成查询（Q）、键（K）、值（V）矩阵：
[
Q = XW^Q, \quad K = XW^K, \quad V = XW^V
]
其中，( W^Q, W^K, W^V \in \mathbb{R}^{d \times d_k} ) 为可学习参数。单头注意力计算如下：
[
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
]
缩放因子 ( \sqrt{d_k} ) 用于防止点积结果过大导致梯度消失。多头注意力将输入分割为h个头（通常h=8），并行计算后拼接结果：

  1. class MultiHeadAttention(nn.Module):
  2.     def __init__(self, d_model, num_heads):
  3.         super().__init__()
  4.         self.d_model = d_model
  5.         self.num_heads = num_heads
  6.         self.d_k = d_model // num_heads
  7.         self.Wq = nn.Linear(d_model, d_model)
  8.         self.Wk = nn.Linear(d_model, d_model)
  9.         self.Wv = nn.Linear(d_model, d_model)
  10.         self.Wo = nn.Linear(d_model, d_model)
  12.     def forward(self, x):
  13.         batch_size = x.size(0)
  14.         Q = self.Wq(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
  15.         K = self.Wk(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
  16.         V = self.Wv(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
  17.         scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
  18.         attn_weights = torch.softmax(scores, dim=-1)
  19.         context = torch.matmul(attn_weights, V)
  20.         context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
  21.         return self.Wo(context)

多头机制允许模型在不同子空间捕捉多样化的注意力模式，例如一个头关注语法关系，另一个头关注语义关联。

2.2 前馈神经网络

每个注意力子层后接一个全连接前馈网络，包含两个线性变换和一个ReLU激活：
[
\text{FFN}(x) = \text{ReLU}(xW1 + b_1)W_2 + b_2
]
其中，( W_1 \in \mathbb{R}^{d \times d{ff}} ), ( W2 \in \mathbb{R}^{d{ff} \times d} )，通常 ( d_{ff} = 4d )（例如d=512时，d_ff=2048）。前馈网络为模型提供非线性变换能力。

2.3 残差连接与层归一化

每子层采用残差连接解决梯度消失问题，并通过层归一化稳定训练：
[
\text{LayerNorm}(x + \text{Sublayer}(x))
]
残差连接要求子层输出维度与输入维度一致（即d_model不变），层归一化则对每个样本的特征进行归一化：

  1. class LayerNorm(nn.Module):
  2.     def __init__(self, features, eps=1e-6):
  3.         super().__init__()
  4.         self.gamma = nn.Parameter(torch.ones(features))
  5.         self.beta = nn.Parameter(torch.zeros(features))
  6.         self.eps = eps
  8.     def forward(self, x):
  9.         mean = x.mean(-1, keepdim=True)
  10.         std = x.std(-1, keepdim=True)
  11.         return self.gamma * (x - mean) / (std + self.eps) + self.beta

三、解码器层：掩码自注意力与编码器-解码器注意力

解码器同样由N个相同层堆叠，但每层包含三个子层：掩码多头自注意力、编码器-解码器多头注意力、前馈网络。

3.1 掩码多头自注意力

解码器在训练时采用掩码机制防止看到未来信息。通过修改注意力分数矩阵，将未来位置的值设为负无穷：

  1. def masked_softmax(scores, mask):
  2.     scores = scores.masked_fill(mask == 0, -1e9)
  3.     return torch.softmax(scores, dim=-1)

掩码矩阵 mask 的形状为 [seq_len, seq_len]，对角线以下为1（允许关注当前及之前位置），以上为0。

3.2 编码器-解码器注意力

该子层的Q矩阵来自解码器上一层的输出，K和V矩阵来自编码器最终层的输出。此机制允许解码器聚焦编码器输出的相关部分，实现跨模态对齐（如机器翻译中源语言与目标语言的对齐）。

四、输出层：线性变换与Softmax

解码器最终通过线性层将输出维度映射至词汇表大小，并通过Softmax生成概率分布：

  1. class Generator(nn.Module):
  2.     def __init__(self, d_model, vocab_size):
  3.         super().__init__()
  4.         self.proj = nn.Linear(d_model, vocab_size)
  6.     def forward(self, x):
  7.         return torch.log_softmax(self.proj(x), dim=-1)

训练时采用交叉熵损失函数，推理时通过贪心搜索或束搜索生成序列。

五、工程实践建议

  1. 超参数选择：典型配置为d_model=512，num_heads=8，d_ff=2048，层数N=6。增大模型规模（如d_model=1024）可提升性能，但需权衡计算成本。
  2. 学习率调度：采用带暖启的逆平方根学习率（如初始lr=0.001，warmup_steps=4000），避免训练初期梯度震荡。
  3. 批处理优化：使用梯度累积模拟大批量训练（如batch_size=256通过4个累积步实现等效效果）。
  4. 混合精度训练：启用FP16可减少内存占用并加速计算，需配合动态损失缩放防止梯度下溢。

Transformer架构通过分层设计实现了高效的并行计算与深层语义捕捉。理解每一层的数学原理与工程实现，对模型调优、压缩及迁移学习具有重要指导意义。在实际应用中，可结合百度智能云等平台的高性能计算资源，进一步探索大规模预训练与领域适配的可能性。

最新文章