[Paper Review] Neumann Optimizer: A Practical Optimization Algorithm for Deep Neural Networks

Shankar Krishnan, Ying Xiao+1 more|arXiv (Cornell University)|Dec 8, 2017

Advanced Neural Network Applications36 references3 citations

TL;DR

The Neumann Optimizer is a practical, second-order stochastic optimization algorithm for deep neural networks that implicitly computes inverse Hessian directions using a Neumann series expansion, avoiding explicit Hessian computation or Hessian-vector products. It achieves state-of-the-art performance by enabling linear scaling up to 32,000 mini-batch sizes with no drop in validation accuracy and improving generalization by 0.8–0.9% at smaller batches, all while requiring only learning rate tuning and computational cost comparable to Adam.

ABSTRACT

Progress in deep learning is slowed by the days or weeks it takes to train large models. The natural solution of using more hardware is limited by diminishing returns, and leads to inefficient use of additional resources. In this paper, we present a large batch, stochastic optimization algorithm that is both faster than widely used algorithms for fixed amounts of computation, and also scales up substantially better as more computational resources become available. Our algorithm implicitly computes the inverse Hessian of each mini-batch to produce descent directions; we do so without either an explicit approximation to the Hessian or Hessian-vector products. We demonstrate the effectiveness of our algorithm by successfully training large ImageNet models (Inception-V3, Resnet-50, Resnet-101 and Inception-Resnet-V2) with mini-batch sizes of up to 32000 with no loss in validation error relative to current baselines, and no increase in the total number of steps. At smaller mini-batch sizes, our optimizer improves the validation error in these models by 0.8-0.9%. Alternatively, we can trade off this accuracy to reduce the number of training steps needed by roughly 10-30%. Our work is practical and easily usable by others -- only one hyperparameter (learning rate) needs tuning, and furthermore, the algorithm is as computationally cheap as the commonly used Adam optimizer.

Motivation & Objective

• To address the growing wall-time bottleneck in training large deep neural networks by developing a scalable, practical optimization algorithm.
• To enable efficient large-batch training without degradation in model generalization, overcoming the typical trade-off between speed and accuracy.
• To incorporate second-order information for improved descent directions without explicit Hessian approximation or costly Hessian-vector products.
• To design an optimizer that is as easy to use as Adam, with only one tunable hyperparameter (learning rate), while outperforming existing methods in both speed and accuracy.

Proposed method

• The algorithm computes descent directions by solving an intermediate optimization problem that implicitly inverts the Hessian of each mini-batch using a Neumann series expansion.
• It avoids explicit Hessian representation by replacing Hessian applications with single gradient evaluations, enabling efficient computation.
• The method leverages the Neumann series to approximate the inverse Hessian without computing Hessian-vector products or storing the Hessian matrix.
• The algorithm is implemented as a differentiable optimization step that integrates seamlessly into standard deep learning training pipelines.
• It uses a single learning rate hyperparameter and maintains computational efficiency comparable to Adam, making it practical for large-scale models.
• The approach is validated through large-scale experiments on ImageNet using ResNet-50, ResNet-101, Inception-V3, and Inception-ResNet-V2 architectures.

Experimental results

Research questions

• RQ1 Can a second-order optimization method be made practical for large-scale deep learning without explicit Hessian computation or high computational cost?
• RQ2 Does implicit Hessian inversion via the Neumann series enable linear scaling of training performance with increasing batch size while preserving model generalization?
• RQ3 Can the Neumann Optimizer improve generalization performance at smaller mini-batch sizes compared to standard optimizers like Adam or RMSProp?
• RQ4 How does the optimizer perform on non-vision models, such as sequence-to-sequence RNNs, where gradient clipping is common?

Key findings

• The Neumann Optimizer achieves linear scaling up to a mini-batch size of 32,000 on ResNet-50 with no degradation in top-1 validation error compared to baseline models.
• At smaller mini-batch sizes (e.g., 4,000), the optimizer improves top-1 validation error by 0.8–0.9% across all evaluated models, despite no improvement in training loss.
• The improvement in generalization can be traded for faster training: the number of steps can be reduced by 10–30% while maintaining baseline accuracy.
• The algorithm maintains computational cost comparable to Adam, requiring only learning rate tuning, and is easily deployable in practice.
• Regularization further improves performance, reducing top-1 error from 23.5% (no regularization) to 23.0% (with regularization) on ResNet-50 at batch size 4,000.
• The optimizer fails to improve performance on a sequence-to-sequence Tacotron model, likely due to breakdown in Hessian approximation under aggressive gradient clipping.