层归一化与残差连接的作用

构建多层神经网络 (neural network)时，像Transformer中堆叠多层自注意力 (self-attention)子层和前馈子层那样，训练过程中会面临挑战。主要是，梯度在通过多层向后传播时可能消失（变得非常小）或爆炸（变得非常大），这使得模型难以有效学习。此外，中间层的激活分布在训练期间可能变化（这种现象有时与内部协变量偏移有关），使学习过程变得复杂。在Transformer架构中，残差连接和层归一化 (normalization)是两种简单但非常有效的方法，它们用于解决这些问题，并使多层模型得以训练。

残差连接

残差连接，也称作跳跃连接，提供了一条备用通路供梯度流经网络。不是简单地将子层的输出传递给下一层，而是我们将子层的输入加到其输出上。

如果一个子层由函数 $\text{SubLayer}(\cdot)$ 表示，其输入为 $x$，那么带有残差连接的块的输出是：

$$\text{输出} = x + \text{子层}(x)$$

这种结构使得网络能够轻松地学习恒等函数，如果某个特定子层没有益处；子层的输出可以简单地趋近于零。更重要的是，在反向传播 (backpropagation)时，梯度可以直接通过加法操作从输出流回到输入 $x$。这绕过了子层内部的变换，提供了一条“捷径”，有助于防止梯度信号在穿过多层时过度减弱。

残差连接将输入 x 加到 子层 的输出上。

在PyTorch中，这实现起来很简单。假设 sublayer 是一个模块（如多头注意力 (multi-head attention)或前馈网络），x 是输入张量：

import torch
import torch.nn as nn

# 假设 'sublayer' 在其他地方定义（例如，MultiHeadAttention, FeedForward）
# class SubLayer(nn.Module):
#     def __init__(self, d_model, ...):
#         super().__init__()
#         # ... 定义层 ...
#     def forward(self, x):
#         # ... 计算子层输出 ...
#         return processed_x

class ResidualConnection(nn.Module):
    def __init__(self, sublayer):
        super().__init__()
        self.sublayer = sublayer

    def forward(self, x):
        """
        对任意子层应用残差连接。
        """
        # 将原始输入 'x' 加到子层的输出上
        return x + self.sublayer(x)

# 使用示例
# d_model = 512
# input_tensor = torch.randn(batch_size, seq_len, d_model)
# attention_layer = MultiHeadAttention(...) # 假设已定义
# residual_block = ResidualConnection(attention_layer)
# output_tensor = residual_block(input_tensor)

层归一化 (normalization)

归一化技术有助于稳定训练过程，通过控制激活的分布。尽管批量归一化在计算机视觉中很常见，但它是在批次维度上归一化统计量（均值和方差）。这对于序列模型来说可能存在问题，因为批次内的序列长度可能不同，并且它引入了批次元素之间不总是期望的依赖关系。

层归一化（LayerNorm）提供了一种替代方法。它独立地跨特征对每个数据点（例如，序列中的每个词元 (token)）的输入进行归一化。它根据单个训练样本中单个层内所有神经元的求和输入，计算用于归一化的均值和方差。

给定输入向量 (vector) $x$（表示单个词元位置在所有特征上的激活，$d_{model}$），层归一化计算归一化输出 $h$ 如下：

1. 计算跨特征维度（$d_{model}$）的均值（$\mu$）和方差（$\sigma^2$）：

   $$\mu = \frac{1}{d_{model}} \sum_{i=1}^{d_{model}} x_i$$ $$\sigma^2 = \frac{1}{d_{model}} \sum_{i=1}^{d_{model}} (x_i - \mu)^2$$

2. 归一化输入 $x$：

   $$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$$

   其中 $\epsilon$ 是为数值稳定性添加的一个小常数。

3. 使用可学习参数 (parameter) $\gamma$（伽马，缩放）和 $\beta$（贝塔，平移）来缩放和移动归一化后的输出，这些参数与 $x$ 具有相同的维度：

   $$h_i = \gamma_i \hat{x}_i + \beta_i$$

这些可学习参数 $\gamma$ 和 $\beta$ 使得网络能够自适应地缩放和移动归一化后的激活，如果这对网络来说最优，甚至可能恢复原始激活。层归一化有助于稳定隐藏状态的动态，减少对初始化尺度的敏感性，甚至可以提供轻微的正则化 (regularization)作用。

在PyTorch中，torch.nn.LayerNorm 实现了这一点：

import torch
import torch.nn as nn

# 示例参数
batch_size = 4
seq_len = 10
d_model = 512
epsilon = 1e-5 # Small value for numerical stability

# 输入张量（批次，序列长度，特征）
input_tensor = torch.randn(batch_size, seq_len, d_model)

# 初始化层归一化
# 默认对最后一个维度 (d_model) 进行归一化
layer_norm = nn.LayerNorm(d_model, eps=epsilon)

# 应用层归一化
normalized_output = layer_norm(input_tensor)

# 检查形状
print("输入形状:", input_tensor.shape)
print("输出形状:", normalized_output.shape)

# 验证批次/序列中一个示例的均值和标准差
# 注意：由于 epsilon 和可学习参数 gamma/beta，
# 除非 gamma=1, beta=0，否则输出的均值/标准差不会正好是 0/1。
# 但归一化是在 gamma/beta 应用*之前*内部进行的。
print(
    "\n归一化输出的均值（示例 0，词元 0）:",
    normalized_output[0, 0, :].mean().item()
)
print(
    "归一化输出的标准差（示例 0，词元 0）:",
    normalized_output[0, 0, :].std().item()
)

# nn.LayerNorm 具有可学习参数 gamma（权重）和 beta（偏置）
print("\nLayerNorm 可学习伽马（权重）:", layer_norm.weight.shape)
print("LayerNorm 可学习贝塔（偏置）:", layer_norm.bias.shape)

结合残差连接和层归一化 (normalization)

在Transformer架构中，层归一化和残差连接通常围绕每个子层（多头注意力 (multi-head attention)和逐位置前馈网络）一起应用。原始论文《Attention Is All You Need》中描述的标准结构是在残差加法之后应用归一化（Post-LN）：

$$\text{输出} = \text{LayerNorm}(x + \text{子层}(x))$$

然而，随后的研究和实践经常发现，在残差分支中，将层归一化应用在子层之前（Pre-LN）可以带来更稳定的训练，特别是对于非常深的Transformer模型：

$$\text{输出} = x + \text{子层}(\text{LayerNorm}(x))$$

我们将在第11章讨论扩展定律和架构选择时，进一步探讨Pre-LN与Post-LN的影响。目前，请认识到这种结合（通常被描述为“添加与归一化”步骤）是基础的。

残差块中Post-LN（加法后归一化）和Pre-LN（子层前归一化）结构的比较。

以下是Transformer编码器层通常如何使用Pre-LN方法结合这些组件的示例：

import torch
import torch.nn as nn
# 假设 MultiHeadAttention 和 PositionwiseFeedForward 是在前面部分/外部模块中定义的类

class EncoderLayer(nn.Module):
    """
    实现一个带有Pre-LN的Transformer编码器层。
    """
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads) # Placeholder
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff) # Placeholder

        # 层归一化实例
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        # 用于正则化的Dropout
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # 1. 层归一化，接着是多头自注意力
        norm_x = self.norm1(x)
        # 假设 self_attn 返回注意力输出
        attn_output = self.self_attn(norm_x, norm_x, norm_x, mask) 
        # 带有Dropout的残差连接
        x = x + self.dropout1(attn_output)

        # 2. 层归一化，接着是逐位置前馈网络
        norm_x = self.norm2(x)
        ff_output = self.feed_forward(norm_x)
        # 带有Dropout的残差连接
        x = x + self.dropout2(ff_output)

        return x

# 使用示例
# d_model = 512
# num_heads = 8
# d_ff = 2048 # Feed-forward inner dimension
# batch_size = 4
# seq_len = 10

# input_tensor = torch.randn(batch_size, seq_len, d_model)
# encoder_layer = EncoderLayer(d_model, num_heads, d_ff)
# output_tensor = encoder_layer(input_tensor)
# print("编码器层输出形状:", output_tensor.shape)

总而言之，残差连接有助于梯度流动和信息在深层网络中的传播，而层归一化则稳定了激活分布。它们的结合使用是成功训练深层Transformer模型的重要因素，构成了编码器和解码器层的核心部分。