Robust Learning Rate Selection for Stochastic Optimization via Splitting Diagnostic
This addresses the challenge of learning rate tuning for practitioners in machine learning, though it is incremental as it builds on existing stochastic optimization methods.
The paper tackles the problem of selecting learning rates in stochastic optimization by proposing SplitSGD, a dynamic schedule that adapts to local geometry using a stationarity detection method, resulting in robust performance and better generalization compared to methods like Adam in experiments.
This paper proposes SplitSGD, a new dynamic learning rate schedule for stochastic optimization. This method decreases the learning rate for better adaptation to the local geometry of the objective function whenever a stationary phase is detected, that is, the iterates are likely to bounce at around a vicinity of a local minimum. The detection is performed by splitting the single thread into two and using the inner product of the gradients from the two threads as a measure of stationarity. Owing to this simple yet provably valid stationarity detection, SplitSGD is easy-to-implement and essentially does not incur additional computational cost than standard SGD. Through a series of extensive experiments, we show that this method is appropriate for both convex problems and training (non-convex) neural networks, with performance compared favorably to other stochastic optimization methods. Importantly, this method is observed to be very robust with a set of default parameters for a wide range of problems and, moreover, can yield better generalization performance than other adaptive gradient methods such as Adam.