Jorge：面向 GPU 高效二阶优化的近似预条件方法
机器学习
2023-10-30 v2 人工智能
分布式、并行与集群计算
摘要
尽管相比一阶优化器具有更好的收敛性质，深度学习中的二阶优化器因其显著的计算成本而较少流行。此类优化器的主要效率瓶颈在于预条件步骤中的矩阵求逆计算，其在 GPU 上计算代价高昂。本文中，我们提出 Jorge，一种兼具两者优势的二阶优化器——兼具二阶方法的快速收敛收益与一阶方法典型的高计算效率。我们通过使用预条件计算的近似完全消除矩阵求逆，解决了计算矩阵逆的主要瓶颈。这使得 Jorge 在 GPU 上就挂钟时间而言极为高效。进一步，我们描述了一种直接从调优良好的 SGD 基线确定 Jorge 超参数的方法，从而显著减少调参工作。我们的实证评估展示了使用 Jorge 的明显优势，在样本效率与挂钟时间上均优于 SGD、AdamW 与 Shampoo 等 SOTA 优化器。
引用
@article{arxiv.2310.12298,
title = {Jorge: Approximate Preconditioning for GPU-efficient Second-order Optimization},
author = {Siddharth Singh and Zachary Sating and Abhinav Bhatele},
journal= {arXiv preprint arXiv:2310.12298},
year = {2023}
}