大语言模型中的归一化技术：LayerNorm与RMSNorm的深入研究

在LLama等大规模Transformer架构的语言模型中，归一化模块是构建网络稳定性的关键组件。本文将系统分析归一化技术的必要性，并详细阐述为何原始Transformer架构中的LayerNorm在LLama模型中被RMSNorm所替代的技术原理。

归一化技术的基础原理

归一化的核心定义

归一化（Normalization）是一种将数据映射到特定数值区间的数学变换技术，旨在提升计算稳定性并优化学习效率。归一化的本质在于调整数据的量纲规模而保持分布形态不变。

归一化前后数据分布

归一化前后按特征划分的数据分布

通过对比归一化前后的数据分布可以明确观察到，尽管数据点的相对位置关系保持恒定，但数值范围发生了显著变化。如上图所示，横轴上原本分布在30至70区间的数据经归一化后映射至0到1区间，然而数据的拓扑分布结构保持不变。这构成了归一化技术的基本理论框架。

归一化技术的技术必要性

归一化技术在深度学习中具有双重技术价值。首要方面是提高数值计算稳定性和优化收敛速度。当数据量级差异显著时，评估指标（如均方误差MSE）可能产生数量级上的波动。在MSE计算过程中，大值输入会导致梯度幅值过大，从而使优化过程不稳定。

以均方误差计算为例，考虑两组输入值a和b的对比情况：

1. 未经归一化处理:

• a = [100, 200, 300], b = [110, 190, 310]
• MSE = (1/3) * [(100–110)² + (200–190)² + (300–310)²]
• MSE = (1/3) * [100 + 100 + 100] = 100

1. 经归一化处理:

• 归一化后的值：a_normalized = [1.0, 2.0, 3.0], b_normalized = [1.1, 1.9, 3.1]
• MSE = (1/3) * [(1.0–1.1)² + (2.0–1.9)² + (3.0–3.1)²]
• MSE = (1/3) * [0.01 + 0.01 + 0.01] = 0.01

从计算结果可见，未归一化的数据由于数值较大，产生了较大的MSE值，这可能导致梯度幅值过大，使模型训练过程不稳定。而归一化处理后，MSE值显著降低，梯度幅值维持在合理区间，有助于参数更新过程的稳定性。

通过将数据规范化到一致的数值范围，归一化能够有效控制梯度幅值，实现更为平稳的参数优化过程。

归一化的另一技术价值在于潜在的内存效率提升。归一化通过缩减数据表示范围间接提高了内存利用效率。特别是在采用8位或16位浮点格式进行训练时，数据范围的收窄使得低精度表示更为精确。这种特性可通过量化技术或FP16计算进一步降低模型训练阶段的内存占用。

需要说明的是，归一化本身并不直接减少内存占用，它需要与量化或低精度计算技术协同应用才能实现有效的内存优化。

归一化的实证效果分析

通过实际数据集分析可验证归一化的技术效果。以加州住房数据集为例，采用归一化技术的模型（蓝线）展现出快速且稳定的学习曲线，而未归一化的模型（红线）则表现出学习不稳定性。这种差异在训练过程中尤为明显，未归一化情况下可能出现数值溢出，导致损失函数出现异常峰值。

归一化前后MSE比较——加州住房数据集

归一化前后准确率性能比较——CIFAR-10

在CIFAR-10图像分类任务中，类似的数值稳定性问题同样存在。应用归一化技术后，学习过程能够迅速达到稳定状态。因此，归一化已成为现代深度学习的基础技术组件，对于提升模型性能至关重要。

层归一化（Layer Normalization）

层归一化最初由Jimmy Lei Ba、Jamie Ryan Kiros和Geoffrey Hinton在2016年发表的论文"Layer Normalization"中提出。该研究提出LayerNorm作为批量归一化的替代方案，以解决其局限性，特别强调了其在序列数据和循环神经网络模型中的适用性。

层归一化通过计算每个样本的均值（μ）和标准差（σ）来实现数据归一化。由于LayerNorm是在每个样本基础上独立应用的，不依赖于批量大小，因此在处理序列数据的模型中表现尤为出色。Transformer模型作为序列数据处理的核心架构，广泛采用了LayerNorm技术。

在Transformer架构中，LayerNorm扮演着至关重要的角色。Transformer处理的序列数据由多个token组成，LayerNorm所归一化的"样本"实际上是每个token位置的嵌入向量或隐藏状态。

例如，对于句子"The cat sat on the mat"，Transformer模型会为每个单词（token）生成特定的向量表示（嵌入）。LayerNorm在此场景下归一化的是该token位置的嵌入向量或隐藏状态。LayerNorm的核心机制是对每个token位置计算的嵌入或隐藏状态向量进行归一化，以维持特定范围内的数值稳定性。

其中：

• x表示每个样本的输入值
• μ表示平均值
• σ表示标准差
• ϵ是一个微小常数，用于确保数值稳定性，防止除零错误，通常设定为0.00001或0.000001

LayerNorm将数据归一化为均值为0，方差为1的分布。

Transformer输入示例和LayerNorm计算过程

LayerNorm应用结果

从上图可观察到，每个token的嵌入维度为5，均值和方差计算在嵌入向量的5维值空间内进行。值得注意的是，这一计算过程独立于其他数据点或token。另一关键特性是数据分布向均值0对齐，这一特性与RMSNorm形成明显对比。

RMS归一化（RMS Normalization）

RMS归一化（RMSNorm）是层归一化的变体，其特点是仅使用均方根（Root Mean Square）进行归一化，省略了均值计算环节。RMSNorm在降低计算成本的同时能够维持模型性能，同时提供更加稳定的学习过程。

RMS(x)的定义如下：

其中RMS代表特征值的均方根。

Transformer输入示例和RMSNorm计算过程

从结果可以观察到，RMSNorm不执行均值对齐操作。由于它没有执行x-μ运算，因此不存在向0对齐的过程。与传统LayerNorm需要两次数据扫描（一次计算均值，一次计算方差）不同，RMSNorm仅需一次处理即可完成RMS计算。这种简化方法在保持性能的同时有效降低了计算资源需求。

LayerNorm与RMSNorm的技术对比

均值对齐特性分析

与LayerNorm显著不同的是，RMSNorm保持非零均值分布，不执行均值对零的对齐操作。

通过对比LayerNorm和RMSNorm应用到前述Transformer输入示例后的均值分布结果，差异变得更为明显。尽管两种方法都将数据分布收敛到特定区域，LayerNorm呈现零均值特性，而RMSNorm则表现为非零均值分布。

从梯度传播视角的比较分析

RMSNorm通过保持非零均值而不强制均值对零的对齐，有效缓解了梯度消失问题。LayerNorm的零均值对齐可能导致梯度在反向传播过程中逐层衰减。

LayerNorm与RMSNorm：跨Epoch的梯度稳定性对比

上图基于简化网络模型分析了LayerNorm和RMSNorm的梯度传播效果。左侧图表展示了按训练epoch划分的LayerNorm和RMSNorm的平均梯度范数变化。梯度范数是衡量网络权重变化速率的指标，具体表示反向传播过程中各层权重梯度的幅值。较大的梯度范数值表示权重变化幅度较大，而过大或过小的值可能导致学习不稳定。

RMSNorm相比LayerNorm维持了更高的梯度范数，尤其在训练初期阶段。这一特性对于防止深度学习模型中的梯度消失问题至关重要。右侧图表显示了RMSNorm与LayerNorm之间的梯度比率，初期RMSNorm的梯度显著高于LayerNorm，随着学习进程的推进，这一差异逐渐减小。

与LayerNorm相比，RMSNorm在计算效率和训练稳定性方面表现卓越。网络结构越深，RMSNorm通过有效防止梯度消失问题实现更稳定的学习过程，同时能够在减少计算资源消耗的条件下达成相似的性能水平。

结论

本文深入分析了原始Transformer架构中采用的层归一化技术与LLama架构中实现的RMS归一化技术之间的技术差异与特性。LayerNorm通过调整数据的均值和方差确保数值稳定性，在各类序列模型中发挥关键作用。而RMSNorm作为一种无需利用均值即可稳定数据尺度的技术，能够有效解决深层网络中的梯度消失问题，同时保持较快的训练收敛速度。

这两种归一化技术在现代深度学习模型中均发挥着不可替代的作用，对于最大化提升模型性能（尤其是在Transformer架构中）至关重要。RMSNorm在计算效率和训练稳定性方面对LayerNorm形成了有效补充，为提升复杂深层模型性能做出了重要贡献。

https://avoid.overfit.cn/post/224e11d8a7d84870b1a3d5e7ea410a35

作者：Hugman Sangkeun Jung