CBN(Cross-Iteration Batch Normalization)论文详解

原文链接：Cross-Iteration Batch Normalization
代码链接：https://github.com/Howal/Cross-iterationBatchNorm

随着BN的提出，现有的网络基本都在使用。但是在显存有限或者某些任务不允许大batch（比如检测或者分割任务相比分类任务训练时的batchsize一般会小很多）的情况下，BN效果就会差很多，如下图所示，当batch小于16时分类准确率急剧下降。

为了改善小batch情况下网络性能下降的问题，有各种新的normalize方法被提出来了（LN、IN、GN），详情请看文章GN-Group Normalization。上述不同的normalize方式适用于不同的任务，其中GN就是为检测任务设计的，但是它会降低推断时的速度。

本文的思想其实很简单，既然从空间维度不好做，那么就从时间维度进行，通过以前算好的BN参数用来计算更新新的BN参数，从而改善网络性能。从上图可以看出使用新提出的CBN，即使batch size小也可以获得较好的分类准确率。

通过过去的BN参数来计算新的BN参数有个问题，就是过去的BN参数是根据过去的网络参数计算出的feature来估算的，新的BN参数计算时，参数已经更新过了，如果直接使用之前的参数来计算新的BN参数会使得参数估计不准，网络性能下降，如上图中的Naive CBN。为了改进这种缺陷，文章使用泰勒多项式来估算新的BN参数。下面来介绍具体的计算方式。

一、Batch Normalization

先来回顾一下原始的BN计算公式。
BN操作的公式如下
${\hat{x}}_{t,i}({\theta }_t)=\frac{x_{t,i}({\theta }_t)-{\mu }_t({\theta }_t)}{\sqrt{{\sigma }_t({\theta }_t{)}^2+ϵ}}$
$y_{t,i}({\theta }_t)=\gamma {\hat{x}}_{t,i}({\theta }_t)+\beta$
上式中${\theta }_t$表示第t个mini-batch训练时的网络参数，$x_{t,i}({\theta }_t)$表示第t个mini-batch中第i个样本经过网络得到的feature map，${\hat{x}}_{t,i}({\theta }_t)$表示bn后均值为0和方差为1的特征，${\mu }_t({\theta }_t)$和${\sigma }_t({\theta }_t)$表示当前mini-batch计算出来的均值和方差，$ϵ$是防除零的系数，$\gamma$和$\beta$表示BN中需要学习的参数。

${\mu }_t({\theta }_t)$和${\sigma }_t({\theta }_t)$计算方式如下
${\mu }_t({\theta }_t)=\frac{1}{m}{\sum }_{i=1}^m{x}_{t,i}({\theta }_t)$
${\sigma }_t({\theta }_t)=\sqrt{\frac{1}{m}{\sum }_{i=1}^m(x_{t,i}({\theta }_t)-{\mu }_t({\theta }_t){)}^2}=\sqrt{{\nu }_t({\theta }_t)-{\mu }_t({\theta }_t{)}^2}$
其中${\nu }_t({\theta }_t)=\frac{1}{m}{\sum }_{i=1}^m{x}_{t,i}({\theta }_t{)}^2$，m表示mini-batch里面有m个样本。

二、估计之前的均值和方差

假设现在是第t次迭代，那t之前的迭代计算时，拿$t-\tau$次迭代来说，已经计算过了对应的均值和方差了，但是之前的计算是使用之前的网络参数得到的，用符号表示为${\mu }_{t-\tau }({\theta }_{t-\tau })$和${\nu }_{t-\tau }({\theta }_{t-\tau })$。现在想要估计的是参数${\mu }_{t-\tau }({\theta }_t)$和${\nu }_{t-\tau }({\theta }_t)$。文章认为连续几次网络参数的变化是平滑的，所以根据泰勒展开式可以估计出上述的两个参数：
${\mu }_{t-\tau }({\theta }_t)={\mu }_{t-\tau }({\theta }_{t-\tau })+\frac{\partial {\mu }_{t-\tau }({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }}({\theta }_t-{\theta }_{t-\tau })+O(\Vert {\theta }_t-{\theta }_{t-\tau }{\Vert }^2)$
${\nu }_{t-\tau }({\theta }_t)={\nu }_{t-\tau }({\theta }_{t-\tau })+\frac{\partial {\nu }_{t-\tau }({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }}({\theta }_t-{\theta }_{t-\tau })+O(\Vert {\theta }_t-{\theta }_{t-\tau }{\Vert }^2)$

其中$\frac{\partial {\mu }_{t-\tau }({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }}$和$\frac{\partial {\nu }_{t-\tau }({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }}$表示对网络参数求偏导数，$O(\Vert {\theta }_t-{\theta }_{t-\tau }{\Vert }^2)$表示泰勒展开式的高阶项，当$({\theta }_t-{\theta }_{t-\tau })$较小时，高阶项可以忽略不计。

注意：要精确的获得上式中的$\frac{\partial {\mu }_{t-\tau }({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }}$和$\frac{\partial {\nu }_{t-\tau }({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }}$的值计算量会很大，因为网络中的第l层的参数${\mu }_{t-\tau }^l({\theta }_{t-\tau })$和${\nu }_{t-\tau }^l({\theta }_{t-\tau })$会依赖之前的层，例如依赖$\frac{\partial {\mu }_{t-\tau }^l({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }^r}\ne 0$和$\frac{\partial {\nu }_{t-\tau }^l({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }^r}\ne 0$，这里$r\le l$,${\theta }_{t-\tau }^r$表示第r层的参数。实际上文章发现当$r\le l$时，偏导数$\frac{\partial {\mu }_{t-\tau }^l({\theta }_t)}{\partial {\theta }_t^r}$和$\frac{\partial {\nu }_t^l({\theta }_t)}{\partial {\theta }_t^r}$减小的非常快

所以上两等式可以近似表示为下两式
${\mu }_{t-\tau }^l({\theta }_t)\approx {\mu }_{t-\tau }^l({\theta }_{t-\tau })+\frac{\partial {\mu }_{t-\tau }^l({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }^l}({\theta }_t^l-{\theta }_{t-\tau }^l)$
${\nu }_{t-\tau }^l({\theta }_t)\approx {\nu }_{t-\tau }^l({\theta }_{t-\tau })+\frac{\partial {\nu }_{t-\tau }^l({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }^l}({\theta }_t^l-{\theta }_{t-\tau }^l)$

三、Cross-Iteration Batch Normalization

CBN的工作方式如下

上一节利用之前的参数估算出来了当前参数下l层在$t-\tau$次迭代的参数值，通过这些估计值可以计算出l层在当前迭代时的BN参数，计算公式如下：
${\bar{\mu }}_{t,k}^l({\theta }_t)=\frac{1}{k}{\sum }_{\tau =0}^{k-1}{\mu }_{t-\tau }^l({\theta }_t)$
${\bar{\nu }}_{t,k}^l({\theta }_t)=\frac{1}{k}{\sum }_{\tau =0}^{k-1}max[{\nu }_{t-\tau }^l({\theta }_t),{\mu }_{t-\tau }^l({\theta }_t{)}^2]$
${\bar{\sigma }}_{t,k}({\theta }_t)=\sqrt{{\bar{\nu }}_{t,k}({\theta }_t)-{\bar{\mu }}_{t,k}({\theta }_t{)}^2}$
注意：在有效的统计时${\nu }_{t-\tau }^l({\theta }_t)\ge {\mu }_{t-\tau }^l({\theta }_t{)}^2$是一直都会满足的，但是利用泰勒展开式估算不一定能满足条件，所以上述使用了max函数来保证这一点。还要说的是，代码实现的时候并没有这样写，实现时k是指所有${\nu }_{t-\tau }^l({\theta }_t)\ge {\mu }_{t-\tau }^l({\theta }_t{)}^2$的元素，即过滤了不满足条件的值

最后CBN更新特征的方式同BN
${\hat{x}}_{t,i}^l({\theta }_t)=\frac{x_{t,i}^l({\theta }_t)-{\bar{\mu }}_t^l({\theta }_t)}{\sqrt{{\bar{\sigma }}_t^l({\theta }_t{)}^2+ϵ}}$

CBN的伪代码如下

四、计算量优化

文章附录还有一个求偏导的优化方式来节省计算量，假设l和l-1层卷积的通道数为$C^l$和$C^{l-1}$,K表示l层卷积的kernelsize，因此${\mu }_{t-\tau }^l$和${\nu }_{t-\tau }^l$通道数为$C^l$，${\theta }_{t-\tau }^l$的维度数为$C^l\times C^{l-1}\times K$，如果直接计算$\frac{\partial {\mu }_{t-\tau }^l({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }^l}$和$\frac{\partial {\nu }_{t-\tau }^l({\theta }_{t-\tau })}{\partial {\theta }_{t-\tau }^l}$，时间复杂度为$O(C^l\times C^l\times C^{l-1}\times K)$。通过下面的推导我们可以知道可以在$O(C^{l-1}\times K)$和$O(C^l\times C^{l-1}\times K)$时间内求出$\mu$和$\nu$

下面拿$\mu$来举例说明，为了简化符号，下面用${\mu }^l$和${\theta }^l$代替${\mu }_{t-\tau }^l({\theta }_{t-\tau })$和${\nu }_{t-\tau }^l({\theta }_{t-\tau })$。
${\mu }_j^l=\frac{1}{m}{\sum }_{i=1}^m{x}_{i,j}^l$
上式中${\mu }_j^l$表示${\mu }^l$中第j个通道的值，$x_{i,j}^l$表示第i个样本的第j个通道值。

$x_{i,j}^l={\sum }_{n=1}^{C^{l-1}}{\sum }_{k=1}^K{\theta }_{j,n,k}^l\cdot y_{i+offset(k),n}^{l-1}$
上式中，n和k分别表示输入feature的通道维度索引和卷积的kernel维度索引，offset表示卷积时的索引值，$y^{l-1}$表示l-1层的输出。

$\frac{\partial {\mu }^l}{\partial {\theta }^l}\in R^{C^l\times C^l\times C^{l-1}\times K}$计算公式如下
$$\begin{gathered} [\frac{\partial {\mu }^l}{\partial {\theta }^l}{]}_{j,q,p,\eta }=\frac{\partial {\mu }_j^l}{\partial {\theta }_{q,p,n}^l} \\ =\frac{\partial \frac{1}{m}{\sum }_{i=1}^m{\sum }_{n=1}^{C^{l-1}}{\sum }_{k=1}^K{\theta }_{j,n,k}^l\cdot y_{i+offset(k),n}^{l-1}}{\partial {\theta }_{q,p,n}^l} \\ =\{\begin{array}{lrc}\frac{1}{m}{\sum }_{i=1}^m{y}_{i+offset(\eta ),p}^{l-1}& & ,j=q\\ 0& & ,j\ne q\end{array} \end{gathered}$$
从上式可以看出当j=q时才要计算，其它情况不用计算，这样就减少了计算量。

到这里CBN就介绍完了，思想很简单，但是公式较多，可以看看代码实现，代码量不多。最后要说的是，这篇文章提出来时为了解决小batch训练的情况，yolov4提出有一点就是方便用户使用一个GPU训练，所以yolov4借用了该算法的思想。