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 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)