当 Batch Size 增大时，学习率该如何随之变化？

引言

随着算力的飞速进步，有越多越多的场景希望能够实现'算力换时间'，即通过堆砌算力来缩短模型训练时间。

理想情况下，我们希望投入 $N$ 倍的算力，那么达到同样效果的时间则缩短为 $1/N$，此时总的算力成本是一致的。这个'希望'看上去很合理和自然，但实际上并不平凡，即便我们不考虑通信之类的瓶颈，当算力超过一定规模或者模型小于一定规模时，增加算力往往只能增大 Batch Size。

然而，增大 Batch Size 一定可以缩短训练时间并保持效果不变吗？

这就是接下来我们要讨论的话题：当 Batch Size 增大时，各种超参数尤其是学习率该如何调整，才能保持原本的训练效果并最大化训练效率？我们也可以称之为 Batch Size 与学习率之间的 Scaling Law。

方差视角

直觉上，当 Batch Size 增大时，每个 Batch 的梯度将会更准，所以步子就可以迈大一点，也就是增大学习率，以求更快达到终点，缩短训练时间，这一点大体上都能想到。问题就是，增大多少才是最合适的呢？

二次方根

这个问题最早的答案可能是平方根缩放，即 Batch Size 扩大到 $N$ 倍，则学习率扩大到 $rac{1}{ ext{sqrt}(N)}$ 倍，出自 2014 年的《One weird trick for parallelizing convolutional neural networks》[1]，推导原理是让 SGD 增量的方差保持不变。

具体来说，我们将随机采样一个样本的梯度记为 $g_i$，其均值和协方差分别记为 $ar{g}$ 和 $ ext{Cov}(g)$，这里的 $ar{g}$ 就是全体样本的梯度。当我们将采样数目增加到 $B$ 个时，有

即增加采样数目不改变均值，而协方差则缩小到 $rac{1}{B} ext{Cov}(g)$。对于 SGD 优化器来说，增量为 $ ext{lr} imes g_{batch}$，其协方差正比于 $ ext{lr}^2 imes rac{1}{B}$，而我们认为优化过程中适量的（不多不少的）噪声是有必要的，所以当 Batch Size 变化时，我们通过调整学习率 $ ext{lr}$ 让增量的噪声强度即协方差矩阵保持不变，从而得出

这就得到了学习率与 Batch Size 的平方根缩放定律，后来的《Train longer, generalize better: closing the generalization gap in large batch training of neural networks》[2] 也认同这个选择。

线性缩放

有意思的是，线性缩放即 $ ext{lr} imes B$ 在实践中的表现往往更好，甚至刚才说的最早提出平方根缩放的《One weird trick for parallelizing convolutional neural networks》[1] 作者也在论文中指出了这一点，并表示他也无法给出合理的解释。

某种程度上来说，线性缩放更符合我们的直观认知，尤其是像《Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour》[3] 那样，假设连续的 $B$ 个 Batch 的梯度方向变化不大的话，那么线性缩放几乎是显然成立的。

不过，这个假设显然过强，放宽这个假设则需要将 SGD 跟 SDE（随机微分方程）联系起来，这由《Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations》[4] 完成，但首先用于指出学习率与 Batch Size 的缩放关系的论文应该是《On the Generalization Benefit of Noise in Stochastic Gradient Descent》[5]。