MaxVA: Fast Adaptation of Step Sizes by Maximizing Observed Variance of Gradients
This work addresses convergence problems in adaptive gradient methods for practitioners in deep learning, particularly for training Transformers and large models, but it is incremental as it builds on existing methods like Adam.
The paper tackles the issue of unstable or extreme adaptive learning rates in Adam by proposing MaxVA, which replaces the running mean of squared gradients with a weighted mean to maximize estimated variance, resulting in faster adaptation and more desirable convergence behaviors. The method demonstrates improved efficacy on tasks like machine translation, natural language understanding, and BERT pretraining.
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.