从零实现Transformer：第 4 部分 - Residual Connection的两种实现 Pre-LN 和 Post-LN

flyfish

Pre-LN = Pre-Layer Normalization
Post-LN = Post-Layer Normalization

Pre = 预先、在前面
Post = 在后、在末尾
Layer Normalization = 层归一化
Pre-LN：前置层归一化
Post-LN：后置层归一化

Pre-LN vs Post-LN

符号定义
$x$：当前模块原始输入
$\text{LN}(\cdot )$：层归一化
$\text{Sublayer}(\cdot )$：子层（自注意力 / FFN前馈网络）
$\text{Dropout}(\cdot )$：随机失活
$+$：残差连接逐元素相加

Post-LN 公式（原始 Transformer）

$$\mathbf{y}=\text{LN}(\mathbf{x}+\text{Dropout}(\text{Sublayer}(\mathbf{x})\left)\right)$$

对应代码

return self.norm(x + self.dropout(sublayer(x)))

Pre-LN 公式（现代大模型 GPT）

$$\mathbf{y}=\mathbf{x}+\text{Dropout}\left(\text{Sublayer}\right(\text{LN}(\mathbf{x})\left)\right)$$

对应代码

return x + self.dropout(sublayer(self.norm(x)))

直接对比

Post-LN：最后归一
Pre-LN：先归一

AI生成的Post-LN 和 Pre-LN

FFN（PositionwiseFeedForward，前馈网络）

import torch
import torch.nn as nn

# ===================== 公共模块 两者完全一致，无任何区别 =====================
class LayerNormalization(nn.Module):
    """层归一化"""
    def __init__(self, features: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = nn.Parameter(torch.zeros(features))
    
    def forward(self, x: torch.Tensor):
        mean = x.mean(dim=-1, keepdim=True)
        var = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
        normalized = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * normalized + self.beta

class PositionwiseFeedForward(nn.Module):
    """Transformer前馈网络"""
    def __init__(self, d_model: int, d_ff: int, dropout: float):
        super().__init__()
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.linear_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.ReLU()
    
    def forward(self, x):
        return self.linear_2(self.dropout(self.activation(self.linear_1(x))))

# ===================== 仅残差连接的 forward 函数不同！ =====================
# 版本1：Post-LN（原始Transformer）
class ResidualConnection_PostLN(nn.Module):
    def __init__(self, features: int, dropout: float):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNormalization(features)

    def forward(self, x, sublayer):
        #  Post-LN 公式：LN(x + Dropout(Sublayer(x)))
        return self.norm(x + self.dropout(sublayer(x)))

# 版本2：Pre-LN（现代大模型 GPT）
class ResidualConnection_PreLN(nn.Module):
    def __init__(self, features: int, dropout: float):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = LayerNormalization(features)

    def forward(self, x, sublayer):
        #  Pre-LN 公式：x + Dropout(Sublayer(LN(x)))
        return x + self.dropout(sublayer(self.norm(x)))

# ===================== 【测试代码】验证两种结构 =====================
if __name__ == "__main__":
    # 固定随机种子，保证结果可复现
    torch.manual_seed(42)
    
    # 超参数配置
    d_model = 512    # 模型维度
    d_ff = 2048     # 前馈网络中间维度
    dropout = 0.1   # Dropout概率
    
    # 构造输入：[batch_size, seq_len, d_model]
    x = torch.randn(2, 10, d_model)  
    print(f"输入张量形状: {x.shape}")
    
    # 初始化子层（前馈网络）
    ffn = PositionwiseFeedForward(d_model, d_ff, dropout)
    
    # 1. 测试 Post-LN 残差连接
    post_ln = ResidualConnection_PostLN(d_model, dropout)
    out_post = post_ln(x, ffn)
    print(f"\nPost-LN 输出形状: {out_post.shape}")
 
    
    # 2. 测试 Pre-LN 残差连接
    pre_ln = ResidualConnection_PreLN(d_model, dropout)
    out_pre = pre_ln(x, ffn)
    print(f"\nPre-LN 输出形状: {out_pre.shape}")

输出

输入张量形状: torch.Size([2, 10, 512])

Post-LN 输出形状: torch.Size([2, 10, 512])

Pre-LN 输出形状: torch.Size([2, 10, 512])