VuePress 站点

🎯 课程主题

层归一化与前馈神经网络的代码实现——Transformer 的两个基础模块。

📝 核心知识点

1. 层归一化（LayerNorm）

概念说明：对每个样本的特征维度做归一化。
关键细节：
- 计算均值和方差
- 归一化：$(x - \mu) / \sigma$
- 加入平滑小数值 $\epsilon$ 避免除零
- 乘以可学习参数 $\gamma$（代码中用 a_2）
- 加上可学习参数 $\beta$（代码中用 b_2）
- $\gamma$ 和 $\beta$ 是可学习参数

2. 前馈神经网络（FFN）

概念说明：两层线性变换 + ReLU 激活。
关键细节：
- 第一层：线性变换 $W_1$（升维，$d \to d_{ff}$）
- 激活函数：ReLU
- Dropout 防止过拟合
- 第二层：线性变换 $W_2$（降维，$d_{ff} \to d$）
- 不带循环，一直向前计算

3. 代码结构

概念说明：两个模块的 PyTorch 实现。
关键细节：
- LayerNorm 继承 nn.Module
- FFN 包含两个 nn.Linear 层
- FFN 的 forward 是一行代码链式调用

🧮 核心公式与推导

层归一化公式：

$$\text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$

其中：

$\mu = \frac{1}{d}\sum_{i=1}^{d} x_i$（均值）
$\sigma^2 = \frac{1}{d}\sum_{i=1}^{d}(x_i - \mu)^2$（方差）
$\epsilon$：平滑小数值，避免除零
$\gamma, \beta$：可学习参数

前馈神经网络公式：

$$\text{FFN}(x) = \max(0, xW_1 + b_1) W_2 + b_2$$

等价于： $$\text{FFN}(x) = \text{Dropout}(\text{ReLU}(xW_1 + b_1)) W_2 + b_2$$

🏗️ 模型架构与数据流向

💻 代码实战

import torch
import torch.nn as nn
import torch.nn.functional as F


class LayerNorm(nn.Module):
    """层归一化"""
    def __init__(self, features, eps=1e-6):
        """
        features: 特征维度 d_model
        eps: 平滑小数值, 避免除零
        """
        super().__init__()
        self.a_2 = nn.Parameter(torch.ones(features))  # γ, 可学习
        self.b_2 = nn.Parameter(torch.zeros(features)) # β, 可学习
        self.eps = eps
    
    def forward(self, x):
        """
        x: [batch, seq_len, d_model]
        return: [batch, seq_len, d_model] 归一化后
        """
        # 计算均值和方差 (在最后一维, 即特征维度)
        mean = x.mean(-1, keepdim=True)      # [batch, seq_len, 1]
        std = x.std(-1, keepdim=True)        # [batch, seq_len, 1]
        
        # 归一化: (x - μ) / (σ + ε)
        normalized = (x - mean) / (std + self.eps)
        
        # 缩放和平移: γ * normalized + β
        return self.a_2 * normalized + self.b_2


class PositionwiseFeedForward(nn.Module):
    """前馈神经网络 (位置式前馈网络)"""
    def __init__(self, d_model, d_ff, dropout=0.1):
        """
        d_model: 模型维度 (如512)
        d_ff: FFN内部维度 (如2048)
        dropout: dropout比例
        """
        super().__init__()
        self.w_1 = nn.Linear(d_model, d_ff)   # 第一层: 升维
        self.w_2 = nn.Linear(d_ff, d_model)   # 第二层: 降维
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, x):
        """
        x: [batch, seq_len, d_model]
        return: [batch, seq_len, d_model]
        """
        # x → Linear1 → ReLU → Dropout → Linear2
        return self.w_2(self.dropout(F.relu(self.w_1(x))))


# 使用示例
if __name__ == "__main__":
    d_model = 512
    d_ff = 2048
    batch_size = 2
    seq_len = 10
    
    # 1. 层归一化
    ln = LayerNorm(d_model)
    x = torch.randn(batch_size, seq_len, d_model)
    output_ln = ln(x)
    print(f"LayerNorm输出: {output_ln.shape}")  # [2, 10, 512]
    print(f"归一化后均值: {output_ln[0, 0].mean():.4f}")  # ≈0
    print(f"归一化后方差: {output_ln[0, 0].std():.4f}")    # ≈1
    
    # 2. 前馈神经网络
    ffn = PositionwiseFeedForward(d_model, d_ff)
    output_ffn = ffn(x)
    print(f"FFN输出: {output_ffn.shape}")  # [2, 10, 512]
    
    # 3. PyTorch内置实现对比
    ln_pytorch = nn.LayerNorm(d_model)
    output_pytorch = ln_pytorch(x)
    print(f"PyTorch LayerNorm输出: {output_pytorch.shape}")

⚠️ 常见问题与避坑指南

LayerNorm 在特征维度（最后一维）做归一化，不是 batch 维度
$\epsilon$ 用于避免除零，通常取 $10^{-6}$
$\gamma$ 初始化为 1，$\beta$ 初始化为 0
FFN 的 $d_{ff}$ 通常是 $d_{model}$ 的 4 倍（如 512→2048）
代码中 std 计算的是标准差（已开方），不是方差

💡 个人总结与延伸

层归一化和前馈神经网络是 Transformer 的基础组件。LayerNorm 稳定训练，FFN 增加非线性表达能力。现代大模型中，RMSNorm（去掉中心化和偏置）比 LayerNorm 更高效；FFN 的激活函数从 ReLU 发展到 GELU、SwiGLU 等，提升模型性能。