AdaSAM: Boosting sharpness-aware minimization with adaptive learning rate and momentum for training deep neural...

Area of Science:

• Deep Learning Optimization
• Machine Learning Theory

Background:

• Sharpness Aware Minimization (SAM) improves deep neural network generalization by flattening loss landscapes.
• Integrating SAM with adaptive learning rates and momentum (AdaSAM) shows empirical promise but lacks theoretical backing.
• Analyzing AdaSAM is challenging due to coupled perturbation, adaptive learning rate, and momentum steps.

Purpose of the Study:

• To theoretically analyze the convergence rate of AdaSAM in stochastic non-convex settings.
• To provide the first non-trivial convergence rate for SAM with adaptive learning rate and momentum acceleration.
• To demonstrate AdaSAM's linear speedup property with respect to mini-batch size.

Main Methods:

• Introduced a delayed second-order momentum term to decouple stochastic gradient steps, adaptive learning rate, and perturbed gradient.
• Analyzed convergence by taking expectations to make components independent.
• Bounded the adaptive learning rate to ensure analytical feasibility.

Main Results:

• AdaSAM achieves a theoretical convergence rate of O(1/bT), demonstrating linear speedup with mini-batch size (b).
• The analysis successfully decouples and bounds the adaptive learning rate and momentum components.
• Experimental results on NLP and synthetic tasks show AdaSAM outperforms SGD, AMSGrad, and SAM.

Conclusions:

• AdaSAM offers a theoretically guaranteed convergence rate and linear speedup, validating its effectiveness.
• The proposed analytical framework provides a foundation for understanding complex optimizers like AdaSAM.
• AdaSAM demonstrates superior performance, making it a valuable optimizer for deep learning tasks.