  Sharpness aware minimization (SAM) optimizer has been extensively explored as
it can generalize better for training deep neural networks via introducing
extra perturbation steps to flatten the landscape of deep learning models.
Integrating SAM with adaptive learning rate and momentum acceleration, dubbed
AdaSAM, has already been explored empirically to train large-scale deep neural
networks without theoretical guarantee due to the triple difficulties in
analyzing the coupled perturbation step, adaptive learning rate and momentum
step. In this paper, we try to analyze the convergence rate of AdaSAM in the
stochastic non-convex setting. We theoretically show that AdaSAM admits a
$\mathcal{O}(1/\sqrt{bT})$ convergence rate, which achieves linear speedup
property with respect to mini-batch size $b$. Specifically, to decouple the
stochastic gradient steps with the adaptive learning rate and perturbed
gradient, we introduce the delayed second-order momentum term to decompose them
to make them independent while taking an expectation during the analysis. Then
we bound them by showing the adaptive learning rate has a limited range, which
makes our analysis feasible. To the best of our knowledge, we are the first to
provide the non-trivial convergence rate of SAM with an adaptive learning rate
and momentum acceleration. At last, we conduct several experiments on several
NLP tasks, which show that AdaSAM could achieve superior performance compared
with SGD, AMSGrad, and SAM optimizers.
