单个GPU也能训练GPT-3！快来看看HP调优新范式吧！

大数据文摘转载自微软研究院AI头条

伟大的科学成就不能仅靠反复试验取得。例如太空计划中的每一次发射都是基于数百年的空气动力学、推进和天体等基础研究。同样，在构建大规模人工智能系统时，基础研究大大减少了试错次数，效益明显。

超参数（Hyperparameter，HP）调优是深度学习的关键，但也是一个昂贵的过程，对于具有数十亿参数的神经网络来说更是如此。假如 HP 选择不当，会导致模型性能不佳、训练不稳定。当训练超大型深度学习模型时，这些问题更加严重。

最近，有研究表明不同的神经网络参数化会导致不同的无限宽度极限（infinite-width limits），他们提出了最大更新参数化（Maximal Update Parametrization，µP），该参数化的无限宽极限是实现“最大”特征学习的。直观地说，它确保每一层激活在训练期间的更新大小保持一致，而不管宽度如何。相比之下，虽然标准参数化 （standard parametrization，SP） 在初始化时保证了激活的一致性，但实际上在训练时，由于每层学习率的不平衡，导致激活在宽模型中爆炸。

来自微软和 OpenAI 的研究者首次提出了基础研究如何调优大型神经网络（这些神经网络过于庞大而无法多次训练）。他们通过展示特定参数化保留不同大小模型的最佳超参数来实现这一点。利用 µP 将 HP 从小型模型迁移到大型模型。也就是说，该研究在大型模型上获得了接近最优的 HP。

本文的想法非常简单，论文中引入了一种特殊参数化 µP，窄和宽的神经网络共享同一组最优超参数。即使宽度→∞也是如此。

论文作者之一、来自微软的 Greg Yang 在推特上表示：“你不能在单个 GPU 上训练 GPT-3，更不用说调优它的超参数了。但是利用新的理论进步，你竟可以在单个 GPU 上调优它的超参数！”

具体而言，该研究证明，在 µP 中，即使模型大小发生变化，许多最优的 HP 仍保持稳定。这导致一种新的 HP 调优范式：µTransfer，即在 µP 中对目标模型进行参数化，并在较小的模型上间接调优 HP，将其零样本迁移到全尺寸模型上，无需调优后者。该研究在 Transformer 和 ResNet 上验证 µTransfer，例如，1）通过从 13M 参数的模型中迁移预训练 HP，该研究优于 BERT-large (350M 参数)，总调优成本相当于一次预训练 BERT-large；2）通过从 40M 参数迁移，该研究的性能优于已公开的 6.7B GPT-3 模型，调优成本仅为总预训练成本的7%。

• 论文地址：https://arxiv.org/pdf/2203.03466.pdf

• 项目地址：https://github.com/microsoft/mup

通过大大减少猜测训练超参数的需要，这种技术可以加快对巨大神经网络的研究，例如 GPT-3 和未来可能更大的继任者。

扩展初始化容易，但扩展训练难

大型神经网络很难训练，部分原因是不了解其行为如何随着规模增加而变化。在深度学习的早期工作中，研究者采用启发式算法。一般来说，启发式方法试图在模型初始化时保持激活大小一致，无论宽度如何。然而，随着训练的开始，这种一致性会在不同的模型宽度处中断，如图1左侧所示。

与随机初始化不同，模型训练期间的行为更难进行数学分析。该研究用 µP 解决，如图1右侧所示，该图显示了网络激活扩展（activation scales）在模型宽度增加的最初几个训练步骤中的稳定性。

图1：在 PyTorch 的默认参数化中，左图，在经过一次 step 训练后，激活扩展的宽度会出现差异。但是在右图的 µP 中，无论训练 step 宽度如何，激活扩展都会发生一致的变化。

事实上，除了在整个训练过程中保持激活扩展一致之外，µP 还确保不同且足够宽的神经网络在训练过程中表现相似，以使它们收敛到一个理想的无穷宽极限，该研究称之为特征学习极限。

如图所示，µP 是唯一在宽度上保持最佳学习率的参数化，在宽度为8192的模型中实现了最佳性能，并且对于给定的学习率，更宽的模型性能更好——即曲线不相交。

图2：左侧，该研究在 CIFAR10 上以不同的学习率（沿 x 轴显示）训练不同宽度（对应于不同颜色和图案的曲线）的多层感知器 (MLP)，并沿 y 轴绘制训练损失。右侧，参数化的 2D 平面由以下插值形成：1）PyTorch 默认值和 µP（x 轴）之间的初始化扩展，以及 2）PyTorch 默认值和 µP（y 轴）之间的学习率扩展。在这个平面上，PyTorch 默认用 (0,0) 表示，µP 默认用 (1,1) 表示。

基于张量程序（Tensor Programs）的理论基础，µTransfer 自动适用于高级架构，例如 Transformer 和 ResNet。此外，它还可以同时迁移各种超参数。

以 Transformer 为例，图3展示了关键超参数如何在宽度上保持稳定。超参数可以包括学习率、学习率 schedule、初始化、参数乘数等，甚至可以单独针对每个参数张量。该研究在最大宽度为4096的 Transformer 上验证了这一点。

图3：在 µP 中参数化并在 WikiText-2 上训练的不同宽度的 transformer。随着模型宽度的增加，最优学习率、交叉熵温度、初始化规模和学习率方案保持稳定。在右下角的图中，该研究尝试了如下学习率方案：(a) 线性衰减，(b) StepLR @ [5k, 8k]，衰减因子为0.1，(c) StepLR @ [4k, 7k]，衰减因子为0.3，(d) 余弦退火，(e) 常数，(f) 逆平方根衰减。

模型深度的实验扩展

现代神经网络扩展不止涉及宽度一个维度。该研究还探索了如何通过将 µP 与非宽度维度的简单启发式算法相结合，将其应用于现实的训练场景。下图4使用相同的 transformer 设置来显示最佳学习率如何在合理的非宽度维度范围内保持稳定。

图4：在 µP 中参数化并在 Wikitext-2 上训练的不同大小的 transformer。如图3所示，最优学习率不仅可以跨宽度迁移，还可在测试范围内实验性地跨其他扩展维度迁移，例如深度、批大小和序列长度。这意味着可以将理论上的跨宽度迁移与实验验证的跨其他扩展维度迁移相结合，以获得能在小模型上间接调整超参数并迁移到大模型的 µTransfer。

除了学习率，其他超参数的情况如下图所示：

测试 µTransfer

在验证完单个超参数的迁移之后，研究者试图将它们组合到更现实的场景中。下图5对比了两种情况，一种是 µTransfer 从一个小型 proxy 模型迁移调整过的超参数，另一种是直接调整大型目标模型。在这两种情况下，调整都是通过随机搜索完成的。

图5：µTransfer 大约将计算效率提高了一个数量级。

由于 proxy 模型即使很小也能有意义地预测最佳超参数（如图3、图4所示），因此随着该研究用数十亿个参数训练更大的目标模型，研究者预计性能差距会扩大。

µP + GPT-3

在这项工作之前，模型越大，调优成本越高，预计调优效果就越差。研究者预计 µTransfer 将给最大的模型带来最大的增益，因此该研究与 OpenAI 合作，在 GPT-3 上评估 µTransfer。

使用 µP 对 GPT-3 的一个相对位置编码版本进行参数化后，该研究调整了一个具有4000万个参数的小型 proxy 模型，然后按照 µTransfer 的方法将最佳超参数组合复制到 GPT-3 的67亿参数变体中。在此调整阶段使用的总计算量仅为67亿模型预训练使用计算量的7%。如下图所示，这个使用 µTransfer 的模型优于 GPT-3 论文中相同大小的模型（绝对位置编码），它的性能与 GPT-3 论文中参数数量翻倍的模型（绝对位置编码）相当。

理论意义

µP 给出了一个扩展规则，该规则在训练损失方面唯一地保留了跨不同宽度模型的最佳超参数组合。相反，其他扩展规则（如 PyTorch 中的默认初始化或 NTK 参数化），随着网络变得越来越宽，超参数空间中的最优值却越来越远。研究者认为：实际使用特征学习神经网络时，为了获得适用的见解，µP 的特征学习极限会比 NTK 极限更自然。因此，过参数化神经网络应该在大型宽度设置中重现 µP 的特征学习极限。

过去几年开发的张量程序（TP）理论使这项进展成为可能。TP 理论使研究人员能够计算任何通用计算图在其矩阵维数变大时的极限。TP 方法产生了基本的理论结果，例如神经网络 - 高斯过程对应的架构普遍性和动态二分定理，并通过推导 µP 和特征学习极限形成了 µTransfer。研究者认为将 TP 理论扩展到深度、批大小等扩展维度是大型模型在宽度之外可靠扩展的关键。

研究者表示：基础研究是对反复试错的一种高成本效益补充，该研究将继续推导出更具原则性的大规模机器学习方法。

点「在看」的人都变好看了哦！