©作者 | 机器之心编辑部
来源 | 机器之心
鉴于语言模型预训练成本巨大,因而研究者一直在寻找减少训练时间和成本的新方向。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 中触发梯度裁剪。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。
📝 稿件基本要求:
• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注
• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题
• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道:
• 投稿邮箱:hr@paperweekly.site
• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿
△长按添加PaperWeekly小编
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
·
·
·