Low Rank Saddle Free Newton: A Scalable Method for Stochastic Nonconvex Optimization
This addresses scalability and stability issues in stochastic optimization for deep learning practitioners, though it is incremental as it builds on existing Newton and saddle-free approaches.
The paper tackles the challenge of using second-order Newton methods in stochastic nonconvex optimization for deep learning by proposing the Low Rank Saddle Free Newton (LRSFN) method, which avoids full Hessian computation and improves stability, showing it can outperform first-order methods in generalizability on large-scale tasks.
In modern deep learning, highly subsampled stochastic approximation (SA) methods are preferred to sample average approximation (SAA) methods because of large data sets as well as generalization properties. Additionally, due to perceived costs of forming and factorizing Hessians, second order methods are not used for these problems. In this work we motivate the extension of Newton methods to the SA regime, and argue for the use of the scalable low rank saddle free Newton (LRSFN) method, which avoids forming the Hessian in favor of making a low rank approximation. Additionally, LRSFN can facilitate fast escape from indefinite regions leading to better optimization solutions. In the SA setting, iterative updates are dominated by stochastic noise, and stability of the method is key. We introduce a continuous time stability analysis framework, and use it to demonstrate that stochastic errors for Newton methods can be greatly amplified by ill-conditioned Hessians. The LRSFN method mitigates this stability issue via Levenberg-Marquardt damping. However, generally the analysis shows that second order methods with stochastic Hessian and gradient information may need to take small steps, unlike in deterministic problems. Numerical results show that LRSFN can escape indefinite regions that other methods have issues with; and even under restrictive step length conditions, LRSFN can outperform popular first order methods on large scale deep learning tasks in terms of generalizability for equivalent computational work.