Paper Reading 《Decorrelated Batch Normalization》

Paper Reading: DBN

Decorrelated Batch Normalization

CVPR2018. Github(Written in Lua). Paper.

1. Intuition

Ioffe & Szegedy(2017) 提出的Batch Normalization:

$\widehat{x_i}=\gamma \frac{x_i-\mu }{\sqrt{{\sigma }^2+ϵ}},\text{where}\mu =\frac{1}{m}{\sum }_{j=1}^m{\mu }_j,{\sigma }^2=\frac{1}{m}{\sum }_{j=1}^m(x_j-\mu {)}^2$

但是最关键的问题是$m$个维度之间关联度较高。

因此本文从ZCA白化入手提出了DBN, i.e. Decorrelated Batch Normalization：

$\widehat{x_i}={\Sigma }^{-\frac{1}{2}}(x_i-\mu )$

实现DBN的四个问题？

1. DBN如何进行反馈传播？
2. 为何选ZCA不选PCA？
3. 如何计算${\Sigma }^{-\frac{1}{2}}$矩阵？
4. 白化操作的样本量如何确定？

2. 算法细节

2.1 Notation

令$\mathbf{X}\in {\mathbb{R}}^{d\times m}$, $d$为维度，$m$为mini-batch的大小，白化变换$\varphi :{\mathbb{R}}^{d\times m}\to {\mathbb{R}}^{d\times m}$可定义为：
$$\varphi (\mathbf{X})={\Sigma }^{-\frac{1}{2}}(\mathbf{X}-\mu \cdot 1^T)$$
其中$\mu =\frac{1}{m}\mathbf{X}\cdot 1,\Sigma =\frac{1}{m}(\mathbf{X}-\mu \cdot 1^T)(\mathbf{X}-\mu \cdot 1^T{)}^T+ϵ\cdot \mathbf{I}$。白化之后目的是$\hat{\mathbf{X}}=\varphi (\mathbf{X})s.t.\hat{\mathbf{X}}{\hat{\mathbf{X}}}^T=I$

那实现DBN该如何解决上述四个问题呢？

2.2 随机轴交换

${\Sigma }_{PCA}^{-\frac{1}{2}}={\Lambda }^{-\frac{1}{2}}\mathbf{D}$

正交特征向量$\mathbf{D}=[{\mathbf{d}}_1,\cdots ,{\mathbf{d}}_d]$是理论上是不能够被唯一确定的，但是PCA就选定了$\Lambda$按照${\sigma }_1,\cdots ,{\sigma }_d$从大到小的顺序进行排列，即旋转坐标轴到方差最大的方向。

但是恰恰是由于PCA总是选取方差最大的方向，而神经网络的激活数值会随权重更新而改变，这使得其在不同的batch、iteration的过程中的所旋转的坐标轴方向也都是不同的，也就是所说的stochastic axis swapping。这也导致了如下的loss震荡、收敛效果差的问题：

2.3 ZCA白化的BP反向传播过程实现

公式推导+解释

鉴于上述情况选取ZCA白化${\Sigma }_{ZCA}^{-1/2}=\mathbf{D}{\Lambda }^{-\frac{1}{2}}{\mathbf{D}}^T$.

那么ZCA白化后续的梯度下降的反向传播就应当如下所示：

2.3.1 Forward Pass (Appendix A.1)

从$x_j\to \widehat{x_j}$

2.3.2 Back Propagation (Appendix A.2)

BP过程从$L\to x_i$

简化版本：
$$\frac{\partial L}{\partial {\mathbf{x}}_i}=\left(\frac{\partial L}{\partial {\tilde{\mathbf{x}}}_i}-\mathbf{f}+{\tilde{\mathbf{x}}}_i^T\mathbf{S}-{\tilde{\mathbf{x}}}_i^T\mathbf{M}\right){\Lambda }^{-1/2}{\mathbf{D}}^T$$

2.4 算法流程

在每一轮训练中的前馈传播和反馈传播的伪算法。

2.4.1 前馈传播算法

白化是发生在每一个mini-batch中的，$\mu$和$\Sigma$在每一个batch中都是迭代更新的。

• 在Forward Pass中先计算PCA，ZCA也是在PCA基础上计算得到。（与直接用$W_{zca}={\Sigma }^{-1/2}$相比，不知数值上有区别否？）
• 步骤10&11的$\lambda$就是个相同的超参，不是严格的moving average迭代更新${\mu }_{n+1}=\frac{n}{n+1}{\mu }_n+\frac{1}{n+1}{X}_{n+1}$

2.4.2 反向传播算法

值得一提的是，在CNN算法中，DBN的输入形如${\mathbf{X}}_C\in {\mathbb{R}}^{h\times w\times d\times m}$，$h,w$分别表示feature map的维度(height$\times$width)，$d,m$分别表示feature maps的数量和batch中样本的数量。

2.5 组白化

为了保证每个batch中有足够的样本数量来做白化操作，我们将激活层数值沿特征维度$d$，划分成$k_G(k_G<d)$个较小的组(防止出现batch样本数$m<<d$的情况)，并在每个组中进行白化。$k_G=1$时，DBN退化为BN。此时计算复杂度从$O(d^2\max (m,d))$降至$O(\frac{d}{k_G}(k_G^2(\max (m,k_G))))$，通常，我们选择$k_G<m$，此时组白化计算复杂度为$O(md{k}_G)$。

3. 实验结果

3.1 消融实验

3.2 CNN和ResNet

在CNN上比较BN与DBN

加在ResNet中加入DBN效果也会提高一部分。

作者在附录B中也讨论了DBN其实比较time costly，主要是ZCA中矩阵求逆过于耗时，所以该组也在CVPR19中提出了采用Iterative Normalization，并在CVPR21上发表了Group Whitening详细探讨了组白化的内容。