【大模型】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]
import torchclass RMSNorm(torch.nn.Module): # nn.Module是所有层的父类，层元素就必须继承nn.Moduledef __init__(self, dim, eps): # 用于储存层的元素super().__init__()self.weight = torch.nn.Parameter(torch.ones(dim)) # 初始化权重参数self.eps = eps # 防止根号下为0def _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)：最后用求得的导数乘以xdef forward(self, x): # 数据流output = self._norm(x.float().type_as(x)) # 将x变成浮点数进行归一化，并保持x原始的数据类型return output * self.weight # 将归一化后的输出乘以可学习的参数 weight,调整每一个维度的缩放if __name__ == '__main__':batch_size = 1dim = 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)
