剖析 Transformer 骨干:编码器层(Encoder Layer)内部结构与核心机制
剖析 Transformer 骨干:编码器层(Encoder Layer)内部结构与核心机制
在 Transformer 的整体架构中,编码器(Encoder)扮演着“全局语义理解引擎”的角色。与受因果掩码约束的解码器不同,编码器拥有无约束的双向可见视野,能够让输入序列中的每一个 Token 充分感知全文上下文。
标准 Transformer 的编码器由 个结构完全相同的编码器层(Encoder Layer,原始论文中 )堆叠而成。每个编码器层在输入与输出之间维持着统一的特征维度 ,形成一条高维向量的流水线。深入拆解单个编码器层,其内部由多头自注意力子层、前馈全连接子层以及贯穿二者的残差连接与层归一化精密耦合而成。
一、 编码器层的整体结构与数据流
单个编码器层接收上一层的输出张量 (其中 为批大小, 为序列长度, 为隐藏特征维度),经过两级特征加工后,输出同等形状的表征张量。
输入张量 X (B, L, d_model)
│
┌───────────────┴───────────────┐
│ ▼
│ [ Multi-Head Self-Attention ]
│ │
▼ ▼
┌─────────────┐ ┌─────────────┐
│ 残差分支 │ │ 子层输出 │
└──────┬──────┘ └──────┬──────┘
│ │
└───────────────► ◄─────────────┘
│ (逐元素相加 Add)
▼
[ LayerNorm ]
│
┌───────────────┴───────────────┐
│ ▼
│ [ Position-wise FFN ]
│ │
▼ ▼
┌─────────────┐ ┌─────────────┐
│ 残差分支 │ │ 子层输出 │
└──────┬──────┘ └──────┬──────┘
│ │
└───────────────► ◄─────────────┘
│ (逐元素相加 Add)
▼
[ LayerNorm ]
│
输出张量 Y (B, L, d_model)
每个编码器层包含两个核心计算单元:
- 多头自注意力子层(Multi-Head Self-Attention): 负责 Token 之间的横向跨位置信息交互(Spatial Routing)。
- 前馈全连接网络子层(Position-wise Feed-Forward Network, FFN): 负责每个 Token 内部特征维度的纵向非线性投影与知识提取(Channel Mixing)。
二、 子层一:全向多头自注意力
在编码器中,多头自注意力机制有三个决定性特征:
- 同源映射(Self-Attention): 查询矩阵 、键矩阵 、值矩阵 全部由上一层的同一个输入矩阵 经过三组不同的可学习线性映射投影生成:
- 无因果遮掩(Bidirectional Context): 编码器不使用下三角下遮掩掩码。这意味着第 个词在计算注意力权重时,可以无障碍地看到序列中第 0 到第 的所有词。这种无偏见的双向注意力使得 BERT 等纯编码器架构在分类、实体识别和阅读理解任务中具备远超自回归模型的语义抽取精度。
- 填充屏蔽(Padding Mask Only): 仅针对因对齐长度而补充的
<pad>占位符进行注意力打分屏蔽(置为 ),防止无效字符污染语义分布。
三、 子层二:基于位置的前馈全连接网络(FFN)
自注意力层虽然擅长建立词与词之间的关联,但它本质上只是对现有向量进行线性加权平均,非线性拟合能力有限。因此,紧随其后的 FFN 子层承担起非线性变换的重任。
1. 数学表述与升降维设计
FFN 独立且相同地应用于序列中的每一个 Token 位置(Position-wise):
- 第一层(升维): 将特征维度从 投影至 (通常为 ,例如 512 升维至 2048)。
- 激活函数: 原始论文使用 ReLU;在现代变体中,通常采用平滑度更高且无硬饱和区的 GELU。
- 第二层(降维): 将维度从 压缩回 ,以便与主干残差向量保持通道一致。
2. 深度认知:FFN 到底在做什么?
学术界近年(如 Geva 等人的研究)发现,FFN 不仅仅是一个简单的感知机,它在大模型中实际充当着键值记忆库(Key-Value Memory)的角色:
- 第一层矩阵 (Memory Keys): 负责检测特定的语义模式或概念触发词。
- 第二层矩阵 (Memory Values): 负责根据第一层激活的置信度,将对应的世界知识或上下文特征补充进 Token 向量中。
简而言之:自注意力层负责“把信息路由到正确的位置”,而 FFN 层负责“提取并赋予该位置更丰富的知识与概念特征”。
四、 连接纽带:残差连接与层归一化(Add & Norm)
编码器每个子层的外部都包裹着残差连接(Residual Connection)和层归一化(Layer Normalization)。这一结构对于深层 Transformer 网络的收敛至关重要。
1. 残差连接(Residual Connection)
残差连接由 He 等人在 ResNet 中提出,计算公式为:
自注意力与 FFN 计算出的量可以视为对原始表征的一种“微小修正增量”()。残差通路为梯度提供了一条无阻碍穿透数十甚至上百层的超高速公路,彻底解决了深层网络中的梯度消失与梯度爆炸问题。
2. 层归一化(Layer Normalization vs Batch Normalization)
在卷积视觉模型中普遍使用 Batch Normalization(BN),但 Transformer 坚定选择了 Layer Normalization(LN):
- BN 的局限: 沿 Batch 维度对同一特征通道进行归一化。在 NLP 场景下,句子长度参差不齐(长尾分布严重),且推理时 Batch Size 经常动态变化甚至为 1,导致均值与方差极不稳定。
- LN 的优势: 针对单个样本在单个时刻的特征向量,在隐藏维度 上独立计算均值 与方差 :
其中 是可学习的缩放与平移参数。LN 不依赖其他样本或不同时序位置,非常适合可变长序列的自适应归一化。
3. 架构范式之争:Post-LN vs Pre-LN
归一化层的位置顺序对深层网络的训练稳定性有决定性影响:
| 对比维度 | Post-LN(原始 Transformer 方案) | Pre-LN(现代主流大模型方案) |
|---|---|---|
| 计算拓扑 | ||
| 残差直连 | 残差分支被 LN 阻断,主干数值随深度被反复重整。 | 残差主干保持纯粹恒等映射,梯度无衰减穿透。 |
| 训练稳定性 | 极度脆弱;深层梯度剧烈震荡,必须依赖严格的 Warmup 策略。 | 极其平稳;无需激进的 Warmup 即可轻松收敛百层网络。 |
| 表征容量 | 理论上特征多样性略高,上限潜力大。 | 深层子层贡献略有微小衰减,但工程容错性极高。 |
原始 Transformer 使用 Post-LN,导致深层堆叠时必须经过繁琐的学习率预热;当前绝大多数主流模型已全面转向 Pre-LN 或其改良形式。
五、 PyTorch 模块化代码实现
以下代码实现了支持 Pre-LN 结构、采用 GELU 激活的生产级 TransformerEncoderLayer 与多层堆叠的 TransformerEncoder:
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
"""标准无掩码多头自注意力实现"""
def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
super().__init__()
assert d_model % num_heads == 0
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_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 forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
B, L, _ = x.shape
Q = self.w_q(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2)
K = self.w_k(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2)
V = self.w_v(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2)
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn_weights = self.dropout(torch.softmax(scores, dim=-1))
context = torch.matmul(attn_weights, V)
context = context.transpose(1, 2).contiguous().view(B, L, self.d_model)
return self.w_o(context)
class PositionwiseFeedForward(nn.Module):
"""基于位置的前馈全连接层 (采用 GELU 激活)"""
def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
super().__init__()
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: torch.Tensor) -> torch.Tensor:
# 升维 -> GELU 非线性激活 -> Dropout -> 降维
return self.linear2(self.dropout(F.gelu(self.linear1(x))))
class TransformerEncoderLayer(nn.Module):
"""单层编码器:采用工业界更稳定的 Pre-LN 结构"""
def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
super().__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
self.ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
# 子层 1:Pre-LN 多头自注意力 + 残差连接
norm_x = self.norm1(x)
attn_out = self.self_attn(norm_x, mask)
x = x + self.dropout(attn_out)
# 子层 2:Pre-LN 前馈网络 + 残差连接
norm_x = self.norm2(x)
ffn_out = self.ffn(norm_x)
x = x + self.dropout(ffn_out)
return x
class TransformerEncoder(nn.Module):
"""N 层编码器堆叠组成的完整 Encoder"""
def __init__(self, num_layers: int, d_model: int, num_heads: int, d_ff: int, dropout: float = 0.1):
super().__init__()
self.layers = nn.ModuleList([
TransformerEncoderLayer(d_model, num_heads, d_ff, dropout)
for _ in range(num_layers)
])
# Pre-LN 结构需要在所有层堆叠完成后补充最后的归一化层
self.final_norm = nn.LayerNorm(d_model)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
for layer in self.layers:
x = layer(x, mask)
return self.final_norm(x)
# ---------------- 验证前向传播 ----------------
if __name__ == "__main__":
batch_size = 4
seq_len = 32
d_model = 512
num_heads = 8
d_ff = 2048
num_layers = 6
# 模拟输入嵌入层输出
dummy_input = torch.randn(batch_size, seq_len, d_model)
# 实例化 6 层编码器
encoder = TransformerEncoder(num_layers, d_model, num_heads, d_ff)
encoded_output = encoder(dummy_input)
print("输入尺寸:", dummy_input.shape) # torch.Size([4, 32, 512])
print("输出尺寸:", encoded_output.shape) # torch.Size([4, 32, 512])
六、 现代大模型对编码器组件的演进与改良
随着模型体量与上下文长度的爆炸式增长,编码器层内部的传统组件逐步演化出更高效率的改良方案:
- RMSNorm(Root Mean Square Layer Normalization):
在 LLaMA、Mistral 等架构中,标准 LN 中的减均值操作被证明对维持梯度的贡献并不显著。RMSNorm 直接舍弃了均值中心化,仅通过均方根进行缩放,减少了 7% 左右的显存读写与核函数调用耗时:
- SwiGLU 激活前馈层:
传统的双层 MLP 逐渐被门控线性单元(GLU)所取代。SwiGLU 引入了带门控分支的三矩阵结构:
虽然增加了单个参数矩阵,但它在参数量相当的情况下展示出了明显更优的梯度流转性质和语义表征能力。
编码器层的精妙之处在于其高度模块化与几何对称性:多头自注意力完成空间维度的动态全局交互,前馈网络实现通道维度的深度非线性提炼,残差连接确保梯度在深层网络中无损回传,层归一化则稳固特征分布。这一套“交互-变换-稳定”的设计范式,不仅让编码器层成为了 BERT、RoBERTa 等语言理解模型的基石,其核心思想更深植于现代所有主流大语言模型的 Backbone 之内。
- 点赞
- 收藏
- 关注作者
评论(0)