封面图片:John Cameron on Unsplash

批量归一化 (Batch normalization) 基本上已经成为卷积神经网络中的固定成员。它是由Sergey Ioffe等人在“Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift” 这篇文章中提出,用于解决神经网络的“内部协变量漂移”问题。你可以在这里找到原文:

Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
Training Deep Neural Networks is complicated by the fact that thedistribution of each layer’s inputs changes during training, as the parametersof the previous layers change. This slows down the training by requiring lowerlearning rates and careful parameter initialization, and makes it notoriousl…

本文是对原文核心内容的摘抄。我尽可能在保持原意清晰的基础上使用更接近中文习惯的表达方式。


随机梯度下降(Stochastic gradient descent,简写为SGD)以及它的各种衍生体是优化神经网络卓有成效的方法。该方法为优化参数 $\Theta$ 需要最小化损失值

$$ \Theta = \text{arg}\min_\Theta \frac{1}{N} \sum^N_{i=1} \ell(x_i, \Theta)$$

其中 $x_{1...N}$ 代表训练数据集。随机梯度下降分步执行,每一步涉及到数量为 $m$ 的小批量数据 $x_{1...m}$。损失函数的梯度值可以通过该小批量数据基于下式近似得到

$$\frac{1}{m} \frac{\partial \ell(x_i, \Theta)}{\partial \Theta}$$

每次使用小批量样本而非单个样本有多个好处。首先,小批量样本的损失梯度是对整个训练集的近似,批尺寸越大、近似程度越高。其次,考虑现代计算平台所提供的并行算力,一次计算 $m$ 个样本的效率要高于单个样本计算 $m$ 次。

随机梯度下降虽然简洁高效,但是优化过程中模型的超参数尤其是学习率、参数初始值需要细心的调教。训练的复杂程度可归因于每一层的输入受该层之前所有层的影响,以至于参数的微小变化会随着网络加深而被放大。

面对输入分布的变化,网络层不得不去持续适应。一个学习系统的输入分布发生改变的现象被称为协变量漂移。然而,协变量漂移的概念可以从学习系统的整体拓展到它的局部,例如一个子网络或者一层。如网络计算

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

其中 $F_1$ 与 $F_2$ 为任意变换,参数 $\Theta_1, \Theta_2$ 可以通过最小化损失 $\ell$ 来学习。学习 $\Theta_2$ 可以看作是将 $x=F_1(u, \Theta_1)$ 当做输入馈入子网络

$$ \ell = F_2(x, \Theta_2)$$

例如当 $m$ 为批尺寸 ,$\alpha$为学习率时,梯度下降的一步

$$\Theta_2 \gets \Theta_2 - \frac{\alpha}{m} \sum^m_{i=1} \frac{\partial F_2(x_i, \Theta_2)}{\partial \Theta_2}$$

等同于输入为 $x$ 的独立网络 $F_2$。因此,让训练数据与测试数据具备相同分布规律这种改善网络训练效率的做法同样适用于子网络。所以训练时保持 $x$ 的分布稳定是有益的。这样 $\Theta_2$ 就不必再去适应 $x$ 的分布变化。

子网络之外的网络层同样会受益于子网络的固定输入分布。例如一个使用Sigmod激活函数 $g(x)=\frac{1}{1+e^{-x}}$ 的网络层 $z=g(Wu+b)$。当 $|x|$ 增大时,$g'(x)$ 趋于0。这意味着除了那些绝对值较小的 $x$,流向 $u$ 的梯度会消失致使训练迟缓。另外,$x$ 受到 $W, b$ 以及之前所有层的影响,训练时这些参数的改变可能会将 $x$ 的诸多维度移动至非线性的饱和区域进而减慢收敛。这些效应都会因为网络深度增加而放大。在实际中,饱和效应造成的梯度消失可以通过采用ReLU激活函数 $ReLU(x) = \max(x, 0)$ 、妥善的初始化以及小学习率来解决。然而,如果保证非线性输入的分布稳定的话,网络就不易被卡在饱和区域,进而可以加速训练。

这种深度神经网络训练时内部节点的分布变化被称为内部协变量漂移。消除它自然可以加速训练过程。批量归一化(Batch Normalization)正是这样做的一种有效方案。它通过添加归一化步骤来修正层输入的均值与方差,还可以降低梯度对于参数数值范围的依赖程度,改善梯度在网络中的传播。如此就可以使用更大的学习率而无需担心网络发散。有了批量归一化就无需使用Dropout,而且可以使用那些具备饱和区域的非线性激活函数。

以消除内部协变量漂移为目标

内部协变量漂移为定义为网络训练过程中由于权重参数变化造成的网络激活分布的改变。将网络的输入通过线性变换使得具备0均值与单位方差以解耦的白化(Whitened)操作可以加速网络训练是早已人尽皆知的事实。由于每一层的输入都是前一层的输出,因此为每一层实现类似的操作应当是有益的。如此迈出了消除内部协变量漂移的第一步。

在训练的每一步添加白化操作可以通过直接修改网络或者改变优化算法的参数以网络激活输出为依据来实现。然而,如果将这些改动散布在优化过程中,那么梯度下降在更新网络参数的同时也会更新归一化参数,减弱梯度的影响。例如输入为 $u$,偏置为 $b$ 的层通过减去训练数据激活数值的平均值来归一化:

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

其中 $x=u+b$,$\chi = {x_{1...N}}$ 为训练数据对应 $x$ 的一系列数值,$E[x]=\frac{1}{N}\sum^N_{i=1}x_i$。如果梯度下降忽略 $E[x]$ 对 $b$ 的依赖,它将会按照 $b \gets b + \Delta b$ 来更新,其中 $\Delta b \propto -\partial\ell/\partial\hat{x}$。那么

$$ u+(b+\Delta b) - E[u+(b+\Delta b)] = u+b-E[u+b]$$

对 $b$ 以及归一化的组合更新造成网络成输出不变,进而loss不变。随着训练继续,$b$ 持续增大而loss不变。这种状况在归一化不仅置中而且缩放激活数值时恶化。在梯度下降外部计算归一化参数会造成模型爆炸。

究其原因是由于梯度下降优化没有考虑归一化造成的。为了纠正这一点,需要保证对于任意参数值,网络的激活输出分布是一定的。再次定义 $x$ 为层输入,$\chi$ 为训练数据集对应 $x$ 的集合。归一化可以写作

$$\hat{x}=\text{Norm}(x, \chi)$$

它的值不仅与当前训练个样本 $x$ 而且与所有训练样本的 $\chi$ 有关,而它们若是网络某一层输出的话都取决于网络参数 $\Theta$。对于反向传播需要计算Jacobians

$$\frac{\partial \text{Norm}(x, \chi)}{\partial x} \text{ 与 } \frac{\partial \text{Norm}(x, \chi)}{\partial \chi} $$

忽略后一项就会造成模型爆炸。在这一构架下,层输入的白化操作是昂贵的。它需要计算协方差矩阵 $\text{Cov}[x]=E_{x\in \chi}[xx^T]-E[x]E[x]^T$ 以及它的平方根以便获得白化激活 $\text{Cov}[x]^{-1/2}(x-E[x])$,还有这些变换的倒数来反向传播。所以消除内部协变量漂移的寻求目标应当是一种可导并且无需训练集全部样本的归一化操作。

基于小批量统计的归一化

由于白化每一层的输入耗费巨大且并非总是可导,问题需要做出必要的简化。第一点,分别而非联合归一化每一层的输入与输出,使得它们的均值为0方差为1。对于一个维度为 $d$ 的输入 $x=(x^{(1)}...x^{(d)})$,每一个维度的归一化

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

其中的平均值与方差由训练集得到。该种方式在不解耦特征的情况下依旧可以加速训练。

单纯的归一化每一层的输入可能改变每一层的行为。例如归一化Sigmoid的输入可能将造成输入被限制在非线性的线性区域。要解决这一问题,需要保证嵌入网络的变换具备恒等变换的能力。为此,对于每一个激活 $x^{(k)}$ 设定参数 $\gamma^{(k)}$ 与 $\beta^{(k)}$ 来缩放与偏移归一化后的数值。

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

这些参数与网络参数一样是可学习的,并可以还原网络的表征能力。如果令 $\gamma ^{(k)}=\sqrt{Var[x ^{(k)}]}$ 且 $\beta^{(k)}=E[x^{(k)}]$,激活函数就与原先相同。

原本的归一化需要使用整个训练集来实现,但是这对于随机梯度下降来说是不现实的。因此需要第二个简化:使用小批量数据来近似每个激活的均值与方差。这样用来归一化的统计值可以全面参与到梯度反向传播中。注意小批量的使用是建立在计算每个维度的方差而非联合协方差的基础上。在联合的情况下,由于小批量的数据量相比白化激活的数量少造成奇异协方差矩阵,而不得不使用正则化。

考虑尺寸为 $m$ 的小批量数据 $\mathcal{B}$。由于归一化是针对每个激活独立执行的,我们仅关注特定激活 $x^{(k)}$ 并忽略 $k$。该小批量数据包含 $m$ 个数值

$$\mathcal{B}=\{x_{1...m}\}$$

记归一化后的数值为 $\hat{x}_{1...m}$,线性变换为 $y_{1...m}。则变换

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

为批量归一化变换。现将其称为算法1。该算法中 $\epsilon$ 是保持小批量样本方差数值稳定的常数。

输入:小批量样本中 $x$ 的数值:$\mathcal{B}=\{x_{1...m}\}$;需要学习的参数:$\gamma, \beta$
输出:$\{y_i = \mathrm{BN}_{\gamma, \beta}(x_i)\}$
$$\mu _{\mathcal{B}} \gets \frac{1}{m}\sum^m_{i=1}x_i \text{ 小批量均值}$$ $$\sigma^2_{\mathcal{B}} \gets \frac{1}{m}\sum^m_{i=1}(x_i-\mu_{\mathcal{B}})^2 \text{ 小批量方差}$$ $$\hat{x_i} \gets \frac{x_i-\mu_{\mathcal{B}}}{\sqrt{\sigma^2_{\mathcal{B}} + \epsilon}} \text{ 归一化}$$ $$y_i \gets \gamma \hat{x}_i + \beta \equiv \mathrm{BN}_{\gamma, \beta}(x_i) \text{ 缩放与偏移}$$

该批量归一化可以用于网络中的任意激活。虽然 $y=\mathrm{BN}{\gamma, \beta}(x)$ 表明 $\gamma$ 与 $\beta$ 是习得的,但是批量归一化并非是孤立的处理每个样本的激活。相反,$\mathrm{BN}{\gamma, \beta}(x)$ 同时依赖当前训练样本以及小批量数据中的其它样本。缩放与偏移后的数值 $y$ 会传递到网络的其它层。归一化之后的激活 $\hat{x}$ 是变换的内部量,但是却很重要。只要小批量数据采样自相同的分布并忽略 $\epsilon$ 的话,任意 $\hat{x}$ 的数值分布的均值为0方差为1。每一个归一化的激活 $\hat{x}^{(k)}$ 可以看作是添加了线性变换 $y^{(k)}=\gamma^{(k)}\hat{x}^{(k)} +\beta^{(k)}$ 之后的子网络的输入。这些子网络输入都具备固定的均值与方差。尽管这些归一化的 $\hat{x}^{(k)}$ 的联合分布可能会随着训练而改变,归一化的输入仍然应当对子网络的训练起加速作用,进而加速整个网络。

在训练时还应当考虑针对该变换损失 $\ell$ 梯度的反向传播,并且计算批量归一化变换中参数的梯度。根据链式法则:

$$\frac{\partial\ell}{\partial\hat{x}_i}=\frac{\partial\ell}{\partial y_i} \cdot \gamma \\ \frac{\partial\ell}{\partial\sigma^2_\mathcal{B}}= \sum^m_{i=1} \frac{\partial\ell}{\partial\hat{x}_i}\cdot(x_i-\mu_{\mathcal{B}})\cdot\frac{-1}{2}(\sigma^2_{\mathcal{B}}+\epsilon)^{-3/2} \\ \frac{\partial\ell}{\partial\mu_\mathcal{B}}= (\sum^m_{i=1} \frac{\partial\ell}{\partial\hat{x}_i}\cdot\frac{-1}{\sqrt{\sigma^2_{\mathcal{B}}+\epsilon}})+\frac{\partial\ell}{\partial\sigma^2_\mathcal{B}}\cdot\frac{\sum^m_{i=1}-2(x_i-\mu_{\mathcal{B}})}{m} \\ \frac{\partial\ell}{\partial{x}_i}= \frac{\partial\ell}{\partial\hat{x}_i}\cdot\frac{1}{\sqrt{\sigma^2_{\mathcal{B}}+\epsilon}}+\frac{\partial\ell}{\partial\sigma^2_\mathcal{B}}\cdot\frac{2(x_i-\mu_{\mathcal{B}})}{m}+\frac{\partial\ell}{\partial\mu_\mathcal{B}}\cdot\frac{1}{m}\\ \frac{\partial\ell}{\partial\gamma}= \sum^m_{i=1} \frac{\partial\ell}{\partial{y}_i}\cdot\hat{x}_i \\ \frac{\partial\ell}{\partial\beta}=\sum^m_{i=1} \frac{\partial\ell}{\partial{y}_i}$$

可见,批量归一化是一种能够在网络中实现激活归一化的可导变换。这确保了网络训练过程中,每一层需要处理的输入仅包含微弱的内部协变量漂移,进而加速训练。另外学习到的仿射变换使得批量归一化具备恒等变换的能力,确保了网络的能力。

批量归一化网络的训练与推演

在激活之后插入批量归一化变换即可构建批量归一化网络。之前以 $x$ 为输入的层现在以 $\mathrm{BN}(x)$ 为输入。批量归一化的后的模型可以使用批量梯度下降、基于小批量$m>1$数据的随机梯度下降或者其它变体如Adagrad来训练。针对小批量数据的激活使用归一化可以提升训练效率,但是对于推演来说却是多余的,因为输出应当只取决于输入。因此,一旦网络训练完成,使用基于全体而非小批量统计的归一化

$$\hat{x}=\frac{x-E[x]}{\sqrt{Var[x]+\epsilon}}$$

忽略 $\epsilon$ 这些归一化的激活在训练时拥有均值0与方差1。使用无偏置的方差估计 $Var[x]=\frac{m}{m-1}\cdot E_{\mathcal{B}}[\sigma^2_{\mathcal{B}}]$,其中 $m$ 为训练小批量数据的数量 $\sigma^2_{\mathcal{B}}$ 是其方差。使用滑动平均可以跟踪模型训练时的准确率。由于推演时均值与方差是固定的,归一化等同于针对每一个激活值的线性变换。

批量归一化的卷积网络

批量归一化可用在网络的任意激活部分。这里关注跟随逐元素非线性激活的仿射变换:

$$z=g(Wu+b)$$

其中 $W$ 与 $b$ 是模型参数,$g(\cdot)$ 是非线性激活如Sigmoid或者ReLU。这种组合方式同时适用于全连接与卷积层。由于批量归一化的对象为 $x=Wu+b$,即插入到非线性激活之前。虽然 $u$ 也可以被归一化,但是由于它是其它非线性激活的输出,它的分布形状会随着训练而变化,限制它前两阶矩不会消除协变量漂移。相反,$Wu+b$ 倾向于具备类高斯对称的、非稀疏的分布。对其归一化可以产生更加稳定分布的激活。

由于归一化对象为 $Wu+b$,偏置 $b$ 在后续求均值时效应微弱因而可被忽略。因此 $z=g(Wu+b)$ 可以替换为

$$z=g(\mathrm{BU}(Wu)$$

批量归一化变换被各自应用在 $x=Wu$ 的每个具备独立 $\gamma^{(k)}, \beta^{(k)}$ 的维度上。

对于卷积层来说,归一化应当具备卷积性质,这样同一张特征图上不同位置的不同元素的归一化方法保持一致。为做到这一点需要联合归一化一组小批量数据中所有位置的激活。算法1中 $\matchcal{B}$ 为一个特征图在小批量数据中跨元素与位置的数值集合。如果批量尺寸为 $m$,特征图尺寸为 $p\times q$,则等效批尺寸为 $m'=|\mathcal{B}|=m\cdot pq$。每一张特征图而非每一个激活有自己的$\gamma^{(k)}, \beta^{(k)}$ 。

批量归一化支持高学习率

传统深度网络中过高的学习率会造成梯度爆炸或者消失,亦或卡在局部最优。批量归一化可以解决这一问题。通过归一化网络的激活可以阻止激活梯度参数中较小的变化被放大到较大的次优变化。例如避免卡在非线性激活的饱和区域。

批量归一化还使得训练对于参数尺度的变化更加有弹性。一般来说大学习率会放大层参数的数值范围,进而放大反向传播时的梯度数值造成模型爆炸。然而,批量归一化之后的层受参数范围的影响变小。例如对于缩放银子 a,

$$\mathrm{BN}(Wu)=\mathrm{BN}((aW)u)$$

可以证明

$$\frac{\partial\mathrm{BN}((aW)u)}{\partial u}=\frac{\partial\mathrm{BN}( Wu)}{\partial u}\\\frac{\partial\mathrm{BN}((aW)u)}{\partial (aW)}=\frac{1}{a}\cdot \frac{\partial\mathrm{BN}( Wu)}{\partial W}$$

缩放因子 $a$ 不影响层的Jacobian以及梯度传播。更进一步,较大的参数会导致较小的梯度,批量归一化能够稳定参数的增长。

另外不难猜测批量归一化会造成层Jacobian的奇异值趋向于1而有助于训练。对于两个具备归一化输入的相邻层,归一化向量的变换为 $\hat{z}=F(\hat{x})$。如果假设 $\hat{x}$ 与 $\hat{z}$ 呈高斯分布且无关,且$F(\hat{x}) \approx J\hat{x}$ 为对应模型参数的线性变换,则 $\hat{x}$ 与 $\hat{z}$ 拥有单位方差,且 $I=\mathrm{Cov}[\hat{z}]=J\mathrm{Cov}[\hat{x}]J^T=JJ^T$。所以 $JJ^T=I$, $J$ 的奇异值均为1,可以保持反向传播时梯度的尺度。实际中,变换并非总是线性的,归一化也并不总是服从高斯分布。但是批量归一化应当会使梯度传播更加顺畅。更加准确的效应需要研究。

批量归一化提供模型正则化

批量归一化模型在训练过程中会同时考虑当前批次的其它样本,进而不会依据单一样本给出决定性的数值。在实验中这体现为网络泛化能力的提升。而Dropout仅仅用来防止过拟合,且在批量归一化网络中可以省去。

加速归一化网络的训练

归一化网络模型还应当采取如下措施加速训练:

  • 增大学习率
  • 去掉Dropout层
  • 减小L2权重正则化
  • 加速学习率衰减
  • 去除局部响应归一化
  • 更加彻底的随机化训练样本
  • 减少样本的人为畸变

总结

我认为这是一篇学术性非常强的文章。大量的公式以及推导需要耐心去解读。但是彻底读懂作者要表达的含义一定会带来相对应的回报。Batch Normalization几乎已经称为卷积神经网络中不可或缺的层类型。在设计网络时要留意作者在文中提出的使用方法与建议,以发挥其最大效能。