深度学习模型的优化算法

本文讲解训练深度学习模型的优化方法，包括优化器与初始化参数的策略。

问题的引入

一个有偏置项的全连接网络可以通过如下方式定义：

\begin{aligned} z_{i + 1} & = σ_{i} (W_{i}^{T} z_{i} + b_{i}), i = 1, . . ., L \\ h_{θ} & \equiv z_{L + 1} \\ z_{1} & \equiv x \end{aligned}

其中， $σ_{i}$ 为非线性激活函数， $z_{i} \in R^{n_{i}}$ ， $W_{i} \in R^{n_{i} \times n_{i + 1}}$ ， $b_{i} \in R^{n_{i + 1}}$ ，需要训练的参数为 $θ = {W_{1 : L}, b_{1 : L}}$

当一个 batch 有多个样本时，则可写成矩阵的形式如下：

Z_{i + 1} = σ_{i} (Z_{i} W_{i} + 1 b_{i}^{T})

注意， $b_{i} \in R^{n_{i + 1}}$ ，而 $Z_{i + 1} \in R^{m \times n_{i + 1}}$ ，为了保证能够正确相加，则需要用 $1 \in R^{m + 1}$ 与之相乘。而在实践中，通常并不这样做，而是利用广播机制将其自动对齐。因此，上式也可以直接写为：

Z_{i + 1} = σ_{i} (Z_{i} W_{i} + b_{i}^{T})

为了去训练一个全连接网络或者任何深度学习模型，我们不得不思考如下几个问题：

如何选择网络的宽度与深度？
如何优化目标函数？
如何初始化网络的权重？
如何保证网络在多次优化迭代后仍能继续训练？

优化算法

这里仅讨论梯度算法，先考虑以下两类：

批量梯度下降法。这是根据所有训练集估计梯度的算法，得到的是标准梯度，一定能保证收敛到极值点。但是缺陷在于计算太慢，且非常容易爆显存，一般只在高精度的任务中用到。
随机梯度下降法。与第一条相反，这种方法每次只随机选取单个样本估计梯度。这种算法的缺陷很明显，梯度估计很不准确，会让参数在极值点附近剧烈抖动，并且无法并行计算。

综合以上算法，使用一个以上而又不是全部的训练样本就有了深度学习中最常用的小批量（minibatch）随机梯度下降法，现在，通常也将其叫作随机梯度下降法，后文皆用该名称指代。

随机梯度下降法

提到优化算法，首先想到的就是随机梯度下降法（SGD）。每次抽取 $m$ 个小批量的样本，通过计算它们的梯度均值，即可得到梯度的估计 $\hat{g}$ ，算法如下：

o60ly

优点：

既有较为精确的梯度估计，又能保证适宜的计算效率

缺点：

学习率难以确定，设置过小会使收敛速度太慢，设置过大会使其剧烈震荡甚至不收敛。

比如，对于 $θ \in R^{2}$ ，考虑函数 $f (θ) = \frac{1}{2} θ^{T} P θ + q^{T} θ$ ，它对于不同学习率的优化情况如图所示：

8h1hj

牛顿法

考虑上面随机梯度下降算法的缺点，一个重要的方面就是震荡问题，如何减轻震荡呢？牛顿法就是这样一种更全局的方法，它在选择下一步的更新方向时，不仅仅考虑一阶导数最大的方向，还会考虑二阶导数，即走完这一步后，导数会不会变得更大。

参数更新步骤如下：

θ \leftarrow θ - α (\nabla_{θ}^{2} f (θ))^{- 1} \nabla_{θ} f (θ)

其中， $\nabla_{θ}^{2} f (θ) \in R^{n \times n}$ 为 Hessian 矩阵。

牛顿法默认 $α = 1$ ，同样求解上面的方程，效果如图：

z9o8l

优点：

可以看到，牛顿法在更新时没走“弯路”，每次都指向极值点的位置，解决了震荡的问题

缺点：

Hessian 矩阵的计算消耗太大，在深度学习问题中是不切实际的
牛顿法仅适用于凸优化问题，对于非凸函数的优化，Hessian 甚至都不是正定的

动量法

受以上两种方法的启发，考虑在梯度下降法中结合一些像牛顿法一样的全局结构，这就是动量法的思想。动量法引入动量项 $u$ ，和折扣因子 $β$ ，更新步骤如下：

\begin{aligned} u_{t + 1} & = β u_{t} + (1 - β) \nabla_{θ} f (θ_{t}) \\ θ_{t + 1} & = θ_{t} - α u_{t + 1} \end{aligned}

$v$ 本质是负梯度的指数衰减平均，保留了前面的梯度信息。当许多连续的梯度与当前梯度方向相同时，步长最大，否则，前面的梯度就起到了平滑的作用，很好地抑制了震荡。

从物理角度上来理解，可以将该算法视为牛顿力学下的粒子运动，当前的负梯度为粒子受到的力，目标函数的值为粒子的位置。那么每一步的负梯度也就是粒子受到的力会改变粒子的动量，当力的方向改变时，则根据矢量加法得到粒子下一步的动量方向，这也解释了该算法求历史负梯度平均的意义。

下图就是普通随机梯度下降法（ $α = 0$ ）与使用动量的随机梯度下降法对比：

gskxb

动量在计算初期很小，参数更新会很慢，通常在更新时给其乘以一个按时间衰减的系数：

θ_{t + 1} = θ_{t} - \frac{α}{1 - β^{t + 1}} v_{t}

Nesterov 动量法

这是动量法的一个改进。它的改变在于动量的更新：它根据超前一步的梯度来更新动量，更新步骤如下：

\begin{aligned} {\tilde{θ}}_{t + 1} & = θ_{t} - α u_{t} \\ g & = \nabla_{\tilde{θ}} f ({\tilde{θ}}_{t + 1}) \\ u_{t + 1} & = β u_{t} + (1 - β) g \\ θ_{t + 1} & = θ_{t} - α u_{t + 1} \end{aligned}

事实上，也可以写成下面的形式：

\begin{aligned} u_{t + 1} & = β u_{t} + (1 - β) \nabla_{θ} f (θ_{t} - α u_{t}) \\ θ_{t + 1} & = θ_{t} - α u_{t + 1} \end{aligned}

这种改进本质上是考虑了目标函数的二阶导信息，所以有了更快的收敛速度。这里不详细推导，可以参考下面这篇知乎文章：

比Momentum更快：揭开Nesterov Accelerated Gradient的真面目 - 知乎 (zhihu.com)

如图是二者的对比：

x3mg9

Adam 算法

为了解决学习率的设置问题，近年来提出了一些自适应学习率算法，本文只介绍 Adam 算法。这个算法也可以看作结合了动量的算法，更新步骤如下：

\begin{aligned} u_{t + 1} & = β_{1} u_{t} + (1 - β_{1}) \nabla_{θ} f (θ_{t}) \\ v_{t + 1} & = β_{2} v_{t} + (1 - β_{2}) (\nabla_{θ} f (θ_{t}))^{2} \\ θ_{t + 1} & = θ_{t} - \frac{α}{ϵ + \sqrt{v_{t + 1}}} u_{t + 1} \end{aligned}

事件中，通常还要做无偏修正：

\begin{aligned} u_{t + 1} & = β_{1} u_{t} + (1 - β_{1}) \nabla_{θ} f (θ_{t}) \\ v_{t + 1} & = β_{2} v_{t} + (1 - β_{2}) (\nabla_{θ} f (θ_{t}))^{2} \\ {\hat{u}}_{t + 1} & = u_{t + 1} / (1 - β_{1}^{t}) (bias correction) \\ {\hat{v}}_{t + 1} & = v_{t + 1} / (1 - β_{2}^{t}) (bias correction) \\ θ_{t + 1} & = θ_{t} - \frac{α}{ϵ + \sqrt{{\hat{v}}_{t + 1}}} {\hat{u}}_{t + 1} \end{aligned}

效果如图所示，非常惊艳！

vjgkv

初始化算法

模型参数的初始化对深度学习模型训练的影响非常大。比如，回顾我们前面的 MLP，假设将 $W_{i}$ 和 $b_{i}$ 都初始化为 0，则反向传播时会计算梯度为 0，整个模型将无法正常训练。我们期望有好的初始化算法，能使神经网络在正向传播和反向传播时，激活范数和梯度范数保持稳定，即方差一致性，这也是所有初始化算法的基本思想。

设 $fan_in$ 为输入维度， $fan_out$ 为输出维度。

Xavier 初始化

这种初始化算法适合激活函数为 sigmoid 或 tanh 的情况。

均匀分布：

$W_{i, j} \sim U (- a, a)$ ，其中

a = gain \times \sqrt{\frac{6}{fan_in + fan_out}}

高斯分布：

$W_{i, j} \sim N (0, {std}^{2})$ ，其中

std = gain \times \sqrt{\frac{2}{fan_in + fan_out}}

这里不做推导，可见论文 Understanding the difficulty of training deep feedforward neural networks

Kaiming 初始化

Kaiming 初始化是针对 Xavier 初始化在 ReLU 这一类整流线性激活函数表现不佳而提出的改进。

均匀分布：

$W_{i, j} \sim U (- b o u n d, b o u n d)$ ，其中

b o u n d = gain \times \sqrt{\frac{3}{fan_in}}

高斯分布：

$W_{i, j} \sim N (0, {std}^{2})$ ，其中

std = \frac{gain}{\sqrt{fan_in}}

对于 ReLU 函数， $gain = \sqrt{2}$

下面对高斯分布的正向过程做一些理论说明：

考虑使用高斯分布初始化参数，即 $W_{i} \sim N (0, σ^{2} I)$ ，在 MNIST 数据集上，用有 50 层隐藏层的 ReLU 为激活函数的神经网络训练，选取不同的的初始化方差得到正向传播时的激活范数和反向传播时的梯度范数如图：

7n8oy

通过上图，我们发现方差的选择对训练有巨大的影响，在 $σ^{2} = \frac{2}{n}$ 时，激活范数和梯度范数在所有层中都保持相近，而 $σ^{2} = \frac{3}{n}$ 时，激活范数却在层间传递时不断上升，这显然不是我们期望的情况。为保证方差基本不变，我们显然应该选择 $σ^{2} = \frac{2}{n}$ 。下面证明：

考虑将网络的中间层变量视为符合高斯分布的相互独立的随机变量， $x \sim N (0, 1)$ ， $W \sim N (0, \frac{1}{n})$ ，则 $E [w^{T} x] = 0$ ， $Var [w^{T} x] = 1$ （中心极限定理）。

所以，如果我们使用线性激活函数， $z_{i} \sim N (0, I)$ ， $W_{i} \sim N (0, \frac{1}{n} I)$ ，则 $z_{i + 1} = W_{i}^{T} z_{i} \sim N (0, I)$ 。由此，可知：如果我们仅仅使用线性层，并且初始化参数的方差 $σ^{2} = \frac{1}{n}$ ，那么下一层的输出将与上一层有着相同的概率分布。换句话说，选择这个方差可以保持后续层范数不变。

回到上面的例子，如果使用 ReLU 激活函数，有一半的 $z_{i}$ 会被置 0，此时为了取得相同的效果，则要两倍的方差，所以选择 $W_{i} \sim N (0, \frac{2}{n} I)$ ，这就解释了为什么上图 $σ^{2} = \frac{2}{n}$ 时，各层范数近乎不变。

这个结果与前面给出的公式 ${std}^{2} = \frac{2}{fan_in}$ 一致。