斯坦福大学研究人员推出了Sophia:一种可扩展的语言模型预训练的二阶优化器

Stanford researchers release Sophia a scalable language model pre-trained with second-order optimizer

鉴于训练语言模型的高昂成本,对优化过程进行非平凡改进将极大地减少完成训练过程所需的时间和金钱。Adam及其变种长时间以来一直是最先进的,而二阶(基于Hessian的)优化器很少被使用,因为它们的每步开销更大。

研究人员提出了对角Hessian的轻量级估计作为二阶优化器Sophia的预处理器。Sophia是一种可以比Adam快两倍的新型优化器,它可以解决LLM的问题。在更新之后进行逐元素剪切,剪切是通过取梯度的平均值除以估计Hessian的平均值找到的。剪切限制了最坏情况下更新的大小,并减轻了轨迹的非凸性和快速Hessian变化的影响。添加一些新的代码可能将200万美元的预算降低到100万美元的范围内(假设适用比例定律)。

平均每步的时间和内存开销很低,因为Sophia只在每几次迭代中估计对角线Hessian。Sophia在模型语言建模中的大小从1.25亿到7.7亿之间的GPT-2模型方面,以步数、计算总量和挂钟时间来衡量,速度是Adam的两倍。研究人员证明了Sophia能够适应潜在的语言建模任务下的大参数变化。运行时限与损失的条件数无关。

关键特点

  • Sophia在PyTorch中实现起来很简单,因为它在梯度之前要求对角线Hessian的轻量级估计(请参见第一张图片中的伪代码),然后逐个剪切元素。
  • Sophia还有助于预训练的稳定性。与Adam和Lion相比,它很少引发梯度剪切。不需要使用重新参数化技巧,其中焦点温度与层索引变化。
  • Sophia通过在锋利的尺寸(具有大Hessian的尺寸)上比在平坦尺寸(具有小Hessian的尺寸)上更严重地惩罚更新,确保所有参数维度上的一致损失减少。在二维空间中,Adam收敛较慢。

这个项目的重要方面

  • 这表明即使资源有限,学术界也可以研究LLM的预训练并开发新颖、有效的算法。
  • 除了回顾以前优化课程的材料外,研究人员在整个研究过程中广泛使用理论推理。

在明天发布的代码中,研究人员使用了略微修改的常见LR定义的版本。虽然对于输入更整洁,但论文中的LR定义对于计算机代码可能更好。