批归一化和层归一化

引言

本文探讨了BN和LN，BN适合于CNN；而LN适合于RNN。虽然我们现在还不知道BN为什么有效，但是重点是它有效，我们就使用它。

Batch Normalization

Batch Normalization(批归一化,以下简称BN)是由Sergey Ioffe​​¹​​等人提出来的，是一个广泛使用的深度神经网络训练的技巧，它不仅可以加快了模型的收敛速度，还可以简化初始化要求，即可以使用较大的学习率。

概念

原文提出了一些概念，用于解释为什么BN有用，但是

covariate shift与 internal covariate shift

在原文 ​​¹​​中，以分布稳定性角度，covariate shift描述的是模型输入分布变化的现象。

internal covariate shift 说的是在深度神经网络隐藏层之间输入分布变化的现象。

从模型的角度看，训练数据和测试数据分布相差较大，也是一种covariate shift。

算法描述

原文中算法描述如下：

其中$x_1,\cdots,x_m$是训练批次$\mathcal{B}$中的所有样本，该算法可学习的参数是$\gamma$和$\beta$；输出是每个样本经过BN后的结果。

共四部，前三步分别是计算批次内样本的均值、方差、进行标准化。

最后一步是反标准化操作，将标准化后的数据再扩展和平移。为了让模型自己去学习是否需要标准化，以及多大程度。其中$\varepsilon$是一个很小的常数，防止分母为零。

上面说的都是针对训练数据的，对于测试数据，或者说线上数据应该怎么做呢？

因为线上数据可能一次只输入一条，因此无法计算均值和方差。原文的做法是保存训练数据每个批次的均值和方差，主要思想是求所有批次得到的均值和方差的期望，使用的是指数移动平均值(EMA)，着重考虑最近迭代的均值和方差。

算法实现

class BatchNorm(nn.Module):
    def __init__(self, num_features, epsilon=1e-05, momentum=0.1, device=None):
        '''
        num_features: 全连接网络的输出大小
        momentum: EMA中使用的参数
        '''
        super(BatchNorm, self).__init__()
        
        self.device = device
        
        # 需要学习的参数，用Parameter生成
        self.beta = nn.Parameter(torch.zeros(1, num_features))
        self.gamma = nn.Parameter(torch.ones(1, num_features))
        
        self.epsilon = epsilon
        self.momentum = momentum
        
        self.moving_mean = torch.zeros(1, num_features)
        self.moving_var = torch.ones(1, num_features)

    
    def forward(self, X):
        '''
        X: [batch_size, num_features]
        '''
        if self.device:
            self.moving_mean = self.moving_mean.to(device)
            self.moving_var = self.moving_var.to(device)
        
        # 如果是训练模式
        if self.training:
            # 当前批次的均值和方差
            mean = X.mean(dim=0) 
            var = ((X - mean)**2).mean(dim=0)
            # 标准化
            X_normalized = (X - mean) / torch.sqrt(var + self.epsilon)
            # 更新移动平均值  和nn.BatchNorm1d的做法一样
            self.moving_mean = (1 - self.momentum) * self.moving_mean + self.momentum * mean
            self.moving_var =  (1 - self.momentum) * self.moving_var + self.momentum * var
        else:
            # 如果是推理模式
            X_normalized = (X - self.moving_mean) / torch.sqrt(self.moving_var + self.epsilon)
        
        # 公式中的y
        Y = self.gamma * X_normalized + self.beta
        
        return Y # [batch_size, num_features]

    def __repr__(self):
        return f'BatchNorm(num_features={self.moving_mean.size(1)}, momentum={self.momentum})'

下面我们用一个回归任务来看一下批归一化的效果。

以下示例参考了莫凡Python​​²​​

# 超参数
N_SAMPLES = 2000
BATCH_SIZE = 64
EPOCH = 12
LR = 0.03
N_HIDDEN = 8
ACTIVATION = torch.relu
B_INIT = -0.2   # 使用一个负值的参数初始化

# training data
x = np.linspace(-7, 10, N_SAMPLES)[:, np.newaxis]
noise = np.random.normal(0, 2, x.shape)
y = np.square(x) - 5 + noise

# test data
test_x = np.linspace(-7, 10, 200)[:, np.newaxis]
noise = np.random.normal(0, 2, test_x.shape)
test_y = np.square(test_x) - 5 + noise

train_x, train_y = torch.from_numpy(x).float(), torch.from_numpy(y).float()
test_x = Variable(torch.from_numpy(test_x).float(), volatile=True)  # not for computing gradients
test_y = Variable(torch.from_numpy(test_y).float(), volatile=True)

train_dataset = Data.TensorDataset(train_x,train_y)
train_loader = Data.DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)

# show data
plt.scatter(train_x.numpy(), train_y.numpy(), c='#FF9359', s=50, alpha=0.2, label='train')
plt.legend(loc='upper left')
plt.show()

以$y= x^2$函数画图，增加了一些噪音。

然后我们构造一个深层的网络来拟合这些数据，用到了上面自定义的批归一化实现。

class Net(nn.Module):
    def __init__(self, batch_normalization=False):
        super(Net, self).__init__()
        self.do_bn = batch_normalization
        self.fcs = []
        self.bns = []
        # 使用自定义的Batch Norm层
        self.bn_input = BatchNorm(1) 

        # 输入层的大小是1，隐藏层的大小是10
        for i in range(N_HIDDEN):             
            input_size = 1 if i == 0 else 10
            fc = nn.Linear(input_size, 10)
            # 通过setattr 动态构建神经网络
            setattr(self, 'fc%i' % i, fc)       
            self._set_init(fc)                 
            self.fcs.append(fc)
            if self.do_bn:
                bn = BatchNorm(10)
                setattr(self, 'bn%i' % i, bn)
                self.bns.append(bn)
        # 输出层的大小也是1，我们做的是回归
        self.predict = nn.Linear(10, 1)         # output layer
        self._set_init(self.predict)            # parameters initialization

    def _set_init(self, layer):
        init.normal_(layer.weight, mean=0., std=.1)
        init.constant_(layer.bias, B_INIT)

    def forward(self, x):
        # 保存激活之前的输入
        pre_activation = [x]
        if self.do_bn: 
            x = self.bn_input(x)     # 输入的BN
        # 每个隐藏层的输入
        layer_input = [x]
        for i in range(N_HIDDEN):
            x = self.fcs[i](x)
            pre_activation.append(x)
            if self.do_bn: 
                x = self.bns[i](x)   # 隐藏层的BN
            x = ACTIVATION(x)
            layer_input.append(x)
            
        out = self.predict(x)
        return out, layer_input,

打印网络结构：

nets = [Net(batch_normalization=False), Net(batch_normalization=True)]
print(*nets)    # print net architecture

我们构建了两个网络实例，一个使用了批归一化，另一个没有使用。作为对比。

训练，并画图。

opts = [torch.optim.Adam(net.parameters(), lr=LR) for net in nets]

loss_func = torch.nn.MSELoss()

f, axs = plt.subplots(4, N_HIDDEN+1, figsize=(10, 5))
plt.ion()   # something about plotting

def plot_histogram(l_in, l_in_bn, pre_ac, pre_ac_bn):
    for i, (ax_pa, ax_pa_bn, ax,  ax_bn) in enumerate(zip(axs[0, :], axs[1, :], axs[2, :], axs[3, :])):
        [a.clear() for a in [ax_pa, ax_pa_bn, ax, ax_bn]]
        if i == 0: p_range = (-7, 10);the_range = (-7, 10)
        else:p_range = (-4, 4);the_range = (-1, 1)
        ax_pa.set_title('L' + str(i))
        ax_pa.hist(pre_ac[i].data.numpy().ravel(), bins=10, range=p_range, color='#FF9359', alpha=0.5);ax_pa_bn.hist(pre_ac_bn[i].data.numpy().ravel(), bins=10, range=p_range, color='#74BCFF', alpha=0.5)
        ax.hist(l_in[i].data.numpy().ravel(), bins=10, range=the_range, color='#FF9359');ax_bn.hist(l_in_bn[i].data.numpy().ravel(), bins=10, range=the_range, color='#74BCFF')
        for a in [ax_pa, ax, ax_pa_bn, ax_bn]: a.set_yticks(());a.set_xticks(())
        ax_pa_bn.set_xticks(p_range);ax_bn.set_xticks(the_range)
        axs[0, 0].set_ylabel('PreAct');axs[1, 0].set_ylabel('BN PreAct');axs[2, 0].set_ylabel('Act');axs[3, 0].set_ylabel('BN Act')
    plt.pause(0.01)
    
# training
losses = [[], []]  # recode loss for two networks
for epoch in range(EPOCH):
    print('Epoch: ', epoch)
    layer_inputs, pre_acts = [], []
    for net, l in zip(nets, losses):
        net.eval（)              # set eval mode to fix moving_mean and moving_var
        pred, layer_input, pre_act = net(test_x)
        l.append(loss_func(pred, test_y).data)
        layer_inputs.append(layer_input)
        pre_acts.append(pre_act)
        net.train()             # free moving_mean and moving_var
    plot_histogram(*layer_inputs, *pre_acts)     # plot histogram

    for step, (b_x, b_y) in enumerate(train_loader):
        b_x, b_y = Variable(b_x), Variable(b_y)
        for net, opt in zip(nets, opts):     # train for each network
            pred, _, _ = net(b_x)
            loss = loss_func(pred, b_y)
            opt.zero_grad()
            loss.backward()
            opt.step()    # it will also learns the parameters in Batch Normalization
            
plt.ioff()

L0是输入层，L1到L8是隐藏层，绘画的是每次迭代各个层输入值的分布情况。

红色是无BN的网络，蓝色的有BN的网络。

PreAct是激活之前的值，Act是激活之后的值。可以看到，无BN的网络，激活函数为ReLU的情况话，所有网络的输出基本不变，看上去像是死掉了。

而使用了BN的网络，每层的分布较为分散，没有集中在某处，经过BN激活之后的值也存在很多大于零的部分。

下面画出两个网络拟合曲线和损失曲线。

# plot training loss
plt.figure(2)
plt.plot(losses[0], c='#FF9359', lw=3, label='Original')
plt.plot(losses[1], c='#74BCFF', lw=3, label='Batch Normalization')
plt.xlabel('step');plt.ylabel('test loss');plt.ylim((0, 2000));plt.legend(loc='best')

# evaluation
# set net to eval mode to freeze the parameters in batch normalization layers
[net.eval（) for net in nets]    # set eval mode to fix moving_mean and moving_var
preds = [net(test_x)[0] for net in nets]
plt.figure(3)
plt.plot(test_x.data.numpy(), preds[0].data.numpy(), c='#FF9359', lw=4, label='Original')
plt.plot(test_x.data.numpy(), preds[1].data.numpy(), c='#74BCFF', lw=4, label='Batch Normalization')
plt.scatter(test_x.data.numpy(), test_y.data.numpy(), c='r', s=50, alpha=0.2, label='train')
plt.legend(loc='best')
plt.show()

为什么BN有用

BN使得深层神经网络更易于训练，但是具体为什么，现在还没有定论，不过存在一些假设。

• 假设1：原文​​¹​​​的作者猜测是因为BN减少了 internal covariate shift(ICS)，使得神经网络更易于训练。
  ❌ Shibani Santurkar​​​³​​等人通过实验证明，这种表现和ICS无关。
• 假设2：BN通过2个可学习的参数调整隐藏层的输入分布来使优化器更好地工作。
  ❓ 这个假设强调是因为参数之间的相互依赖性，让优化任务更加困难，但是没有确凿的证据。
• 假设3：BN重新定制了底层的优化问题，使之更加平滑且稳定。
  ❓ 这是最新的研究，并且还未有人提出异议。他们提供了一部分理论支持，但是一些基本问题仍然未得到解答，比如BN是如何帮助泛化的。

Layer Normalization

Jimmy​​⁴​​基于Batch Normalization提出了Layer Normalization(层归一化，以下简称LN)。

原文指出，如果将BN应用到RNN中会出现一些问题，由于NLP任务中句子的长度是不固定的，如果使用BN，会导致每个时间步的统计量不同。可能某个时间步，某个句子没有输入了；更糟糕的是，BN无法适应于在线学习(每个批次只有一个样本)和批次数量过小的情况。

如果说BN是针对整个批次计算的，那么LN就是针对一个样本所有特征计算的。

或者说BN是对一个隐藏层的所有神经元进行归一化。

算法描述

类似BN，但是是对每个样本自身进行计算，因此训练时和测试时是一样的，不需要计算EMA。

令$x=(x_1,x_2,\cdots,x_H)$是某个时间步LN层的H大小的输入向量表示，LN通过下面的公式将$x$进行归一化：
$\mu = \frac{1}{H}\sum_{i=1}^H x_i,\quad \sigma = \sqrt{\frac{1}{H}\sum_{i=1}^H (x_i - \mu)^2}, \quad N(x) = \frac{x-\mu}{\sigma},\quad h = g \,\odot N(x) + b$
其中$h$就是LN层的输出，$\odot$是点乘操作，$\mu$和$\sigma$是输入各个维度的均值和方差，$g$和$b$是两个可学习的参数，和$h$的维度相同。

在标准的RNN中，在每个时间步中，RNN单元的总输入的平均值要么增加，要么减小，导致梯度爆炸或消失。应用LN后，归一化操作重新缩放所有的总输入到某个层中，从而保证更稳定的隐藏层到隐藏层的传递。

算法实现

class LayerNorm(nn.Module):
    def __init__(self, normalized_shape, epsilon=1e-05):
        '''
        normalized_shape: 输入tensor的shape或输入tensor最后一个维度大小
        '''
        super(LayerNorm, self).__init__()
        
        if isinstance(normalized_shape, int):
            normalized_shape = (normalized_shape, )
        else:
            normalized_shape = (normalized_shape[-1], )
        
        self.normalized_shape = torch.Size(normalized_shape)
                
        # 需要学习的参数，用Parameter生成
        self.beta = nn.Parameter(torch.zeros(*normalized_shape))
        self.gamma = nn.Parameter(torch.ones(*normalized_shape))
        
        self.epsilon = epsilon


    
    def forward(self, X):
        '''
        X: [batch_size, *]
        '''

        # 计算每个样本的均值和方差
        mean = X.mean(dim=-1, keepdim = True) 
        var = ((X - mean)**2).mean(dim=-1, keepdim = True)
        # 标准化
        X_normalized = (X - mean) / torch.sqrt(var + self.epsilon)
     
        
        # 公式中的h
        Y = self.gamma * X_normalized + self.beta
        
        return Y # [batch_size, num_features]

    def __repr__(self):
        return f'LayerNorm(normalized_shape={self.normalized_shape})'

参考

---

1. ​​Batch Normalization: Accelerating Deep Network Training b y Reducing Internal Covariate Shift​​​ ​​↩︎​​ ​​↩︎​​ ​​↩︎​​
2. ​​莫凡Python​​​ ​​↩︎​​
3. ​​How Does Batch Normalization Help Optimization? ​​​ ​​↩︎​​
4. ​​Layer Normalization​​​ ​​↩︎​​