残差、LayerNorm 与前馈层

残差、LayerNorm 与前馈层：Transformer 的稳定器

小象实战讲义 · 深度学习基础：从神经网络到 Transformer

在上一节，我们学习了 Transformer 的核心机制——多头注意力。然而，一个只有注意力堆叠的网络在训练时极易不稳定，梯度会消失或爆炸。本节我们将学习 Transformer 中另外三个关键组件：残差连接（Residual Connection）、层归一化（Layer Normalization） 和位置前馈网络（Position-wise Feed-Forward Network），它们与注意力一起构成了 Transformer 的稳定骨架。学完本节，你将能回答：为什么 Transformer 能堆叠数十层而不崩溃？LayerNorm 和 BatchNorm 有何本质区别？代码中如何实现“Add & Norm”操作？

先修：缩放点积注意力与多头注意力的形状（本章 5.1、5.2 节）、nn.Linear 与 nn.LayerNorm 的基本用法。后续：下一节用本节这几块零件，从零拼出一个最小 Transformer block，并做参数量对账。

💡 核心导读

1. 残差连接：将子层（如注意力、前馈网络）的输入直接加到其输出上，形成 output = input + sublayer(input) 的短路路径，确保梯度能有效回传，缓解深层网络梯度消失问题。
2. 层归一化（LayerNorm）：在每个样本、每个时间步上，沿所有特征维度计算均值和方差进行归一化，稳定激活值的分布，使训练过程更平滑。它与 BatchNorm 的计算维度不同，尤其适合变长序列和在线学习。
3. 位置前馈网络（FFN）：一个应用于序列中每个位置（Position-wise）的两层全连接网络，通常中间维度扩大（如 4 倍），用于引入非线性变换和增加模型容量。
4. 子层连接范式：Transformer 的标准做法是“后归一化（Post-LN）”，即 LayerNorm(x + Sublayer(x))。我们将通过代码对比其与“前归一化（Pre-LN）”变体的区别。

🧭 前置知识

1. 矩阵加法与广播机制（形状需匹配）。
2. 均值和方差的计算。
3. 梯度消失与爆炸的基本概念。
4. 全连接层（线性变换）与激活函数（如 ReLU、GELU）。

残差连接：为梯度铺设高速公路

残差连接最早在 ResNet（He et al., 2016）中被提出，用于解决极深卷积网络的退化问题。其核心思想异常简单：让网络层学习输入与输出之间的残差（差值），而非直接学习完整的输出映射。

在 Transformer 中，每个子层（注意力层或前馈层）都被一个残差连接所包裹。具体操作如下：

子层输出 = 输入 + Dropout(子层(输入))      # 残差连接本身
Add & Norm = LayerNorm(子层输出)           # 之后再做一次层归一化

这里的“输入”指的是进入该子层的张量。加法操作要求 输入 与 子层(输入) 的形状完全相同。在 Transformer 的原始设计中，所有子层和嵌入层的输出维度都统一为 d_model（例如 512），这确保了形状匹配。

为什么有效？ 假设没有残差连接，梯度需要流经整个子层的计算图。在深层堆叠时，连续的矩阵乘法可能导致梯度值急剧缩小（消失）或放大（爆炸）。残差连接引入了一条从输入直接到输出的“短路”路径。在反向传播时，梯度可以通过这条路径几乎无损地回传，极大地缓解了梯度消失问题，使得训练数十层的深度网络成为可能。

下面用一个刻意做“坏”的例子验证这一点：4 层线性层的权重全部固定成 0.1（4×4 矩阵、输入梯度全是 1，所以每层沿梯度方向实际缩到 0.4 倍），再看第一层拿到的梯度有多大。

import torch
import torch.nn as nn

dim, n_layers, batch = 4, 4, 2
torch.manual_seed(0)
x = torch.randn(batch, dim)

def build_stack():
    torch.manual_seed(1)
    layers = [nn.Linear(dim, dim) for _ in range(n_layers)]
    for layer in layers:
        nn.init.constant_(layer.weight, 0.1)   # 4x4 全 0.1：每层把梯度约缩到 0.4 倍
        nn.init.constant_(layer.bias, 0.0)
    return layers

for use_residual in (False, True):
    layers = build_stack()
    h = x
    for layer in layers:
        out = layer(h)
        h = out + h if use_residual else out
    h.sum().backward()
    # 只看最靠近输入的第一层：没有残差时它几乎收不到梯度
    g = layers[0].weight.grad.norm().item()
    tag = "有残差" if use_residual else "无残差"
    print(f"{tag}｜第一层权重梯度范数 = {g:.6f}｜输出均值 = {h.mean().item():.4f}")

无残差｜第一层权重梯度范数 = 0.366573｜输出均值 = -0.0051
有残差｜第一层权重梯度范数 = 15.716805｜输出均值 = -0.7705

差了 40 多倍。原因可以直接算出来：无残差时，第一层的梯度要连乘后面三层，每层沿梯度方向缩到 $0.1 \times 4 = 0.4$ 倍，$0.4^3 \approx 0.064$，只剩约 6%；有残差时，每层的雅可比从 $W^\top$ 变成 $W^\top + I$，同一方向上的乘数从 0.4 变成 1.4，$1.4^3 \approx 2.74$，比值 $2.74/0.064 \approx 43$，正好对上打印结果。这也解释了为什么深层网络一旦加上残差，就敢往几十层堆。

层归一化：稳定每层激活的分布

层归一化（Layer Normalization）由 Ba 等人于 2016 年提出（arXiv:1607.06450）。它的目标是稳定训练过程中每层激活值的分布，从而加速收敛。

需要说明一句：论文给出的动机是减少层内激活的“协变量偏移”（Covariate Shift），但后续工作（Santurkar et al., 2018, arXiv:1805.11604）做了对照实验，认为“减少内部协变量偏移”并不是归一化起作用的主要原因——真正管用的是它把优化曲面变平滑了。这个争论至今没有完全定论，你只要记住结论：归一化确实让深层网络更好训，机制解释仍在讨论中。

LayerNorm 与 BatchNorm 的本质区别

理解 LayerNorm 的关键在于明确其计算均值和方差所沿着的维度。

• BatchNorm：沿着 Batch 维度进行归一化。对于形状为 (batch, seq, feature) 的张量，它计算每个特征（feature）在所有 batch 样本上的均值和方差。这导致其效果依赖于 batch size，并且在序列长度可变或 batch size 为 1 时效果不佳。
• LayerNorm：沿着 Feature 维度进行归一化。它计算每个样本、每个时间步上，所有特征维度的均值和方差。因此，其统计量与 batch 内其他样本无关，非常适合 RNN、Transformer 这类处理变长序列的模型，也适用于在线学习。

根据 PyTorch 官方文档（nn.LayerNorm），其计算公式为： $$y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} \odot \gamma + \beta$$ 其中，$\mathrm{E}[x]$ 和 $\mathrm{Var}[x]$ 是在 normalized_shape 指定的维度上计算的均值和方差。$\gamma$ 和 $\beta$ 是可学习的缩放（weight）和平移（bias）参数，维度与 normalized_shape 相同。

在 Transformer 中的具体应用

在原始 Transformer 论文（Vaswani et al., 2017, arXiv:1706.03762）中，每个子层（自注意力、编码器-解码器注意力、前馈网络）的输出都经过“Add & Norm”操作：LayerNorm(x + Sublayer(x))。这种设计被称为 后归一化（Post-LayerNorm）。

让我们通过代码直观感受 LayerNorm 的效果，并对比 Post-LN 与 Pre-LN。

import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(0)

# 1. 演示 LayerNorm 的基本计算
batch_size, seq_len, d_model = 2, 4, 8
# 构造输入，模拟某个子层（如注意力）计算后的输出，其分布可能已经改变
x = torch.randn(batch_size, seq_len, d_model) * 0.5 + 2.0  # 生成分布均值2、标准差0.5
sub_output = torch.randn(batch_size, seq_len, d_model) * 2.0 - 1.0  # 模拟子层输出，分布不同
# 注意：下面只打印 8 个元素，经验统计量会与生成参数有偏差

# 应用 Post-LN: LayerNorm(x + sub_output)
layer_norm = nn.LayerNorm(d_model)
post_ln_output = layer_norm(x + sub_output)

# 计算原始输入、子层输出及 Post-LN 输出的统计量（针对第一个样本的第一个位置）
sample_input = x[0, 0]
sample_sub = sub_output[0, 0]
sample_post_ln = post_ln_output[0, 0]

print("原始输入 x[0,0] 的均值与标准差:")
print(f"  均值: {sample_input.mean().item():.4f}, 标准差: {sample_input.std(unbiased=False).item():.4f}")
print("模拟子层输出 sub[0,0] 的均值与标准差:")
print(f"  均值: {sample_sub.mean().item():.4f}, 标准差: {sample_sub.std(unbiased=False).item():.4f}")
print("Post-LN 输出 post_ln[0,0] 的均值与标准差:")
print(f"  均值: {sample_post_ln.mean().item():.4f}, 标准差: {sample_post_ln.std(unbiased=False).item():.4f}")
print(f"Post-LN 输出的缩放参数 gamma (前3个): {layer_norm.weight[:3].tolist()}")
print(f"Post-LN 输出的偏置参数 beta (前3个): {layer_norm.bias[:3].tolist()}")

原始输入 x[0,0] 的均值与标准差:
  均值: 1.7592, 标准差: 0.4597
模拟子层输出 sub[0,0] 的均值与标准差:
  均值: -0.0743, 标准差: 2.3944
Post-LN 输出 post_ln[0,0] 的均值与标准差:
  均值: 0.0000, 标准差: 1.0000
Post-LN 输出的缩放参数 gamma (前3个): [1.0, 1.0, 1.0]
Post-LN 输出的偏置参数 beta (前3个): [0.0, 0.0, 0.0]

从输出可以看到，不管子层输出把分布拉成什么样（这里是均值 −0.07、标准差 2.39），经过 LayerNorm(x + sub_output) 之后这个 8 维向量的均值都被拉回 0、标准差都被拉回 1。注意这只说明“统计量对齐”，不是说向量服从标准正态分布；另外 $\gamma=1,\beta=0$ 是初始状态，训练中它们会学出各维度需要的尺度和偏移。

位置前馈网络：每个位置的独立非线性变换

位置前馈网络（Position-wise Feed-Forward Network, FFN）是 Transformer 编码器和解码器每个层中都包含的另一个关键子层。它对序列中的每个位置独立、相同地进行变换。

其结构非常简单，是一个两层全连接网络：

FFN(x) = max(0, x W1 + b1) W2 + b2

在原始论文中，中间层的维度被扩大：d_ff = 4 * d_model。例如，当 d_model=512 时，d_ff=2048。第一层使用 ReLU 激活函数（现代实现中常用 GELU）。由于它对每个位置独立操作，因此可以并行计算，效率很高。

FFN 的作用是为模型引入非线性变换能力和额外的可学习参数，增强其表示能力。注意力机制主要负责聚合上下文信息（“在哪里看”），而 FFN 则负责对聚合后的信息进行深度加工（“怎么看”）。

import torch
import torch.nn as nn
torch.manual_seed(0)

# 2. 实现并演示位置前馈网络 (FFN)
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)  # 扩大维度
        self.w2 = nn.Linear(d_ff, d_model)  # 投影回 d_model
        self.dropout = nn.Dropout(dropout)
        # 使用 GELU 作为现代更常用的激活函数
        self.activation = nn.GELU()

    def forward(self, x):
        # x 形状: (batch, seq, d_model)
        return self.w2(self.dropout(self.activation(self.w1(x))))

# 实例化并运行
d_model, d_ff = 8, 32
ffn = PositionwiseFeedForward(d_model, d_ff, dropout=0.0)  # 暂不用 dropout 以便观察
x = torch.randn(2, 4, d_model)
output = ffn(x)

print(f"输入 x 形状: {tuple(x.shape)}")
print(f"FFN 中间层 (w1) 输出形状: {tuple(ffn.w1(x).shape)}")  # (batch, seq, d_ff)
print(f"FFN 最终输出形状: {tuple(output.shape)}")  # 应与输入相同 (batch, seq, d_model)
print(f"输入与输出形状相同: {tuple(x.shape) == tuple(output.shape)}")

输入 x 形状: (2, 4, 8)
FFN 中间层 (w1) 输出形状: (2, 4, 32)
FFN 最终输出形状: (2, 4, 8)
输入与输出形状相同: True

组装起来：Transformer 子层连接标准实现

现在，我们将残差连接、LayerNorm 和前馈网络（或注意力）组装成一个完整的子层连接模块。我们同时实现标准的 Post-LN 和另一种常见的变体 Pre-LN 进行对比。

import torch
import torch.nn as nn
torch.manual_seed(0)

# 3. 实现并对比 Post-LN 与 Pre-LN 子层连接
class SublayerConnectionPostLN(nn.Module):
    """标准 Post-LN: LayerNorm(x + Sublayer(x))"""
    def __init__(self, size, dropout):
        super().__init__()
        self.norm = nn.LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        # sublayer 是一个函数，例如一个注意力层或 FFN 的前向传播
        return self.norm(x + self.dropout(sublayer(x)))

class SublayerConnectionPreLN(nn.Module):
    """Pre-LN 变体：先归一化，再子层计算，最后残差连接。"""
    def __init__(self, size, dropout):
        super().__init__()
        self.norm = nn.LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        # 先对输入进行归一化，再送入子层
        return x + self.dropout(sublayer(self.norm(x)))

# 模拟一个前馈子层（块自包含：这里重新定义一份）
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()

    def forward(self, x):
        return self.w2(self.dropout(self.activation(self.w1(x))))

d_model, d_ff = 8, 32
ffn = PositionwiseFeedForward(d_model, d_ff, dropout=0.1)

# 创建输入
x = torch.randn(2, 4, d_model)

# 应用两种连接方式
post_ln_conn = SublayerConnectionPostLN(d_model, dropout=0.1)
pre_ln_conn = SublayerConnectionPreLN(d_model, dropout=0.1)

output_post = post_ln_conn(x, ffn)
output_pre = pre_ln_conn(x, ffn)

print("输入 x 的统计（第一个样本的第一个位置）:")
print(f"  均值: {x[0,0].mean().item():.4f}, 标准差: {x[0,0].std().item():.4f}")
print("\nPost-LN 输出的统计（同一位置）:")
print(f"  均值: {output_post[0,0].mean().item():.4f}, 标准差: {output_post[0,0].std().item():.4f}")
print("\nPre-LN 输出的统计（同一位置）:")
print(f"  均值: {output_pre[0,0].mean().item():.4f}, 标准差: {output_pre[0,0].std().item():.4f}")
print(f"\n两种方式输出是否相同？ {torch.allclose(output_post, output_pre, atol=1e-6)}")

输入 x 的统计（第一个样本的第一个位置）:
  均值: 0.2064, 标准差: 1.4183

Post-LN 输出的统计（同一位置）:
  均值: 0.0000, 标准差: 1.0690

Pre-LN 输出的统计（同一位置）:
  均值: 0.1652, 标准差: 1.5048

两种方式输出是否相同？ False

可以看到，Post-LN 和 Pre-LN 的输出分布不同。原始 Transformer 使用的是 Post-LN；BERT 沿用 Post-LN，而 GPT-2 起把 LayerNorm 挪到每个子块的输入侧，也就是 Pre-LN，此后多数大模型跟着用 Pre-LN。两者的取舍是：Pre-LN 的梯度尺度更稳定，不依赖学习率 warmup；在相同 lr_max 下，Xiong 等（2020，arXiv:2002.04745）实测 Pre-LN 收敛得比 Post-LN 更快（IWSLT14 上第 9 个 checkpoint 就接近 Post-LN 第 15 个的效果）。Post-LN 当然不是不能用——配上 warmup 依然很强，原始论文和 BERT 都是这么训的。选哪个取决于你要复现哪篇论文、以及训练深度有多大。

📝 动手练一练

问题：假设我们有一个形状为 (batch=3, seq_len=5, d_model=6) 的张量 x。请手动计算（不直接调用 nn.LayerNorm）对最后一个维度（d_model）进行 LayerNorm 的结果。计算第一个样本第一个位置向量的归一化结果，并与 PyTorch 官方实现对比。忽略可学习的 $\gamma$ 和 $\beta$（即假设它们为 1 和 0），eps=1e-5。

import torch
import torch.nn as nn
torch.manual_seed(0)

batch, seq, d_model = 3, 5, 6
x = torch.randn(batch, seq, d_model)

# 手动实现 LayerNorm (针对最后一个维度 dim=2)
eps = 1e-5
mean = x.mean(dim=-1, keepdim=True)  # 形状 (3,5,1)
var = x.var(dim=-1, keepdim=True, correction=0)  # 无偏估计设为0，与PyTorch默认一致
x_manual = (x - mean) / torch.sqrt(var + eps)

# 使用 PyTorch 官方实现
layer_norm = nn.LayerNorm(d_model, elementwise_affine=False)  # 关闭可学习参数
x_pytorch = layer_norm(x)

# 比较第一个样本第一个位置的结果
pos_result_manual = x_manual[0, 0].tolist()
pos_result_pytorch = x_pytorch[0, 0].tolist()

print("第一个样本第一个位置的原始向量:")
print([round(v, 6) for v in x[0,0].tolist()])
print("\n手动 LayerNorm 结果:")
print([round(v, 6) for v in pos_result_manual])
print("\nPyTorch LayerNorm 结果:")
print([round(v, 6) for v in pos_result_pytorch])
print(f"\n两者是否接近？ {torch.allclose(x_manual, x_pytorch, atol=1e-7)}")

第一个样本第一个位置的原始向量:
[-1.12584, -1.15236, -0.250579, -0.433879, 0.84871, 0.692009]

手动 LayerNorm 结果:
[-1.130551, -1.164283, -0.017284, -0.250428, 1.38093, 1.181617]

PyTorch LayerNorm 结果:
[-1.130551, -1.164283, -0.017284, -0.250428, 1.38093, 1.181618]

两者是否接近？ True

常见错误

本章小结

1. 残差连接通过引入 output = input + sublayer(input) 的短路路径，有效缓解了深层 Transformer 中的梯度消失问题，是模型能够堆叠深度的关键。
2. 层归一化（LayerNorm） 沿着特征维度对每个样本独立进行归一化，稳定了层间激活值的分布，其效果不依赖于 batch size，特别适合序列模型。
3. 位置前馈网络（FFN） 是一个应用于每个位置的两层全连接网络，通常中间扩大维度，用于引入非线性并增强模型表示能力。
4. 子层连接范式主要有 Post-LN（原始 Transformer 使用）和 Pre-LN（现代变体常用）两种，区别在于 LayerNorm 与残差加法的顺序，它们影响训练的稳定性和收敛速度。
5. “Add & Norm” 是 Transformer 每个子层的标准操作，它将残差连接、Dropout 和 LayerNorm 组合在一起，构成了模型稳定的基础骨架。

📋 行动清单

• 在代码示例中，将 SublayerConnectionPostLN 和 SublayerConnectionPreLN 中的 dropout 率从 0.1 改为 0.3，重新运行并观察输出值的变化幅度。
• 修改 FFN 的实现，将激活函数从 GELU 换为原始论文的 ReLU，比较两种激活函数下输出值的分布差异（例如打印输出值的均值和标准差）。
• 尝试构造一个形状不匹配的场景：将 FFN 的 self.w2 输出维度故意设错（如改为 d_model+1），运行代码并理解产生的错误信息。

延伸阅读

本段 PyTorch 官方文档/教程条目的框架出处：PyTorch: An Imperative Style, High-Performance Deep Learning Library（Paszke et al., 2019，arXiv:1912.01703）；页面内容以 PyTorch 2.14.1 为准。

1. Attention Is All You Need（Vaswani et al., 2017，arXiv:1706.03762）—— 3.1 节给出 LayerNorm(x + Sublayer(x)) 的原式，3.3 节描述位置前馈网络（$d_{ff}=2048$），3.5 节是位置编码。
2. Layer Normalization（Ba et al., 2016，arXiv:1607.06450）—— 层归一化原理的原始论文。
3. On Layer Normalization in the Transformer Architecture（Xiong et al., 2020，arXiv:2002.04745）—— 解释 Post-LN 为什么依赖 warmup、Pre-LN 为什么更稳。
4. PyTorch 官方文档 nn.LayerNorm（https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html）—— 层归一化的 API 说明与计算公式。

—— 小象教研组