引言
在学习完GPT2之后，从本文开始进入Llama模型系列。
本文介绍Llama模型的改进之RMSNorm(均方根层归一化)。它是由Root Mean Square Layer Normalization论文提出来的，可以参阅其论文笔记1。
LayerNorm
层归一化(LayerNorm)对Transformer等模型来说非常重要，它可以帮助稳定训练并提升模型收敛性。LayerNorm针对一个样本所有特征计算均值和方差，然后使用这些来对样本进行归一化：
μ = 1 H ∑ i = 1 H x i , σ = 1 H ∑ i = 1 H ( x i − μ ) 2 , N ( x ) = x − μ σ , h = g ⊙ N ( x ) + b (1) \mu = \frac{1}{H}\sum_{i=1}^H x_i,\quad \sigma = \sqrt{\frac{1}{H}\sum_{i=1}^H (x_i - \mu)^2}, \quad N(\pmb x) = \frac{\pmb x-\mu}{\sigma},\quad \pmb h = \pmb g \,\odot N(\pmb x) + \pmb b \tag 1 μ=H1i=1∑Hxi,σ=H1i=1∑H(xi−μ)2 ,N(x)=σx−μ,h=g⊙N(x)+b(1)
这里 x = ( x 1 , x 2 , ⋯ , x H ) \pmb x = (x_1,x_2,\cdots, x_H) x=(x1,x2,⋯,xH)表示某个时间步LN层的输入向量表示，向量维度为 H H H； h \pmb h h实LN层的输出； g , b \pmb g,\pmb b g,b实两个可学习的参数。
为什么层归一化有用？一些解释如下2：
减少内部协变量偏移（Internal Covariate Shift）： 内部协变量偏移是指在深度神经网络的训练过程中，每一层输入的分布会发生变化，导致网络的训练变得困难。层归一化通过对每一层的输入进行归一化处理，可以减少内部协变量偏移，使得每一层的输入分布更加稳定。 稳定化梯度： 层归一化有助于保持每一层输出的均值和方差稳定，从而使得梯度的传播更加稳定。这有助于减少梯度消失或梯度爆炸的问题，提高梯度在网络中的流动性，加快训练速度。 更好的参数初始化和学习率调整： 通过层归一化，每一层的输入分布被归一化到均值为0、方差为1的标准正态分布，这有助于更好地初始化网络参数和调整学习率。参数初始化与学习率调整的稳定性对模型的训练效果至关重要。 增强模型的泛化能力： 层归一化可以减少网络对训练数据分布的依赖，降低了过拟合的风险，从而提高模型的泛化能力。稳定的输入分布有助于模型更好地适应不同数据集和任务。RMSNorm
虽然LayerNorm很好，但是它每次需要计算均值和方差。RMSNorm的思想就是移除(1)式中 μ \mu μ的计算部分1：
x ˉ i = x i RMS ( x ) g i RMS ( x ) = 1 H ∑ i = 1 H x i 2 (2) \bar x_i = \frac{x_i }{ \text{RMS}(\pmb x)} g_i \quad \text{RMS}(\pmb x) =\sqrt{\frac{1}{H} \sum_{i=1}^H x_i^2} \tag 2 xˉi=RMS(x)xigiRMS(x)=H1i=1∑Hxi2 (2)
同时在实现也可以移除平移偏置 b \pmb b b。
单看(2)式的话，相当于仅使用 x \pmb x x的均方根来对输入进行归一化，它简化了层归一化的计算，变得更加高效，同时还有可能带来性能上的提升。
实现
RMSNorm的实现很简单：
import torch import torch.nn as nn from torch import Tensor class RMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(hidden_size)) def _norm(self, hidden_states: Tensor) -> Tensor: variance = hidden_states.pow(2).mean(-1, keepdim=True) return hidden_states * torch.rsqrt(variance + self.eps) def forward(self, hidden_states: Tensor) -> Tensor: return self.weight * self._norm(hidden_states.float()).type_as(hidden_states) torch.rsqrt是torch.sqrt的倒数；eps是一个很小的数，防止除零；hidden_states.float()确保了标准差计算的精确度和稳定性，然后在forward方法中，通过.type_as(hidden_states)将结果转换回原来的数据类型，以保持与输入张量相同的数据类型，使得归一化处理后的结果与输入数据类型一致。
下面通过一个简单的网络来测试一下：
import torch import torch.nn as nn from torch import Tensor class SimpleNet(nn.Module): def __init__(self): super(SimpleNet, self).__init__() self.linear = nn.Linear(in_features=10, out_features=5) self.rmsnorm = RMSNorm(hidden_size=5) def forward(self, x): x = self.linear(x) x = self.rmsnorm(x) return x net = SimpleNet() input_data = torch.randn(2, 10) # 2个样本，每个样本包含10个特征 output = net(input_data) print("Input Shape:", input_data.shape) print("Output Shape:", output.shape) Input Shape: torch.Size([2, 10]) Output Shape: torch.Size([2, 5]) 参考
[论文笔记]Root Mean Square Layer Normalization ↩︎ ↩︎
批归一化和层归一化 ↩︎
tpu泛化能力llama稳定性gpt数据集神经网络transformer训练数据样本包
