Scalable Second Order Optimization for Deep Learning

Abstract

Optimization in machine learning, both theoretical and applied, is presently dominated by first-order gradient methods such as stochastic gradient descent. Second-order optimization methods, that involve second derivatives and/or second order statistics of the data, are far less prevalent despite strong theoretical properties, due to their prohibitive computation, memory and communication costs. In an attempt to bridge this gap between theoretical and practical optimization, we present a scalable implementation of a second-order preconditioned method (concretely, a variant of full-matrix Adagrad), that along with several critical algorithmic and numerical improvements, provides significant convergence and wall-clock time improvements compared to conventional first-order methods on state-of-the-art deep models. Our novel design effectively utilizes the prevalent heterogeneous hardware architecture for training deep models, consisting of a multicore CPU coupled with multiple accelerator units. We demonstrate superior performance compared to state-of-the-art on very large learning tasks such as machine translation with Transformers, language modeling with BERT, click-through rate prediction on Criteo, and image classification on ImageNet with ResNet-50.

arXiv comment24 pages, Code available here: https://bit.ly/3uXXtKy

References in corpus (14)

• Optimizing Neural Networks with Kronecker-factored Approximate Curvature
• Deep Learning Recommendation Model for Personalization and Recommendation Systems
• Lingvo: a Modular and Scalable Framework for Sequence-to-Sequence Modeling
• Adaptive Bound Optimization for Online Convex Optimization
• Convergence rates of sub-sampled Newton methods
• Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis
• Shampoo: Preconditioned Stochastic Tensor Optimization
• Practical Quasi-Newton Methods for Training Deep Neural Networks
• Limitations of the Empirical Fisher Approximation for Natural Gradient Descent
• Developing a Recommendation Benchmark for MLPerf Training and Inference
• Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training
• Faster SGD Using Sketched Conditioning
• A Large Batch Optimizer Reality Check: Traditional, Generic Optimizers Suffice Across Batch Sizes
• Disentangling Adaptive Gradient Methods from Learning Rates

Cited by in corpus (10)

• On the Opportunities and Risks of Foundation Models
• $m^{\ast }$ of two-dimensional electron gas: a neural canonical transformation study
• Large-Scale Deep Learning Optimizations: A Comprehensive Survey
• LocoProp: Enhancing BackProp via Local Loss Optimization
• NG+ : A Multi-Step Matrix-Product Natural Gradient Method for Deep Learning
• Noether's Learning Dynamics: Role of Symmetry Breaking in Neural Networks
• Tensor Normal Training for Deep Learning Models
• Bridging the Gap between Deep Learning and Frustrated Quantum Spin System for Extreme-scale Simulations on New Generation of Sunway Supercomputer
• Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping
• Training Efficiency and Robustness in Deep Learning