Achieving small-batch accuracy with large-batch scalability via Hessian-aware learning rate adjustment.

Area of Science:

• Machine Learning
• Deep Learning Optimization
• Computer Vision

Background:

• Synchronous data-parallel neural network training often uses large batch sizes for increased parallelism.
• Large batch sizes can lead to poor generalization performance due to reduced gradient noise.
• Generalization is crucial for model performance on unseen data.

Purpose of the Study:

• To address the generalization degradation issue in large-batch neural network training.
• To propose a novel learning rate adjustment framework to improve generalization.
• To encourage convergence to flat minima known for better generalization.

Main Methods:

• Developed a general learning rate adjustment framework.
• Incorporated three critical heuristics focusing on loss landscape geometry.
• Implemented a Hessian-aware learning rate schedule.

Main Results:

• The Hessian-aware learning rate schedule significantly improves generalization in large-batch training.
• Achieved 92.31% accuracy on CIFAR-10 with ResNet20 using a batch size of 16,384.
• Performance is comparable to 92.83% accuracy with a batch size of 128, at negligible computational overhead.

Conclusions:

• The proposed Hessian-aware learning rate adjustment framework effectively mitigates generalization issues in large-batch training.
• This method enables high parallelism without sacrificing model performance on unseen data.
• The approach offers a computationally efficient solution for improving deep learning model generalization.