【大模型】RMS Normalization原理及实现

1.RMS Normalization的原理

说RMS Normalization之前，先讲Layer Normalization 和 Batch Normalization。

BN和LN是两种常见的归一化方法。它们的目的是帮助模型训练更稳定、收敛更快。BN是在Batch维度上进行归一化，即对同一batch中每个特征维度的值进行归一化。LN则是在层的维度上进行归一化，即对每一个样本的特征进行归一化。

RMS Normalization属于LN。

再来说RMS Normalization和Layer Normalization。

Layer Normalization：利用均值和方差对特征进行归一化。

RMS Normalization：利用均方根对特征进行归一化。

LLaMA架构中采用RMS Normalization的原因是通过只计算均方根，从而减少计算量，同时在实验中也确实获得了更加稳定的训练。

在这里插入一点NLP任务中，对于将特征进行"归一化"目的的一些个人小理解：在NLP中，使用Layer Normalization进行归一化是为了使输入特征在每一层的神经元中保持稳定的分布，避免特征值之间出现过大的波动。通过归一化，Layer Normalization 将特征重新调整为均值为 0、方差为 1 的分布，从而让模型的训练更加稳定和高效，使得数据变得更加**"平滑"** 。这里的"平滑"是指数值的尺度更一致、更稳定 ，不会有特别大的数值差异，能够防止特征值在网络层中传递时变得过大或过小。这种一致性有助于缓解模型训练中的一些问题，如梯度爆炸 或梯度消失 ，并能让模型更容易优化。在使用RMS Normalization进行归一化则是直接使特征本身的数值变得更加"平滑"。

2.RMS Normalization公式

2.RMS Normalization的实现

该函数在神经网络中需要对输入的数据进行处理，再输出相应的处理好的数据，对应的实现方式就用层来实现。

因为RMS Normalization属于LN，所以，x-->[batch_size, hidden_states]

python 复制代码

import torch


class RMSNorm(torch.nn.Module):  # nn.Module是所有层的父类，层元素就必须继承nn.Module
    def __init__(self, dim, eps):  # 用于储存层的元素
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(dim))  # 初始化权重参数
        self.eps = eps  # 防止根号下为0

    def _norm(self, x):  # 定义类函数里的方法（"_"表示只在该类的内部调用）
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        # x.pow(2)：求平方
        # x.pow(2).mean(-1, keepdim=True)：所有的平方求一个均值
        # x.pow(2).mean(-1, keepdim=True) + self.eps：加上一个防止根号下为0的元素
        # torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)：开平方再求导
        # rsqrt(x) = 1 / sqrt(x)
        # x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)：最后用求得的导数乘以x

    def forward(self, x):  # 数据流
        output = self._norm(x.float().type_as(x))  # 将x变成浮点数进行归一化，并保持x原始的数据类型
        return output * self.weight  # 将归一化后的输出乘以可学习的参数 weight,调整每一个维度的缩放


if __name__ == '__main__':

    batch_size = 1
    dim = 4  # 特征维度
    x = torch.Tensor([0.1, 0.1, 0.2, 0.3])
    # 初始化RMSNorm对象
    rms_norm = RMSNorm(dim=dim, eps=0)
    output = rms_norm(x)

    print("输入数据: \n", x)
    print("RMSNorm输出数据: \n", output)