1行代码改进大模型训练,Llama训练速度升至1.47倍,华人团队出品

全文1514字,阅读约需5分钟,帮我划重点

划重点

01四名华人学者提出了大模型训练优化器Cautious Optimizers,可提升大模型训练效率至1.47倍。

02该优化器以哈密顿量和下降动力学为理论基础,在加速的同时不影响收敛特性。

03作者在600M到1B不同参数规模的Llama模型上进行了试验,获得了最高47%的加速率。

04为此,研究相关代码已经开源,在GitHub上有使用方法的详细讲解。

05实验结果显示,Cautious Optimizer能有效减少训练震荡,使收敛更平稳高效。

以上内容由腾讯混元大模型生成,仅供参考

克雷西 发自 凹非寺

量子位 | 公众号 QbitAI

只要改一行代码,就能让大模型训练效率提升至1.47倍。

图片

拥有得州大学奥斯汀分校背景四名华人学者,提出了大模型训练优化器Cautious Optimizers。

在提速的同时,Cautious能够保证训练效果不出现损失,而且语言和视觉模型都适用。

图片

该优化器以哈密顿量和下降动力学为理论基础,在加速的同时不影响收敛特性。

作者在600M到1B不同参数规模的Llama模型上进行了试验,获得了最高47%的加速率。

该研究相关代码已经开源,在GitHub上有使用方法的详细讲解。

图片

一行代码改进大模型训练

Cautious Optimizers在PyTorch当中增加的一行代码,核心思路是引入实现一种掩蔽机制,从而避免参数更新的方向与当前梯度方向相悖

图片

因为这两个方向一旦不一致,就有可能导致损失函数暂时增加,造成收敛速度的减缓。

不过作者并未在方向不一致的来源问题上过度纠结,而是引入了一种判断机制,在参数更新之前增加一步计算,从而过滤掉方向不一致的情形。

这也正是上面代码的直接作用。

图片△GD:梯度下降,GDM:带动量的梯度下降,C-GDM:本项目

具体来说,加入的两行代会对u和g两个向量求内积,u向量对应优化器给出的参数更新方向,而g向量对应当前时刻的梯度方向。

作者设计了一个对齐掩码函数ϕ,当u和g的内积小于0时(即方向不一致),ϕ的输出为0向量;当内积大于等于0时,ϕ的输出为全1向量。

而一旦ϕ为零向量时,w_t计算式中含u的项也会变为零向量,导致此项更新被跳过。

图片

这样就可以判断参数更新和梯度方向是否一致,如果不一致则不会用于参数更新,避免了训练过程中损失函数的回升。

训练效率提升47%

为了评估Cautious Optimizers的具体效果,作者分别在语言模型Llama和视觉模型MAE上进行了试验。

作者选取了60M、100M、350M和1B四种参数规模的Llama模型,在C4语料库上进行预训练。

优化器选用了AdamW和Lion,以及它们对应的Cautious版本:C-AdamW和C-Lion,每个实验中进行1万步迭代。

结果C-AdamW和C-Lion在所有规模上都表现出明显的收敛加速效果。

尤其是在1B规模上,相比原版的AdamW和Lion,它们的样本效率分别提高了47%和28%,这表明Cautious Optimizer能有效减少训练震荡,使收敛更平稳高效。

图片

并且,Cautious Optimizer在所有情况下都取得了更低的困惑度,印证了其出色的泛化性能。

图片

为了评估模型的实际效果,研究者在语句匹配、文本蕴含、情感分类等6个GLUE下游任务上测试了AdamW和C-AdamW优化后1B模型的表现,

结果表明,C-AdamW的平均得分比AdamW高出2%,在大多数任务上都取得了进步,说明Cautious跳过部分参数更新的方式不会引起模型性能下降。

图片

对于视觉模型,作者以ViT为骨干网络,在ImageNet-1K数据集上预训练了MAE模型。

由于视觉任务的特殊性,训练过程采用了随机遮挡图像块并重建的范式,因此优化目标是最小化重建误差,而非通常的分类损失。

作者对比了AdamW和C-AdamW的表现,即训练50轮后的最终重建误差,结果C-AdamW的误差为0.5926,低于AdamW的0.6085。

图片

一作曾在一周内复刻o1

本项目是由四名华人学者共同打造的。

第一作者Kaizhao Liang,是AI推理加速服务商SambaNova公司的一名高级ML工程师。

在o1模型发布一周内,该公司就推出了一个类似o1模型思考过程的开源平替,主要作者正是Liang。

图片

其他三名作者是得州大学奥斯汀分校CS助理教授Qiang Liu,以及他的两名博士生,Lizhang Chen和Bo Liu。

此外,Liang的人工智能硕士学位也是从该校获得。

论文地址:
https://arxiv.org/abs/2411.16085

GitHub:
https://github.com/kyleliang919/C-Optim