AEGD: adaptive gradient descent with energy

Abstract

We propose AEGD, a new algorithm for optimization of non-convex objective functions, based on a dynamically updated 'energy' variable. The method is shown to be unconditionally energy stable, irrespective of the base step size. We prove energy-dependent convergence rates of AEGD for both non-convex and convex objectives, which for a suitably small step size recovers desired convergence rates for the batch gradient descent. We also provide an energy-dependent bound on the stationary convergence of AEGD in the stochastic non-convex setting. The method is straightforward to implement and requires little tuning of hyper-parameters. Experimental results demonstrate that AEGD works well for a large variety of optimization problems. Specifically, it is robust with respect to initial data, capable of making rapid initial progress. The stochastic AEGD shows comparable and often better generalization performance than SGD with momentum for deep neural networks. The code is available at https://github.com/txping/AEGD.