2023/05/25 19:49

马腾宇团队新出大模型预训练优化器，比Adam快2倍，成本减半

鉴于语言模型预训练成本巨大，因而研究者一直在寻找减少训练时间和成本的新方向。Adam 及其变体多年来一直被奉为最先进的优化器，但其会产生过多的开销。本文提出了一种简单的可扩展的二阶优化器 Sophia，在与 Adam 比较中，Sophia 在减少了 50% step 数量的情况下实现了与 Adam 相同的验证预训练损失。

大语言模型（LLM）的能力随着其规模的增长而取得了显著的进展。然而，由于庞大的数据集和模型规模，预训练 LLM 非常耗时，需要进行数十万次的模型参数更新。例如，PaLM 在 6144 个 TPU 上进行了为期两个月的训练，总共耗费大约 1000 万美元。因此，提高预训练效率是扩展 LLM 规模的一个主要瓶颈。

本文来自斯坦福大学的研究者撰文《 Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training 》，文中提出了 Sophia（Second-order Clipped Stochastic Optimization）轻量级二阶优化器，旨在通过更快的优化器提高预训练效率，从而减少达到相同预训练损失所需的时间和成本，或者在相同预算下实现更好的预训练损失。

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

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) 的梯度，即

。设 m_t 为 EMA（ exponential moving average ）的梯度，则更新的分子为 m_t ← β_1m_t−1 + (1 − β_1) g_t 。

Sophia 使用基于对角 Hessian 的 pre-conditioner，根据参数维度的曲率直接调整更新的大小。为了减少开销，该研究仅在每 k 个step内（实现中 k = 10）估计一次 Hessian。在时间步 t 上，估计器返回小批次损失的 Hessian 对角线的估计

。每 k 个step更新一次 EMA，得到对角 Hessian 估计的以下更新规则：

该研究只考虑对角 Hessian 的正项，并在更新中引入按坐标裁剪，更新规则改写为：

对角 Hessian 估计器

该研究引入了两个对角 Hessian 估计器，它们的内存和运行时间成本都与计算梯度相似。估计器分别为 Hutchinson 无偏估计器以及 GNB（ Gauss-Newton-Bartlett ）估计器。伪代码如下所示：

实验

研究将使用 Hutchinson 估计器和 GNB 估计器的算法分别称为 Sophia-H 和 SophiaG。本文用 GPT-2 评估了 Sophia 的自回归语言建模，模型尺寸从 125M 到 770M 不等。结果表明，Sophia 在 step、总计算量和所有模型大小的 wall-clock 时间方面比 AdamW 和 Lion 快 2 倍。此外，扩展定律更有利于 Sophia 而不是 AdamW。

实验语言建模设置

该实验在 OpenWebText 上训练自回归模型。遵循 GPT-2 的标准协议，将上下文长度设置为 1024。使用只有解码器的 Transformer，模型参数量分别为 125M (小型)、355M (中型) 和 770M (大型)。

基线：研究主要比较 Sophia 和 Adam。Adam 采用解耦权重衰减 (AdamW)，这是语言建模任务中主要使用的优化器，而 Lion 是通过符号搜索发现的一阶自适应优化器。所有优化器都进行了很好的调整。权重衰减被设置为 0.1，β_1 = 0.9， β_2 = 0.95。对于 Lion，使用 β_1 = 0.95 和 β_2 = 0.98。125M 和 355M 模型是在 10 个 A5000 GPU 上训练的，而 770M 模型是在 8 个 A100 GPU 上训练的。

评估：研究人员使用每个优化器对模型进行 100K、200K 或 400K 个step的预训练评估，以比较速度。值得注意的是，与标准一样，LR 调度取决于预先指定的总目标step 数，如图 5 (a) 所示。这使得同一优化器的损失曲线在不同的 step 数下是不同的，因为总 step 数较少的 LR 调度会更早地衰减 LR。本文主要评估了这些模型在 OpenWebText 上的 log 困惑度，并绘制了损失曲线。此外，该研究还报告了 SuperGLUE 上的上下文学习结果，然后对 5 个提示的结果取平均值。

实验结果

图 4 展示了相同 step 数 (100K) 下 OpenWebText 上的验证损失曲线 (token 级 log 困惑度)。与 AdamW 和 Lion 方法相比，本文所提出来的方法获得了更好的验证损失。随着模型大小的增加，Sophia 和基线之间的差距也变得更大。Sophia-H 和 Sophia-G 在 355M 模型上的验证损失都小 0.04 (图 4 (b))。

同样 100K step，Sophia-H 在 770M 模型上的验证损失小了 0.05 (图 4，(c))。可以看出，这是个明显的改进，因为根据该机制中的扩展定律和图 5 中的结果，损失 0.05 的改进相当于实现相同验证损失的 step 数或总计算量的改进的双倍。

Sophia 在 step 数、总计算时间和 wall-clock 时间方面快了两倍。Sophia 对验证损失的改进在于减少 step 数或总计算量。在图 1 (a) 和 (b) 和图 5 中，通过比较达到相同验证损失水平所需的 step 数或总计算量来评估优化器。从图 1 (a) 和 (b) 中可以看出，与 AdamW 和 Lion 相比，Sophia-H 和 Sophia-G 在不同的模型尺寸下实现了 2 倍的加速。

扩展定律更有利于 Sophia-H 而不是 AdamW。在图 1 (c) 中，该研究绘制了预训练 100K step 的不同大小模型的验证损失。Sophia 和 AdamW 之间的差距随着模型的扩大而增大。此外，Sophia-H 训练的 540M 模型比 AdamW 训练的 770M 模型的损失更小。Sophia-H 训练的 355M 模型与 AdamW 训练的 540M 模型的损失相当。

评估下游任务的小样本 (SuperGLUE)。如图 6 所示，验证损失的改善也使得下游任务准确率的提高。在预训练 step 数相同的情况下，使用 Sophia 预训练的 GPT-2 medium 和 GPT-2 large 在大多数子任务上具有更好的少样本准确率。此外，用 Sophia-H 预训练的模型与用 AdamW 预训练的模型具有相当的小样本准确率。

分析

比较 wall-clock 时间与计算量。表 1 比较了每一个 step 的总计算量 (TFLOPs) 和 A100 GPU 上的 wall-clock 时间。本文报告了每个 step 的平均时间，Hessian 计算花费的时间的总计算。较小的批量大小，即每 10 个 step 以计算对角 Hessian 估计，Hessian 计算占总计算量的 6%，与 AdamW 相比，整体 wall-clock 时间开销小于 5%。在内存使用方面，优化器 m 和 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优化器

相关数据

网格搜索技术

网格搜索是一项模型超参数优化技术，常用于优化三个或者更少数量的超参数，本质是一种穷举法。对于每个超参数，使用者选择一个较小的有限集去探索。然后，这些超参数笛卡尔乘积得到若干组超参数。网格搜索使用每组超参数训练模型，挑选验证集误差最小的超参数作为最好的超参数。

来源：Deep Learning Book

权重技术

线性模型中特征的系数，或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0，则相应的特征对模型来说没有任何贡献。

来源：Google AI Glossary

调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源，如线程、进程或数据流；也可以指硬件资源，如处理器、网络连接或扩展卡。进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态，允许多位用户有效地同时共享系统资源，或达到指定的服务质量。 see planning for more details

来源：维基百科

人工智能技术

在学术研究领域，人工智能通常指能够感知周围环境并采取行动以实现最优的可能结果的智能体（intelligent agent）

来源：Russell, S., & Norvig, P. (2003). Artificial Intelligence: A Modern Approach.

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

超参数技术

在机器学习中，超参数是在学习过程开始之前设置其值的参数。相反，其他参数的值是通过训练得出的。不同的模型训练算法需要不同的超参数，一些简单的算法（如普通最小二乘回归）不需要。给定这些超参数，训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式，并且必须对其进行调整以便模型能够最优地解决机器学习问题。在实际应用中一般需要对超参数进行优化，以找到一个超参数元组（tuple），由这些超参数元组形成一个最优化模型，该模型可以将在给定的独立数据上预定义的损失函数最小化。

来源：Wikipedia

伪代码技术

伪代码，又称为虚拟代码，是高层次描述算法的一种方法。它不是一种现实存在的编程语言；它可能综合使用多种编程语言的语法、保留字，甚至会用到自然语言。它以编程语言的书写形式指明算法的职能。相比于程序语言它更类似自然语言。它是半形式化、不标准的语言。

来源：维基百科

准确率技术

分类模型的正确预测所占的比例。在多类别分类中，准确率的定义为：正确的预测数/样本总数。在二元分类中，准确率的定义为：(真正例数+真负例数)/样本总数

来源：Google ML Glossary

参数模型技术

在统计学中，参数模型是可以使用有限数量的参数来描述的分布类型。这些参数通常被收集在一起以形成单个k维参数矢量θ=（θ1，θ2，...，θk）。

来源：Wikipedia

自回归模型技术

自回归模型，是统计上一种处理时间序列的方法，自回归模型被广泛运用在经济学、资讯学、自然现象的预测上。

来源：维基

批次技术

模型训练的一次迭代（即一次梯度更新）中使用的样本集。

来源：Google ML glossary

优化器技术

优化器基类提供了计算梯度loss的方法，并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法，如梯度下降和Adagrad。优化器是提供了一个可以使用各种优化算法的接口，可以让用户直接调用一些经典的优化算法，如梯度下降法等等。优化器（optimizers）类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类，但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer（tensorflow下的优化器包）等等这些算法。

来源：维基百科

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用，如语音识别，机器翻译，词性标注，句法分析和资讯检索。由于字词与句子都是任意组合的长度，因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题)，也使得在语料库中估算字串的机率变得很困难，这也是要使用近似的平滑n元语法(N-gram)模型之原因。

来源：维基百科

GPT-2技术

GPT-2是OpenAI于2019年2月发布的基于 transformer 的大型语言模型，包含 15 亿参数、在一个 800 万网页数据集上训练而成。据介绍，该模型是对 GPT 模型的直接扩展，在超出 10 倍的数据量上进行训练，参数量也多出了 10 倍。在性能方面，该模型能够生产连贯的文本段落，在许多语言建模基准上取得了 SOTA 表现。而且该模型在没有任务特定训练的情况下，能够做到初步的阅读理解、机器翻译、问答和自动摘要。

来源：OpenAI博客