Mu-scaling：大模型科研新范式，无需训练即可预知表现

智源研究院认知与数据团队提出了Mu-scaling：一种基于最大更新参数化(Maximal Update Parametrization, MUP)的损失预测方法。该方法通过MUP放缩规则，生成一组宽度不同的模型序列，借助小模型的损失(loss)值直接预测大模型的损失。实验证明，对于超参数空间中的损失低谷(loss basin)，模型损失与参数规模的关系可以精确拟合为幂律。该方法可直接确定大模型训练所应使用的部分超参数，不需要重新搜索。基于该发现，在大参数规模上，不同模型结构和算法可以通过损失预测直接进行有意义的比较，而不需要实际训练。这一特性预示了大模型时代一种新的科研范式。论文地址：https://arxiv.org/abs/2304.06875文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

1. 背景：损失预测（Loss Prediction）

2023年3月，OpenAI在GPT-4技术报告[1]中展示了一种能够实现可预测扩展(Predictable Scaling)的AI底层框架，在模型开始训练之前或开始之后短时间内预测最终的损失(loss)以及部分其他指标，如下图所示。本文所公开的Mu-scaling正是一种实现损失预测的方法。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

GPT-4的预测文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

2. 实现损失预测的主要难点

较多现有工作已经通过实验证明，模型训练损失与运算量、模型参数规模、训练数据规模之间存在幂律关系(Scaling Law)[2]。然而，该关系必须通过在所有规模的模型上进行超参数(例如学习率)的调参才能够成立，而在大模型开始训练之前，往往并不知道最优的超参数选择。即使能够确定最优超参数，如何设计中介模型序列对其最终训练损失进行预测，也是悬而未决的问题。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

3. Mu-scaling方法

本文提出Mu-scaling方法，通过最大更新参数化(MUP)[3]解决上述问题，只需要在小模型上进行一次最优参数搜索，即可通过修正的幂律精确预测大模型的损失。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

3.1 最大更新参数化(MUP)

神经网络的无穷宽度极限理论证明：在（学习率、初始化方差、乘子）这一组超参数满足某种随神经网络宽度放缩的规则的前提下，不同宽度的模型的最优超参数值是对齐的，这组放缩规则被称为MUP（下图）。通过MUP规则，可以在小宽度模型上对（学习率、初始化方差、乘子）进行调参，并直接计算得到该组超参数在大宽度模型上的最优值。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

使用uP的不同宽度 loss 是对齐的（右），不使用的是散乱的（左）文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

然而，MUP本身不具有直接预测loss值的能力，对于其他的超参数，以及不同模型之间而言，其最终训练性能仍然无法比较。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

3.2 本文方法：Mu-scaling

Mu-scaling（下图）利用MUP随模型宽度参数化的过程，生成有且只有宽度不同的模型序列。经过小模型上的单次超参数搜索后，该序列中所有模型的最优（学习率、初始化方差、乘子）值都可以直接计算得到。本文通过实验证明了：对于在上述三维超参数空间中位于（统一的）最优超参数值附近的点(即对序列中所有模型都具有较低训练损失的点，也就是超参数空间中的损失低谷(loss basin))，模型的训练/预测损失值与模型参数量之间满足高精确度的幂律。以这一模型序列为媒介，任意大模型的最优损失可以通过小模型拟合幂律直接计算得到，而对应该最优损失的（学习率、初始化方差、乘子）值则直接通过MUP计算得到，从而在大模型训练开始之前即可预测loss。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

对于其他的超参数，以及所有模型结构的设计和改进而言，只需分别通过Mu-scaling生成模型序列并预测其损失值，即可在大参数规模上进行可信的比较。更多细节请参阅原文。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

本文方法文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

4. 实验结果

本文以12层GPT-2模型在OpenWebText数据集上从头开始预训练20k步作为完整训练过程，在如下宽度序列上实现Mu-scaling：文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

实验设置文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

其中128～1024宽度的训练损失用于拟合幂律，2048和3072宽度作为预测目标。我们在256宽度上对（学习率、初始化方差、乘子）进行单次最优超参数搜索，发现以下三组参数：(7.5e-4, 0.04, 6.0) / (1e-3, 0.05, 5.0) / (2e-3, 0.06, 4.0) 在MUP放缩之后，对于序列中所有模型都具有较低loss，即处于loss basin附近。而另有两组超参数(3e-3, 0.02, 2.0)以及(1e-4, 0.01, 0.5)具有较高的loss。对于这五组超参数，幂律拟合预测结果如下：文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

实验结果文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

可以发现，对于loss basin中的超参数点，只需要训练8M～200M参数的模型，即可通过Mu-scaling准确预测10倍以上参数规模的训练损失；而距离loss basin越远，拟合效果相应下降。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

5. 大模型科研新范式

基于Mu-scaling损失预测，不同模型在任意参数规模下都能够直接比较，这预示着大模型时代一种全新的科研范式。在该范式下，计算资源不足的研究者也可以对较大参数规模的模型产出较为可信的结论。具体流程如下：文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

(1) 对于每种模型设计，在小参数上搜索最优（学习率、初始化方差、乘子）组。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

(2) 生成一组仅宽度不同的模型序列，每个模型通过MUP计算其对应的最优（学习率、初始化方差、乘子）值。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

(3) 训练上述序列中的若干小宽度模型，拟合幂律，预测大宽度下的损失。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

(4) 对所有模型设计，重复(1) – (3)，横向比较，找出大宽度下损失最低的模型。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

上述过程中，训练大宽度模型是非必要的，只作为算力充足条件下的可选项。更多细节请参考原文。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

我们希望该工作能够启发建立一种新的科研范式，减少由于小模型与大模型之间科研结论无法推广所导致的资源浪费，从而在学术和工业界中形成高效的迭代和合作。文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

参考文献

[1] OpenAI. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

[2] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361.文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html

[3] Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. 2021. Tuning large neural networks via zero-shot hyperparameter transfer. In Advances in Neural Information Processing Systems.文章源自菜鸟学院-https://www.cainiaoxueyuan.com/ai/44659.html