Improving generalization performance of adaptive gradient method via bounded step sizes

Abstract

While adaptive gradient methods such as Adam have been widely used in the training of deep neural networks, a recent study has provided a synthetic function that shows the non-convergence problem of Adam. This issue stems from the existence of extreme gradients and the mismatch between the first and second moments. Several adaptive optimizers have been continuously developed. However, designing a fast optimizer with excellent generalization capability is still challenging. We propose an adaptive method with bounded step sizes, named AdaBS, which removes the extreme step sizes and ensures that it appropriately adjusts adaptive step sizes to mitigate the over-adaptation of step sizes in Adam. In particular, AdaBS effectively clips step sizes that are too large or too small by using two static bounds with a predetermined boundary to control updates. When determining the step size, static bound clipping will be used if the preconditioner is outside the modest boundary, and vanilla Adam will be used if the preconditioner is inside the boundary. AdaBS establishes a trust region around the basic step size and obtains benefits of both Adam and SGD, i.e. fast convergence and better generalization. Finally, we conduct extensive experiments on a variety of practical tasks with benchmark datasets, including image classification and modeling language tasks. Empirical results demonstrate AdaBS’s promising performance with remarkably fast convergence, superior generalization, and robustness.