总结2: Batch Normalization反向传播公式推导及其向量化

1. 背景介绍

上周学习了吴恩达的Deep Learning专项课程的第二门课，其中讲到了Batch Normalization（BN）。但在课程视频中，吴恩达并没有详细地给出Batch Normalization反向传播的公式推导，而是从high level的角度解释了为什么Batch Normalization会work，以及如何在TensorFlow等framework中使用Batch Norm。

我自己首先尝试了一下推导BN反向传播的公式，但是用代码实现后跑的结果都不甚理想，收敛速度比不使用BN还要慢甚至有时候无法收敛，应该是公式推错了。接着我在网上搜索到Google最开始提出BN的论文， 里面给出了反向传播的公式，但是不是以向量化的形式给出的。众所周知，我们实现深度网络应该依赖于向量计算。我看着公式以自己的理解写出了向量的形式，但是实现后结果仍旧不正常。

接着在网上搜索其他人介绍BN的博客文章，绝大多数文章都是前面讲一大堆BN的好处，消除Internal Convariate Shift，加快收敛，减少Dropout的使用，起到部分正则化等等，然后涉及到核心的公式部分时，话锋一转，说BN反向传播部分的推导很简单，就是利用了Chain Rule，接着就给出了与论文中一模一样的公式。看着让人很是头疼。

就这样停滞了大概三四天的时间，但是我实在不甘心仅仅会使用TensorFlow中提供的BN模块，而搞不懂BN的详细推导。终于，我下定决心抽出一整天的时间拿出纸笔一步步的演算，最终心静下来花了大概一个小时算出来，然后代码实现之后跑起来结果就正常了。

2. BN反向传播的详细推导

2.1 单个activation进行batch norm的情况

假设神经网络的第L层是BatchNorm层，其输入数据为 ZL ，其中 ZL 的维度是(n, dL )，n是样本个数， dL 代表第L层神经元的个数。

对 ZL 进行normalize得到：

Znorm,L=ZL−1∗mean(ZL)T1∗var(ZL)T‾‾‾‾‾‾‾‾‾√

其中， 1 代表元素全为1的列向量，如不特殊说明，其长度为n； ∗ 代表矩阵乘法； mean(ZL) 是长度为 dL 的列向量； (√var(ZL) 也是长度为 dL 的列向量。

然后对 Znorm,L 加上可被学习的scale参数 γL 和shift参数 βL ：

YL=Znorm,L×(1∗γT)+1∗βT

其中， × 代表矩阵对应元素相乘； γT 代表 γ 的转置。 γ 和 β 都是长度为 dL 的向量。

以上是Batch Norm层正向传播的严谨的公式，很多文章里都习惯于使用具有broadcasting功能的公式，如用一个矩阵减去一个向量等操作，虽然python里的numpy支持这种运算，但是公式如果也用这种方式写则很不严谨，也对我们的求导造成很大的困扰。

接下来是反向传播部分的推导。因为 Znorm,L 是对 ZL 按列normalize得到的，每两列之间是完全独立的，求导的时候也不会互相干扰，所以为了推导的清晰简便起见，我们取 ZL 的第j列进行推导，以下将 ZLj 记为小写的 x 。

现在我们有如下数据：

x=[x1,x2,...,xn]T

mean(x)=∑ni=1xin

var(x)=∑np=1(xp−mean(x))2n

stddev(x)=var(x)‾‾‾‾‾‾√

xnorm=⎡⎣⎢⎢⎢⎢xnorm1xnorm2...xnormn⎤⎦⎥⎥⎥⎥=⎡⎣⎢⎢⎢⎢⎢(x1−mean(x))/var(x)‾‾‾‾‾‾√(x2−mean(x))/var(x)‾‾‾‾‾‾√...(xn−mean(x))/var(x)‾‾‾‾‾‾√⎤⎦⎥⎥⎥⎥⎥

通过chain rule，我们有：

(αLαxi)1×1=(αLαxnorm)T1×n∗(αxnormαxi)n×1

其中，

αLαxnorm=⎡⎣⎢⎢⎢⎢⎢⎢⎢⎢αLαxnorm1αLαxnorm2...αLαxnormn⎤⎦⎥⎥⎥⎥⎥⎥⎥⎥n×1

αxnormαxi=⎡⎣⎢⎢⎢⎢⎢⎢⎢⎢αxnorm1αxiαxnorm2αxi...αxnormnαxi⎤⎦⎥⎥⎥⎥⎥⎥⎥⎥n×1

因为 xnormj=xj−mean(x)var(x)√ ，通过分式求导法则有：

αxnormjαxi=(I{j=i}−1n)var(x)‾‾‾‾‾‾√−αvar(x)√αxi(xj−mean(x))var(x)

下面求 αvar(x)√αxi ：

αvar(x)‾‾‾‾‾‾√αxi=12var(x)‾‾‾‾‾‾√αvar(x)αxi

因为 var(x)=∑np=1(xp−mean(x))2n ，我们有：

αvar(x)αxi=1n∑p=1n2(xp−mean(x))(I{p=i}−1n)=2n∑p=1nxp(I{p=i}−1n)−2n∑p=1nmean(x)(I{p=i}−1n)=2n(xi−mean(x))−2n(mean(x)−mean(x))=2n(xi−mean(x))

所以：

αvar(x)‾‾‾‾‾‾√αxi=12var(x)‾‾‾‾‾‾√αvar(x)αxi=12var(x)‾‾‾‾‾‾√2n(xi−mean(x))=xi−mean(x)nvar(x)‾‾‾‾‾‾√

所以：

αxnormjαxi=(I{j=i}−1n)var(x)‾‾‾‾‾‾√−αvar(x)√αxi(xj−mean(x))var(x)=(I{j=i}−1n)var(x)‾‾‾‾‾‾√−(xi−mean(x))nvar(x)√(xj−mean(x))var(x)=I{j=i}−1nvar(x)‾‾‾‾‾‾√−(xi−mean(x))(xj−mean(x))nvar(x)3‾‾‾‾‾‾‾√

因为

αxnormαx=⎡⎣⎢⎢⎢⎢⎢⎢⎢⎢αxnorm1αx1αxnorm2αx1...αxnormnαx1αxnorm1αx2αxnorm2αx2...αxnormnαx2............αxnorm1αxnαxnorm2αxn...αxnormnαxn⎤⎦⎥⎥⎥⎥⎥⎥⎥⎥=In×nvar(x)‾‾‾‾‾‾√−1∗1Tnvar(x)‾‾‾‾‾‾√−(x−mean(x))(x−mean(x))Tnvar(x)‾‾‾‾‾‾√3

So:

αLαx=(αxnormαx)TαLαxnorm

以下是我用Scala实现的Batch Norm反向传播方法的向量化版本，值得注意的是，由于向量对向量求导的结果是一个矩阵，多个向量对向量求导的结果是一个三维张量，而Scala中的Breeze数值计算库现并不支持张量运算，所以在我用Scala实现的Batch Normalization的反向传播版本中，不可避免地使用的for循环对 ZL 的每一列循环计算；而如果读者使用的Python，Numpy支持三位数组，故可以把此for循环也替换成张量运算，使得运算速度更快！

  private def backWithBatchNorm(dYCurrent: DenseMatrix[Double], yPrevious: DenseMatrix[Double]): (DenseMatrix[Double], DenseMatrix[Double]) = {
    val numExamples = dYCurrent.rows
    val oneVector = DenseVector.ones[Double](numExamples)

    val dZDelta = dYCurrent *:* this.activationFuncEval(zDelta)
    val dZNorm = dZDelta *:* (oneVector * beta.t)
    val dAlpha = dZDelta.t * oneVector / numExamples.toDouble
    val dBeta = (dZDelta *:* zNorm).t * oneVector / numExamples.toDouble

    val dZ = DenseMatrix.zeros[Double](z.rows, z.cols)
    for (j <- 0 until z.cols) {
      val dZNormJ = dZNorm(::, j)
      val dZJ = (DenseMatrix.eye[Double](dZNormJ.length) / currentStddevZ(j) - DenseMatrix.ones[Double](dZNormJ.length, dZNormJ.length) / (numExamples.toDouble * currentStddevZ(j)) - (z(::, j) - currentMeanZ(j)) * (z(::, j) - currentMeanZ(j)).t / (numExamples.toDouble * pow(currentStddevZ(j), 3.0))) * dZNormJ
      dZ(::, j) := dZJ
    }

    val dWCurrent = yPrevious.t * dZ / numExamples.toDouble
    val dYPrevious = dZ * w.t

    val grads = DenseMatrix.vertcat(dWCurrent, dAlpha.toDenseMatrix, dBeta.toDenseMatrix)

    (dYPrevious, grads)
  }