Scaling Laws

交叉熵，梯度下降，自注意力那些东西就不在这里赘述了。稍微挑一点 AIST5050 里我觉得有意思的东西记录一下。

$C=6ND$

设：

• $N$：模型的 trainable parameter count。
• $D$：训练过程中处理的 token occurrences 数量。
• $C$：总 training compute，单位为 FLOPs。

对于 dense neural networks，可以近似估计：

$$\boxed{C_{\mathrm{train}}\approx6ND}.$$

推导过程如下，考虑一个 dense layer：

$$Y=WH,$$

其中：

$$W\in\mathbb R^{d_{\mathrm{out}}\times d_{\mathrm{in}}}, \qquad H\in\mathbb R^{d_{\mathrm{in}}\times b}.$$

该层具有 $N_W=d_{\mathrm{out}}d_{\mathrm{in}}$ 个权重。

Forward pass 的主要矩阵乘法需要：

$$2d_{\mathrm{out}}d_{\mathrm{in}}b =2N_Wb$$

次 FLOPs（结果矩阵有 $d_{\mathrm{out}}b$ 个元素，得到每个元素需要 $d_{\mathrm{in}}$ 次乘法与 $d_{\mathrm{in}}-1$ 次加法）。

Backward pass 主要计算两个矩阵乘法：

$$\frac{\partial J}{\partial W}=GH^\top, \qquad \frac{\partial J}{\partial H}=W^\top G.$$

它们各自也需要大约 $2N_Wb$ FLOPs。

因此：

$$C_{\mathrm{forward+backward}} \approx2N_Wb+2N_Wb+2N_Wb =6N_Wb.$$

对所有 dense layers 和训练 tokens 累加：

$$C_{\mathrm{train}}\approx6ND.$$

这个近似模型忽略了部分 attention computation、optimizer overhead、bias operations 和 recomputation 等开销。

Chinchilla Scaling Law

Chinchilla（Hoffmann et al., 2022）研究了固定 compute budget 下的 optimal model/data allocation。其经验结论之一是，compute-optimal 模型应当在模型规模增加时同步增加训练 tokens。

在 Chinchilla 的训练设置下，这一比例为：

$$\frac{D^{*}}{N^{*}}\approx20.$$

即每个参数对应约 20 个 training tokens。

此外，Chinchilla 的一种经验 loss model（通常关注在 held-out 数据上的 loss）为：

$$\boxed{ \hat f(N,D)= f_{\mathrm{floor}}+\frac{A}{N^\alpha} +\frac{B}{D^\beta} }$$

其中 $A,B,\alpha,\beta>0$。

三项分别表示：

• $f_{\mathrm{floor}}$：不可消除的基准 loss。
• $A/N^\alpha$：有限模型容量带来的 loss。
• $B/D^\beta$：有限训练数据量带来的 loss。

模型参数 $\theta$ 是在每次训练运行中学习的；这里的 $A,B,\alpha,\beta,f_{\mathrm{floor}}$ 是跨多次实验拟合出的 scaling-law 常数。也就是说，模型的 Loss 可以近似分解为一个基准值，以及两种由资源不足导致的额外 Loss。

给定预算：

$$6ND=C,$$

代入 $D=C/(6N)$：

$$\hat f_C(N) =f_{\mathrm{floor}} +\frac{A}{N^\alpha} +B\left(\frac{6N}{C}\right)^\beta.$$

对 $N$ 求导：

$$\hat f_C'(N)= -\alpha A N^{-\alpha-1} +\beta B\left(\frac6C\right)^\beta N^{\beta-1}.$$

令导数为零并整理：

$$\frac{\alpha A}{N^\alpha} =\beta B\left(\frac{6N}{C}\right)^\beta.$$

因此：

$$N^{\alpha+\beta} =\frac{\alpha A}{\beta B} \left(\frac C6\right)^\beta.$$

最终：

$$\boxed{ N^{*}= \left(\frac{\alpha A}{\beta B}\right)^{\frac1{\alpha+\beta}} \left(\frac C6\right)^{\frac{\beta}{\alpha+\beta}} }$$

并且：

$$D^{*}=\frac{C^{*}}{6N^{*}}.$$

所以：

$$N^{*}\propto C^{\frac{\beta}{\alpha+\beta}}, \qquad D^{*}\propto C^{\frac{\alpha}{\alpha+\beta}}.$$

这里的最优解是 fitted loss model 下的最优 compute allocation，即，在所有满足 $6ND=C$ 的方案中，使得预测 Loss $\hat f(N,D)$ 最小的那组模型参数量 $N^{*}$ 与训练数据量 $D^{*}$。

如果进一步采用 Chinchilla 根据经验得出的近似常数比例，可以得到：

$$N^{*}\propto\sqrt C,\qquad D^{*}\propto\sqrt C.$$

这意味着，如果 compute budget 增加四倍，推荐的 model size 和 training tokens 都增加约两倍。