经典回顾：OpenAI Critical Batch Size

这篇博客对 OpenAI 在 critical batch size（CBS）方向上的经典工作 An Empirical Model of Large-Batch Training 做一个回顾，总结其核心的思想与方法。

CBS 是什么？

Batch size（B）是影响模型训练结果和效率的重要超参数，因此它的挑选特别重要。从经典随机优化的角度可以对其有一个直观的理解：

1. 当 B 很小的时候，梯度噪声很大，梯度估计很不准，此时学习率需要设置的比较小，而且要迭代相当多的步数才能达到预期的目标。在这种情况下，虽然对卡数的要求很小，但是需要训练相当长的时间，在实际中是不被期待的。
2. 当 B 在一定范围内成倍增大的时候，梯度噪声得到有效控制，我们可以从容地增大学习率，并期待能用成倍少的步数达到目标的效果。此时，我们能够充分利用卡的数量，把训练时间有效减少，同时保持较高的数据效率，这在实践中是比较可行的。
3. 当 B 非常大的时候，梯度估计更准确带来的增益几乎消失，我们没法再有效增大学习率，也不能期待进一步减少迭代步数。此时，我们需要花费更多的卡，但运行时间不能再有效减小，不仅算力耗费大，而且也是非常数据低效的。

因此，选择不同的 B 对应着不同的时间-数据/计算的权衡，如下图所示：

CBS 正是一个反映第二阶段与第三阶段临界点的一个量。当 B 小于 CBS 时，我们视卡的数量，可以通过增大 B 来实现更高效的并行，缩短训练时间，保持较高的数据效率；而当 B 大于 CBS 时，进一步增大 B 难以进一步缩短训练时间，但会显著增加卡的数量和计算成本，数据效率降低。因此，CBS 能够科学地指导实际中 B 的选择，实现卡数和训练时间的高效分配，是一个重要的、值得估计的量。

单点二次模型与 CBS 信噪比估计量

从前面的初步分析可以看出，CBS 和梯度的信噪比存在深刻的关联。从这个角度出发，OpenAI 在原论文中，基于简单的二次模型，提出了一个 gradient noise scale 估计量，能够较为简单地估计 CBS 的数量级趋势，并且能够指导实际中 CBS 的经验拟合。这一节就对这部分的核心机制进行说明。

从降低 loss 的角度来说，我们期待梯度下降（$B \to \infty$）的效果是最好的，但这需要及其多的卡数才能实现，对算力的要求过大。CBS 的核心就是找到一个 B，使得梯度噪声不大，估计仍然比较准确，此时达到目标性能的迭代步数和梯度下降差不多，但总体的算力耗费会少很多。因此，我们接下来在简单的单点二次模型上，先分析 GD（$B \to \infty$）的理想效果，然后研究 SGD（B）相较于 GD 的效果减幅，找到一个合适的 B 实现性能和算力的平衡。

GD 的最优学习率和理想效果

我们目前只考虑训练过程中的某一个时刻（单点），设网络当前的权重为 $\theta$，对应的 loss function 为 $L(\theta) = \mathbb{E}_{x \sim \rho}[L_x(\theta)]$，其中 $x$ 为样本，$\rho$ 为对应的样本分布。为了实现 GD 场景下的理想效果，我们首先需要找到一个最优的学习率。由于神经网络的 $L(\theta)$ 本身非常复杂，所以为了简化分析，我们在临域内对 loss 采用二阶近似，且把这个叫做单点二次模型。设当前的梯度为 $g = \nabla L(\theta)$，当前的 Hessian 矩阵为$H = \nabla^2 L(\theta)$ ，那么当前点附近 loss 的二次近似可表示为：

\[L(\theta + \Delta \theta) \approx L(\theta) + g^\top \Delta\theta + \frac{1}{2} \Delta\theta^\top H \Delta\theta.\]

在 GD 的语境下，有 $\Delta\theta = -\eta g$，因此可得

\[L(\theta + \Delta \theta) \approx L(\theta) - \eta \|g\|_2^2 + \frac{\eta^2}{2} g^\top H g.\]

对 lr 求导，可以得到最优的 lr 为 $\eta_* = \frac{\Vert g\Vert_2^2}{g^\top Hg}$，对应的最优的 loss 下降幅度为

\[\Delta L_* = L(\theta) - L(\theta + \Delta \theta) \approx \eta_* \|g\|_2^2 - \frac{\eta_*^2}{2} g^\top H g = \frac{1}{2} \frac{\Vert g\Vert_2^4}{g^\top Hg}.\]

题外话，可以看到，最优 lr 和 loss 降幅都和 Hessian 矩阵 H 有着密切关系，现代的 adaptive optimizer 或多或少也有着近似 Hessian 二阶信息的思想存在。

SGD 的最优学习率和理想效果

在实际情况中，B 是有限的，我们只能使用 SGD 对模型进行优化。接下来我们就分析 SGD 对应的最优学习率和 loss 降幅，看看有限 B 带来的性能退化，并找到一个折中的平衡点。

首先，此时我们不再能得到精确的梯度信息，而是一个近似的估计，为

\[\hat{g} = \frac{1}{B} \sum_{i=1}^B \nabla L_{x_i}(\theta), \quad x_i \sim \rho.\]

我们先看看随机梯度的统计性质，包括期望和协方差。期望上，我们可以得出随机梯度是真实梯度的一个无偏估计：

\[\mathbb{E}[\hat{g}] = \frac{1}{B} \sum_{i=1}^B \mathbb{E}_{x_i\sim \rho} [\nabla L_{x_i}(\theta)] = \frac{1}{B} \sum_{i=1}^B g(\theta) = g.\]

我们也可以类似计算其协方差：

\[\Sigma(\hat{g}) = \frac{1}{B^2} \sum_{i=1}^B \Sigma_{x_i \sim \rho}(\nabla L_{x_i}(\theta)) = \frac{1}{B^2} \sum_{i=1}^B \Sigma(\theta) = \frac{1}{B} \Sigma.\]

其中 $\Sigma$ 为样本梯度的协方差，为

\[\begin{aligned} \Sigma &= \Sigma_{x \sim \rho}(\nabla L_{x}(\theta)) = \mathbb{E}_x[(\nabla L_{x}(\theta) - g)(\nabla L_{x}(\theta) - g)^\top] \\ &= \mathbb{E}_x [\nabla L_{x}(\theta) \nabla L_{x}^\top(\theta)] - gg^\top. \end{aligned}\]

基于这些统计性质，我们现在就可以研究 SGD 的最优学习率和 loss 降幅了。和 GD 的情况不同，此时的梯度是随机的，所以我们更新后得到的 loss $L(\theta + \Delta \theta)$ 也是随机的。因此，我们研究其期望，并找到合适的学习率，使 loss 的期望下降最快。和之前类似，我们有

\[\begin{aligned} \mathbb{E}[L(\theta + \Delta \theta)] &\approx L(\theta) + \mathbb{E}[g^\top \Delta\theta] + \frac{1}{2} \mathbb{E}[\Delta\theta^\top H \Delta\theta] \\ &= L(\theta) - \eta \mathbb{E}[g^\top \hat{g}] + \frac{\eta^2}{2} \mathbb{E}[\hat{g}^\top H \hat{g}] \\ &= L(\theta) - \eta \|g\|_2^2 + \frac{\eta^2}{2} \mathbb{E}[\hat{g}^\top H \hat{g}]. \end{aligned}\]

最后一项 $\mathbb{E}[\hat{g}^\top H \hat{g}]$ 的推导会稍微复杂一些：

\[\begin{aligned} \mathbb{E}[\hat{g}^\top H \hat{g}] &= \mathbb{E}[tr(\hat{g}^\top H \hat{g})] = \mathbb{E}[tr(H \hat{g} \hat{g}^\top)] = tr(H \mathbb{E}[\hat{g} \hat{g}^\top]) \\ &= tr(H (\Sigma(\hat{g}) + \mathbb{E}[\hat{g}] \mathbb{E}[\hat{g}^\top])) \\ &= tr(H (\frac{\Sigma}{B} + {g}{g}^\top)) \\ &= \frac{tr(H\Sigma)}{B} + {g}^\top H g. \end{aligned}\]

总结一下，更新一步之后的 loss 期望为：

\[\begin{aligned} \mathbb{E}[L(\theta + \Delta \theta)] \approx L(\theta) - \eta \|g\|_2^2 + \frac{\eta^2}{2} (\frac{tr(H\Sigma)}{B} + {g}^\top H g). \end{aligned}\]

对 lr 求导，可以得到 B 对应最优的 lr 为

\[\eta_*(B) = \frac{\Vert g\Vert_2^2}{g^\top Hg + \frac{1}{B} tr(H\Sigma)} = \frac{\eta_*}{1 + \frac{1}{B} \frac{tr(H\Sigma)}{g^\top Hg}}.\]

将其带入 loss 的表达式，可以得到 B 对应的最优 loss 降幅为

\[\Delta L_*(B) = L(\theta) - L(\theta + \Delta \theta) \approx \frac{\|g\|_2^4}{2 (g^\top Hg + \frac{1}{B} tr(H\Sigma))} = \frac{\Delta L_*}{1 + \frac{1}{B} \frac{tr(H\Sigma)}{g^\top Hg} }.\]

可以很明显看出，随着 B 的增大，$\eta_\ast(B)$ 和 $\Delta L_\ast(B)$ 都在增大，并逐渐趋近于 GD 对应的 $\eta_\ast$ 和 $\Delta L_\ast$ 。细致一看，可以发现，当 B 远小于 $\frac{tr(H\Sigma)}{g^\top Hg}$ 时，增大 B 近似对应着最优学习率和 loss 降幅的线性增加。此时，我们可以线性增加学习率，并且有效（线性）减小训练步数，这对应着我们最开始提到的阶段二。当 B 远大于 $\frac{tr(H\Sigma)}{g^\top Hg}$ 时，再增加 B 带来的 loss 降幅收益几乎消失，训练时间不会再明显减少，反而需要更多的卡，这对应着我们最开始提到的阶段三。两阶段的临界点自然引出了 CBS 的估计量，如下图所示：

Local CBS 的估计量：Gradient Noise Scale

从上一节可以看出，阶段二和阶段三的临界点发生在 $\frac{tr(H\Sigma)}{g^\top Hg}$ 附近，我们将其叫做 Gradient Noise Scale，并明确定义为

\[B_{noise} = \frac{tr(H\Sigma)}{g^\top Hg}.\]

可以看到其分子部分反映的是梯度的噪声水平（按照曲率大小加权，曲率越大权重越大），分母部分反映的是梯度的信号水平。当 loss 较低或者任务很复杂的时候，信噪比较低，$B_{noise}$ 比较大，更适用于使用大 B 来更准地估计梯度，加速训练。反之，用大 B 不能带来明显的训练加速，反而会造成算力和数据的浪费。这和我们开始对 CBS 的感知非常一致，原论文正是用这个量来估计 local 的数量级，并预测其走势。

在实际的神经网络训练中，Hessian 矩阵通常是比较难计算的，所以文章又进一步假设 Hessian 阵为单位阵（各向同性），提出了一个简化版的 local CBS 估计量，适用于平时的实际估计：

\[B_{simple} = \frac{tr(\Sigma)}{\|g\|^2_2}.\]

在实际中，可以在训练的不同时间点估计该统计量，来判断当前的 local CBS，并视情况动态调节 B。

全程推广与时间-算力权衡

上一节，我们考虑并分析了训练过程中单点附近 GD/SGD 的理论性质，并提出了单点附近 local CBS 的估计量。现在，我们将思想拓展到训练的全程，并回答一个问题：为了训练到同一个目标 loss，不同 B 分别需要多少优化步、处理多少数据（或花费多少算力）？文章最终得到一个关于训练时间-数据/计算效率的双曲线 Pareto frontier，并基于此提出了 global CBS 的估计量。

我们记 $B_{noise}(s)$ 为 s 时刻对应的 gradient noise scale。记 S 为达到目标 loss 的优化步数，其衡量了串行时间效率。记 E 为达到目标 loss 的所需数据量，其衡量了数据效率或者计算成本。我们也假设每一步的学习率都已经按上一节调至最优。

小试牛刀：Gradient Noise Scale 不变

我们先考虑最简单的情况，假设训练全程 $B_{noise}(s)$ 都不变，恒为 $B_{noise}$。当 bs 为 B 时，从上一节的推导我们可以知道，当参数处于同一个点时，SGD 每一步的 loss 降幅与 GD 对应降幅的比值为

\[\frac{\Delta L_*(B)}{\Delta L_*} \approx \frac{1}{1 + B_{noise}/B}.\]

如果进一步假设不同 B 的训练可以沿同一条有效进度轨迹对齐，且该比例能够沿训练过程累积，我们可以近似认为 GD 一步得到 loss 降幅需要 SGD $1 + B_{noise}/B$ 步来实现。设 GD 达到目标 loss 的步数为 $S_{min}$（从当前推导思想出发，GD 是最快的），则 SGD 大致需要 $S(B) = S_{min}(1 + B_{noise}/B)$ 步，对应数据量为 $E(B) = S_{min}(B + B_{noise})$。从这里出发，我们可以得到时间 S 和数据 E 的双曲线 Pareto frontier。

我们先考虑时间和数据两种效率的极端情况，即双曲线的两侧渐近线。为了让训练时间 $S(B)$ 最短，我们需要让 B 趋近于无穷，此时步数趋近于 $S(\infty) = S_{min}$，但数据/计算的耗费是无穷大的。为了让数据/计算耗费最小，我们需要让 $B=1$，此时数据耗费为 $E_{min} = E(1) = S_{min}(1 + B_{noise}) \approx S_{min} B_{noise}$，但训练步数也是最长的。

我们再考虑一般的情况。基于 $S_{min}$ 和 $E_{min} \approx S_{min} B_{noise}$ 的定义，我们可以得到 S 与 E 的双曲线关系：

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

在实际中，我们可以在不同 B 训练的时候收集数据，并拟合出 $S_{min}$、$E_{min}$ 及其对应的双曲线。如下图所示：

最后，我们看看我们提出的 local CBS 的估计量 $B_{noise}$ 在时间-数据权衡的视角里意味着什么。一，一旦我们完成双曲线的拟合，我们就可以得到

\[B_{noise} \approx E_{min}/S_{min},\]

这给了 Gradient Noise Scale 之外的另一种估计方法，在实际中也是更准的一种选择 CBS 的方法。

二，将 $B_{noise}$ 带入 S 和 E 的表达式，我们可以得到

\[S(B_{noise}) = 2S_{min}, \quad E(B_{noise}) \approx 2E_{min},\]

这意味着 $B_{noise}$ 在时间和数据效率的两种极限之间取了一个折中之处。自然地，原论文把这种情况的 global CBS 估计量就定义为：

\[B_{crit} = \frac{E_{min}}{S_{min}} \approx B_{noise}.\]

继续深入：Gradient Noise Scale 变化

实际的训练过程中，$B_{noise}$ 是不断变化的（先小后大），所以上面的 global CBS 估计方法并不一定成立。我们将上一小节的思想继续推广至 $B_{noise}(s)$ 随时间变化的场景，推导原则是一样的。我们仍然假设不同 B 的训练可以沿同一条有效进度轨迹对齐，且该比例能够沿训练过程累积。以 GD 的训练时间为基准，设在时间 $s$ 处，SGD 的 bs 为 $B(s)$，那么为了完成 GD 此时 $ds$ 时间的 loss 降幅，SGD 近似需要 $(1 + B_{noise}(s)/B(s))ds$ 的时间来实现。和之前推导类似，我们可以得出，为了达到目标效果，SGD 需要的时间和数据量分别为：

\[\begin{aligned} & S(B(s)) = \int (1 + \frac{B_{noise}(s)}{B(s)})ds, \\ & E(B(s)) = \int B(s)(1 + \frac{B_{noise}(s)}{B(s)})ds = \int (B(s) + {B_{noise}(s)})ds. \end{aligned}\]

和上一小节类似，我们也先考虑时间和数据两种效率的极端情况。为了让训练时间 S 最小，我们仍然是需要 $B(s) = \infty$，此时有 $S_{min} = \int ds$，即 GD 所需要的时间。为了让训练数据 E 最小，我们仍然是需要 $B(s)=1$，此时 $E_{min} = \int (1 + {B_{noise}(s)})ds \approx \int B_{noise}(s)ds$。

为了分析简便，我们接下来暂时假设 SGD 的 bs 不随时间变化，即 $B(s) = B$，这也是实际中的默认情况。我们看到，此时上一节的双曲线规律仍然成立。具体来说，我们有：

\[\begin{aligned} & \frac{S(B)}{S_{min}} - 1 = \frac{1}{B} \frac{\int B_{noise}(s) ds}{S_{min}} \approx \frac{1}{B} \frac{E_{min}}{S_{min}}, \\ & \frac{E(B)}{E_{min}} - 1 \approx B \frac{\int ds}{E_{min}} = B \frac{S_{min}}{E_{min}}, \\ & \implies \frac{E(B)}{E_{min}} - 1 \approx \left(\frac{S(B)}{S_{min}} - 1\right)^{-1}. \end{aligned}\]

所以，如果我们用随时间恒定的 bs schedule 来训练的话，通过改变不同的 B 收集训练数据，实际中仍然能拟合到和上一节一样的的双曲线性质。

现在，和上一节一样，我们也可以类似定义 CBS 的估计量，即

\[B_{crit} = \frac{E_{min}}{S_{min}} \approx \frac{\int B_{noise}(s)ds}{\int ds}.\]

不难看出，这就是 gradient noise scale 在时间上的平均。这意味着，如果用随时间恒定的 B schedule 来训练的话，CBS 是可以用 gradient noise scale 在训练过程上的平均来反映的。事实上，和上一节一样，我们有 $S(B_{crit}) = 2S_{min},\ E(B_{crit}) \approx 2E_{min}$，这意味着 $B_{crit}$ 在时间和数据效率的两种极限之间取了一个折中之处。

Batch Size Schedule

在上一节，我们通过假设 $B(s) = B$，得到了时间 S 和数据 E 的双曲线函数，并提出了基于双曲线拟合和 gradient noise scale 的 global CBS 估计。不过，这看起来并不是一个最优的选项。既然 $B_{noise}$ 是随时间变化的，那我们为什么不随时间动态调整我们的 B 呢？例如，在训练初期，梯度信号比较大，信噪比较高，$B_{noise}$ 比较低，我们可以用比较小的 B 来实现有效的 loss 下降；在训练后期，梯度信号比较少，信噪比较低，$B_{noise}$ 比较高，我们可以用比较大的 B 来实现有效的 loss 下降；这样子看起来无论是时间还是数据效率都是更优的。在这一节，我们就做一个理论上的探讨。

首先把问题 formulate 一下，我们希望找到一个合理的 bs schedule $B(s)$，尽可能地减小训练时间 S 和数据总量 E，即

\[\min_{B(s)} S(B(s)) + \frac{1}{\lambda} E(B(s)) = \min_{B(s)} \int (1 + \frac{B_{noise}(s)}{B(s)})ds + \frac{1}{\lambda} \int (B(s) + {B_{noise}(s)})ds,\]

其中 $\lambda$ 表示的是训练时间与数据耗费的权衡，为一个自由变量。把问题中与 $B(s)$ 无关的项去掉，我们也可以等价地转化为研究

\[\min_{B(s)} \int \frac{B_{noise}(s)}{B(s)}ds + \frac{1}{\lambda} \int B(s)ds.\]

为了得到最优的 bs schedule，我们对 $B(s)$ 做变分，并令其为 0，可以得到

\[- \frac{B_{noise}(s)}{B^2(s)} + \frac{1}{\lambda} = 0.\]

这轻松导出了我们最优的 bs schedule：

\[B(s) = \sqrt{\lambda B_{noise}(s)}.\]

其实这里和直观猜测有些不同，直观感觉既然 $B_{noise}(s)$ 是 local CBS 的一个良好估计，那就直接让 $B(s) = B_{noise}(s)$ 就行了。但理论告诉我们，从全局的考量来看，更优的选择应该是根号阶数的。现在，我们将 bs schedule 代回 S 和 E 的表达式，并看看有没有类似先前的双曲线规律。我们有训练时间为

\[S(B(s)) = \int (1 + \frac{B_{noise}(s)}{B(s)})ds = S_{min} + \frac{\int \sqrt{B_{noise}(s)} ds}{\sqrt{\lambda}} := S_{min} + \frac{\Phi}{\sqrt{\lambda}},\]

其中 $\Phi = \int \sqrt{B_{noise}(s)} ds$。我们也可以得到数据耗费为

\[E(B(s)) = \int (B(s) + {B_{noise}(s)})ds \approx E_{min} + \sqrt{\lambda} \Phi.\]

可以推得和前文类似的双曲线规律：

\[\frac{E(B(s))}{E_{min}} - 1 \approx \frac{\Phi^2}{E_{min} S_{min}} \left(\frac{S(B(s))}{S_{min}} - 1\right)^{-1} := \gamma \left(\frac{S(B(s))}{S_{min}} - 1\right)^{-1},\]

其中 $\gamma = {\Phi^2}/{E_{min} S_{min}} = (\int \sqrt{B_{noise}(s)} ds)^2 / E_{min} S_{min}$。由于使用 bs schedule 进行训练的时间和数据耗费理论上肯定不会差于固定 bs，我们猜测 $\gamma \leq 1$。事实上，这是成立的。基于柯西不等式，我们有：

\[\gamma = \frac{(\int \sqrt{B_{noise}(s)} ds)^2}{E_{min} S_{min}} \leq \frac{(\int {B_{noise}(s)} ds) (\int ds) }{E_{min} S_{min}} = \frac{E_{min} S_{min}}{E_{min} S_{min}} = 1,\]

当 $B_{noise}(s)$ 不随时间变化时取等，此时回退到之前的情况，采用 bs schedule 并不能带来效率的提升。而对于一般的 $B_{noise}(s)$ 随时间变化的情况有 $\gamma < 1$，这意味着引入 bs schedule 对应的双曲线更低，有着更好的时间和数据效率。一般来说，gradient noise scale 变化越剧烈，$\gamma$ 通常越小，引入 bs schedule 的理论收益越大。如下图所示，用了 bs schedule 的点都落在 const bs 的曲线内部，说明效率更高。

对于权衡系数 $\lambda$，原文也给出了一个自然但非唯一的选择：

\[\lambda = B_{crit} = \frac{E_{min}}{S_{min}}.\]

此时对应的训练时间为：

\[S(B(s)) = S_{min} + \frac{\Phi}{\sqrt{\lambda}} = S_{min} + \Phi \frac{\sqrt{S_{min}}}{\sqrt{E_{min}}} = S_{min}(1 + \frac{\Phi}{\sqrt{S_{min} E_{min}}}) = S_{min}(1 + \sqrt{\gamma}),\]

对应的样本耗费为：

\[E(B(s)) \approx E_{min} + \sqrt{\lambda} \Phi = E_{min}(1 + \sqrt{\gamma}).\]

和之前类似，这个选择相对于纯时间最优和纯数据最优的极限恶化相同的倍数，达到了某种平衡折中的阶段。

CBS 的经验估计

最后，我们基于前面的理论，总结一下实际中 CBS 的估计方法。一种是基于 gradient noise scale 的，一种是基于拟合双曲线的。我们默认不考虑 bs schedule 的情况。

基于 gradient noise scale 的 CBS 估计

回顾前面分析训练全程的理论，我们已知 $B_{crit} \approx \frac{\int B_{noise}(s)ds}{\int ds}$ 是 global CBS 的一个比较好的估计量，所以在实际中可以跑一次基准训练，在不同时间估计一下 $B_{noise}(s)$ 或者计算更简单的 $B_{simple}(s)$，然后取平均就可以。由于我们的理论是基于最优学习率的，所以估计 CBS 的基准训练的学习率等超参也需要比较好。如文章提到了学习率太小会抬高 noise scale 的估计。实际中可以用公认比较安全的超参进行实验。

在实践中，这个方法计算比较简单，也可以比较好地估计真实 CBS 的数量级和趋势，但距离精准预测 CBS 还有较大的差距，所以不太适合精确确定 bs 的场景。

基于拟合双曲线的估计

从前面的分析中，我们也知道，可以先拟合时间-数据双曲线，然后再用 $B_{crit} = \frac{E_{min}}{S_{min}}$ 来获得 CBS 的估计。具体的实操流程为：

1. 先固定目标性能。我们前面的分析都是假设了达到某个目标 loss 的。CBS 一般是和 loss 相关的，loss 越小，梯度信号通常越小，CBS 往往越大。因此，我们需要先固定一个目标性能。
2. 扫描不同 bs，并对不同的 bs 扫描学习率等超参。
3. 对于每个 batch size，选择最快达到目标性能的 run，记录其训练步数 S 和数据耗费 E。
4. 拟合时间-数据双曲线，然后用 $B_{crit} = \frac{E_{min}}{S_{min}}$ 来获得 CBS 的估计。

在实践中，这个方法计算成本更高，但能够更为精准地预测 CBS。两种方法的对比如下图所示：

总结

本文对 OpenAI 在 CBS 上的经典工作 An Empirical Model of Large-Batch Training 做了一个回顾，总结其核心的理论思想与对应的 CBS 估计方法。