详解大模型中的LayerNorm和RMSNorm
详解大模型中的LayerNorm和RMSNorm
在深度学习领域,归一化技术是提高模型训练效率和稳定性的关键手段。本文将深入探讨Layer Normalization和RMSNorm这两种归一化方法在Transformer模型中的应用,帮助读者理解它们的原理、差异和实际效果。
一、背景
在像LLama这样的基于大型Transformer的语言模型中,Layer Normalization(LayerNorm)模块嵌入在整个网络中。本文将检查为什么需要归一化,并详细说明为什么原始Transformer的LayerNorm被替换为RMSNorm。
二、什么是归一化?
首先,让我们澄清什么是归一化。归一化是一种技术,它将数据归一化到一定范围内,以提高数值稳定性和学习效率。关键点是归一化调整数据的尺度而不改变其形状。
我们看应用归一化前后的数据,可以看到虽然每个数据点的相对位置得以保持,但只有尺度发生了变化。例如,在上图中,水平轴上介于30到70之间的数据在归一化后变为介于0和1之间,但数据的相对分布保持不变。这就是归一化的基本概念。
三、为什么我们需要归一化?
归一化的第一个原因是提高数值稳定性和学习速度。例如,当两个值之间的差异较小时与较大时,评估值(例如MSE,均方误差)可能会有很大差异。在MSE中,如果出现较大的值,梯度可能会变得极大,从而使学习变得不稳定。
示例:
假设我们有两个输入值a和b,并计算它们的MSE。
- 未经归一化:
- a = [100, 200, 300], b = [110, 190, 310]
- MSE = (1/3) * [(100–110)^2 + (200–190)^2 + (300–310)^2]
- MSE = (1/3) * [100 + 100 + 100] = 100
- 经过归一化: 首先对值进行归一化:
- a_normalized = [1.0, 2.0, 3.0], b_normalized = [1.1, 1.9, 3.1]
- MSE = (1/3) * [(1.0–1.1)^2 + (2.0–1.9)^2 + (3.0–3.1)^2]
- MSE = (1/3) * [0.01 + 0.01 + 0.01] = 0.01
从上述结果可以看出,未经归一化时,较大的值会导致较大的MSE值,这可能会导致较大的梯度,从而使学习变得不稳定。然而,当应用归一化时,MSE值变小,梯度保持稳定。当通过归一化在一致的范围内处理数据时,梯度的大小也保持稳定,这使得在学习过程中可以更稳定地更新参数。
第二个原因是可能节省内存。归一化可以通过缩小数据范围间接提高内存效率。例如,当使用8位或16位浮点格式进行训练时,如果数据在较小的范围内,它可以使用较低的位宽表示。这可以通过量化或FP16操作减少模型训练期间的内存使用。(然而,仅归一化本身并不能立即节省内存。它需要与量化或低比特操作技术结合使用,才能看到实际的内存减少效果。)
四、归一化的现实世界示例
我们可以通过像加利福尼亚州住房数据集这样的真实数据来验证归一化的效果。应用了归一化的模型(蓝色)能够快速稳定学习,而没有应用归一化的模型(红色)学习不稳定。这种差异在训练过程中尤为明显,没有归一化时,可能会出现过大的值,导致损失峰值。
在处理图像的CIFAR-10数据集中也出现了类似的数值稳定性问题。当应用归一化时,学习能够非常快速地稳定下来。因此,归一化作为现代深度学习中的基本技术,必须应用以提高性能。
在这次讲座中,我们将重点介绍Layer Normalization和RMS归一化这两种归一化技术。特别是,我们将检查为什么这些归一化技术在像Transformer这样的深度神经网络结构中很重要,以及它们各自是如何工作的。
五、Layer Normalization
Layer Normalization最初是在Jimmy Lei Ba、Jamie Ryan Kiros和Geoffrey Hinton于2016年发表的论文“Layer Normalization”中提出的。这篇论文提出了LayerNorm作为克服Batch Normalization限制的替代方案,强调了其适用于序列数据和RNN模型。
Layer Normalization是一种通过计算每个样本的均值(μ)和标准差(σ)来归一化数据的技术。LayerNorm特别适用于处理序列数据的模型,因为它是在每个样本的基础上应用的,与批量大小无关。由于Transformer模型处理序列数据,LayerNorm在Transformer中得到了广泛的应用。
LayerNorm在Transformer模型中扮演着关键角色。Transformer是处理序列数据的结构,其中输入数据由多个标记组成。在Transformer模型中,LayerNorm处理的“样本”是每个标记位置的嵌入向量或隐藏状态。
例如,当一个句子是“The cat sat on the mat”时,Transformer模型为每个词(标记)生成一个独特的向量(嵌入)。在这里,LayerNorm归一化该标记位置的嵌入向量或隐藏状态。换句话说,LayerNorm的基本概念是归一化在每个标记位置计算的嵌入或隐藏状态向量,以保持在一定范围内的稳定性。
μ是均值。
σ是标准差。
ϵ是一个小的常数,用于数值稳定性,以防止除以零的错误。通常使用0.00001或0.000001。
LayerNorm将数据归一化为均值为0和方差为1。
从上图可以看出,每个标记的嵌入维度为5,并且在嵌入向量的5维向量值内计算均值和方差。换句话说,计算是独立于其他数据或标记进行的。另一个重要特征是数据移动以与均值0对齐。这是一个重要特征,与RMSNorm形成鲜明对比。
六、RMS归一化
RMS归一化(RMSNorm)是Layer Normalization的一个变体,它仅使用均方根(RMS)进行归一化,而不计算均值。RMSNorm的计算开销比LayerNorm小,并且可以在学习更稳定的同时保持性能。
这里,RMS(x)定义如下:
其中的
计算得到
和
从上述结果可以看出,RMSNorm不与0对齐。由于它不执行x-μ,因此没有与0对齐的过程。与传统的LayerNorm相比,后者需要一次数据扫描来计算均值,另一次来计算方差,RMSNorm只需要一次RMS计算处理。这是一种高效的可以节省计算资源同时保持性能的方法。
七、LayerNorm和RMSNorm的比较
- 零均值与非零均值
与LayerNorm不同,RMSNorm维持非零均值而不与0对齐均值。
LayerNorm和RMSNorm上述Transformer输入示例后的结果的均值分布,结果更加清晰。尽管这两种方法都将数据分布汇聚到一个区域,但LayerNorm显示零均值,而RMSNorm显示非零均值。
- 从梯度角度进行比较
RMSNorm通过维持非零均值而不与0对齐均值,帮助解决梯度消失问题。虽然LayerNorm将均值与0对齐,但这可能导致在反向传播过程中梯度消失。 LayerNorm与RMSNorm:按周期比较梯度稳定性
上图使用一个非常简单的网络分析了LayerNorm和RMSNorm的梯度流动程度。左侧图表显示了按训练周期的LayerNorm和RMSNorm的平均梯度范数。在这里,梯度范数简单地衡量网络权重变化的速度。具体来说,梯度范数代表在反向传播过程中为每层权重计算的梯度的大小。较大的值意味着权重变化更大,如果它太大或太小,学习可能会变得不稳定。
RMSNorm在早期周期中比LayerNorm维持更高的梯度范数。这在防止深度学习模型中的梯度消失问题中很重要。右侧图表显示了RMSNorm和LayerNorm之间的梯度比率,其中RMSNorm的梯度最初要大得多,但随着学习的进行,差异减小。
RMSNorm在效率和稳定性方面优于LayerNorm。网络越深,使用RMSNorm越能稳定学习,防止梯度消失问题,并且可以在减少计算资源的同时实现类似的性能。
八、结论
在这次讲座中,我们检查了应用于原始Transformer的Layer Normalization和应用于LLama架构的RMS Normalization之间的差异和特征。LayerNorm通过调整数据的均值和方差来确保数值稳定性,并在各种序列模型中发挥重要作用。另一方面,RMSNorm是一种无需使用均值即可稳定数据尺度的技术,即使在深度网络中也能解决梯度消失问题,同时保持快速的学习速度。这两种技术都是现代深度学习模型中必不可少的,并且对于最大化性能,特别是在Transformer模型中是必要的。RMSNorm在效率和稳定性方面补充了LayerNorm,并在提高更复杂和更深模型的性能方面发挥着重要作用。