【Transformer】Normalization

嗜睡的篠龙

已于 2024-09-02 22:04:10 修改

阅读量668

点赞数 20

分类专栏： AI大模型技术积累技术分享文章标签： transformer pytorch 人工智能

于 2024-09-02 00:40:49 首次发布

本文链接：https://blog.csdn.net/weixin_43799388/article/details/141792358

版权

AI大模型技术积累同时被 2 个专栏收录

3 篇文章 0 订阅

订阅专栏

技术分享

3 篇文章 0 订阅

订阅专栏

文章目录

Normalization概述

Normalization，也就是归一化（或者叫标准化）操作，已经广泛应用于神经网络模型中，成为必不可少的组件。各种Normalization操作，其主要的作用如下 [1]：

保证训练稳定性：当训练更深的神经网络时，计算反向传播的参数往往是指数级变化，太大或者太小的参数送入激活函数后很容易造成梯度爆炸或者梯度消失，因此Normalization就使得每一层激活函数的输入，可以控制在近似理想的数值区间中，使得训练更加稳定。
加速模型收敛：神经网络的每一层都在拟合一个数据分布，如果不进行 Normalization，那么每次的输入分布可能随时都在变化，这样学习起来就很困难。在 Normalization之后，绝大部分数值都集中在了一个可接受的范围内，每一层的参数就安心拟合这个分布就好了，这个范围正好又是激活函数的“舒适区”，所以模型收敛速度会更快。
减少模型参数初始化的影响：这是Normalization的额外好处，因为早期神经网络参数初始化对模型效果的影响还是挺大的，使用Normalization之后，可以让模型的训练不再那么依赖权重的初始化。

主流Norm的统一公式：
$\rm{Norm}(\bold{X}) = \frac{\bold{X}-\mu}{\sigma} * \gamma + \beta$
其中， $\mu$ 是 $\bold{X}$ 按照某一维度的均值， $\sigma$ 是归一化的分母，在不同Norm方法中，他的形式不同，比如在LayerNorm中使用的是标准差，而在WeightNorm中使用的是L2范数。 $\gamma$ 和 $\beta$ 是可学习的参数，对归一化后的数据进行仿射变换（缩放和平移），使得归一化后的数据分布，更符合激活函数中的非线性变化，而不是在近似线性变换区域，如下图所示（仅做理解帮助）[2] ：
在这里插入图片描述

三种活跃在面试中的normalization方法

几个思考题

为了更深入理解不同Norm方法的实现方式，让我们先来思考几个问题，【】里面给出了对应的答案，具体分析放到后面各部分的讲解中：

对于形状为[2, 3, 16, 16]的图片输入（每个维度分别表示batch，channel，height，width）进行BN操作时：

计算的 $\mu$ 和 $\sigma$ 个数是多少？【3】
可训练参数 $\gamma$ 和 $\beta$ 的维度是多少？【3】

对于形状为[2, 10, 4]的文本序列输入（每个维度分别表示batch，sequence，dimension）进行LN操作时：

计算的 $\mu$ 和 $\sigma$ 个数是多少？【2*10=20】
可训练参数 $\gamma$ 和 $\beta$ 的维度是多少？【4】

在这里插入图片描述

BN（Batch Norm）

BN方法是指批归一化（Batch Normalization），这是一种用于训练深度神经网络的技术，尤其是应用于计算机视觉领域，旨在提高训练速度、稳定性以及减少对模型参数初始化的敏感度。BN通过减少内部协变量偏移（internal covariate shift）的影响来改善模型的表现，即在训练过程中减少各层输入分布的变化。

BN计算步骤

在每个特征维度上，计算batch内所有数据上的均值（mean）：
$\mu_B=\frac{1}{m} \sum_{i=1}^m x_i$
其中， $m$ 是批次内样本的数量， $x_i$ 是样本数据
同样在每个特征维度上，计算batch内所有数据上的方差（variance）：
$\sigma_B^2=\frac{1}{m} \sum_{i=1}^m\left(x_i-\mu_B\right)^2$
对每个特征做标准化（normalization）：
$\hat{x}_i=\frac{x_i-\mu_B}{\sqrt{\sigma_B^2+\epsilon}}$
其中， $\epsilon$ 是为了防止除以零的微小常数
最后，对每个特征进行缩放（scaling）和平移（shifting），以保持网络表示的多样性：
$y_i=\gamma \hat{x}_i+\beta$
其中， $\gamma$ $和 $ $\beta$ 是可学习的参数，分别代表缩放和平移

分析一

从上述计算步骤可以看出，假设我们的输入图像有3个特征通道，BN首先计算第一个特征通道上，batch内所有数据的均值，然后再依次计算第二个和第三个特征通道上的均值，方差同理。也就是说，当前BN层有几个特征通道，就会计算几个均值和方差，所以torch中的BN层初始化时，我们通常只传入一个参数，就是num_features，也就是当前输入数据中的特征通道数。
在这里插入图片描述

所以思考题1中的第一个问题，对于形状为[2, 3, 16, 16]的图片输入（每个维度分别表示batch，channel，height，width）进行BN操作时，计算的 $\mu$ 和 $\sigma$ 个数都是3

分析二

从上述计算步骤以及分析一可以看出，BN操作是对每个特征通道上的所有数据进行归一化，那么在进行仿射变换时，也应该对每个特征通道上的数据进行缩放和平移，因此思考题1中的第二个问题，答案也是3。从torch.nn.BatchNorm2d的实现中也可以看出，两个可学习参数 $\gamma$ 和 $\beta$ 初始化时，只传入了参数num_features
在这里插入图片描述
但是要注意，在最后真正计算时，要把维度view到和输入数据一致的shape，然后进行点乘和点加计算。

训练和推理的不同处理

由于BN计算方式的局限性，训练的时候有大批的数据可以组成 batch，然后计算所有batch数据的均值和方差，但是推理的时候，如果只有一个样本，那BN操作的意义就不大了。所以在推理的时候实际上是采用了训练时候的数据分布来进行归一化的。代码层面就是运行时均值和方差的累积保存了。这样的话就必须要保证训练和预测的分布必须一致，也导致BN的泛化能力没那么强。

同步BN

单机多卡分布式训练时，不同GPU上的数据不同，计算得到的均值和方差也不一致，越大的batch统计得到的均值方差越符合整体数据集的均值方差，如果进行多卡实验时（这里以两张卡为例）整体的batch是4，那么单卡的batch是2，就很有必要同步bn。pytorch中的实现为SyncBatchNorm类，如果模型包含bn层，就可以调用下面函数把所有bn层修改成syncbn：model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)。

BN的Pytorch代码实现

import torch
import torch.nn as nn

class BatchNorm2d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super(BatchNorm2d, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum

        # 初始化可训练参数
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))

        # 初始化运行时均值和方差
        self.running_mean = torch.zeros(num_features)
        self.running_var = torch.ones(num_features)

        self.training = True

    def forward(self, x):
        # x.shape = [bs, c, h, w]
        if self.training:
            mean = x.mean(dim=(0, 2, 3), keepdim=True)  # [16]
            var = x.var(dim=(0, 2, 3), keepdim=True)

            # 更新运行时均值和方差
            self.running_mean = (1 - self.momentum) * \
                self.running_mean + self.momentum * mean
            self.running_var = (1 - self.momentum) * \
                self.running_var + self.momentum * var
        else:
            mean = x.mean(dim=(0, 2, 3))
            var = x.var(dim=(0, 2, 3))

        # 归一化
        x_norm = (x - mean) / \
            (torch.sqrt(var) + self.eps)
        output = self.gamma[None, ..., None, None] * \
            x_norm + self.beta[None, ..., None, None]

        return output

if __name__ == '__main__':
    bn = BatchNorm2d(num_features=16)

    x = torch.rand(32, 16, 8, 8)

    res = bn(x)
    print(res.shape)

LN（Layer Norm）

LN是层归一化（Layer Norm）方法，起初主要用于处理NLP问题，尤其是应用在基于transformer的模型中。根据LN所处的位置不同，也可以分为pre-LN和post-LN，这个我们到后面再详细分析。

分析一

LN的计算形式与BN一样，只是计算的维度不同。LN是对每个样本计算所有特征的均值和方差，所以有多少个样本，就会有多少个均值和方差需要计算。为了直观理解，我们直接来看pytorch源码中的示例：

在这里插入图片描述
可以看到对于文本序列和图像，LN的样本定义是不同的：

对于文本来说，每个词，是一个样本，比如上面这个NLP Example，LN会对所有的词（batch * sentence_length = 100个），在embedding_dim维度上计算均值和方差
对于图像来说，每张图像，是一个样本，比如上面这个Image Example，LN会对所有的图片（batch = 20个），在[C, H, W]维度上计算均值和方差

所以思考题2中的第一个问题，对于形状为[2, 10, 4]的文本序列输入（每个维度分别表示batch，sequence，dimension）进行LN操作时，计算的 $\mu$ 和 $\sigma$ 个数都是 2*10=20 个。

分析二

LN中的仿射变换和BN类似，是对所有的特征（比如上图这个例子，每个词有10个特征）进行缩放和平移的，对于每一个特征通道，都要有一组缩放和平移的参数，因此思考题2中的第二个问题，对于形状为[2, 10, 4]的文本序列输入（每个维度分别表示batch，sequence，dimension）进行LN操作时，可训练参数 $\gamma$ 和 $\beta$ 的维度就是特征通道数，也就是4。

LN的Pytorch代码实现

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-5):
        super(LayerNorm, self).__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(normalized_shape))
        self.beta = nn.Parameter(torch.zeros(normalized_shape))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        std = x.std(dim=-1, keepdim=True)
        y = (x - mean) / (std + self.eps)
        output = y * self.gamma + self.beta
        return output

RMSNorm（RMS Norm）

RMSNorm（Root Mean Square Normalization）是一种较新的归一化技术，它基于层归一化（Layer Normalization）的变体。RMSNorm的主要特点是它只使用激活函数的均方根（RMS）来进行归一化，而不是使用均值和标准差。

RMSNorm的具体实现步骤

计算均方根：对于输入张量 x，计算其每个样本在所有特征维度上的平方和的均值，然后取平方根
$\operatorname{RMS}(x)=\sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2}$
其中d是特征维度
进行归一化：使用计算得到的RMS值对输入张量进行归一化
$\hat{x}_i=\frac{x_i}{\operatorname{RMS}(x)}$
缩放和平移：通常，RMSNorm还会包括一个可学习的缩放参数 γ和一个平移参数 β（可以去掉），以增加模型的表达能力
$y_i=\gamma \hat{x}_i+\beta$

RMSNorm的Pytorch代码实现

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-8):
        super(RMSNorm, self).__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        # 计算均方根RMS
        rms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        # 归一化
        x_norm = x * rms
        # 缩放和平移
        return x_norm * self.gamma

三种Norm的Pytorch实现方式

这里给出了BN和LN的两种实现方式：手动实现以及调用torch.nn库中的函数，并且通过实验证明，手动实现方式输出结果与调用函数库输出结果一致：

import torch
import torch.nn as nn


def batchnorm2d(x: torch.Tensor, eps=1e-5):
    '''
    x_bn = (x - mean) / sqrt(var + eps)
    y = x_bn * gamma + beta
    1. bn操作，对batch中所有数据，按照特征通道维度求平均/方差，有几个特征通道，就有几个平均值/方差值
    2. bn操作后，需要对结果再进行缩放和平移，对应为可学习参数gamma和beta，维度是输入向量的特征通道数
    '''
    gamma = nn.Parameter(torch.ones(
        x.shape[1]))  # 可学习参数gamma，维度为输入向量中特征维度，beta同
    beta = nn.Parameter(torch.zeros(x.shape[1]))

    # 按照特征维度进行求平均，同时保持输入shape
    mean = torch.mean(x, dim=(0, 2, 3), keepdim=True)
    # 这里设置unbiased为false，表示不使用无偏估计
    var = torch.var(x, dim=(0, 2, 3), keepdim=True, unbiased=False)

    y = (x - mean) / torch.sqrt(var + eps)
    # 注意这里是对特征进行仿射变换，所以gamma和beta的维度和特征维度一致
    output = y * gamma.view(1, -1, 1, 1) + beta.view(1, -1, 1, 1)
    return output


def layernorm(x: torch.Tensor, eps=1e-5):
    gamma = nn.Parameter(torch.ones(x.shape[-1]))
    beta = nn.Parameter(torch.zeros(x.shape[-1]))
    mean = torch.mean(x, dim=-1, keepdim=True)
    var = torch.var(x, dim=-1, keepdim=True, unbiased=False)
    y = (x - mean) / torch.sqrt(var + eps)
    output = y * gamma + beta  # 注意这里是对特征进行仿射变换，所以gamma和beta的维度和特征维度一致
    return output


class rmsnorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super(rmsnorm, self).__init__()

        self.w = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        # rms = torch.mean(x ** 2, dim=-1, keepdim=True)
        # y = x / torch.sqrt(rms + self.eps)
        y = x * torch.rsqrt(torch.mean(x ** 2, dim=-1,
                            keepdim=True) + self.eps)
        output = y * self.w
        return output


if __name__ == '__main__':
    bn_input = torch.rand((2, 3, 4, 4))  # [bs, c, h, w]
    ln_input = torch.rand((2, 3, 4))  # [bs, s, d]

    print('*' * 40)
    manual_bn_output = batchnorm2d(bn_input)
    print(manual_bn_output)
    bn = nn.BatchNorm2d(3)
    torch_bn_output = bn(bn_input)
    print(torch_bn_output)

    print('*' * 40)
    manual_ln_output = layernorm(ln_input)
    print(manual_ln_output)
    ln = nn.LayerNorm(4)
    torch_ln_output = ln(ln_input)
    print(torch_ln_output)

    print('*' * 40)
    rms = rmsnorm(dim=4)
    manual_rmsn_output = rms(ln_input)
    print(manual_rmsn_output)

参考资料

[1] https://note.mowen.cn/note/detail?noteUuid=6zX7ET1-uPTxMOLXOasIh
[2] https://www.bilibili.com/video/BV1L2421N7jQ/?share_source=copy_web&vd_source=79b1ab42a5b1cccc2807bc14de489fa7
[3] https://blog.csdn.net/qq_37668436/article/details/121997884

嗜睡的篠龙

关注

20
点赞
踩
4

收藏

觉得还不错? 一键收藏
打赏
0
评论
【Transformer】Normalization

Normalization，也就是归一化（或者叫标准化）操作，已经广泛应用于神经网络模型中，成为必不可少的组件。保证训练稳定性：当训练更深的神经网络时，计算反向传播的参数往往是指数级变化，太大或者太小的参数送入激活函数后很容易造成梯度爆炸或者梯度消失，因此Normalization就使得每一层激活函数的输入，可以控制在近似理想的数值区间中，使得训练更加稳定。
复制链接

扫一扫