Low Rank Saddle Free Newton: Scalable Stochastic Nonconvex Optimization

Newton methods have fallen out of favor for modern optimization problems (e.g. deep learning) because of concerns about per-iteration computational complexity. In this setting highly subsampled first order methods are preferred. In this work we motivate the extension of Newton methods to the highly stochastic regime, and argue for the use of the scalable low rank saddle free Newton (LRSFN) method. In this setting, iterative updates are dominated by stochastic noise, and stability of the method is key. In stability analysis, we demonstrate that stochastic errors for Newton methods can be greatly amplified by ill-conditioned matrix operators. The LRSFN algorithm mitigates this issue by the use of Levenberg-Marquardt damping, but generally second order methods with stochastic Hessian and gradient information may need to take small steps, unlike in deterministic problems. Numerical results show that even under restrictive step-length conditions, LRSFN can outperform popular first order methods on nontrivial deep learning tasks in terms of generalizability for equivalent computational work.