本文约5700字，建议阅读15分钟本文从多个视角讨论了“BatchSize与学习率之间的ScalingLaw”这一经典炼丹问题。
随着算力的飞速进步，有越多越多的场景希望能够实现“算力换时间”，即通过堆砌算力来缩短模型训练时间。
理想情况下，我们希望投入倍的算力，那么达到同样效果的时间则缩短为，此时总的算力成本是一致的。这个“希望”看上去很合理和自然，但实际上并不平凡，即便我们不考虑通信之类的瓶颈，当算力超过一定规模或者模型小于一定规模时，增加算力往往只能增大BatchSize。
然而，增大BatchSize一定可以缩短训练时间并保持效果不变吗？
这就是接下来我们要讨论的话题：当BatchSize增大时，各种超参数尤其是学习率该如何调整，才能保持原本的训练效果并最大化训练效率？我们也可以称之为BatchSize与学习率之间的ScalingLaw。
01方差视角
直觉上，当BatchSize增大时，每个Batch的梯度将会更准，所以步子就可以迈大一点，也就是增大学习率，以求更快达到终点，缩短训练时间，这一点大体上都能想到。问题就是，增大多少才是最合适的呢？
02二次方根
这个问题最早的答案可能是平方根缩放，即BatchSize扩大到倍，则学习率扩大到倍，出自2014年的《Oneweirdtrickforparallelizingconvolutionalneuralnetworks》[1]，推导原理是让SGD增量的方差保持不变。
具体来说，我们将随机采样一个样本的梯度记为，其均值和协方差分别记为和，这里的就是全体样本的梯度。当我们将采样数目增加到个时，有
即增加采样数目不改变均值，而协方差则缩小到。对于SGD优化器来说，增量为，其协方差正比于，而我们认为优化过程中适量的（不多不少的）噪声是有必要的，所以当BatchSize变化时，我们通过调整学习率让增量的噪声强度即协方差矩阵保持不变，从得出
这就得到了学习率与BatchSize的平方根缩放定律，后来的《Trainlonger,generalizebetter:closingthegeneralizationgapinlargebatchtrainingofneuralnetworks》[2]也认同这个选择。
03线性缩放
有意思的是，线性缩放即在实践中的表现往往更好，甚至刚才说的最早提出平方根缩放的《Oneweirdtrickforparallelizingconvolutionalneuralnetworks》[1]作者也在论文中指出了这一点，并表示他也无法给出合理的解释。
某种程度上来说，线性缩放更符合我们的直观认知，尤其是像《Accurate,LargeMinibatchSGD:TrainingImageNetin1Hour》[3]那样，假设连续的个Batch的梯度方向变化不大的话，那么线性缩放几乎是显然成立的。
不过，这个假设显然过强，放宽这个假设则需要将SGD跟SDE（随机微分方程）联系起来，这由《StochasticModifiedEquationsandDynamicsofStochasticGradientAlgorithmsI:MathematicalFoundations》[4]完成，但首先用于指出学习率与BatchSize的缩放关系的论文应该是《OntheGeneralizationBenefitofNoiseinStochasticGradientDescent》[5]。
事后来看，这个联系的建立其实并不难理解，设模型参数为，那么SGD的更新规则可以改写成：
其中即为梯度的噪声，到目前为止，我们还没有对这个噪声的分布做任何假设，只知道它的均值为，协方差为。接下来我们假设这个噪声的分布是正态分布，那么上述迭代可以进一步改写成
这就意味着SGD的迭代格式实际上在近似地求解SDE：
因此，要想在发生变化时，运行结果不产生明显变化，上述SDE的形式应该不变，这就得到了线性缩放。这个过程中最关键的一步是，SDE的噪声项步长是非噪声项的平方根，从而分离出一项来。
这一点我们在《生成扩散模型漫谈：一般框架之SDE篇》也有过评析，简单来说就是零均值的高斯噪声长期会有一定的抵消作用，所以必须增大步长才能将噪声效应体现出来。
以上结论都是基于SGD优化器得出的，论文《OntheSDEsandScalingRulesforAdaptiveGradientAlgorithms》[6]将它推广到了RMSProp、Adam等优化器上，结果是平方根缩放。
无独有偶，稍早一点的《LargeBatchOptimizationforDeepLearning:TrainingBERTin76minutes》[7]在测试Adam及其变体LAMB时，也应用了平方根缩放。更多内容还可以参考博客《HowtoScaleHyperparametersasBatchSizeIncreases》[8]。
04直面损失
可以肯定的是，不管是平方根缩放还是线性缩放，它们都只能在局部范围内近似成立，因为它们都包含了“只要BatchSize足够大，那么学习率就可以任意大”的结论，这显然是不可能的。此外，前面两节的工作都围绕着方差做文章，但我们的根本任务是降低损失函数，因此以损失函数为导向或许更为本质。
05单调有界
这个视角下的经典工作是OpenAI的《AnEmpiricalModelofLarge-BatchTraining》[9]，它通过损失函数的二阶近似来分析SGD的最优学习率，得出“学习率随着BatchSize的增加而单调递增但有上界”的结论。
整个推导过程最关键的思想是将学习率也视作优化参数：设损失函数是，当前Batch的梯度是，那么SGD后的损失函数则是，我们将最优学习率的求解视为优化问题：
这个目标显然很直观，就是选择学习率使得平均而言训练效率最高（损失函数下降得最快）。为了求解这个问题，我们将损失函数近似地展开到二阶：
这里的就是Hessian矩阵，而是损失函数的梯度，理想的目标函数是基于全量样本来求的，这也就是为什么它的梯度就是的均值。接着求期望，我们得到：
最后一项有少许技巧：
变换过程主要利用到了。现在只要假定的正定性，那么问题就变成了二次函数的最小值，容易解得：
这就得出了“随着单调递增有上界“的结果，其中：
06实践分析
当时，，所以，即线性缩放，这再次体现了线性缩放只是小BatchSize时的局部近似；当时，逐渐趋于饱和值，这意味着训练成本的增加远大于训练效率的提升。
所以，相当于一个分水岭，当BatchSize超过这个数值时，就没必要继续投入算力去增大BatchSize了。
对于实践来说，最关键的问题无疑就是如何估计和了，尤其是直接关系到学习率的缩放规律和训练效率的饱和问题，二者的直接计算涉及到Hessian矩阵，其计算量正比于参数量的平方，在数亿参数量都算小模型的今天，计算Hessian矩阵显然是不现实的事情，所以必须寻找更有效的计算方式。
我们先来看，它的式子是，分子分母都有一个，这无疑有一种让我们将它们“约掉”的冲动。事实上简化的思路也是如此，假设近似于单位阵的若干倍，那么得到：
在计算上更为可行，并且实验发现它通常是的一个良好近似，因此我们选择估计而不是。注意只需要对角线上的元素，因此不用算出完整的协方差矩阵，只需要将每个梯度分量单独算方差然后求和。在数据并行场景，可以直接利用每个设备上算出来梯度来估计梯度方差。
需要指出的是，式（10）等结果实际上是动态的，也就是说理论上每一步训练的、、都是不同的，所以如果我们希望得到一个静态的规律，需要持续训练一段时间，等到模型的训练进入“正轨”后计算的才可靠的，或者也可以在训练过程中持续监控，以便判断当前设置与最优的差距。
至于，其实就没必要根据公式来估计了，直接在某个小BatchSize下对学习率进行网格搜索，搜出一个近似的，然后结合估计的就可以反推出了。
07数据效率
从上述结果出发，我们还可以推导关于训练数据量和训练步数的一个渐近关系。推导过程也很简单，将（10）代入到损失函数中可以算得，在最优学习率下每一步迭代带来的损失函数减少量是：
其中。接下来的重点是对这个结果的解读。
当也就是全量SGD时，每一步损失函数减少量达到了最大的，这时候可以用最少的训练步数（记为）达到目标点。
当有限时，每一步的损失下降量平均只有，这意味我们需要步才能达到全量SGD单步的下降量，所以训练的总步数大致上就是。
由于BatchSize为，所以训练过程消耗的样本总数则是，这是的增函数，且当，这表明只要我们使用足够小的BatchSize去训练模型，那么所需要的总训练样本数也会相应地减少，代价是训练步数非常多。
进一步地，利用这些记号我们可以写出它们之间的关系是：
这就是训练数据量和训练步数之间的缩放规律，表明数据量越小，那么应该缩小BatchSize，让训练步数更多，才能更有机会达到更优的解。
这里的推导是经过笔者简化的，假设了和在整个训练过程的不变性，如果有必要也可以按照原论文附录用积分更精细地处理动态变化的情形（但需要引入假设），这里就不展开了。
此外，由于，所以上式也提供了估计的另一个方案：通过多次实验加网格搜索得到多个对，然后拟合上式就可以估计出，继而计算。
08自适应版
不得不说，OpenAI不愧为各种ScalingLaw的先驱之一，前述分析可谓相当精彩，并且结果也相当丰富，更难得的是，整个推导过程并不复杂，给人一种大道至简的本质感。
不过，目前的结论都是基于SGD来推的，对于Adam等自适应学习率优化器的适用性还不明朗，这部分内容由《SurgePhenomenoninOptimalLearningRateandBatchSizeScaling》[10]完成。
09符号近似
分析Adam的思路跟SGD一样，都是基于二阶展开，不同的是方向向量由换成了一般的向量，此时我们有：
现在需要确定以及计算相应的和。由于只需要一个渐近关系，所以跟《配置不同的学习率，LoRA还能再涨一点？》一样，我们选择SignSGD即作为Adam的近似。这个近似的合理性体现在两点：
1.无论取何值，Adam第一步的更新向量都是；
2.当时，Adam的更新向量始终为。
为了计算和，我们还需要跟“线性缩放”[11]一节一样，假设服从分布，而为了简化计算，我们还要进一步假设是对角阵，即假设分量之间是相互独立的，这样一来我们可以独立地处理每一个分量。
由重参数得等价于，因此：
这里的是误差函数[12]，它是跟类似的值域为的型函数，可以作为的光滑近似。
但本身没有初等函数表达式，所以我们最好找一个初等函数近似，才能更直观地观察变化规律，之前我们在《GELU的两个初等函数近似是怎么来的》就讨论过这个话题，不过那里的近似还是太复杂了（都涉及到指数运算），这里我们整个简单点的：
我们选择，使得这个近似在处的一阶近似跟的一阶近似相等。当然，都做了这么多重近似了，这个c的值其实已经不大重要，我们只需要知道存在这么个就行了。基于这个近似，我们得到：
可以发现，Adam跟SGD的一个明显区别是这一步就已经跟相关了。不过好在，此时的二阶矩更简单了，因为的平方必然是1，所以：
利用这些结果，我们就可以求得：
10两个特例
相比SGD的式（10），Adam的式（20）更为复杂，以至于无法直观看出它对的依赖规律，所以我们从几个特殊例子入手。
首先考虑，此时，所以：
它跟SGD的的区别是它关于梯度并不是齐次的，而是正比于梯度的scale。
接着我们考虑是对角阵的例子，即时，此时：
这里求和的每一项关于都是单调递增有上界的，所以总的结果也是如此。为了捕捉最本质的规律，我们可以考虑进一步简化（这里开始跟原论文不一样）：
这里的假设是存在某个跟无关的常数【比如可以考虑取全体的某种均值，其实这里的类似前面的，按照的定义来估计也可以】，使得对任意来说把换成都是一个良好近似，于是：
当即时，可以进一步写出近似：
这表明在BatchSize本身较小时，Adam确实适用于平方根缩放定律。
11涌现行为
如果我们将近似（24）应用到原始的式（20），会发现它存在一些全新的特性，具体来说我们有：
其中，以及：
注意是的单调递增函数，但式（27）最后的近似并不是的单调递增函数，它是先增后减的，最大值在取到。这意味着存在一个相应的，当BatchSize超过这个后，最佳学习率不应该增大反而要减小！这便是原论文标题所说的“Surge现象”。
当然这里还有一个限制，是始终小于1的，如果，那么最优学习率与BatchSize的关系依旧是单调递增的。
关于Adam的，其实OpenAI在论文附录中曾不加证明地“猜测”Adam的最优学习率应该是：
其中。现在看来，这个形式只是Hessian矩阵对角线元素占主导时的近似结果，当非对角线元素的作用不可忽略时，则有可能涌现出“BatchSize足够大时学习率反而应该减小”的Surge现象。
如何直观地理解Surge现象呢？笔者认为，这本质上是自适应学习率策略的次优性的体现。仍以近似为例，越大就越准，则是，然而是最科学的更新方向吗？
不一定，尤其是训练后期这种自适应策略可能还有负面作用。因此，当取适当值时，的噪声反而可能修正这种次优性，而继续增大时噪声减少，反而减少了修正的机会，从而需要更谨慎地降低学习率。
12效率关系
同SGD的分析一样，最后我们还可以考虑，将式（27）代入式（21），恢复记号然后化简（化简过程不需要任何近似）得到：
其中：
注意这里是一个新的记号，它不是，后者是由反解出来的理论最优BatchSize，结果是：
它们之间的关系是：
由于式（30）形式上跟SGD的式（13）是一样的，所以那一节的分析同样适用，因此同样可以导出式（14）：
只不过现在。这样一来，我们就有得到一种估计和的方案：通过多次实验得到多个对，实验过程中还可以顺便估计，然后拟合上式得到，继而估计，最后由（31）式解出。
如果，那么不存在最优的，如果则说明Hessian矩阵对角线元素占主导，此时适用于缩放规律（25），增大BatchSize总可以适当增大学习率；当时，可以由（33）解出最优的，BatchSize超出这个值学习率反而应该下降。
13补充说明
需要指出的是，上面几节分析的出发点和最终结论，其实跟原论文《SurgePhenomenoninOptimalLearningRateandBatchSizeScaling》[10]大同小异，但中间过程的近似处理有所不同。
原论文得到的大部分结论，都是在假设下的近似结果，所以得到Surge现象几乎总会出现的结论，这其实是不大科学的。
最明显的是这个假设的形式本身就有点问题，它右端是跟i相关的，我们总不能给每个分量都配一个单独的BatchSize，所以为了得到一个全局的结果就只能是，但这未免有点苛刻了。
本文的做法则是引入近似（24），这可以看成是平均场近似，直觉上比逐点的假设更为合理一些，所以原则上结论会更为精准，比如可以得到“即使Hessian矩阵的非对角线元素不可忽略，Surge现象也不一定会出现”的结论（取决于）。
特别地，这种精准性并没有牺牲简洁性，比如式（27）同样很简明清晰，式（30）形式也跟原论文一致，并且不需要额外的近似假设，等等。
最后，稍微感慨一下，OpenAI对SGD的分析其实已经是2018年的工作了，而Surge现象这篇论文则是今年中才发布的，从SGD到Adam居然花了6年时间，这是让人比较意外的，大体是OpenAI的“威望”以及猜测（29），让大家觉得Adam已经没什么好做了，没想到Adam可能会有一些新的特性。
当然，作为Adam的近似究竟有多合理、能多大程度上代表实际情况等问题，笔者认为还值得进一步思考。
文章小结
本文从多个视角讨论了“BatchSize与学习率之间的ScalingLaw”这一经典炼丹问题，其中着重介绍了OpenAI基于损失函数的二阶近似的推导和结论，以及后续利用同样的思想来分析Adam优化器的工作。
参考文献
[1]
[2]
[3]
[4]
[5]
[6]
[7]
[8]
[9]
[10]
[11]
[12]
