Scalable and Practical Natural Gradient for Large-Scale Deep Learning

Area of Science:

• Deep Learning
• Machine Learning Optimization
• Computer Vision

Background:

• Large-scale distributed training of deep neural networks often leads to decreased generalization performance due to increased effective mini-batch sizes.
• Existing methods to mitigate this issue involve complex adjustments to learning rates, batch sizes, and batch normalization techniques.

Purpose of the Study:

• To introduce a scalable and practical natural gradient descent (SP-NGD) method for training deep neural networks.
• To enable models to achieve generalization performance comparable to first-order optimization methods while accelerating convergence.
• To demonstrate the scalability of SP-NGD with large mini-batch sizes and minimal computational overhead.

Main Methods:

• Implementation of Scalable and Practical Natural Gradient Descent (SP-NGD).
• Evaluation on a ResNet-50 image classification task on the ImageNet dataset.
• Comparison with highly optimized first-order optimization methods.

Main Results:

• SP-NGD achieved a top-1 validation accuracy of 75.4% in 5.5 minutes with a mini-batch size of 32,768 using 1,024 GPUs.
• Demonstrated convergence to 74.9% accuracy with an extremely large mini-batch size of 131,072 in just 873 SP-NGD steps.
• SP-NGD exhibited negligible computational overhead compared to first-order methods.

Conclusions:

• SP-NGD is a principled and effective approach for large-scale distributed deep learning.
• The method successfully addresses the generalization gap associated with large mini-batch sizes.
• SP-NGD offers accelerated convergence and scalability, making it a practical choice for training deep neural networks efficiently.