Small steps and giant leaps: Minimal Newton solvers for Deep Learning

TL;DR

The paper introduces CurveBall, a minimal second-order optimizer for deep learning that offers faster convergence than SGD without the high computational cost of traditional second-order methods.

Contribution

It presents a novel second-order method that efficiently estimates the inverse Hessian effect with minimal overhead, improving convergence in deep learning training.

Findings

Faster convergence on CIFAR and ImageNet models.

Comparable computational cost to standard forward passes.

Effective on large-scale deep learning models.

Abstract

We propose a fast second-order method that can be used as a drop-in replacement for current deep learning solvers. Compared to stochastic gradient descent (SGD), it only requires two additional forward-mode automatic differentiation operations per iteration, which has a computational cost comparable to two standard forward passes and is easy to implement. Our method addresses long-standing issues with current second-order solvers, which invert an approximate Hessian matrix every iteration exactly or by conjugate-gradient methods, a procedure that is both costly and sensitive to noise. Instead, we propose to keep a single estimate of the gradient projected by the inverse Hessian matrix, and update it once per iteration. This estimate has the same size and is similar to the momentum variable that is commonly used in SGD. No estimate of the Hessian is maintained. We first validate our method, called CurveBall, on small problems with known closed-form solutions (noisy Rosenbrock function and degenerate 2-layer linear networks), where current deep learning solvers seem to struggle. We then train several large models on CIFAR and ImageNet, including ResNet and VGG-f networks, where we demonstrate faster convergence with no hyperparameter tuning. Code is available.