Despite their better convergence properties compared to first-order optimizers, second-order optimizers for deep learning have been less popular due to their significant computational costs. The primary efficiency bottleneck in such optimizers is matrix inverse calculations in the preconditioning step, which are expensive to compute on GPUs. In this paper, we introduce Jorge, a second-order optimizer that promises the best of both worlds -- rapid convergence benefits of second-order methods, and high computational efficiency typical of first-order methods. We address the primary computational bottleneck of computing matrix inverses by completely eliminating them using an approximation of the preconditioner computation. This makes Jorge extremely efficient on GPUs in terms of wall-clock time. Further, we describe an approach to determine Jorge's hyperparameters directly from a well-tuned SGD baseline, thereby significantly minimizing tuning efforts. Our empirical evaluations demonstrate the distinct advantages of using Jorge, outperforming state-of-the-art optimizers such as SGD, AdamW, and Shampoo across multiple deep learning models, both in terms of sample efficiency and wall-clock time.
翻译:尽管二阶优化器相较一阶优化器具有更优的收敛特性,但由于其巨大的计算成本,在深度学习领域尚未得到广泛应用。这类优化器的核心效率瓶颈在于预处理步骤中的矩阵逆运算,该运算在GPU上计算成本高昂。本文提出Jorge——一种兼具二阶方法快速收敛优势与一阶方法高计算效率的二阶优化器。我们通过完全消除矩阵逆计算的方法解决主要计算瓶颈,采用预处理计算的近似方案。这使得Jorge在GPU上的运行时间效率极高。此外,我们描述了直接从良好调优的SGD基线确定Jorge超参数的方法,从而大幅减少调参工作量。实证评估表明,使用Jorge具有显著优势,在多个深度学习模型上,无论是样本效率还是运行时间,均优于SGD、AdamW和Shampoo等最先进优化器。