BatchNorm 论文笔记

ICML 2015
训练网络通常来说是复杂的，因为每一层的分布都会随着每一次的训练而改变。输入分布的改变叫做 covariate shift 。

首先对于 SGD 优化来讲，其目标为：

$\Theta=\argmin_\Theta\frac{1}{N}\sum_{i=1}^N \ell(\mathrm{x}_i,\Theta)$

其中 $N$ 是整个训练集的大小，当选取 batch-size 大小为 $m$ 时，每次训练得到的梯度为（原文似乎少了一个求和符号？）：

$\frac{1}{m}\sum_{i=1}^m\ell(\mathrm{x}_i,\Theta)$

现在来考虑损失函数是：

$\ell=F_2(F_1(u,\Theta_1),\Theta_2)$

$F_1,F_2$ 可以为任意的变换（$\ell$ 为最终的 loss，那么其实一般 $F_2$ 就是最终的损失函数，比如交叉熵，$F_1$ 则是网络的最后一层，这里探讨的是更一般的情况）。这个式子也可以看成是将输入 $x=F_1(u,\Theta_1)$ 喂给 $\ell=F_2(x,\Theta_2)$ 得到的结果，所以假如说存在一个输入分布是能够使网络快速学习的，那么这样的分布也能使子网络进行快速学习。基于此，最好使每个子网络的输入分布不变。

现在重点讨论使子网络输入分布不变，这样做有什么好处呢？考虑一个 sigmoid 激活函数：

$\mathrm{z}=g(W\mathrm{u}+\mathrm{b})$

其中 $\mathrm{u}$ 是网络输入，$W$ 是权重矩阵，$\mathrm{b}$ 是偏置，$g(x)=\dfrac{1}{1+\exp{(-x)}}$，当 $|x|$ 变大时，$g'(x)$ 将趋近于 $0$（这便是饱和非线性 saturating nonlinearity，意思是函数的值域有界，当输入趋于正负无穷时，函数值不趋于无穷），这样的话只有在 $x$ 较小的时候，梯度才比较大，在剩下的情况下网络学习得都很慢。随着网络得加深，这种现象也越来越严重，通常采用 $ReLU(x)=\max(x,0)$ 、更好的初始化、较小的学习率来解决这一问题。现在假如能够让输入分布保持相对稳定，那么优化器便更容易脱离饱和状态，加速网络训练。

文章中把在训练过程中内部结点的输入数据分布不断改变的行为叫做 Internal Covariate Shift，基于此提出了 Batch Normalization，可以加快网络的训练，甚至对网络有正则化作用，减少了对 Dropout 的需要。

历史上已经研究过，当网络的输入“白化”（whitened，线性变换成 $0$ 均值，单位方差，并去相关）之后，网络的训练速度会加快。所以文章的基本思路就是对网络每一层的输入做白化操作。

但是并没有这么简单，首先考虑对激活值进行标准化，假设一个带参数的网络是这个样子的：

$x=u+b$

其中 $u$ 是输入，$x$ 是输出，$b$ 是参数，假如对于所有的训练样本得到的输出结果的集合为 $\mathcal{X}=\{x_{1\cdots N}\}$，进行的标准化操作为：

$\hat{x}=x-E[x]$

其中 $E[x]=\frac{1}{N}\sum_{i=1}^Nx_i$，假如梯度下降的时候不考虑 $E[x]$ 对 $b$ 的依赖，那么进行的更新为：

$b\gets b+\Delta b$

其中 $b\propto -\partial\ell/\partial\hat{x}$，那么更新完之后：

$\begin{aligned} \hat{x}'&=x'-E[x']\\ \hat{x}'&=u+(b+\Delta b)-E[u+(b+\Delta b)]\\ \hat{x}'&=u+b-E[u+b]=\hat{x} \end{aligned}$

会发现梯度更新并不会改变激活值，这不就白学了。问题在于梯度下降的时候没有把 Normalization 考虑进去。那么为了解决这一问题，可以要求网络在训练的时候，给定参数得到的激活值总是呈现想要的分布。也就是让参数更新时能够保证激活的输出直接就是 normalize 好的。

令向量 $\mathrm{x}$ 为一层的输入，$\mathcal{X}$ 为训练集在这一层的输入集合，normalization 的过程即为：

$\hat{\rm{x}}=\text{Norm}(\rm{x},\mathcal{X})$

norm 之后的输入与 $\rm{x}$ 有关，也和 $\mathcal{X}$ 有关（而之前考虑的情况是认为 $E[x]$ 不参与反向传播的计算，即之前并没有考虑 $\mathcal{X}$ 这一部分），对于这样的 norm 计算反向传播的话，需要计算：

$\frac{\partial\rm{Norm}(\rm{x},\mathcal{X})}{\partial\rm{x}},\frac{\partial\rm{Norm}(\rm{x},\mathcal{X})}{\partial\rm\mathcal{X}}$

但是假如使用这样一种方法的话，对该输入进行白化是很昂贵的，因为需要计算

$\rm{Cov}[x]=\rm{E}_{\rm{x}\in\mathcal{X}}[\rm{xx}^T]-\rm{E}[\rm{x}]\rm{E}[\rm{x}]^T$

然后再计算

$\rm{Cov[x]}^{-\frac{1}{2}}(\rm{x-E[x]})$

并且还要计算这些变换对应的导数。由于在整个输入上进行 normalization 十分昂贵，并且并不是完全可导的，因此考虑在每一个维度上分别进行 normalization。考虑一个 $d$ 维输入 $\mathrm{x}=(x^{(1)},\cdots,x^{(d)})$，对每个维度进行 normalization：

$\hat{x}^{(k)}=\frac{x^{(k)}-\mathrm{E}[x^{(k)}]}{\sqrt{\mathrm{Var}[x^{(k)}]}}$

其中期望和方差是在整个训练集上进行计算的，可以这样理解，训练集中每个数据作为当前层的输入有 $d$ 维，normalization 对每个维度单独进行，单个维度计算时用的个数是整个数据集的大小。

但是仅仅进行 normalization 操作可能会改变每层能够表示的东西，为了解决这一问题，文章使得增加的变换可以表示恒等映射，方法是增加缩放和平移参数 $\gamma^{(k)},\beta^{(k)}$：

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

假如学出来的 $\gamma^{(k)}=\sqrt{\mathrm{Var}[x^{(k)}]},\beta^{(k)}=\mathrm{E}[x^{(k)}]$，那么整体的效果便是一个恒等映射。

对于 mini-batch 的情况，考虑大小为 $m$ 的 mini-batch $\mathcal{B}$，先略去 $x^{(k)}$ 的上标，用 $x$ 来表示单个维度的情况，$\mathcal{B}=\{x_{1...m}\}$，定义 normalize 之后的值为 $\hat{x}_{1...m}$ ，整体变换之后为 $y_{1...m}$ ，称以下变换为 Batch Normalizing Transform：

$\mathrm{BN}_{\gamma,\beta}:x_{1...m}\to y_{1...m}$

整体算法如下：

好的，现在让我们来推一下反向传播，主要使用链式法则，我重新用比较易于理解的方式写了一下。

$\begin{aligned} \frac{\partial\ell}{\partial\hat{x}_i}&=\gamma\frac{\partial\ell}{\partial y_i}\\ \frac{\partial\ell}{\partial\sigma_\mathcal{B}^2}&=\sum_{i=1}^m(x_i-\mu_{\mathcal{B}})(-\frac{1}{2}(\sigma_\mathcal{B}^2+\epsilon)^{-\frac{3}{2}})\\ \frac{\partial\ell}{\partial\mu_{\mathcal{B}}}&=\sum_{i=1}^m(-\frac{1}{\sqrt{\sigma_{\mathcal{B}}^2+\epsilon}}\frac{\partial\ell}{\partial\hat{x}_i})+\frac{\partial\ell}{\partial\sigma_\mathcal{B}^2}\frac{1}{m}\sum_{i=1}^m(-2(x_i-\mu_\mathcal{B}))\\ \frac{\partial\ell}{\partial x_i}&=\frac{1}{\sqrt{\sigma_\mathcal{B}^2+\epsilon}}\frac{\partial\ell}{\partial\hat{x}_i}+\frac{2(x_i-\mu_\mathcal{B})}{m}\frac{\partial\ell}{\partial\sigma_\mathcal{B}^2}+\frac{1}{m}\frac{\partial\ell}{\partial\mu_\mathcal{B}}\\ \frac{\partial\ell}{\partial\gamma}&=\sum_{i=1}^m\hat{x}_i\frac{\partial\ell}{\partial y_i}\\ \frac{\partial\ell}{\partial\beta}&=\sum_{i=1}^m\frac{\partial\ell}{\partial y_i} \end{aligned}$

现在来看 inference 的情况，BN 在训练的时候需要计算一个 batch 中每个维度的均值和方差，但是在 inference 的时候我们希望对于一个样本就能够得到一个确定的结果，所以 BN 在训练和测试的时候的设置是不同的（我觉得这是一个比较重要的点）。操作如下：

1. $N_{\mathrm{BN}}^{\mathrm{tr}}$ 为训练好的网络权重，首先将其赋值给 inference 网络：$N_{\mathrm{BN}}^{\mathrm{inf}}\gets N_{\mathrm{BN}}^{\mathrm{tr}}$
2. 然后对于每个维度，计算：$\mathrm{E}[x]\gets \mathrm{E}_\mathcal{B}[\mu_\mathcal{B}],\mathrm{Var}[x]\gets\frac{m}{m-1}\mathrm{E}_\mathcal{B}[\sigma_\mathcal{B}^2]$
3. 然后在 $N_{\mathrm{BN}}^\mathrm{inf}$ 中，用 $\frac{x-\mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+\epsilon}}\gamma+\beta$ 来代替 $y=\mathrm{BN}_{\gamma,\beta}(x)$

文章中算法流程如下：

对于卷积神经网络而言，我们希望 normalization 可以保持住卷积网络的性质——在同一个 feature map 中的不同元素按照同一种方式进行 normalize，假如 feature map 的大小是 $p\times q$，batch-size 是 $m$，那么 batch norm 的有效 batch-size 为 $m'=mpq$，用 group normalization 论文中的图会更清晰一些：

Batchnorm 能够带来的其他好处：

（1）能够增大学习率：传统的深度网络中，过高的学习率可能会导致梯度消失或者梯度爆炸，并陷于局部最优解。使用 batch-norm 之后，可以防止参数微小的改变引起的过大梯度传播。BN 使得训练能够更好地适应较大 scale 的参数，当学习率比较大时，一般会导致参数的 scale 较大，这样会导致反向传播时得到的梯度过大，但是有了 BN 之后，参数的 scale 并不影响反向传播，首先证明，对于一个标量 $a$，有 $\mathrm{BN}(W\mathrm{u})=\mathrm{BN}((aW)\mathrm{u})$：

$\begin{aligned} \mathrm{BN}(W\mathrm{u})&=\gamma\frac{W\mathrm{u}-W\mathrm{E[u]}}{\sqrt{\mathrm{Var}[W\mathrm{u}]+\epsilon}}+\beta\\ \mathrm{BN}((aW)\mathrm{u})&=\gamma\frac{aW\mathrm{u}-aW\mathrm{E[u]}}{a\sqrt{\mathrm{Var}[W\mathrm{u}]+\epsilon/a^2}}+\beta\approx \mathrm{BN}(W\mathrm{u}) \end{aligned}$

所以自然的有 $\frac{\partial\mathrm{BN}((aW)\mathrm{u})}{\partial\mathrm{u}}=\frac{\partial\mathrm{BN}(W\mathrm{u})}{\partial\mathrm{u}}$，而并且有：

$\begin{aligned} \frac{\partial\mathrm{BN}((aW)\mathrm{u})}{\partial(aW)}&=\frac{1}{a}\frac{\partial\mathrm{BN}(W\mathrm{u})}{\partial W} \end{aligned}$

可以看到，较大的权重会导致较小的梯度。

（2）能够对模型进行正则化：在训练的过程中，在处理一个训练样本时，一个 mini-batch 中的其他样本也会被考虑其中，因此可以提高模型的泛化能力，可以替代 dropout

文章在 mnist 和 imagenet 上进行了一系列的实验，验证了以上的性质，具体请参见 paper 原文。

关于实际框架使用中 BN，我决定再开一篇文章来写（待填坑）