CBN论文笔记

论文:Cross-Iteration Batch Normalization.Zhuliang Yao,Yue Cao,Shuxin Zheng,Gao Huang,Stephen Lin 


摘要

    当mini-batch的大小设置的比较小时,BN的作用大大减小。因为太少的样本估计的统计特征无法代表整个数据集的统计特征。为了解决这个问题,本文提出了Cross-Iteration Batch Normalization(CBN),可以联合最近几次迭代的样本来提高估计的质量。这其中的难点在于由于每次迭代后模型权重会改变,因此不同迭代之间的网络激活值无法比较。由此本文提出一种基于泰勒多项式的方法来补偿网络权重的改变,从而能准确地估计统计特征,并应用到BN中。当目标检测和图像分类模型使用小mini-batch时,使用CBN超过使用BN或不加入补偿的CBN的模型。代码开源:https://github.com/Howal/Cross-iterationBatchNorm 


引言

    BN对深度学习的成功起到了非常重要的作用,它可以解决网络训练时由于参数改变而出现的内部协方差平移(internal covariate shift)的问题。这种平移会使得网络收敛的速度变慢,BN通过计算每个mini-batch的均值和方差将网络激活值进行归一化,来解决这个问题。加入BN后,网络训练可以使用更高的学习率、以及更随意的网络初始权重。

    BN假设mini-batch里面的统计特征(分布)可以反映整个数据集的统计特征(分布)。当mini-batch很小时这个假设会失效,少量的样本无法代表整个分布。BN的这个缺点导致它无法广泛应用到内存消耗量很大的任务,比如:目标检测、语义分割、动作识别等,因为这个这些任务的mini-batch不能设置的很大。

    同时也出现了很多替代方法,包括:Layer Normalization(LN)、Instance Normalization(IN)、Group Normalization(GN)等。然后这些方法并不能像BN一样具有通用性,比如GN用于目标检测,但是又会拖慢推断速度。SyncBN通过跨GPUs处理更好的mini-batch,但是还是需要多个GPUs。

    然而很少人会利用多个最近的mini-batch中计算统计特征,虽然可以增加统计参数的置信度,但是缺点是由于网络参数的改变导致无法互相比较。如下图所示,直接计算多个mini-batch的置信度会得到一个很低的精度,这种方法称为Naive CBN。

图片1.jpg

    本文提出了一种补偿网络权重改变的方法,称为Cross-Iteration Batch Normalization(CBN)。相邻迭代之间的网络权重是逐渐变化的,因此当前迭代的平均值和方差可以通过一个低阶泰勒多项式逼近当前网络权重,该多项式定义在相对于网络权重的统计特征的梯度上。最近多次迭代的补偿的均值和方差,与当前迭代的均值和方差进行平均,以产生更好的统计估计。

    如上图所示,CBN在小batch_size可以得到较好的性能提升。本文还将CBN应用到图像分类和目标检测任务上进行实验。CBN带来的计算量可以忽略不计,因为先前迭代的统计参数已经计算出来,泰勒多项式也很容易计算。


BN回顾

    BN通过mini-batch内计算得到的统计参数来白化每层的激活值。设θt为网络权重,xt,it)为某层中第t个mini-batch中第i个样本的激活值。BN对该层的输出做如下归一化:

图片2.jpg

    xt,i^(θt)是白化后的均值为0方差为1的输出特征图,开方里后面一项是一个小常数。µtt)和σtt) 是当前mini-batch中计算的均值和方差:

图片3.jpg

    上式中m是样本数量。白化后的特征图xt,i^(θt)经过一个线性变换后可以提高表示能力:

图片4.jpg

    其中γ和β是可学习的参数,初始值为1和0。

    当m很小时,当前mini-batch计算得到的均值和方差受到严重的噪声影响,会大大影响BN的效果。对于图像分类来说,batch size可以设置为32,然后对于显存占用高的模型,比如目标检测、语义分割,batch size会被设置为1或2,此时BN就失效了。


跨迭代

    为了解决BN出现的问题,一个自然的想法是联合计算当前迭代和前几个迭代的统计参数。然而第(t-τ)个迭代的统计参数μt-τt-τ)和νt-τt-τ)是在权重θt-τ下计算得到的。因此直接从多个迭代中计算统计参数会得到一个不精确的估计。

    由于梯度的特性,相邻迭代间的网络权重改变是平滑的,因此可以通过泰勒多项式从μt-τt-τ)和νt-τt-τ)逼近μt-τt)和νt-τt),如下:

图片5.jpg

    其中dμt-τt-τ) / dθt-τ和dνt-τt-τ) / dθt-τ是均值和方差分别相对于网络权重的梯度,O(||θt - θt-τ||2)表示泰勒展开的二阶无穷小,当θt - θt-τ很小时该项省略。上式中的梯度无法以低计算量获得,因为要计算l层之前所有层的梯度。本文发现,随着网络层索引r减小,上面计算的梯度也会减小。网络前面层的权值变化对后层激活分布的影响减弱,可能是BN内部协变量漂移减少的原因。经过研究,这里在l层截断这些偏微分:

图片6.jpg


CBN

    补偿网络权重的改变后,CBN计算当前迭代和最近k-1个迭代来获得统计特征:

图片7.jpg

    然后将得到的统计参数用于白化输出特征图:

图片8.jpg

    因此CBN用于获得统计参数的样本数是BN的k倍。由于每次计算CBN时不需要重复计算前面迭代的样本,因此CBN的计算量、内存占用仅比BN大一点点。

    CBN的一个重要参数是时间窗口大小:k,即用于计算统计参数的迭代数。更大的k可以增加统计的样本数量,但是泰勒展开逼近的精度会下降。本文认为k=8比较合适。但是刚开始梯度变化比较快,此时为burn-in阶段,分类任务的前25epoch和检测任务的前3轮设置k=1。

    图形描述整个计算过程如下:

图片9.jpg

    整个过程的伪代码如下:

图片10.jpg




上一篇:
下一篇:

首页 所有文章 机器人 计算机视觉 自然语言处理 机器学习 编程随笔 关于