第4章前向传播

第4章：前向传播：LayerNorm的完整计算图

好，咱们进入正题。前向传播这部分，说白了就是搞清楚数据在LayerNorm里到底是怎么流动的。我见过不少同学，背公式背得滚瓜烂熟，一让画计算图就懵了。嗯，今天咱们就把这张图彻底画明白。

4.1 从输入到输出的完整路径

先看一个典型的LayerNorm计算流程。假设输入是一个向量 x = [x₁, x₂, ..., xₙ]，我们要做的是沿着特征维度做归一化。我个人习惯把整个过程拆成5步：

计算均值：μ = (1/n) * Σxᵢ
计算方差：σ² = (1/n) * Σ(xᵢ - μ)²
归一化：x̂ᵢ = (xᵢ - μ) / √(σ² + ε)
缩放：yᵢ = γᵢ * x̂ᵢ + βᵢ

你想想看，这其实就是一个「先标准化，再学一个线性变换」的过程。我在项目中遇到过有人把LayerNorm和BatchNorm搞混，觉得都是归一化，随便用就行。结果模型训练到一半，loss死活降不下去。后来一查，原来是归一化的维度选错了。

关键区别：LayerNorm是对每个样本的特征维度做归一化，BatchNorm是对整个batch的同一特征维度做归一化。方向完全不同。

4.2 计算图的节点与边

咱们把计算图画出来，每个节点就是一个操作，每条边就是数据的流向。我习惯用这种结构来理解：

输入 x (shape: [batch, seq_len, hidden_dim])
    │
    ▼
[计算均值 μ]  ← 沿着hidden_dim维度
    │
    ▼
[计算中心化]  x - μ
    │
    ├──→ [计算方差 σ²]  ← 也是沿着hidden_dim维度
    │         │
    │         ▼
    │    [加ε开根号]  √(σ² + ε)
    │         │
    │         ▼
    └──→ [归一化]  x̂ = (x - μ) / √(σ² + ε)
              │
              ▼
         [缩放+平移]  y = γ * x̂ + β
              │
              ▼
           输出 y

这张图里，最容易被忽略的是那个ε。我曾经在写自定义算子时忘了加ε，结果方差为0的时候直接除零，整个训练崩了。ε一般取1e-5或1e-6，就是个数值稳定的小技巧。

4.3 维度视角下的计算图

咱们用具体数字来感受一下。假设输入形状是 [2, 3, 4]，也就是batch_size=2，seq_len=3，hidden_dim=4。LayerNorm通常对最后一个维度做归一化。

步骤	操作	输出形状	说明
1	计算均值 μ	[2, 3, 1]	沿着hidden_dim求平均，保留了batch和seq维度
2	计算方差 σ²	[2, 3, 1]	同样沿着hidden_dim，注意这里用了中心化后的值
3	归一化 x̂	[2, 3, 4]	形状恢复，每个元素都做了标准化
4	缩放 y	[2, 3, 4]	γ和β形状都是[4]，通过广播机制作用到每个位置

这里有个细节要注意：μ和σ²的形状是[batch, seq_len, 1]，而不是[2, 3]。为什么？因为我们要保留维度信息，方便后续的广播运算。如果你用numpy实现，记得用keepdims=True。

我的小技巧：在写代码时，我习惯在每一步都打印一下中间变量的shape。这样能快速定位维度不匹配的问题。尤其是当你自己实现LayerNorm时，这一步能省下大量调试时间。

4.4 计算图中的依赖关系

为什么我们要把计算图画这么细？因为反向传播时，每个节点的梯度都依赖于前向传播中的中间结果。举个例子：

γ的梯度：依赖于x̂和输出y的梯度
β的梯度：直接等于输出y的梯度求和
x的梯度：最复杂，要经过归一化、方差、均值三条路径

说白了，前向传播中保存的中间变量，就是反向传播的「燃料」。我在做Transformer推理优化时，发现有些框架会刻意丢弃一些中间变量来省显存，但代价是反向传播时要重新计算。这就是典型的「时间换空间」策略。

避坑指南：我曾经在实现自定义LayerNorm时，为了省事只保存了x̂和σ²，结果反向传播时发现μ的梯度算不出来。后来老老实实把x - μ也保存了。记住：反向传播需要什么，前向就保存什么。

4.5 代码实现与计算图对应

咱们用PyTorch风格写一段伪代码，把计算图对应起来：

def layer_norm_forward(x, gamma, beta, eps=1e-5):
    # 步骤1: 计算均值
    mu = x.mean(dim=-1, keepdim=True)  # shape: [batch, seq, 1]
    
    # 步骤2: 中心化
    x_centered = x - mu  # shape: [batch, seq, hidden]
    
    # 步骤3: 计算方差
    var = (x_centered ** 2).mean(dim=-1, keepdim=True)  # shape: [batch, seq, 1]
    
    # 步骤4: 归一化
    x_norm = x_centered / torch.sqrt(var + eps)  # shape: [batch, seq, hidden]
    
    # 步骤5: 缩放和平移
    out = gamma * x_norm + beta  # gamma和beta通过广播
    
    # 保存中间变量供反向传播使用
    cache = (x_norm, x_centered, var, mu, eps)
    
    return out, cache

你看，每一步都对应计算图中的一个节点。我个人习惯把cache打包成一个元组，这样反向传播时直接解包就行，代码看起来清爽很多。

4.6 为什么计算图这么重要？

我经常跟团队里的新人说：前向传播的计算图，就是反向传播的导航地图。你如果连前向的路径都画不清楚，反向传播的梯度流就一定会迷路。

举个例子，在LayerNorm中，x̂的梯度要经过两条路径：一条来自γ的梯度，一条来自β的梯度（虽然β的梯度不直接依赖x̂）。而x的梯度更复杂，要经过归一化、方差、均值三条路径的链式法则。这些路径，全都在计算图里标得清清楚楚。

嗯，这一章的内容就到这儿。下一章咱们就拿着这张计算图，开始推导反向传播的梯度公式。到时候你会发现，只要前向的图画对了，反向的推导就是顺着路径走一遍的事。

📚 课程目录