Layer/Batch/Instance Normalization

总览

图中N表示batch，C表示CV中的通道(NLP中的序列长度、时间步)，如果是图像则【H,W】表示每个通道下二维像素矩阵的高和宽，NLP中就只有一维特征向量。Batch Norm依赖Batch，对【Batch, H, W】三个维度做标准化；Layer Norm不依赖Batch，对【C，H，W】三个维度做标准化。Instance Norm既不受Batch也不受其它通道的影响，只对【H，W】两个维度做标准化。

三种标准化的表示式形式都相同，其区别在于$x$的表示不同，其公式如下：

$$y=\frac{x-\mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+ϵ}}\ast \gamma +\beta$$
其中，对于不同Norm方法$x$不同，$\mathrm{E}[x]$表示期望(均值)，$\mathrm{Var}[x]$表示方差，$\gamma$和$\beta$表示可学习参数，向量大小和输入维度一致。

Batch Norm

$$y_{tilm}=\frac{x_{tilm}-{\mu }_i}{\sqrt{{\sigma }_i^2+ϵ}}\ast \gamma +\beta ,{\mu }_i=\frac{1}{HWT}\sum _{t=1}^T\sum _{l=1}^W\sum _{m=1}^H{x}_{tilm},{\sigma }_i^2=\frac{1}{HWT}\sum _{t=1}^T\sum _{l=1}^W\sum _{m=1}^H{\left(x_{tilm}-u_i\right)}^2$$

Layer Norm

• 针对4维输入：
  $$y_{tilm}=\frac{x_{tilm}-{\mu }_t}{\sqrt{{\sigma }_t^2+ϵ}}\ast \gamma +\beta ,{\mu }_t=\frac{1}{HWC}\sum _{i=1}^C\sum _{l=1}^W\sum _{m=1}^H{x}_{tilm},{\sigma }_t^2=\frac{1}{HWC}\sum _{i=1}^C\sum _{l=1}^W\sum _{m=1}^H{\left(x_{tilm}-u_t\right)}^2$$

• 针对3维输入：
  $$y_{til}=\frac{x_{til}-{\mu }_t}{\sqrt{{\sigma }_t^2+ϵ}}\ast \gamma +\beta ,{\mu }_t=\frac{1}{NS}\sum _{i=1}^S\sum _{l=1}^N{x}_{til},{\sigma }_t^2=\frac{1}{NS}\sum _{i=1}^S\sum _{l=1}^N{\left(x_{til}-u_t\right)}^2$$

• 3维输入的代码测试样例：

import torch
from torch import nn

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-5):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        sizes = x.size()
        mean = x.view(sizes[0], -1).mean(-1, keepdim=True)
        std = x.view(sizes[0], -1).std(-1, keepdim=True)
        mean = torch.repeat_interleave(mean, repeats=9, dim=1).view(2, 3, 3)
        std = torch.repeat_interleave(std, repeats=9, dim=1).view(2, 3, 3)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

norm_2 = LayerNorm(features=[2, 3, 3])
norm = nn.LayerNorm(normalized_shape=(3, 3), elementwise_affine=True)
a = torch.tensor([[[1, 2, 3], 
                   [4, 5, 6], 
                   [7, 8, 9]], 
                  [[4, 5, 6], 
                   [7, 8, 9], 
                   [10, 11, 12]]], dtype=torch.float32)

print(norm(a))
print(norm_2(a))

• 代码运行结果：

tensor([[[-1.5492, -1.1619, -0.7746],
         [-0.3873,  0.0000,  0.3873],
         [ 0.7746,  1.1619,  1.5492]],

        [[-1.5492, -1.1619, -0.7746],
         [-0.3873,  0.0000,  0.3873],
         [ 0.7746,  1.1619,  1.5492]]], grad_fn=<NativeLayerNormBackward>)
tensor([[[-1.4606, -1.0954, -0.7303],
         [-0.3651,  0.0000,  0.3651],
         [ 0.7303,  1.0954,  1.4606]],

        [[-1.4606, -1.0954, -0.7303],
         [-0.3651,  0.0000,  0.3651],
         [ 0.7303,  1.0954,  1.4606]]], grad_fn=<AddBackward0>)

结果分析：从结果可以看出，自己实现的和库函数实现在数值大小稍有些不同，但计算方法应该是一致的，可能在具体实现的某个小细节上有所差异

Layer Norm原论文还给出了一个应用在RNN上案例，这个就针对三维的输入[batch_size, seq_len, num_features]。令当前时间步的输入为 $x^t$，上一时间步的隐层状态为$h^{t-1}$，因此可以计算得到$a^t=W_{hh}{h}^{t-1}+W_{xh}{x}^t$。注意这儿$a^t$其实是未经过激活函数的$h^t$。然后直接对$a^t$进行归一化，如下式所示：
$${\mathbf{h}}^t=f\left[\frac{\mathbf{g}}{{\sigma }^t}\odot \left({\mathbf{a}}^t-{\mu }^t\right)+\mathbf{b}\right]{\mu }^t=\frac{1}{H}\sum _{i=1}^H{a}_i^t{\sigma }^t=\sqrt{\frac{1}{H}\sum _{i=1}^H{\left(a_i^t-{\mu }^t\right)}^2}$$
上式中 $\mathbf{g}$ 和 $\mathbf{b}$ 相当于 $\gamma$ 和 $\beta$，$f$ 表示激活函数 $\tanh$。注意：这儿实际上是对每个样本的每一个时间步的未经过激活的 $h^t$ 做标准化，和上上式是不一样的。

• Harwardnlp给出了一份基于每个时间步上Layer Norm实现：

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

Instance Norm

$$y_{tilm}=\frac{x_{tilm}-{\mu }_{ti}}{\sqrt{{\sigma }_{ti}^2+ϵ}}\ast \gamma +\beta ,{\mu }_{ti}=\frac{1}{HW}\sum _{l=1}^W\sum _{m=1}^H{x}_{tilm},{\sigma }_{ti}^2=\frac{1}{HW}\sum _{l=1}^W\sum _{m=1}^H{\left(x_{tilm}-u_{ti}\right)}^2$$
如有不妥指出，欢迎纠正指出！

References

[1] Ioffe, Sergey, and Christian Szegedy. “Batch normalization: Accelerating deep network training by reducing internal covariate shift.” arXiv preprint arXiv:1502.03167 (2015).
[2] Ba, Jimmy Lei, Jamie Ryan Kiros, and Geoffrey E. Hinton. “Layer normalization.” arXiv preprint arXiv:1607.06450 (2016).
[3] Ulyanov, Dmitry, Andrea Vedaldi, and Victor Lempitsky. “Instance normalization: The missing ingredient for fast stylization.” arXiv preprint arXiv:1607.08022 (2016).
[4] http://nlp.seas.harvard.edu/2018/04/03/attention.html