基于随机划分的小批次动量梯度下降法

点击“蓝字”关注我们吧！

高原，博士毕业于华东师范大学统计学院，现为北京大学光华管理学院博士后，主要研究方向包括统计机器学习、统计计算的方法与理论。

（封面为今年刚去世的俄罗斯数学家Boris Polyak（1935—2023）以及他在著作Introduction to Optimization[4]中为动量法绘制的示意图）

今天和大家分享的主题是小批次动量梯度下降法（Minibatch Gradient Descent with Momentum, MGDM）。内容主要包括小批次动量梯度下降法的应用背景、算法介绍以及理论性质。

一、 背景介绍

在大多数的统计学习问题中，我们通常需要通过优化一个形如的损失函数来得到未知参数的估计值。这里的表示第个观测值上的损失函数值。为了求解最优值点，我们常常会使用牛顿迭代法：

其中表示参数第步的迭代值，和则分别表示损失函数在处的一阶导数（即梯度）和二阶导数（即黑塞矩阵）。直观上，牛顿法利用损失函数的前两阶导数对原始的损失函数做了二次逼近，然后找到了该二次函数的极小值点作为参数下一步的迭代值；见图1左侧示意图。牛顿法在一定条件下可以达到二阶收敛速度：，这里的表示某个常数。这意味着牛顿法往往只需要迭代几步就能达到极高的数值收敛精度。但遗憾的是，计算黑塞矩阵并对其求逆所需的计算复杂度将近。在大规模优化问题中，我们面对的样本量和参数量可能都十分巨大，这将使得牛顿法在计算上变得不再可行。

图 1 牛顿法和梯度下降法示意图

为了免于黑塞矩阵相关的计算，各类一阶优化算法走进大家的视野，其中最具代表性的就是梯度下降法（Gradient Descent, GD）：

其中称为学习率。从迭代公式可以看到，GD算法只使用了损失函数的一阶导数信息（即梯度），沿着负梯度方向不断进行迭代，而学习率则控制了沿负梯度方向上走的步长；见图1右侧示意图。由于忽视了损失函数二阶导数的信息，GD算法一般只能达到线性收敛速度：，这里的称作收敛因子。注意到，GD算法每次迭代都需要计算梯度，对应的计算复杂度为。当和都很大时，这依然是个难以接受的计算量。例如，著名的ImageNet数据集包含了超过100万张图片，针对ImageNet的典型深度学习模型（例如ResNet50）通常也有千万级别的待训练参数。很显然在这种场景下，GD算法的每一步迭代依然过于“沉重”。

为了进一步提高GD算法的迭代效率，我们可以从全样本抽取一个较小的子样本（即小批次，Minibatch），并使用小批次上的梯度来代替全样本梯度。这就是所谓的小批次梯度下降法（Minibatch Gradient Descent，MGD）：

其中表示样本量为n的小批次上的损失函数。Polyak (1964)[3]注意到由于GD算法没有使用二阶导数信息，面对条件数（Condition Number）较大优化问题时，GD算法的收敛因子会十分接近1，这会导致极慢的数值收敛速度。因此，他提出了利用算法上一步的迭代信息来改善GD算法，这就是所谓的动量（Momentum）法。将小批次梯度下降方法与动量结合起来，便得到了小批次动量梯度下降法（Minibatch Gradient Descent with Momentum，MGDM）：这里的称为动量参数，它控制了对上一步迭代信息的利用程度。大量的理论和实践都表明，动量可以显著地提高GD算法的收敛速度。

二、模型和方法

现有关于MGD类算法的理论研究中，一般需要所谓的鞅差（Martingale Difference）条件(Polyak and Juditsky, 1992; Chen et al., 2020; Moulines and Bach, 2011)[5][1][2]。直观上来说，这个条件保证了小批次上计算得到的梯度应当是全样本梯度的无偏估计。为了满足这个条件，现有的很多文献会假设小批次是从全样本中通过简单随机有放回抽取得到的。然而在实际中（例如TensorFlow和PyTorch）更常见的做法是：预先将全样本随机划分成若干个互不相交的小批次，然后让算法在这些小批次上依次迭代。我们称以这种方式产生的小批次为基于随机划分的小批次。通过这种方式产生的小批次一般不再满足所谓的鞅差条件，但它能保证在每一轮（Epoch）迭代中全样本的每个样本点都可以被使用一次。然而，在这种设定下，MGDM算法的相关理论性质并没有得到很好的研究，例如数值收敛性质和所得到估计量的统计性质。本文从最简单的线性回归问题出发，希望得到一些有启发的结果。

我们考虑的模型为最简单的线性回归模型，损失函数为二次损失。此外，我们将样本被随机划分成个互不相交的小批次，其中每个小批次都包含个样本点，即。我们假设小批次划分完成后便固定（Fixed）下来，然后让MGDM算法在这M个小批次上依次迭代：

这里的表示参数在经过第t轮第m个小批次迭代后得到的估计值，为第z4m$个小批次上的损失函数；算法的迭代示意图可见图2。

图 2 基于固定小批次的MGDM算法迭代示意图将个小批次上迭代公式写在一起，便得到了如下的线性动力系统（Linear Dynamical System）：其中，且。对于这个线性动力系统，一个自然的问题便是它是否存在稳定解（Stable Solution）。不妨假设每个小批次对应的稳定解是存在的，并记堆叠后的稳定解为。将上述线性系统写成矩阵形式后，我们可以发现该稳定解应当满足以下线性方程组：其中这意味着“稳定解是否唯一存在”等价于“矩阵是否可逆”。下面，我们将从该线性方程组出发，研究稳定解的存在性，算法的数值收敛性，以及稳定解的统计性质。在此之前，我们引入一些符号。记为零均值随机向量的协方差矩阵，并假设它的个特征值从大到小分别为。此外，假设随机噪声均值为0方差为。

三、理论性质

从上节的讨论我们知道，稳定解的存在性主要由矩阵的可逆性决定的。下面定理给出了可逆的条件。

从定理1可以看到，只要学习率，动量参数，以及的特征值之间没有恰好满足某些交互关系，那么大概率都是可逆的。此时，稳定解也是唯一存在的，并可以表示为。但是，稳定的存在并不意味着MGDM算法能够收敛到这个稳定解。例如在实践中，一旦动量参数超过1，那么算法基本上就会发散。但根据定理1我们可以知道，即使稳定解依然是可以存在的。为此，我们进一步研究了MGDM数值收敛性质，结论如下。

定理2给出了MGDM算法收敛到稳定解的几乎充分且必要的条件。总结来说，若想要算法收敛，我们不仅需要动量参数，还需要学习率不能太大。这也与我们实际经验比较吻合。

前面提到，动量法的提出是为了加速GD算法的数值收敛速度。那么在MGDM算法中，动量是如何影响数值收敛速度的呢？相关结论总结在下面的定理中。

定理3给出了MGDM算法的线性收敛形式。我们可以看到，算法的数值收敛速度主要由收敛因子决定。结论（a）给出了收敛因子能够达到的最优值（最小值），并给出了调节参数和所需的取值。可以看到，的大小主要由条件数所决定：条件数越大，则越靠近1，对应的数值收敛速度也就越慢；反之亦然。结论（b）则给出了MGD算法（即，不使用动量）能达到的最优收敛因子，它的大小同样由所决定。对比结论（a）和（b）可以发现，只要条件数，那么MGDM算法总可以达到比MGD算法更快的数值收敛速度，这也进一步证实了动量的加速效应。结论（c）考虑了学习率取得较小的场景，此时动量的存在依然可以有效加速数值收敛。

以上结论考虑的都是MGDM算法的数值收敛性质，然而我们最终关心的是所得到估计量的统计性质。如果算法收敛很快，但最终得到的估计量具有很差的性质，这当然也不是我们想要的结果。我们知道，在线性回归问题中，最优的估计量一般为为全样本最小二乘（OLS）估计，其中和。因此，我们需要考察所得到稳定解与OLS估计量之间的差距。记堆叠后的OLS估计量为，下面的定理刻画了与之间的差距。

定理（4）的结论（a）告诉我们，当动量参数固定时，只要学习率足够小，那么稳定解也能够足够靠近OLS估计。结论（b）刻画了学习率固定时的情况，此时只要足够大，和也可以充分接近。需要注意的是，当时，虽然稳定解依然可能存在，但这个稳定解是没法用MGDM算法逼近的（见定理2）。因此，为了让稳定具有更好的统计有效性（即离OLS估计更接近），我们只能选取一个较小的学习率。根据定理3结论（c）我们知道，当学习率很小时，算法的数值收敛速度可能会变得很慢；但适当调大动量参数，我们仍然可以得到比MGD算法（即）更快的数值收敛速度。

四、总结

本文研究了基于随机划分小批次的动量梯度下降（MGDM）算法，以最小二乘问题为例，我们初步探索了算法的数值收敛性质以及统计学性质。我们发现，动量的存在确实可以有效加速算法的数值收敛。此外，为了达到更好的统计学效率，我们需要选择较小的学习率，并可以选择合适的动量参数来达到更快的数值收敛速度。

参考文献

[1] Chen, X., Lee, J. D., Tong, X. T., Zhang, Y., et al. (2020). Statistical inference for model parameters in stochastic gradient descent. The Annals of Statistics, 48(1):251273.

[2] Moulines, E. and Bach, F. (2011). Non-asymptotic analysis of stochastic approximation algorithms for machine learning. In Shawe-Taylor, J., Zemel, R., Bartlett, P., Pereira, F., and Weinberger, K., editors, Advances in Neural Information Processing Systems, volume 24. Curran Associates, Inc.

[3] Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. Ussr Computational Mathematics and Mathematical Physics, 4(5):1–17.

[4] Polyak, B. T. (1987). Introduction to optimization. New York : Optimization Software, Inc.

[5] Polyak, B. T. and Juditsky, A. B. (1992). Acceleration of stochastic approximation by averaging. SIAM Journal on Control and Optimization, 30(4):838–855.