Notes on Batch Normalization

在训练深层神经网络的过程中, 由于输入层的参数在不停的变化, 因此, 导致了当前层的分布在不停的变化, 这就导致了在训练的过程中, 要求 learning rate 要设置的非常小, 另外, 对参数的初始化的要求也很高. 作者把这种现象称为 internal convariate shift. Batch Normalization 的提出就是为了解决这个问题的. BN 在每一个 training mini-batch 中对每一个 feature 进行 normalize. 通过这种方法, 使得网络可以使用较大的 learning rate, 而且, BN 具有一定的 regularization 作用.

为什么需要 Batch Normalization

在神经网络的优化中最常用最进本的方法是 SGD, 其目标是寻找最小化 loss function 的参数:
\[ \theta = \underset{\theta}{\mathrm{arg min}} \dfrac{1}{N} \sum_{i=1}^N{\mathcal L\left(x_i, \theta\right)} \]
在求解的过程中, 一般是使用 minibatch 的方法, 简单来说, 就是计算下面的梯度:
\[ \dfrac{1}{m}\sum \dfrac{\partial{\mathcal L \left(x_i, \theta\right)}}{\partial \theta} \]
使用 minibatch 的方法有两个好处:

  1. minibatch 计算出来的 loss 可以看做是整个 trainset 的 loss 的近似值.
  2. minibatch 中, 可以并行地计算 m 各样本, 因此, 使用 minibatch 的方法比原始的 SGD 方法速度更快.

然而, SGD 算法有其固有缺点:

  1. 对初始值要求很高, 如果参数的初始化不好, 经常不能收敛
  2. 学习率比较难设置, 由于每一层 input 数据的 scale 不同, 导致 backward 的梯度的scale 也不同, 为了保证不会 gradient vanish, 只能设置较小的 learning rate, 而, 较小的 learning rate 使得整个学习过程很慢
  3. 第 N 层的输入受前面 N-1 层的影响, 在深度学习中, 网络层数很多, 因此, 及时前面 layer 的很小的影响, 当到达第 N 层的时候, 会被放大很多倍.

在深度神经网络中, 每一层输入数据的分布都不同, 因此, 每一层的参数都要去学习不同的分布. 而主要由于上述 #3 的原因, 使得这个过程比较困难. 为了说明这个问题, 使用一个简单的例子. 考虑如下的一个两层的神经网络:
\[ F_2\left(F_1\left(\theta_1, x\right), \theta_2\right) \]
\(F_1\) 的输出作为 \(F_2\) 的输入, 因此, 在学习的过程中, 当 \(F_1\) 的输入 \(x\) 的分布变化的时候(在选择 minibatch 的使用根本无法保证每次选择的 minibatch 的分布是相同的), \(F_2\) 的参数在向最优解收敛的过程中就会产生偏差, 因此, 导致了收敛速度变慢. 而 BN 就是要解决这个问题.
想象一下, 如果 \(F_2\) 的学习不受 \(F_1\) 的输入 \(x\) 的影响, 即不管 \(x\) 输入的是什么, \(F_2\) 的输入都是相同的分布, 这样, \(F_2\) 就不用调整去适应由于 \(x\) 输入变化带来的影响, 那么, 是否就可以解决这个问题了呢? 因此, 这就是 BN 的提出.

Batch Normalization 是什么

Input: Values of \(x\) over a mini-batch, \(\mathcal{B}={x_{1\dots m}}\)
Output: \(\gamma, \beta\)
\[ \begin{align} {\mu}_{\beta} &=\frac{1}{m}\sum_{i=1}^{m}x_i\\ {\sigma}_{\beta}^{2} &=\frac{1}{m}\sum_{i=1}^{m}\left(x_i-{\mu}_{\beta}\right)^2\\ \hat{x}_i &=\frac{x_i-{\mu}_{\beta}}{\sqrt{{\sigma}_{\beta}^2+\epsilon}}\\ y_i&={\gamma}\hat{x}_i+\beta \end{align} \]

为什么 Batch Normalization 可以加速训练

  1. 允许网络使用较高的 learning rate. 在传统的深度网络训练中, 如果使用较大的 learning rate 很容易导致 gradient vanish 或者 gradient explode. 通过在整个网络中 normalize activations, 可以防止参数的较小的改变被应用到较大的或者次优的activation 中. 另一方面, BN 使得网络对于 parameter 的 scale 更加鲁棒. 通常情况下, large learning rate 会 increase the scale of layer parameters, 进而会放大 BP 的梯度, 导致了 model explosion. BN 的使用使得网络在 BP 的时候不会受到parameter scale 的影响. 这是因为: $$ \begin{align} & BN\left(Wu\right)=BN\left(\left(aW\right)u\right)\\ \\ & \frac{\partial{BN((aW)u)}}{\partial{u}}=\frac{\partial{Wu}}{\partial{u}} \\ \\ & \frac{\partial{BN((aW)u)}}{\partial{aW}}=\frac{1}{a} \cdot \frac{\partial{Wu}}{\partial{W}} \end{align} $$
  2. 具有一定的regularization 作用, 可以减少 Dropout 的使用. dropout 的作用是方法 overfitting, 实验发现, BN 可以 reduce overfitting.
  3. 降低 \(L_2\) 权重衰减系数.
  4. 取消 LRN(Local Response Normalization).
  5. Reduce the photomatric distortions. 因为 BN 使得训练过程更快, 能 observe 到的 sample 次数变少, 所以, 减少 distorting 使得网络 focus 在真实的图片上面.
  6. BN 不仅仅限定在 ReLU 上, 而且, 对其它的 activation 也同样适用.

MXNet 源码

Forward 阶段

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
virtual void Forward(const OpContext &ctx,
const std::vector<TBlob> &in_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &out_data,
const std::vector<TBlob> &aux_states) {
using namespace mshadow;
using namespace mshadow::expr;
CHECK_EQ(in_data.size(), 3);
CHECK_EQ(aux_states.size(), 2);
if (ctx.is_train) {
CHECK_EQ(out_data.size(), 3);
CHECK_EQ(req.size(), 3);
} else {
CHECK_GE(out_data.size(), 1);
CHECK_GE(req.size(), 1);
CHECK_EQ(req[batchnorm::kOut], kWriteTo);
}
Stream<xpu> *s = ctx.get_stream<xpu>();
// 1/m
const real_t scale = static_cast<real_t>(in_data[batchnorm::kData].shape_[1]) /
static_cast<real_t>(in_data[batchnorm::kData].shape_.Size());
Tensor<xpu, 4> data;
Tensor<xpu, 4> out;
if (in_data[batchnorm::kData].ndim() == 2) {
Shape<4> dshape = Shape4(in_data[batchnorm::kData].shape_[0],
in_data[batchnorm::kData].shape_[1], 1, 1);
// 输入数据
data = in_data[batchnorm::kData].get_with_shape<xpu, 4, real_t>(dshape, s);
// 输出数据
out = out_data[batchnorm::kOut].get_with_shape<xpu, 4, real_t>(dshape, s);
} else {
data = in_data[batchnorm::kData].get<xpu, 4, real_t>(s);
out = out_data[batchnorm::kOut].get<xpu, 4, real_t>(s);
}
// slope 是原文中的 \gamma
Tensor<xpu, 1> slope = in_data[batchnorm::kGamma].get<xpu, 1, real_t>(s);
// bias 是算法中的 \beta
Tensor<xpu, 1> bias = in_data[batchnorm::kBeta].get<xpu, 1, real_t>(s);
// 获取moving average, 如果 use_global_stats, 那么就要使用 moving average
Tensor<xpu, 1> moving_mean = aux_states[batchnorm::kMovingMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> moving_var = aux_states[batchnorm::kMovingVar].get<xpu, 1, real_t>(s);
// whether use global statistics
if (ctx.is_train && !param_.use_global_stats) {
// 不使用 moving average, 那么, 要针对当前 minibatch 计算期望和方差
Tensor<xpu, 1> mean = out_data[batchnorm::kMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> var = out_data[batchnorm::kVar].get<xpu, 1, real_t>(s);
CHECK(req[batchnorm::kMean] == kNullOp || req[batchnorm::kMean] == kWriteTo);
CHECK(req[batchnorm::kVar] == kNullOp || req[batchnorm::kVar] == kWriteTo);
// The first three steps must be enforced.
// 计算当前 minibatch 的均值
mean = scale * sumall_except_dim<1>(data);
// 计算当前 minibatch 的方差
var = scale * sumall_except_dim<1>(F<mshadow_op::square>(
data - broadcast<1>(mean, data.shape_)));
if (param_.fix_gamma) {
// 减均值, 除以方差, 其中 param_.eps 是一个非常小的值, 防止出现除 0 的错误
Assign(out, req[batchnorm::kOut], (data - broadcast<1>(mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(var + param_.eps, data.shape_)) +
broadcast<1>(bias, out.shape_));
} else {
Assign(out, req[batchnorm::kOut], broadcast<1>(slope, out.shape_) *
(data - broadcast<1>(mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(var + param_.eps, data.shape_)) +
broadcast<1>(bias, out.shape_));
}
} else {
// test/predict 或者 use_global_stats, 直接使用 moving average
Assign(out, req[batchnorm::kOut], broadcast<1>(slope /
F<mshadow_op::square_root>(moving_var + param_.eps),
data.shape_) * data +
broadcast<1>(bias - (slope * moving_mean) /
F<mshadow_op::square_root>(moving_var + param_.eps), data.shape_));
}
}

Backward 阶段

先给出梯度公式
$$ \begin{align} \frac{\partial{l}}{\partial{\hat{x}_i}} &= \frac{\partial{l}}{\partial{y_i}} \cdot \gamma \\ \frac{\partial{l}}{\partial{\sigma_{\mathcal{B}}^2}} &= \sum_{i=1}^m{\frac{\partial{l}}{\partial{\hat{x}_i}}\cdot\left(x_i-\mu_{\mathcal{B}}\right)} \cdot \frac{-1}{2}\left(\sigma_{\mathcal{B}^2}+\epsilon\right)^{-3/2} \\ \frac{\partial{l}}{\partial{\mu_{\mathcal{B}}}} &= \left( \sum_{i=1}^m{\frac{\partial{l}}{\partial{\hat{x}_i}} \cdot \frac{-1}{\sqrt{\sigma_{\mathcal{B}}^2+\epsilon}}} \right) + \frac{\partial{l}}{\partial{\sigma_{\mathcal{B}}^2}} \cdot \frac{\sum_{i=1}^m{-2(x_i-\mu_{\mathcal{B}})}}{m} \\ \frac{\partial{l}}{\partial{x_i}} &= \frac{\partial{l}}{\partial{\hat{x}_i}} \cdot \frac{1}{\sqrt{\sigma_{\mathcal{B}}^2+\epsilon}} + \frac{\partial{l}}{\partial{\sigma_{\mathcal{B}}^2}} \cdot \frac{2(x_i-\mu_{\mathcal{B}})}{m} + \frac{\partial{l}}{\partial{\mu_{\mathcal{B}}}} \cdot \frac{1}{m} \\ \frac{\partial{l}}{\partial{\gamma}} &= \sum_{i=1}^m{\frac{\partial{l}}{\partial{y_i}}} \cdot \hat{x}_i \\ \frac{\partial{l}}{\partial{\beta}} &= \sum_{i=1}^m{\frac{\partial{l}}{\partial{y_i}}} \end{align} $$

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
virtual void Backward(const OpContext &ctx,
const std::vector<TBlob> &out_grad,
const std::vector<TBlob> &in_data,
const std::vector<TBlob> &out_data,
const std::vector<OpReqType> &req,
const std::vector<TBlob> &in_grad,
const std::vector<TBlob> &aux_states) {
using namespace mshadow;
using namespace mshadow::expr;
CHECK_EQ(out_grad.size(), 1);
CHECK_EQ(in_data.size(), 3);
CHECK_EQ(out_data.size(), 3);
CHECK_EQ(in_grad.size(), 3);
Stream<xpu> *s = ctx.get_stream<xpu>();
Tensor<xpu, 4> data, grad, grad_in;
const real_t scale = static_cast<real_t>(out_grad[batchnorm::kOut].shape_[1]) /
static_cast<real_t>(out_grad[batchnorm::kOut].shape_.Size());
if (in_data[batchnorm::kData].ndim() == 2) {
Shape<4> dshape = Shape4(out_grad[batchnorm::kOut].shape_[0],
out_grad[batchnorm::kOut].shape_[1], 1, 1);
data = in_data[batchnorm::kData].get_with_shape<xpu, 4, real_t>(dshape, s);
grad = out_grad[batchnorm::kOut].get_with_shape<xpu, 4, real_t>(dshape, s);
grad_in = in_grad[batchnorm::kData].get_with_shape<xpu, 4, real_t>(dshape, s);
} else {
data = in_data[batchnorm::kData].get<xpu, 4, real_t>(s);
grad = out_grad[batchnorm::kOut].get<xpu, 4, real_t>(s);
grad_in = in_grad[batchnorm::kData].get<xpu, 4, real_t>(s);
}
Tensor<xpu, 1> mean = out_data[batchnorm::kMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> var = out_data[batchnorm::kVar].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> slope = in_data[batchnorm::kGamma].get<xpu, 1, real_t>(s);
// Tensor<xpu, 1> bias = in_data[kBeta].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> gslope = in_grad[batchnorm::kGamma].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> gbias = in_grad[batchnorm::kBeta].get<xpu, 1, real_t>(s);
// update moving avg
Tensor<xpu, 1> moving_mean = aux_states[batchnorm::kMovingMean].get<xpu, 1, real_t>(s);
Tensor<xpu, 1> moving_var = aux_states[batchnorm::kMovingVar].get<xpu, 1, real_t>(s);
if (ctx.is_train && !param_.use_global_stats) {
// get requested temp space
Tensor<xpu, 2> workspace = ctx.requested[batchnorm::kTempSpace].get_space<xpu>(
mshadow::Shape2(3, mean.shape_[0]), s);
Tensor<xpu, 1> gmean = workspace[0];
Tensor<xpu, 1> gvar = workspace[1];
Tensor<xpu, 1> tmp = workspace[2];
// 使用 moving average 算法, 更新 mean 和 var, 用于test/predict
moving_mean = moving_mean * param_.momentum + mean * (1 - param_.momentum);
moving_var = moving_var * param_.momentum + var * (1 - param_.momentum);
// cal
// 计算方差的梯度
gvar = sumall_except_dim<1>((grad * broadcast<1>(slope, data.shape_)) *
(data - broadcast<1>(mean, data.shape_)) *
-0.5f *
F<mshadow_op::power>(broadcast<1>(var + param_.eps, data.shape_),
-1.5f));
// 计算均值的梯度的第一部分
gmean = sumall_except_dim<1>(grad * broadcast<1>(slope, data.shape_));
gmean *= -1.0f / F<mshadow_op::square_root>(var + param_.eps);
tmp = scale * sumall_except_dim<1>(-2.0f * (data - broadcast<1>(mean, data.shape_)));
tmp *= gvar;
// 均值梯度的第二部分
gmean += tmp;
// assign
if (!param_.fix_gamma) {
//\gamma 的梯度
Assign(gslope, req[batchnorm::kGamma],
sumall_except_dim<1>(
grad * (data - broadcast<1>(mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(var + param_.eps, data.shape_))));
// 输入 x 的梯度
Assign(grad_in, req[batchnorm::kData],
(grad * broadcast<1>(slope, data.shape_)) *
broadcast<1>(1.0f / F<mshadow_op::square_root>(var + param_.eps), data.shape_) +
broadcast<1>(gvar, data.shape_) * scale * 2.0f * (data - broadcast<1>(mean,
data.shape_)) +
broadcast<1>(gmean, data.shape_) * scale);
} else {
Assign(grad_in, req[batchnorm::kData], grad *
broadcast<1>(1.0f / F<mshadow_op::square_root>(var + param_.eps), data.shape_) +
broadcast<1>(gvar, data.shape_) * scale * 2.0f * (data - broadcast<1>(mean,
data.shape_)) +
broadcast<1>(gmean, data.shape_) * scale);
}
// 偏置的梯度
Assign(gbias, req[batchnorm::kBeta], sumall_except_dim<1>(grad));
} else {
// use global statistics with freeze moving mean and var.
if (!param_.fix_gamma) {
Assign(gslope, req[batchnorm::kGamma],
sumall_except_dim<1>(
grad * (data - broadcast<1>(moving_mean, data.shape_)) /
F<mshadow_op::square_root>(broadcast<1>(moving_var + param_.eps, data.shape_))));
Assign(grad_in, req[batchnorm::kData], (grad * broadcast<1>(slope, data.shape_)) *
broadcast<1>(
1.0f / F<mshadow_op::square_root>(moving_var + param_.eps), data.shape_));
} else {
Assign(grad_in, req[batchnorm::kData], grad *
broadcast<1>(
1.0f / F<mshadow_op::square_root>(moving_var + param_.eps), data.shape_));
}
}
}

Batch Normalization 使用方法

  1. 一般BN layer 放在 FC 和 conv 的后面
  2. 在 train 阶段, 对每一个 minibatch 使用 BN, 那么, 在 test 或者 predict 的时候怎么使用 BN 呢? 常见的做法是使用整个 train set 计算出 \(\mu\). 由于 train set 的数据量非常大, 计算 \(\mu\) 计算量非常大, 所以, 经常采用的技术是使用 moving average 算法, 在 train 阶段计算出 \(\mu\)

CAFFE 中的使用方法

使用 residual net 网络作为示例

1
2
3
4
5
6
7
8
9
10
11
12
layer {
bottom: "res2a_branch2b"
top: "res2a_branch2b"
name: "bn2a_branch2b"
type: "BatchNorm"
# use_global_stats 表示是否使用全部数据的统计值(该数据实在train 阶段通过
# moving average 方法计算得到)训练阶段设置为 fasle, 表示通过当前的
# minibatch 数据计算得到, test/predict 阶段使用 通过全部数据计算得到的统计值
batch_norm_param {
use_global_stats: true
}
}

MXNet 中的使用方法

1
2
3
4
# data 是输出数据, eps 是为了防止出现除 0 二废除 0 而给出的一个很小的值, momentum
# 是moving average 的动量项, fix_gamma 参考上面的 C++ 代码, use_global_stats 作用同
# CAFFE 中的解释.
bn = mxnet.symbol.BatchNorm(data, eps, momentum, fix_gamma, use_global_stats)

Reference

Batch Normalization