6.4 层归一化#
在上一节内容中，我们详细介绍了批归一化的动机原理及实现过程，总体来讲批归一化的核心思想是以一个小批量数据样本为单位在对应维度上进行标准化。但也正是由于这一特性使得批量归一化会受到小批量样本数量的影响，同时，显而易见批归一化也不能直接用于循环神经网络。在这样的背景下，层归一化（Layer Normalization）[1]便应运而生了。
6.4.1 层归一化动机#
根据「第6.3节 BatchNorm原理：批归一化为什么能加速训练」内容可知，批归一化需要先计算一个小批量数据中每个通道上所有样本特征图的均值和方差，然后再根据对应的公式进行标准化。可见，此时小批量样本的数量便会影响均值和方差的估计结果。同时，在循环神经网络中由于每个样本的序列长度并不相同，如果使用批归一化进行标准化那么当模型推理时只要出现样本序列长度大于训练时最长的样本序列长度时，那么归一化过程将无法进行，因为此时需要对每个时间片的输出结果进行标准化并输入到下一个时刻中。除此之外，对于训练数据来说其长度越长则对应的样本数量将会越少，如果仍旧使用批归一化这类跨样本的归一化方法那么极端情况下将会出现归一化失真的情况。
如图6-15所示为文本序列批归一化序列图，其$x$轴、$y$轴和$z$轴分别表示小批量样本数量、词向量维度和样本序列长度，即图6-15中有5个样本，从左到右其长度分别的5、2、4、3和6，且每个词的维度为4维。如果此时对于整个小批量样本同时在$z$轴这个维度上采用批归一化进行标准化，那么在对进行第3个词的位置进行标准化时第2个样本便不会参与，进一步在对第6个词的位置进行标准化时便只有第5个词参与，而这便会影响整个归一化结果。退一步来说，即便通过这样的方式进行标准化，那么当测试样本的长度大于6时，该位置便没有对应的批归一化模型参数。
由于跨样本间进行标准化会受到上述因素的影响，因此层归一化提出了以每个样本为独立单位进行标准化的思想。具体地，对于普通的前馈神经网络来说，以每个样本当前层的所有神经元为整体进行标准化；对于卷积神经网络来说，以每个样本在当前层的所有特征通道为整体进行标准化；对于循环神经网络来说，则是以每个样本在当前时刻的输出向量为整体进行标准化，即图6-15中$y$˙轴所在的维度。
6.4.2 层归一化原理#
1. 归一化原理
从归一化的计算过程来看，层归一化同批归一化的计算公式类似，即式(6-15)~式(6-18)所示先计算均值和方差，然后进行标准化，最后再进行缩放和平移，唯一的差别在于两者选择归一化的维度不同，但也有研究表明缩放和平移操作反而会对于层归一化有害 [3]。如图6-16所示便是两者在对卷积特征进行归一化时维度选择的差异之处。
如图6-16所示，左右两边为分别是批归一化和层归一化的归一化维度示意图，其一共包含有2个样本和3个特征通道。从图6-16(a)可以看出，批归一化在进行标准化时是将所有小批量样本看成的一个整体，并逐一对每个通道进行标准化，且每个通道有各自独立的均值$\mu_i$、方差$\sigma^2_i$和权重参数$\gamma_i,\beta_i$（4个值均为标量）；对于图6-16(b)中的层归一化来说，它在进行标准化时是将每个样本看做一个整体，并同时对所有通道进行标准化，且每个样本均有各自独立的均值$\mu_i$和方差$\sigma^2_i$（2个值均为标量）但共享权重参数$\gamma,\beta$（2个值均为向量，维度为$p\times q\times c$，分别表示特征图的长、宽和通道数）。
2. CNN中的归一化
下面仍旧以第6.3.2节中特征图用层归一化的方式来进行计算示例。如图6-17所示左右两边分别为一个样本，且均有3个特征通道，则层归一化在进行标准化时以每个样本（图中虚线框部分）为单位进行处理，同时可知$\gamma$和$\beta$的维度均为75，即将3个通道拉伸后的总维度。
根据图6-17可知，则此时对于左右两个样本来说其均值和方差分别为
$$ \begin{aligned} \mu_1=&\frac{1}{75}(0+2+0+1+0+,...,+0+1+0+0+0)\approx0.4933\\[2ex] \sigma^2_1=&\frac{1}{75}((0-0.4933)^2+(2-0.4933)^2+,...,+(0-0.4933)^2)\approx0.3566\\[2ex] \mu_2=&\frac{1}{75}(1+2+1+1+1+,...,+0+1+1+0+0)=\approx0.7733\\[2ex] \sigma^2_2=&\frac{1}{75}((1-0.7733)^2+(2-0.7733)^2+,...,+(0-0.7733)^2)\approx0.5486\\[2ex] \end{aligned} \tag{6-23} $$进一步根据式(6-17)，此时假定$\epsilon=0$，则此时对于左侧的样本点有
$$ \frac{0-0.4933}{\sqrt{0.3566+0}}\approx-0.8261,\;\;\;\;\frac{2-0.4933}{\sqrt{0.3566+0}}\approx2.5229,\;\;\;\;\frac{1-0.4933}{\sqrt{0.3566+0}}\approx0.8484\tag{6-24} $$对于右侧的样本点有
$$ \frac{1-0.7733}{\sqrt{0.5486+0}}\approx0.3060,\;\;\;\;\frac{2-0.7733}{\sqrt{0.5486+0}}\approx1.6561,\;\;\;\;\frac{3-0.7733}{\sqrt{0.5486+0}}\approx3.0062\tag{6-25} $$根据式(6-18)，此时假定$\gamma=[1,1,...,1]_{1\times75},\beta=[0,0,...,0]_{1\times75}$，则有
$$ \begin{aligned} -0.8261\times1+0&=-0.8261\\[1ex] 2.5229\times1+0&=2.5229\\[1ex] 0.8484\times1+0&=0.8484\\[2ex] 0.3060\times1+0&=0.3060\\[1ex] 1.6561\times1+0&=1.6561\\[1ex] 3.0062\times1+0&=3.0062\\[1ex] \end{aligned}\tag{6-26} $$这里需要再次提醒的是，所有样本在进行平移和缩放时共享参数$\gamma$和$\beta$。
最后，图6-17中每个样本最后一个通道层归一化后的结果（可以同图6-13的结果进行对比）为
1 [[ 0.8484, 0.8484, -0.8261, -0.8261, -0.8261],
2 [-0.8261, 0.8484, -0.8261, -0.8261, 0.8484],
3 [-0.8261, 0.8484, 0.8484, 0.8484, 0.8484],
4 [-0.8261, 0.8484, -0.8261, -0.8261, -0.8261],
5 [-0.8261, 0.8484, -0.8261, -0.8261, -0.8261]]
6
7 [[-1.0441, 0.3060, 0.3060, -1.0441, -1.0441],
8 [-1.0441, 0.3060, -1.0441, 1.6561, 0.3060],
9 [ 0.3060, 1.6561, 0.3060, 0.3060, -1.0441],
10 [ 0.3060, 0.3060, -1.0441, -1.0441, -1.0441],
11 [-1.0441, 0.3060, 0.3060, -1.0441, -1.0441]]上述完整计算示例代码可参见Code/Chapter06/C04_LN/ln_compute.py文件。
3. RNN中的归一化
在理解了CNN中的批归一化计算过程后RNN便相对简单多了。不过这里需要知道的一点是，虽然论文[1]中作者提出的原始动机是对RNN中每个时刻的输出进行归一化然后再将归一化后的结果输入到下一个时刻并以此类推，但是在目前的实际使用中往往只会对RNN最后一层所有时刻计算结束后的输出进行标准化，因此前者在PyTorch中也并没有相关实现，不过在TensorFlow的addons模块中可以通过调用 LayerNormSimpleRNNCell进行实现。
如图6-18所示为某RNN网络模型的输出结果，其中序列长度为4、输出向量的维度为3、样本数量为2，那么此时层归一化将分别计算每个样本在每个时刻所有维度的均值和方差进行标准化并以同一组权重参数进行平移和缩放，同时可以知道$\gamma$和$\beta$的维度均为3。
在图6-18中，对于第1个样本的4个输出来说，其均值为$1.3333,0.6667,0.3333,0.3333$，方差为$0.2222,0.8889,0.2222,0.2222$。根据式(6-17)，此时假定$\epsilon=0$，则此时对于$t_1$时刻的标准化结果为
$$ \frac{1-1.3333}{\sqrt{0.2222+0}}\approx-0.7071,\;\;\;\;\frac{1-1.3333}{\sqrt{0.2222+0}}\approx-0.7071,\;\;\;\;\frac{2-1.3333}{\sqrt{0.2222+0}}\approx1.4142\tag{6-27} $$对于$t_2$时刻的标准化结果为
$$ \frac{2-0.6667}{\sqrt{0.8889+0}}\approx1.4142,\;\;\;\;\frac{0-0.6667}{\sqrt{0.8889+0}}\approx-0.7071,\;\;\;\;\frac{0-0.6667}{\sqrt{0.8889+0}}\approx-0.7071\tag{6-28} $$假定此时$\gamma=[1,1,1],\beta=[0,0,0]$则第1个样本$t_1,t_2$时刻缩放和平移后的结果便仍旧是上述结果。
最后，对于上述两个样本，其层归一化后的结果为
1 [[[-0.7071, -0.7071, 1.4142],
2 [ 1.4142, -0.7071, -0.7071],
3 [ -0.7071, -0.7071, 1.4142],
4 [ -0.7071, 1.4142, -0.7071]],
5
6 [[ 0.7071, -1.4142, 0.7071],
7 [-0.7071, -0.7071, 1.4142],
8 [ 0.0000, -1.2247, 1.2247],
9 [-1.4142, 0.7071, 0.7071]]]上述完整计算示例代码可参见Code/Chapter06/C04_LN/layer_normalization.py文件。
4. 测试时的归一化
由于层归一化是以每个样本为整体估算均值和方差对各维度进行标准化，因此层归一化并不需要区分当前是测试状态还是训练状态，其计算过程同训练时保持一致不再赘述。
到此为止对于层归一化的原理及计算过程就介绍完了。下面，我们再来进一步介绍如何从零实现层归一化及其相关的特性。
6.4.3 层归一化实现#
在介绍完整个层归一化操作的原理后，我们再来看如何通过PyTorch进行实现。以下完整示例代码可以参见Code/Chapter06/C04_LN/layer_normalization.py文件
1. 前向传播
根据6.4.2节内容的介绍可知，首先需要定义整个层归一化中的相关参数及维度信息，示例代码如下所示：
1 class LayerNormalization(nn.Module):
2 def __init__(self, normalized_shape=None, dim=-1, eps=1e-5):
3 super(LayerNormalization, self).__init__()
4 self.dim = dim
5 self.eps = eps
6 self.gamma = nn.Parameter(torch.ones(normalized_shape))
7 self.beta = nn.Parameter(torch.zeros(normalized_shape))在上述代码中，第2行normalized_shape用来指定参数$\gamma$和$\beta$的维度，对于RNN中的特征来说便是隐含向量的维度，对于CNN中的特征来说则是传入长、宽和通道数这3个维度，dim用来指定需要进行标准化的维度，默认-1即RNN中特征输出的最后一个维度，也可通过列表来指定相应的维度。第6~7行便是初始化两个权重参数，分别全为1和0的两个向量。
在完成上述初始化工作后便可以进一步实现批归一化的整个前向传播过程，示例代码如下所示：
1 def forward(self, X):
2 mean = torch.mean(X, dim=self.dim, keepdim=True)
3 var = torch.mean((X - mean) ** 2, dim=self.dim, keepdim=True)
4 X_hat = (X - mean) / torch.sqrt(var + self.eps)
5 Y = self.gamma * X_hat + self.beta
6 return Y在上述代码中，第2~3行分别是计算指定维度上的均值和方差，同时由于下一步计算时需要利用到PyTorch中的广播机制，所以设定了keepdim=True。第4行则是进行初始的层归一化操作。第5~6行是进行缩放与平移，并返回最后的结果。
在实现完上述代码后还可以通过如下方式来进行检验，示例代码如下所示：
1 if __name__ == '__main__':
2 batch, sentence_length, embedding_dim = 2, 4, 3
3 embedding = torch.tensor([
4 [[1, 1, 2],[2, 0, 0],[0, 0, 1.],[0, 1, 0]],
5 [[1, 0, 1],[1, 1, 2],[1, 0, 2], [0, 2, 2]]])
6 layer_norm = nn.LayerNorm(embedding_dim)
7 my_layer_norm = LayerNormalization(embedding_dim)
8 print(layer_norm(embedding))
9 print(my_layer_norm(embedding))
10
11 N, C, H, W = 2, 3, 5, 5
12 embedding = torch.randn(N, C, H, W)
13 layer_norm = nn.LayerNorm([C, H, W])
14 print(layer_norm(embedding))
15 my_layer_norm = LayerNormalization([C, H, W], dim=[1, 2, 3])
16 print(my_layer_norm(embedding))在上述代码中，第2~5行是定义一个RNN的输入样本。第6~7行是分别实例化一个层归一化对象，一个来自于PyTorch一个来自于手动实现。第8~9行则是层归一化后的结果，部分输出如下：
1 tensor([[[-0.7071, -0.7071, 1.4142],
2 [ 1.4142, -0.7071, -0.7071],
3 [-0.7071, -0.7071, 1.4142],
4 [-0.7071, 1.4142, -0.7071]],
5 ......], grad_fn=<NativeLayerNormBackward0>)
6 tensor([[[-0.7071, -0.7071, 1.4142],
7 [ 1.4142, -0.7071, -0.7071],
8 [-0.7071, -0.7071, 1.4142],
9 [-0.7071, 1.4142, -0.7071]],
10 ......], grad_fn=<AddBackward0>)第11~12行则是随机定义一个特征图。第13~16则是分别对其进行标准化并输出最终的结果，部分输出如下：
1 tensor([[[[-0.7794, 1.5095, 0.0447, -1.1159, -0.1546],
2 [ 0.3652, -0.3946, -1.0248, 0.8011, 1.5210],
3 [ 1.2956, 0.1933, -1.2409, -1.5562, 0.1174],
4 [-0.5513, -0.2401, 1.0143, 0.4252, 1.6937],
5 [ 2.2888, -0.7463, 1.1979, -0.7542, -0.3546]],
6 ......],grad_fn=<NativeLayerNormBackward0>)
7 tensor([[[[-0.7794, 1.5095, 0.0447, -1.1159, -0.1546],
8 [ 0.3652, -0.3946, -1.0248, 0.8011, 1.5210],
9 [ 1.2956, 0.1933, -1.2409, -1.5562, 0.1174],
10 [-0.5513, -0.2401, 1.0143, 0.4252, 1.6937],
11 [ 2.2888, -0.7463, 1.1979, -0.7542, -0.3546]],
12 ......],grad_fn=<AddBackward0>)2. 使用示例
由于LayerNormalization同样是继承自类nn.Module，所以仍旧可以将它作为一个网络层进行使用，这里以「第7.2节 时序数据建模：RNN 适合处理什么样的序列任务」中介绍的RNN网络模型为例，层归一化的使用方法如下所示：
1 class FashionMNISTRNN(nn.Module):
2 def __init__(self, input_size=28, hidden_size=128,
3 num_layers=1, num_classes=10):
4 super(FashionMNISTRNN, self).__init__()
5 self.rnn = nn.RNN(input_size, hidden_size,num_layers, batch_first=True)
6 self.layer_norm = LayerNormalization(hidden_size)
7 self.fc = nn.Linear(hidden_size, num_classes)
8 def forward(self, x, labels=None):
9 x = x.squeeze(1)
10 x, _ = self.rnn(x)
11 x = self.layer_norm(x)
12 logits = self.fc(x[:, -1].squeeze(1))
13 if labels is not None:
14 loss_fct = nn.CrossEntropyLoss(reduction='mean')
15 loss = loss_fct(logits, labels)
16 return loss, logits
17 else:
18 return logits在上述代码中，第2~3行分别是定义RNN模型中的各个超参数。第5行用于实例化得到RNN类对象。第6行则是定义一个层归一化层。第7行是定义最后的分类层。这里需要注意的是，由于RNN模型前向传播后的输出结果不止一个，所以不能直接将第5~6行放入一个nn.Sequential里面。第8~18行则是整个前向传播的计算过程，详细介绍请参见「第7.1节 RNN原理：循环神经网络的结构与序列建模」内容。
最后，输入结果如下所示：
1 Epochs[1/5]--batch[0/938]--Acc: 0.0781--loss: 2.4072
2 Epochs[1/5]--batch[50/938]--Acc: 0.5312--loss: 1.2232
3 Epochs[1/5]--batch[100/938]--Acc: 0.5469--loss: 1.0935
4 ......
5 Epochs[1/5]--batch[900/938]--Acc: 0.8594--loss: 0.4901
6 Epochs[1/5]--Acc on test 0.81386.4.4 小结#
在本节内容中，我们首先介绍了层归一化算法提出的原因和动机；然后详细介绍了层归一化的原理及过程，包括CNN及RNN中的层归一化计算过程等；进一步，介绍了如何从零开始在PyTorch框架中实现层归一化算法的计算过程；最后，以RNN模型为例对层归一化层的使用进行了示例和验证。
引用#
[1] Ba J L, Kiros J R, Hinton G E. Layer normalization[J]. arXiv preprint, 2016, arXiv:1607.06450.
[2] Paszke A, Gross S, Massa F, et al. Pytorch: An imperative style, high-performance deep learning library[J]. Advances in neural information processing systems, 2019, 32.
[3] Xu J, Sun X, Zhang Z, et al. Understanding and improving layer normalization[J]. Advances in neural information processing systems, 2019, 32.