  Adaptive gradient methods such as RMSProp and Adam use exponential moving
estimate of the squared gradient to compute adaptive step sizes, achieving
better convergence than SGD in face of noisy objectives. However, Adam can have
undesirable convergence behaviors due to unstable or extreme adaptive learning
rates. Methods such as AMSGrad and AdaBound have been proposed to stabilize the
adaptive learning rates of Adam in the later stage of training, but they do not
outperform Adam in some practical tasks such as training Transformers
\cite{transformer}. In this paper, we propose an adaptive learning rate
principle, in which the running mean of squared gradient in Adam is replaced by
a weighted mean, with weights chosen to maximize the estimated variance of each
coordinate. This results in a faster adaptation to the local gradient variance,
which leads to more desirable empirical convergence behaviors than Adam. We
prove the proposed algorithm converges under mild assumptions for nonconvex
stochastic optimization problems, and demonstrate the improved efficacy of our
adaptive averaging approach on machine translation, natural language
understanding and large-batch pretraining of BERT. The code is available at
https://github.com/zhuchen03/MaxVA.
