详解三种常用标准化 Batch Norm & Layer Norm & RMSNorm

云端FFF

已于 2024-06-30 19:13:48 修改

阅读量926

点赞数 9

分类专栏：机器学习 ========================= # LLM专题文章标签： Batch Norm Layer Norm RMSNorm 归一化

于 2024-06-24 22:35:17 首次发布

本文链接：https://blog.csdn.net/wxc971231/article/details/139925707

版权

机器学习 ========================= 同时被 2 个专栏收录

11 篇文章 7 订阅

订阅专栏

LLM专题

6 篇文章 0 订阅

订阅专栏

参考：
- BN究竟起了什么作用？一个闭门造车的分析
- 《动手学深度学习》7.5 节

深度学习中，归一化是常用的稳定训练的手段，CV 中常用 Batch Norm； Transformer 类模型中常用 layer norm，而 RMSNorm 是近期很流行的 LaMMa 模型使用的标准化方法，它是 Layer Norm 的一个变体
值得注意的是，这里所谓的归一化严格讲应该称为 标准化Standardization ，有时也称为 白化whitening。它描述一种把样本调整到均值为 0，方差为 1 的缩放平移操作。使用这种方法可以消除输入数据的量纲，有利于随机初始化的网络训练
- 假如输入一个绝对值很大的数字 -10000，那么神经网络的 wx+b 的参数 w、b 需要经过很多步的更新才能从 0 → 1000，然后才能让这一层的输出也接近均值0方差1，使激活函数ReLU、Tanh等正常工作
- 归一化、标准化、正则化等术语常常被混用，可以看标准化、归一化概念梳理（附代码）这篇文章理清
详细讨论前，先粗略看一下 Batch Norm 和 Layer Norm 的区别
1. BatchNorm是对整个 batch 样本内的每个特征做归一化，这消除了不同特征之间的大小关系，但是保留了不同样本间的大小关系。BatchNorm 适用于 CV 领域，这时输入尺寸为 $b\times c\times h\times w$ (批量大小x通道x长x宽)，图像的每个通道 $c$ 看作一个特征，BN 可以把各通道特征图的数量级调整到差不多，同时保持不同图片相同通道特征图间的相对大小关系
2. LayerNorm是对每个样本的所有特征做归一化，这消除了不同样本间的大小关系，但是保留了一个样本内不同特征之间的大小关系。LayerNorm 适用于 NLP 领域，这时输入尺寸为 $b\times l\times d$ (批量大小x序列长度x嵌入维度)，如下图所示
  
  注意这时长 $l$ 的 token 序列中，每个 token 对应一个长为 $d$ 的特征向量，LayerNorm 会对各个 token 执行 $l$ 次归一化计算，保留每个 token $d$ 维嵌入内部的相对大小关系，同时拉近了不同 token 对应特征向量间的距离。与之相比，BN 会消除 $d$ 维特征向量各维度之间的大小关系，破坏了 token 的特征（以下第 2 节会进一步说明这一点）

1. Batch Normalization

BN 对同一 batch 内同一通道的所有数据进行归一化，设输入的 batch data 为 $\pmb{x}$ ，BN 运算如下
$\mathrm{BN}(\mathbf{x})=\boldsymbol{\gamma} \odot \frac{\mathbf{x}-\hat{\boldsymbol{\mu}}_{\mathcal{B}}}{\hat{\boldsymbol{\sigma}}_{\mathcal{B}}}+\boldsymbol{\beta} .$ 其中 $\odot$ 表示按位置乘， $\pmb{\gamma}$ 和 $\pmb{\beta}$ 是 拉伸参数scale 和 偏移参数shift，这两个参数的 size 和特征维数相同，代表着把第 $i$ 个特征的 batch 分布的均值和方差移动到 $\beta^i, \gamma^i$ ， $\pmb{\gamma}$ 和 $\pmb{\beta}$ 是需要与其他模型参数一起学习的参数。 $\hat{\boldsymbol{\mu}}_{\mathcal{B}}$ 和 $\hat{\boldsymbol{\sigma}}_{\mathcal{B}}$ 表示 batch data 中各特征的均值和方差，如下计算
$\begin{aligned} \hat{\boldsymbol{\mu}}_{\mathcal{B}}&=\frac{1}{|\mathcal{B}|} \sum_{\mathbf{x} \in \mathcal{B}} \mathbf{x} \\ \hat{\boldsymbol{\sigma}}_{\mathcal{B}}^{2}&=\frac{1}{|\mathcal{B}|} \sum_{\mathbf{x} \in \mathcal{B}}\left(\mathbf{x}-\hat{\boldsymbol{\mu}}_{\mathcal{B}}\right)^{2}+\epsilon \end{aligned}$ 注意我们在方差估计值中添加一个小的常量 $\epsilon$ ，以确保我们永远不会尝试除以零
注意一些细节
1. 在 MLP 中应用 BN 时，均值和方差的计算发生在各个特征维度上。此时输入数据形式通常为 $\pmb{x}\in\mathbb{R}^{b\times n}$ ，其中 $b=|\mathcal{B}|$ 为 batch size， $n$ 为特征维度，有 $\hat{\boldsymbol{\mu}}_{\mathcal{B}},\hat{\boldsymbol{\sigma}}_{\mathcal{B}}^{2},\pmb{\gamma},\pmb{\beta} \in \mathbb{R}^{1\times n}$
2. 在 CNN 中应用 BN 时，均值和方差的计算发生在各个通道上。此时输入数据形式通常为 $\pmb{x}\in\mathbb{R}^{b\times c\times h\times w}$ ，其中 $b=|\mathcal{B}|$ 为 batch size， $c, h, w$ 分别为为通道数量和图像长宽尺寸，有 $\hat{\boldsymbol{\mu}}_{\mathcal{B}},\hat{\boldsymbol{\sigma}}_{\mathcal{B}}^{2},\pmb{\gamma},\pmb{\beta} \in \mathbb{R}^{1\times c \times 1\times 1}$ ，如下图所示
3. BN 层在”训练模式“（通过小批量统计数据规范化）和“预测模式”（通过数据集统计规范化）中的功能不同。训练过程中，我们无法得知使用整个数据集来估计平均值和方差，所以只能根据每个小批次的平均值和方差不断训练模型；预测模式下，可以根据整个数据集精确计算批量规范化所需的平均值和方差
BatchNorm是一种在深度学习训练中广泛使用的归一化技术，有很多好处，包括正则化效应、减少过拟合、减少对权重初始值的依赖、允许使用更高的学习率等
- 一方面，BN 使每一层隐藏值分布主动居中，并将它们重新调整为学习到的最佳均值和方差，这种操作可能将参数的量级进行了统一，因此直觉上往往被认为可以使优化更加平滑
- 另一方面，BN 有效性的科学性解释一度存在争议。15 年提出BN的论文声称 BN 减小了所谓的 内部协变量偏移internal covariate shift，因此可以提高模型性能，但其分析中假设了每层隐变量值都服从某种正态分布，这个假设过强了，很多后续研究指出了其问题。18 年的论文 How Does Batch Normalization Help Optimization? 认为 BN 的主要作用是使得整个损失函数的 landscape 更为平滑，从而使得我们可以更平稳地进行训练。相关分析可以参考苏神的博文

示例代码参考自《动手学深度学习》7.5 节，适用于全连接层和卷积层，训练过程中使用滑动平均法计算 batch 数据的均值和方差；评估过程中使用最新的均值和方差结果

class BatchNorm(nn.Module):
    # num_features：完全连接层的输出数量或卷积层的输出通道数。
    def __init__(self, num_features, num_dims):
        super().__init__()
        if num_dims == 2: # 全连接层
            shape = (1, num_features)
        else:             # 卷积层
            shape = (1, num_features, 1, 1)
        
        # 参与求梯度和迭代的拉伸和偏移参数，分别初始化成1和0
        self.gamma = nn.Parameter(torch.ones(shape))
        self.beta = nn.Parameter(torch.zeros(shape))
        
        # 非模型参数的变量初始化为0和1
        self.moving_mean = torch.zeros(shape)
        self.moving_var = torch.ones(shape)

    def batch_norm(self, X, gamma, beta, moving_mean, moving_var, eps, momentum):
        if not torch.is_grad_enabled():
            # 如果是在预测模式下，直接使用传入的移动平均所得的均值和方差
            X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps)
        else:
            assert len(X.shape) in (2, 4)
            if len(X.shape) == 2:
                # 使用全连接层的情况，计算特征维上的均值和方差
                mean = X.mean(dim=0)                                       # (num_features,)
                var = ((X - mean) ** 2).mean(dim=0)                        # (num_features,)
            else:
                # 使用二维卷积层的情况，计算通道维上（axis=1）的均值和方差。
                mean = X.mean(dim=(0, 2, 3), keepdim=True)                # (1,num_features,1,1) 保持X的形状，以便后面可以做广播运算
                var = ((X - mean) ** 2).mean(dim=(0, 2, 3), keepdim=True) # (1,num_features,1,1)
                
            # 训练模式下，用当前的均值和方差做标准化
            X_hat = (X - mean) / torch.sqrt(var + eps)
            
            # 更新移动平均的均值和方差
            moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
            moving_var = momentum * moving_var + (1.0 - momentum) * var
            
        Y = gamma * X_hat + beta  # 缩放和移位
        return Y, moving_mean.data, moving_var.data

    def forward(self, X):
        # 如果X不在内存上，将moving_mean和moving_var，复制到X所在显存上
        if self.moving_mean.device != X.device:
            self.moving_mean = self.moving_mean.to(X.device)
            self.moving_var = self.moving_var.to(X.device)
       
        # 保存更新过的moving_mean和moving_var
        Y, self.moving_mean, self.moving_var = self.batch_norm(
            X, self.gamma, self.beta, self.moving_mean,
            self.moving_var, eps=1e-5, momentum=0.9
        )
        return Y

2. Layer Normalization

LN 主要用于 NLP 领域，它对每个 token 的特征向量进行归一化计算。设某个 token 的特征向量为 $\pmb{x}\in \mathbb{R}^d$ ，LN 运算如下
$\mathrm{LN}(\mathbf{x})=\boldsymbol{\gamma} \odot \frac{\mathbf{x}-\hat{\boldsymbol{\mu}}}{\hat{\boldsymbol{\sigma}}}+\boldsymbol{\beta} .$ 其中 $\odot$ 表示按位置乘， $\pmb{\gamma}, \pmb{\beta}\in \mathbb{R}^d$ 和是 拉伸参数scale 和 偏移参数shift，代表着把第 $i$ 个特征的 batch 分布的均值和方差移动到 $\beta^i, \gamma^i$ ， $\pmb{\gamma}$ 和 $\pmb{\beta}$ 是需要与其他模型参数一起学习的参数。 $\hat{\boldsymbol{\mu}}$ 和 $\hat{\boldsymbol{\sigma}}$ 表示特征向量所有元素的均值和方差，如下计算
$\begin{aligned} \hat{\boldsymbol{\mu}}&=\frac{1}{d} \sum_{x^i \in \mathbf{x}} x^i \\ \hat{\boldsymbol{\sigma}}^{2}&=\frac{1}{d} \sum_{x^i \in \mathbf{x}}\left(x^i-\hat{\boldsymbol{\mu}}\right)^{2}+\epsilon \end{aligned}$ 注意我们在方差估计值中添加一个小的常量 $\epsilon$ ，以确保我们永远不会尝试除以零
给定一个长 $l$ 的句子，LN 要进行 $l$ 次归一化计算，之后对每个特征维度施加统一的拉伸和偏移，如下图所示
为什么 LN 比 BN 更适用于 Transformer 类模型呢，这是因为 transformer 模型是基于相似度的，把序列中的每个 token 的特征向量进行归一化有利于模型学习语义，第一步调整均值方差时，相当于对把各个 token 的特征向量缩放到统一的尺度，第二步施加 $γ,β \pmb{\gamma, \beta}$ 时，相当于对所有 token 的特征向量进行了统一的 transfer，这不会破坏 token 特征向量间的相对角度，因此不会破坏学到的语义信息。与之相对的，BN 沿着特征维度进行归一化，这时对序列中各个 token 施加的 transfer 是不同的，破坏了 token 特征向量间的相对角度关系
Transformer 类模型中，LayerNorm 层有两种放置方式
$\text{Pre Norm:} \quad \boldsymbol{x}_{t+1}=\boldsymbol{x}_{t}+F_{t}\left(\operatorname{Norm}\left(\boldsymbol{x}_{t}\right)\right) \\ \text{Post Norm:} \quad \boldsymbol{x}_{t+1}=\operatorname{Norm}\left(\boldsymbol{x}_{t}+F_{t}\left(\boldsymbol{x}_{t}\right)\right)$ 如下图所示

目前比较明确的结论是：同一设置之下，Pre Norm结构往往更容易训练，但最终效果通常不如Post Norm
1. Pre Norm 更容易训练好理解，因为它的恒等路径更突出
2. Pre Norm 中多层叠加的结果更多是增加宽度而不是深度，层数越多，这个层就越“虚”，这是因为 Pre Norm 结构无形地增加了模型的宽度而降低了模型的深度，而我们知道深度通常比宽度更重要，所以是无形之中的降低深度导致最终效果变差了。而 Post Norm 刚刚相反，它每Norm一次就削弱一次恒等分支的权重，所以 Post Norm 反而是更突出残差分支的，因此Post Norm中的层数更加有分量，起到了作用，一旦训练好之后效果更优。详细说明参考为什么Pre Norm的效果不如Post Norm？
过去 BERT 主流的时代往往使用 Post Norm，现在 GPT 时代模型规模都很大，因此大多用 Pre Norm 来稳定训练

3. RMSNorm

RMSNorm 是 LayerNorm 的一个简单变体，来自 2019 年的论文 Root Mean Square Layer Normalization，被 T5 和当前流行 lamma 模型所使用。其提出的动机是 LayerNorm 运算量比较大，所提出的RMSNorm 性能和 LayerNorm 相当，但是可以节省7%到64%的运算
RMSNorm和LayerNorm的主要区别在于RMSNorm不需要同时计算均值和方差两个统计量，而只需要计算均方根 Root Mean Square 这一个统计量，公式如下
$\text{RMSNorm}(\pmb{x})=\boldsymbol{\gamma} \odot\frac{\pmb{x}}{\operatorname{RMS}(x)} \quad \text{where \quad}\operatorname{RMS}(x)=\sqrt{\frac{1}{d} \sum_{x^i \in \mathbf{x}} x_{i}^{2} + \epsilon}$
论文 Do Transformer Modifications Transfer Across Implementations and Applications? 中做了比较充分的对比实验，显示出RMS Norm的优越性。一个直观的猜测是，计算均值所代表的 center 操作类似于全连接层的 bias 项，储存到的是关于预训练任务的一种先验分布信息，而把这种先验分布信息直接储存在模型中，反而可能会导致模型的迁移能力下降

下面给出 Transformer Lamma 源码中实现的 RMSNorm

class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        LlamaRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)