从零开始实现大语言模型（八）：Layer Normalization

1. 前言

Layer Normalization是深度学习实践中已经被证明非常有效的一种解决梯度消失或梯度爆炸问题，以提升神经网络训练效率及稳定性的方法。OpenAI的GPT系列大语言模型使用Layer Normalization对多头注意力模块，前馈神经网络模块以及最后的输出层的输入张量做变换，使shape为[batch_size, num_tokens, embedding_dim]的输入张量的embedding_dim维度数据的均值为0，方差为1。

本文介绍Layer Normalization的基本原理及其对输入张量的embedding_dim维度数据均值及方差做变换的方法，并实现继承自torch.nn.Module的神经网络模块LayerNorm。后续三篇文章将分别介绍前馈神经网络(feed forward network)与GELU激活函数，残差连接(shortcut connection)，Transformer Block，并最终构建出OpenAI的GPT系列大语言模型GPTModel。

2. Layer Normalization

如下图所示，对神经网络模块输出的均值为0.13，方差为0.39的6维向量做Layer Normalizaition，可以使输出向量的均值变为0，方差变为1。

可以使用torch.nn.Sequential(torch.nn.Linear(5, 6), torch.nn.ReLU())创建上图所示输入向量为度为5，输出向量维度为6的神经网络模块layer，并使用torch.randn(2, 5)创建输入张量batch_example。张量batch_example的shape为[2, 5]，其中第一个维度2表明该batch_example包含两个样本数据，第二个维度5表示每个样本数据是一个5维向量。将batch_example输入神经网络模块layer，输出张量out的shape为[2, 6]。其中第一个维度2对应batch_example中输入的样本数量，第二个维度6表示该神经网络模块输出向量的维度为6：

import torch

torch.manual_seed(123)

batch_example = torch.randn(2, 5)
layer = torch.nn.Sequential(torch.nn.Linear(5, 6), torch.nn.ReLU())
out = layer(batch_example)
print(out)

执行上面代码，打印结果如下：

tensor([[0.2260, 0.3470, 0.0000, 0.2216, 0.0000, 0.0000],
        [0.2133, 0.2394, 0.0000, 0.5198, 0.3297, 0.0000]],
       grad_fn=<ReluBackward0>)

可以使用如下代码计算神经网络模块layer各个输出向量的均值和方差：

mean = out.mean(dim=-1, keepdim=True)
var = out.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)

执行上面代码，打印结果如下：

Mean:
 tensor([[0.1324],
        [0.2170]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[0.0231],
        [0.0398]], grad_fn=<VarBackward0>)

在计算神经网络模块layer各个输出向量的均值及方差时，将参数keepdim的值设置为True，可以使均值张量mean及方差张量var的维数与输出张量out保持一致。如果将参数keepdim的值设置为False，则均值张量mean将会是一个2维的向量[0.1324, 0.2170]，而不会是一个shape为[2, 1]的矩阵[[0.1324], [0.2170]]。

参数dim用于指定计算均值和方差的维度。如下图所示，如果参数dim的值设置为0，将计算每列数据的均值及方差。如果参数dim的值设置为1或-1，将会计算每行数据的均值及方差。

Layer Normalization会将神经网络模块输出向量的各个分量减去其均值并除以其标准差（方差的平方根），如此即可使输出向量的均值为0，方差为1：

out_norm = (out - mean) / torch.sqrt(var)
mean = out_norm.mean(dim=-1, keepdim=True)
var = out_norm.var(dim=-1, keepdim=True)
print("Mean:\n", mean)
print("Variance:\n", var)

执行上面代码，打印结果如下：

Mean:
 tensor([[9.9341e-09],
        [5.9605e-08]], grad_fn=<MeanBackward1>)
Variance:
 tensor([[1.0000],
        [1.0000]], grad_fn=<VarBackward0>)

均值张量mean中的两个均值分别为9.9341e-09、5.9605e-08。9.9341e-09表示$9.9341\times 10^{-9}$，即0.0000000099341。由于计算机表示数字的精度有限，在进行浮点数运算时会造成非常小的数值误差，因此均值并不精确地等于0。

3. 构建神经网络模块LayerNorm

构建神经网络模块LayerNorm须实现一个继承自torch.nn.Modeul的类，并在__init__方法中初始化变量eps，创建神经网络参数scale及shift。eps是一个非常小的常数，对输入张量做变换时，除以方差与eps的和，可以避免方差为0时的除0异常。scale的初始值全为1，shift的初始值全为0，神经网络在训练过程中可以根据训练数据适当调整参数scale及shift的值。将经过变换后的输入张量乘以scale并加上shift，可以让神经网络模块LayerNorm具备根据训练数据动态调整输入张量数据分布的能力。

forward方法分别使用.mean及.var函数计算输入张量x最后一个维度数据的均值及方差，将输入张量x减去其均值并除以其方差与eps和的平方根，使输入张量x最后一个维度数据的均值为0，方差为1，得到norm_x。最后将norm_x乘以scale并加上shift，得到LayerNorm模块的输出。具体代码如下所示：

class LayerNorm(torch.nn.Module):
    def __init__(self, embedding_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = torch.nn.Parameter(torch.ones(embedding_dim))
        self.shift = torch.nn.Parameter(torch.zeros(embedding_dim))
 
    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift

在上述神经网络模块LayerNorm的forward方法中，使用.var函数计算输入张量x最后一个维度数据的方差时，将参数unbiased的值设置成了False。当unbiased=False，会使用公式$\frac{{\sum }_i(x_i-\bar{x}{)}^2}{n}$计算输入张量x最后一个维度数据的方差，其中$n$表示输入张量x最后一个维度数据中各个向量的维度。

一个数据集中所有数据的方差计算公式为$\frac{{\sum }_i(x_i-\bar{x}{)}^2}{n}$。如果从数据集中随机抽样得到一个样本数据集，使用样本数据集的方差估计数据集整体的方差，则方差计算公式为$\frac{{\sum }_i(x_i-\bar{x}{)}^2}{n-1}$。

在大语言模型中，Embedding向量的维度$n$一般会特别大，计算方差时不论除以$n$还是$n-1$，做Layer Normalization之后的结果基本不会有什么区别。PyTorch及TensorFolw中内置的LayerNorm模块计算方差时均默认除以$n$，OpenAI的GPT系列大语言模型做Layer Normalization时，同样使用了这种除以$n$的方差计算方法。

可以使用如下代码分别实例化上面构建的LayerNorm类对象以及PyTorch内置的LayerNorm类对象，并将张量batch_example分别输入两个LayerNorm类对象：

layernorm = LayerNorm(5)
layernorm_torch = torch.nn.LayerNorm(5)
print("LayerNorm:\n", layernorm(batch_example))
print("LayerNorm_torch:\n", layernorm_torch(batch_example))

执行上面代码，打印结果如下：

LayerNorm:
 tensor([[ 0.5528,  1.0693, -0.0223,  0.2656, -1.8654],
        [ 0.9087, -1.3767, -0.9564,  1.1304,  0.2940]],
        grad_fn=<AddBackward0>)
LayerNorm_torch:
 tensor([[ 0.5528,  1.0693, -0.0223,  0.2656, -1.8654],
        [ 0.9087, -1.3767, -0.9564,  1.1304,  0.2940]],
       grad_fn=<NativeLayerNormBackward0>)

根据打印结果可知，上面构建的神经网络模块LayerNorm与PyTorch内置的LayerNorm模块对输入张量的变换完全相同。

4. 结束语

OpenAI的GPT系列大语言模型使用Layer Normalization对多头注意力模块，前馈神经网络模块以及最后的输出层的输入张量做变换，将shape为[batch_size, num_tokens, embedding_dim]的输入张量的embedding_dim维度数据减去其均值并除以其方差与eps和的平方根，使embedding_dim维度数据的均值为0，方差为1。

Layer Normalization独立地对输入张量中每个embedding_dim维度数据做变换，不会受到训练大语言模型时每个batch训练数据的batch_size大小，以及部署大语言模型执行线上推理任务时输入Prompt中的num_tokens大小的影响。

OpenAI的GPT系列大语言模型结构中，只有自注意力机制相对复杂。Layer Normalization，GELU激活函数，前馈神经网络，残差链接的原理都非常简单。接下来，让我们去看看前馈神经网络与GELU激活函数吧！