#

层规范化

批量标准化的局限性

你需要保持跑步手段。
对于 RNN 来说很棘手。每个步骤都需要不同的规范化吗？
不适用于小批量；大型 NLP 模型通常使用小批量进行训练。
需要在分布式训练中计算设备间的均值和方差。

层规范化

图层归一化是一种更简单的归一化方法，适用于更广泛的设置。图层归一化会将输入变换为各要素的均值和单位方差为零。请注意，批量归一化修复了每个元素的零均值和单位方差。层归一化对所有元素的每个批次执行此操作。

层归一化通常用于 NLP 任务。

我们在大多数变压器实现中都使用了层归一化。

35from typing import Union, List
36
37import torch
38from torch import nn, Size
39
40from labml_helpers.module import Module

#

层规范化

图层 $L N$ 归一化将输入归一化 $X$ ，如下所示：

当输入 $X \in R^{B \times C}$ 是一批嵌入时，其中 $B$ 是批次大小， $C$ 是要素的数量。 $γ \in R^{C}$ 和 $β \in R^{C}$ 。 $L N (X) = γ \frac{X - C E [ X ]}{C Va r [ X ] + ϵ} + β$

当 input $X \in R^{L \times B \times C}$ 是嵌入序列中的一批时，其中 $B$ $C$ 是批次大小， $L$ 是通道数，是顺序。 $γ \in R^{C}$ 和 $β \in R^{C}$ 。 $L N (X) = γ \frac{X - C E [ X ]}{C Va r [ X ] + ϵ} + β$

当输入 $X \in R^{B \times C \times H \times W}$ 是一批图像表示时，其中 $B$ 是批次大小， $C$ 是通道数， $H$ 是高度和 $W$ 是宽度。这不是一个广泛使用的场景。 $γ \in R^{C \times H \times W}$ 和 $β \in R^{C \times H \times W}$ 。 $L N (X) = γ \frac{X - C , H , W E [ X ]}{C , H , W Va r [ X ] + ϵ} + β$

43class LayerNorm(Module):

#

normalized_shape $S$ 是元素的形状（批次除外）。那么输入应该是 $X \in R^{*\times S [0] \times S [1] \times ... \times S [n]}$
eps 是 $ϵ$ ， $Va r [X] + ϵ$ 用于数值稳定性
elementwise_affine 是否缩放和移动归一化值

我们已经尝试使用与 PyTorchLayerNorm 实现相同的参数名称。

72    def __init__(self, normalized_shape: Union[int, List[int], Size], *,
73                 eps: float = 1e-5,
74                 elementwise_affine: bool = True):

#

84        super().__init__()

#

转换normalized_shape 为torch.Size

87        if isinstance(normalized_shape, int):
88            normalized_shape = torch.Size([normalized_shape])
89        elif isinstance(normalized_shape, list):
90            normalized_shape = torch.Size(normalized_shape)
91        assert isinstance(normalized_shape, torch.Size)

#

94        self.normalized_shape = normalized_shape
95        self.eps = eps
96        self.elementwise_affine = elementwise_affine

#

为增益 $γ$ 和 $β$ 偏置创建参数和参数

98        if self.elementwise_affine:
99            self.gain = nn.Parameter(torch.ones(normalized_shape))
100            self.bias = nn.Parameter(torch.zeros(normalized_shape))

#

x 是形状张量[*, S[0], S[1], ..., S[n]] 。* 可以是任意数量的维度。例如，在 NLP 任务中，这将是[seq_len, batch_size, features]

102    def forward(self, x: torch.Tensor):

#

进行健全性检查以确保形状匹配

110        assert self.normalized_shape == x.shape[-len(self.normalized_shape):]

#

用于计算均值和方差的维度

113        dims = [-(i + 1) for i in range(len(self.normalized_shape))]

#

计算所有元素的均值；即每个元素的均值 $E [X]$

117        mean = x.mean(dim=dims, keepdim=True)

#

计算所有元素的均方值；即每个元素的均值 $E [X^{2}]$

120        mean_x2 = (x ** 2).mean(dim=dims, keepdim=True)

#

所有元素的方差 $Va r [X] = E [X^{2}] - E [X]^{2}$

122        var = mean_x2 - mean ** 2

#

规范化 $\hat{X} = \frac{X - E [ X ]}{Va r [ X ] + ϵ}$

125        x_norm = (x - mean) / torch.sqrt(var + self.eps)

#

缩放和移动 $L N (x) = γ \hat{X} + β$

127        if self.elementwise_affine:
128            x_norm = self.gain * x_norm + self.bias

#

131        return x_norm

#

简单测试

134def _test():

#

138    from labml.logger import inspect
139
140    x = torch.zeros([2, 3, 2, 4])
141    inspect(x.shape)
142    ln = LayerNorm(x.shape[2:])
143
144    x = ln(x)
145    inspect(x.shape)
146    inspect(ln.gain.shape)

#

150if __name__ == '__main__':
151    _test()