Batch normalization with TensorFlow

ResNet的坑还没填，先挖Batch Normalization (BN)的坑吧。读了Ioffe & Szegedy 2015年的论文,做个笔记。参考了几篇相关的博客，还是很有帮助的。如此文所述，BN也是神经网络中的一层，并且包含待训练的参数，而这些参数也是其精髓之一。下面我将分两个部分来说，(1) BN的优势、原理及推导, 以及(2) BN在TensorFlow下的实现。

BN的优势、原理及推导

Motivations

首先，Ioffe和Szegedy在摘要中说明了提出BN的动机，其原文为，
The deep neural network is complicated by the fact that the distribution of each layer’s inputs changes during training, as the parameters of the previous layers change. This slows down the training by requiring lower learning rates and careful parameter initialization, and makes it notoriously hard to train models with saturating nonlinearities.

即在网络的训练过程中，受到上一层参数调整的影响，该层输入的分布会发生变化，使得参数训练要重新适应这一分布变化，导致网络的训练时间和迭代次数增加。为了避免这些问题，传统方法只能采用小的学习率以及适当的参数初始化策略; 并且无法避免大数值输出在经过激活函数后进入饱和区 (saturating nonlinearities) 的问题。Ioffe和Szegedy将这一现象命名为internal covariate shift。

LeCun在1998年的论文中提出过这样一个观点，The newtork training converges faster if its inputs are whitened, i.e., linearily transformed to have zero means and unit variances, and decorrelated. 即对网络的输入进行归一化，有利于加速网络的收敛。这一点也是比较好理解的，参考传统机器学习算法，通常分类器的输入是人工提取的特征，这些特征在量纲和数量级上会有相当大的差异性，如果直接输入网络，数值较大的特征便会影响训练结果。为了解决这一问题，通常对特征进行去量纲的操作，最好的方式便是归一化，这种针对每一维特征进行归一化的处理方法也是BN的核心。

那么，在具有多层结构的网络内部，是否可以对每一层的输入做归一化处理，既稳定了样本的分布，也能加速网络的收敛。

BN的优势

根据论文的Abstract和Intro, BN的优势可以概括为以下四点，

解决internal convariate shift问题，加速网络训练;
避免了梯度传递过程原始数据大小对参数的影响，使得可以选择大的学习率;
正则化了网络，可以不使用dropout;
将输入集中在小区间内，避免经过激活函数的输出进入饱和区。

BN的原理与推导

BN文的第二节对normalization进行了讨论，并在第三节通过两个简化对Batch Normalization进行定义，以解决Full whitenning of each layer’s inputs is cotly的问题。

Nomalize each scalar feature indepently, by making whitened
Computation over a mini-batch instead of m individuals.

定义 $\mathbf{x}=[x^{(1)}, x^{(2)}, ... , x^{(K)}]$ , $x\in \chi,\chi={x_1,x_2,...,x_m}$ 为该层的输入 (上一层的输出)，其中m表示mini-batch中的样本数。$\hat{x}^{(k)}$表示$\mathbf{x}$第$k$维特征的whitening量，由下式给出

$\begin{align} \hat{x}^{(k)} = \frac{x^{(k)} - E[x^{(k)}]}{\sqrt{Var[x^{(k)}]}}, \notag \end{align}$

其中$E[\cdot]$和$Var[\cdot]$分别对应$x_i^{(k)}, i=1,2,…,m$的期望和方差。

以上虽然对输入做了归一话，但如BN文中所述Simply normalizing each input of a layer may change what the layer can represent。以激活函数sigmoid为例，如果$x^{(k)}$本来的数值较大，经过sigmoid函数以后会分布在接近饱和区的两端，而白化以后则剧集在sigmoid(x)=0附近，使得原有的分布消失。为了解决这一问题，需要对$\hat{x}$进行尺度(scale)和位移(shift)变换，从而有，

$y^{(k)} = \gamma^{(k)} \cdot \hat{x}^{(k)} + \beta^{(k)},$

其中$y^{(k)}$即为BN的输出，$\gamma^{(k)}$和$\beta^{(k)}$分别对应scale和shift，是Batch Normalization layer的参数，需要进行初始化，并在网络训练的过程中进行学习和调整。

BN和activation层的关系

我们知道，通常在Conv layer和FC layer之后，会加上非线性的激活函数层对输出进行调整，并作为下一层的输入。而BN层也是在Conv layer和FC layer之后对输出进行归一化的，二者的先后关系如何？参考BN文3.2节，定义$z$为activation layer的输出，则有，

$z = g[\mathrm{BN}(x)],$

即先进行Batch Normalization,后通过激活函数。

BN的训练和测试(inference)

BN的训练和测试可以参考下图，截取自Ioffe和Szegedy的论文，这里要注意的是inference的理解，即新样本需要经过怎样的预处理，以利用训练好的网络进行测试？

网络的训练过程采用了batch training，每个batch会有对应的期望$\muB$和标准差$\sigma{B}$，用于测试网络的样本在每一个layer依然需要归一化，此时对应的$E[x]$和$\mathrm{Var}[x]$用如下式进行估计，

$\begin{align} E[x] &= E[\mu_B], \notag \\ \mathrm{Var}[x] &= \frac{m}{m-1} \cdot E[{\sigma_B}^2], \notag \end{align}$

进而，BN layer的输出定义为，

$y = \frac{\gamma}{\sqrt{\mathrm{Var}[x]+\epsilon}}\cdot x + (\beta - \frac{\gamma E[x]}{\sqrt{\mathrm{Var}[x]+\epsilon}}).$

BN的TensofFlow实现

白化本身是很简单的，关键问题是如何训练每个$x^{(k)}$对应的$\gamma^{(k)}$和$\beta^{(k)}$。TensofFlow提供了多个batch_normalization相关的类或方法，比较简单的有tf.nn.batch_normalization，其参数如下

batch_normalization(
    x, # inputs
    mean, # 均值
    variance, # 方差
    offset, # shift beta
    scale, # scale gamma
    variance_epsilon, # epision
    name=None
)

简单实现一个DNN分类手写体的样例，这里基于tf.nn.moments生成mean和varianciences两个tensor

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data", one_hot=True)
x = tf.placeholder(tf.float32, shape=[None, 784])
y_ = tf.placeholder(tf.float32, shape=[None, 10]) # one-hot 10-dimensional vector
l_FC1 = 512
W_FC1 = tf.Variable(tf.truncated_normal(shape=[784,l_FC1], stddev=0.1))
b_FC1 = tf.Variable(tf.constant(0.1,shape=[l_FC1]))
# batch_norm_l1
x_FC1 = tf.matmul(x, W_FC1)+b_FC1
axis = list(range(len(x_FC1.get_shape()) - 1))
mean_FC1,var_FC1 = tf.nn.moments(x_FC1, axis)
gamma_FC1 = tf.Variable(tf.constant(0.1, shape=mean_FC1.get_shape()))
beta_FC1 = tf.Variable(tf.constant(0.1, shape=mean_FC1.get_shape()))
y_bn1 = tf.nn.batch_normalization(
    x = x_FC1,
    mean = mean_FC1,
    variance = var_FC1,
    offset = beta_FC1,
    scale = gamma_FC1,
    variance_epsilon = 1e-5,
    name= 'BN_FC1')
y_FC1 = tf.nn.relu(y_bn1)
...

测试部分单独开一个坑，见How to apply the batch-normalized net.

References

[1] Ioffe, S. and Szegedy, C., 2015, June. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning (pp. 448-456).
[2] Batch Normalization 学习笔记
[3] 谈谈Tensorflow的Batch Normalization