Jorge: Approximate Preconditioning for GPU-efficient Second-order Optimization
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.
Code (0)
등록된 구현이 없습니다.
Tasks
Computational EfficiencyGPUSecond-order methodsMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
NysAct: A Scalable Preconditioned Gradient Descent using Nystrom Approximation
Adaptive gradient methods are computationally efficient and converge quickly, but they often suffer from poor generalization. In contrast, second-order methods enhance convergence and generalization but typically incur h…
Second-order methodsScalable K-FAC Training for Deep Neural Networks with Distributed Preconditioning
The second-order optimization methods, notably the D-KFAC (Distributed Kronecker Factored Approximate Curvature) algorithms, have gained traction on accelerating deep neural network (DNN) training on GPU clusters. Howeve…
GPUAmortized Proximal Optimization
We propose a framework for online meta-optimization of parameters that govern optimization, called Amortized Proximal Optimization (APO). We first interpret various existing neural network optimizers as approximate stoch…
image-classificationImage ClassificationImage ReconstructionMeta-Learning+1Adaptive Learning Rate via Covariance Matrix Based Preconditioning for Deep Neural Networks
Adaptive learning rate algorithms such as RMSProp are widely used for training deep neural networks. RMSProp offers efficient training since it uses first order gradients to approximate Hessian-based preconditioning. How…
Stochastic OptimizationThe Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
Recent efforts to accelerate LLM pretraining have focused on computationally-efficient approximations that exploit second-order structure. This raises a key question for large-scale training: how much performance is forf…