鉴于语言模型预训练成本巨大，因而研究者一直在寻找减少训练时间和成本的新方向。Adam 及其变体多年来一直被奉为最先进的优化器，但其会产生过多的开销。本文提出了一种简单的可扩展的二阶优化器 Sophia，在与 Adam 比较中，Sophia 在减少了 50% step 数量的情况下实现了与 Adam 相同的验证预训练损失。
大语言模型（LLM）的能力随着其规模的增长而取得了显著的进展。然而，由于庞大的数据集和模型规模，预训练 LLM 非常耗时，需要进行数十万次的模型参数更新。例如，PaLM 在 6144 个 TPU 上进行了为期两个月的训练，总共耗费大约 1000 万美元。因此，提高预训练效率是扩展 LLM 规模的一个主要瓶颈。
Sophia 优化器使用随机估计作为 Hessian 矩阵对角线的 pre-conditioner，并采用剪切（clipping）机制来控制最坏情况下的参数大小更新。在像 GPT-2 这样的预训练语言模型上，Sophia 与 Adam 相比，在减少了 50% step 数量的情况下实现了相同的验证预训练损失。
由于 Sophia 可以维持每个 step 内的的内存和所用时间，这相当于总计算量减少了 50%，wall-clock 时间减少了 50%(参见图 1 (a) 和 (b))。此外，根据扩展定律（从模型大小的角度），在 125M 到 770M 的模型上，Sophia 相对于 Adam 更具优势，并且随着模型大小的增加，Sophia 和 Adam 差距也在增加（在 100K step 内）（图 1（c））。特别的，在 540M 参数模型上（100K step），Sophia 和具有 770M 参数模型上的 Adam 实现了相同的验证损失。需要注意的是，后者需要多 达40％的训练时间和 40％的推理成本。
这项研究也得到了大家的认可。英伟达人工智能科学家 Jim Fan 表示：「多年来，有无数论文推翻了 Adam 优化器，不知道 Sophia 会不会是保留到最后的那个，这是一个可扩展的二阶优化器， 其伪代码只有 13 行，在 GPT-2 （预训练）上比 Adam 快了 2 倍，我很想试试这个优化器！」
论文作者之一、斯坦福大学助理教授马腾宇表示：「（从发布之初，）Adam 可以说是一个 9 岁的优化器，是训练 LLM 的首选，如 GPT-3、OPT、 LLAMA 等。而我们的新研究 Sophia 是一个崭新的优化器，在 LLM 上比 Adam 快了 2 倍。只需要你多写几行代码，你的成本就能从 200 万美元降至 100 万美元（如果扩展定律成立的话）。」
至于该研究的动机，作者表示 Adam 对于异构曲率（heterogeneous curvatures）的适应性不足。另一方面，vanilla Newton 方法在凸函数中具有最优的 pre-conditioner，但对于负曲率和 Hessian 的快速变化容易受到影响。基于这些见解，该研究设计了一种新的优化器 Sophia，它比 Adam 更适应异构曲率，比 Newton 方法更能抵抗非凸性和 Hessian 的快速变化，并且还使用了成本较低的 pre-conditioner。
方法理论方面，在时间步长 t 上，该研究用 θ_t 表示参数。在每个 step 上，该研究从数据分布中采样一个小批次，计算小批次损失，并用 L_t (θ_t) 表示。g_t 表示 L_t (θ_t) 的梯度，即
Sophia 使用基于对角 Hessian 的 pre-conditioner，根据参数维度的曲率直接调整更新的大小。为了减少开销，该研究仅在每 k 个step内（实现中 k = 10）估计一次 Hessian。在时间步 t 上，估计器返回小批次损失的 Hessian 对角线的估计
。每 k 个step更新一次 EMA，得到对角 Hessian 估计的以下更新规则：
该研究只考虑对角 Hessian 的正项，并在更新中引入按坐标裁剪，更新规则改写为：
该研究引入了两个对角 Hessian 估计器，它们的内存和运行时间成本都与计算梯度相似。估计器分别为 Hutchinson 无偏估计器以及 GNB（ Gauss-Newton-Bartlett ） 估计器。伪代码如下所示：
评估：研究人员使用每个优化器对模型进行 100K、200K 或 400K 个step的预训练评估，以比较速度。值得注意的是，与标准一样，LR 调度取决于预先指定的总目标step 数，如图 5 (a) 所示。这使得同一优化器的损失曲线在不同的 step 数下是不同的，因为总 step 数较少的 LR 调度会更早地衰减 LR。本文主要评估了这些模型在 OpenWebText 上的 log 困惑度，并绘制了损失曲线。此外，该研究还报告了 SuperGLUE 上的上下文学习结果，然后对 5 个提示的结果取平均值。
同样 100K step，Sophia-H 在 770M 模型上的验证损失小了 0.05 (图 4，(c))。可以看出，这是个明显的改进，因为根据该机制中的扩展定律和图 5 中的结果，损失 0.05 的改进相当于实现相同验证损失的 step 数或总计算量的改进的双倍。
评估下游任务的小样本 (SuperGLUE)。如图 6 所示，验证损失的改善也使得下游任务准确率的提高。在预训练 step 数相同的情况下，使用 Sophia 预训练的 GPT-2 medium 和 GPT-2 large 在大多数子任务上具有更好的少样本准确率。此外，用 Sophia-H 预训练的模型与用 AdamW 预训练的模型具有相当的小样本准确率。
在 30M 模型上，执行网格搜索来测试 Sophia-H 对超参数的敏感性 (图 7 (c))。所有组合的性能相近，但 β_2 = 0.99 和 ρ = 0.1 的性能最好。此外，这种超参数选择可以跨模型大小迁移。对于 125M、355M 和 770M 的所有实验，都使用了 30M 模型上搜索超参数 ρ = 0.01， β_2 = 0.99。
训练稳定性。与 AdamW 和 Lion 相比，Sophia-H 在预训练中具有更好的稳定性。梯度裁剪 (by norm) 是语言模型预训练中的一项重要技术。在实践中，梯度裁剪触发的频率与训练的稳定性有关 —— 如果梯度被频繁裁剪，迭代可能处于非常不稳定的状态。图 7 (a) 比较了 GPT-2 (125M) 触发梯度裁剪的 step 比例。尽管所有方法都使用相同的裁剪阈值 1.0，但 Sophia-H 很少触发梯度裁剪，而 AdamW 和 Lion 在超过 10% 的 step 中触发梯度裁剪。
原标题：《马腾宇团队新出大模型预训练优化器，比Adam快2倍，成本减半》
本文为澎湃号作者或机构在澎湃新闻上传并发布，仅代表该作者或机构观点，不代表澎湃新闻的观点或立场，澎湃新闻仅提供信息发布平台。申请澎湃号请用电脑访问。