论文阅读：一些实用的 Critical Batch Size 估计方法

这篇博客提出 OpenAI 原始 critical batch size（CBS）估计方法的两个问题，然后选取介绍了一些近年出现的更加方便实用的 CBS 估计方法。

前情回顾

在上一篇文章 [1] 中，我们详细介绍了 OpenAI 在 critical batch size（CBS）上的经典工作 [2]。针对实际训练中 CBS 的估计，该文章提出了一个基于训练时间-数据消耗的双曲线估计方法，我们先对其实操流程做一个回顾：

固定一个目标性能。CBS 是和 loss 相关的。通常来说，loss 越小，梯度信号通常相对越小，噪声相对越大，CBS 往往越大。因此，我们需要先固定一个目标性能，再来谈 CBS 的估计问题。
扫描不同 bs，并对不同的 bs 扫描超参数（如 lr、lr schedule、动量、weight decay 等），目标是对每一个 bs 尽可能把优化过程调到最优，尽快达到目标 loss。
对于每个 bs $B$，选择最快达到目标性能的 run，记录其训练步数 $S(B)$ 和数据耗费 $E(B)$。
拟合形如 $(\frac{E(B)}{E_{min}} - 1) = (\frac{S(B)}{S_{min}} - 1)^{-1}$ 的时间-数据双曲线，然后用 $B_{crit} = \frac{E_{min}}{S_{min}}$ 来获得 CBS 的估计。

然而，代入实际的训练场景中一想，就会发现这个方法实操起来难度还是比较大的，比如存在下面一些问题：

目标性能的选取并不直观。在面对一个新任务的时候，我们可能并不清楚目标 loss 要达到多少才是最合理的，这会直接导致测出来的 CBS 可能并不是我们真正期待的。有时候，更贴合实际需求的是：给定一个模型大小 $N$ 和和数据量 $E$，我们如何确定一个 CBS？
超参调优的难度非常大，且这里面问题最大的看起来是 lr schedule。因为我们需要针对一个目标 loss，选取超参使得达到目标 loss 的训练步数最小。而 lr schedule 是直接影响到训练步数的，且本质上直接蕴涵了多个子超参数，如 warmup 步数，stable 步数，decay 的步数和方式。如果对于每一个 bs，我们都要重新调这些步数直接相关的超参，是非常难的。

这篇博客暂时不提出新的 CBS 估计策略，而是选取一些近年出现的更加方便实用的 CBS 估计方法进行介绍，这些方法都或多或少地缓解了所提到的两个问题。

How Does Critical Batch Size Scale in Pre-training? [3]

这篇文章考虑了在给定模型大小 $N$ 和和数据量 $E$ 的情况下，如何对 CBS 进行估计。相对于 OpenAI 原论文 [2]，虽然这篇文章对 CBS 的定义和估计方法都有一些改变，但和原来本质上是相似的。我们先介绍这篇文章的定义和方法，最后再来说明和 [2] 的关系。

文章所提方法

首先，回顾上面的问题一，我们还是得需要选取一个目标 loss，不过此时我们可以让目标 loss 和 $N$、$E$ 两者挂钩，从而避免盲目的选择。一个自然的选择是给定 N、E 的情况下，模型能够达到的最优 loss，但是完美实现这一点需要大量的调参（bs、优化器的超参等）。从 [1] 的推导中可以知道，在 bs 达到 CBS 之前，相当范围内不同 bs 对应的调优效果是差不多的，所以实际中我们选取一个相对安全、远离 CBS 的 bs 作为基准 $B_{opt}$，然后在 $B_{opt}$ 下调其它超参，获得对应的最优 loss 作为目标 loss 就行了。

回应第二个调参的问题，这篇文章经过比较系统的消融发现 constant lr + EWA 能够达到和常用的 lr schedule 类似的效果，所以文章采用 constant lr + EWA，从而丢弃了 warmup，stable，decay 等 lr schedule 超参数，只额外引入了一个 EWA 系数，降低了调参压力。当然，lr、动量、wd 等超参的调整仍然是无法避免的。

此时，确认完问题一和问题二之后，我们其实已经可以按照 OpenAI 的方法做 CBS 的估计了，但文章换了一个 CBS 的定义，从而也在具体的估计环节和原来有所差异。回顾 [1] 中的理论模型推导，我们已经知道

\[\frac{S(B)}{S_{min}} - 1 = \frac{1}{B} \frac{E_{min}}{S_{min}} \implies S(B) = S_{min} + \frac{E_{min}}{B}.\]

当 B 较小，${E_{min}}/{B} \gg S_{min}$ 的时候，我们有 $S(B) \approx {E_{min}}/{B}$。此时，为了达到目标 loss，线性增大 B 可以带来近似线性的步数减小，消耗的总数据量近似不变，可以放心增大 B 来利用卡数，我们把这个称为 linear scaling regime。当 B 较大，${E_{min}}/{B} \ll S_{min}$ 的时候，增大 B 难以再有效降低迭代步数，数据效率降低。CBS 正是落在 linear scaling regime 与后者的临界点，文章也是基于这个想法提出了 CBS 的定义，后续我们看到，这个和 OpenAI 的思路其实本质上是相似的。

设基准 $B_{opt}$ 达到目标 loss 的步数为 $S(B_{opt})$。我们为了得到 CBS，可以在此基础上继续增加 B 并调超参，如果仍然处于 linear scaling regime 中，我们期待达到目标 loss 的步数为 $S(B) \approx S(B_{opt}) \frac{B_{opt}}{B}$。当逐渐脱离 linear scaling regime 后，实际得到的 $S(B)$ 将会变得大于 $S(B_{opt}) \frac{B_{opt}}{B}$。为了反映这个临界点，我们设置一个容忍界限 $\beta$（原论文为 0.2），并定义 CBS 满足

\[S(B_{crit}) = (1+\beta)S(B_{opt}) \frac{B_{opt}}{B_{crit}}, \quad or\ i.e. \quad E(B_{crit}) = (1 + \beta)E(B_{opt}).\]

即 CBS 达到目标 loss 的训练步数是理想 linear scaling 的 $1 + \beta$ 倍。在实际中，比起一个个 B 试，更方便的是先拟合 $S(B) = S_{min} + \frac{E_{min}}{B}$（这和 [2] 中的双曲线拟合本质上一样！），然后代入 CBS 的定义中可以化简得到

\[B_{crit} = (1 + \beta)B_{opt} + \beta \frac{E_{min}}{S_{min}}.\]

这样我们就得到了给定模型大小 $N$ 和和数据量 $E$ 的情况下的一个 CBS 估计。注意到结论和 $B_{opt}$ 直接相关，且我们的推导都是基于 $B_{opt}$ 在 linear scaling regime 内的，所以实际中要选取一个安全的、比较远离 CBS 的值。最后，文章基于这个方法，可以探索 CBS 关于模型大小和训练数据量的关系，最终的结果为 CBS 和模型大小关系不大，但是会随着训练数据量的增大而增大，如下图所示：

图 1：CBS 主要和训练数据量相关。

和 OpenAI 原论文的关系

仔细思考一下，其实这篇论文的方法本质上和 OpenAI 论文 [2] 的方法是一样的，两者都是通过拟合 $(S_{min}, E_{min})$，只不过选取拟合的曲线不同，以及 CBS 对时间-数据的权衡不同。我们接下来简单说明一下。

首先，从我们上一篇博客 [1] 中的理论模型推导中可以知道，训练步数 $S(B)$ 和数据耗费 $E(B)$ 满足如下关系：

\[\begin{aligned} & {S(B)} = {S_{min}} + \frac{E_{min}}{B}, \\ & {E(B)} = E_{min} + B {S_{min}}, \\ & \frac{E(B)}{E_{min}} - 1 = \left(\frac{S(B)}{S_{min}} - 1\right)^{-1}. \end{aligned}\]

因此，要获得 $(S_{min}, E_{min})$ 的估计，这三种方式都可以被采用。[3] 这篇论文选用的是第一种，OpenAI 原论文 [2] 选用的是第三种。

其次，观察当前的 CBS 估计结果，记 [2] 中的估计结果为 $B_{crit-oai} = E_{min}/S_{min}$，则当前的结果可写为：

\[B_{crit} = (1 + \beta)B_{opt} + \beta B_{crit-oai} \approx \beta B_{crit-oai},\]

两者之间近似只差一个倍数。这本质上反映的是两种方法在训练时间-训练数据上的权衡不同，具体来说，对于训练时间，我们有

\[S(B_{crit}) \approx (1 + \frac{1}{\beta}) S_{min}, \quad S(B_{crit-oai}) = 2S_{min}.\]

对于训练数据，代入 [1] 中的 ${E(B)} = E_{min} + B S_{min}$，我们也可以类似得到

\[E(B_{crit}) \approx (1 + \beta) E_{min}, \quad S(B_{crit-oai}) = 2 E_{min}.\]

所以两种方法本质上差不多，只不过是时间-训练的权衡有差别。具体到论文 [3] 中，$\beta=0.2$ 意味着更重视数据耗费，而不是训练时间。总而言之，这篇论文在 CBS 估计上的创新主要是确定目标 loss 的选取，以及通过消融定位来缓解调参的压力。

Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training [4]

这篇文章的 CBS 方法基本也是 follow OpenAI 原论文 [2] 的方法的，但针对超参调整困难的问题进行了一些改良。从 [2] 中理论的原则上来说，为了达到目标 loss，对于不同 B，我们都应该把各个超参调到最优，但前文已经提到这是非常困难的，尤其是对于 lr schedule。[2] 中直接规避了这个问题，采用 const lr。[3] 中声称 const lr + EWA 可以实现常用的 lr schedule 的效果，所以采用 const lr + EWA 来降低调参的难度，但仍需要频繁地评测来判断是否达到了目标 loss。除此之外，在实际中，我们还是希望能够在真实使用的 lr schedule 上估计 CBS，[4] 这篇文章提出的方法就缓解了这个问题。

首先，[4] 这篇文章把考虑范围收紧，假设实际中我们已经选好了一个 lr schedule，例如 10% warmup + 90% cosine decay。注意到，哪怕是在这种情况下，要得到 B 对应最优的训练步数/数据量也是不容易的，因为设置不同的步数就意味着不同的 lr schedule，没法复用已有的 trail。[4] 这篇文章就是缓解了这个问题。

解决的方法也非常直接，就是尝试拟合一个 loss-token 的 scaling law。对于每个模型大小 N 和 bs B，使用确定好的 lr schedule，我们可以进行不同数据量 E 的训练，收集最后的 loss L，最后拟合一个如下的 scaling law：

\[L_{N,B}(E) = L_N + C_{N,B} E^{-\beta_{N,B}},\]

其中 $L_N$ 是对应大小模型的不可约 loss，第二项反映了数据越多，loss 越小。这个形式本质上就是 Chinchilla loss scaling law 的形式。此时给定一个目标 loss $L$，我们可以直接反解出对应的训练数据 $E(B)$ 和训练步数 $S(B)$：

\[E(B) = \left(\frac{C_{N,B}}{L-L_N}\right)^{\frac{1}{\beta_{N,B}}}, \quad S(B) = \frac{E(B)}{B}.\]

之后我们就可以利用之前提到的方法进行双曲线拟合，从而得到 CBS 的估计了。这篇文章也在不同的模型大小和数据量下估计了 CBS，结论和 [3] 类似：CBS 和模型大小关系不大，但是会随着训练数据量的增大而增大。

抛弃目标 loss

近期也有一些文章直接抛弃了“对齐目标 loss”这一条件，采用了更为简单直观，但没那么严谨的方式来估计给定模型大小和训练数据量时的 CBS。该方法就是用不同的 B 进行训练，然后找到 loss 开始明显变差的拐点，基本上可以总结为：

固定模型大小和数据量，更改不同的 B 进行训练，收集最后的 loss；
获得最优的 $B_{opt}$ 和对应的最优 loss $L(B_{opt})$；
根据某种规则选择 $L(B)$ 的拐点确定为 $B_{crit}$。

拐点的选择十分灵活。[5] 直接靠肉眼评估拐点；[6] 定义拐点为 $L(B_{crit}) - L(B_{opt}) = 0.01$，0.01 为一个主观的经验值；[7] 定义拐点为 $L(B_{crit}) = 1.0025 L(B_{opt})$。以 [6] 为例，示意图如下所示：

图 2：拐点确定，loss 与最优 loss 相差 0.01 以内。

可以看到，这种方法简单、易于理解，且和原方法是某种对偶关系。原来的方法是固定一个目标 loss，然后找到训练时间和数据耗费的权衡点；现在的方法本质上是固定训练数据，找到训练时间和 loss 的权衡点。相对来说，当前方法的理论支撑比原方法还是弱一些。未来工作可以适当考虑一下这一点。

总结