[Paper Review] Better SGD using Second-order Momentum

Hoang Vu Tran, Ashok Cutkosky|arXiv (Cornell University)|Mar 4, 2021

Stochastic Gradient Optimization Techniques25 references4 citations

TL;DR

This paper proposes SGDHess, a novel stochastic optimization algorithm that enhances SGD by incorporating Hessian-vector products to correct momentum bias, achieving an optimal $O(\epsilon^{-3})$ convergence rate for finding $\epsilon$-critical points in non-convex settings. Unlike prior second-order methods, it avoids large batch sizes and enables adaptive learning rates, improving practicality and performance on deep learning benchmarks.

ABSTRACT

We develop a new algorithm for non-convex stochastic optimization that finds an $ε$-critical point in the optimal $O(ε^{-3})$ stochastic gradient and Hessian-vector product computations. Our algorithm uses Hessian-vector products to "correct" a bias term in the momentum of SGD with momentum. This leads to better gradient estimates in a manner analogous to variance reduction methods. In contrast to prior work, we do not require excessively large batch sizes, and are able to provide an adaptive algorithm whose convergence rate automatically improves with decreasing variance in the gradient estimates. We validate our results on a variety of large-scale deep learning architectures and benchmarks tasks.

Motivation & Objective

• To develop a practical second-order optimization method that improves convergence beyond standard SGD in non-convex deep learning settings.
• To reduce the number of stochastic gradient and Hessian-vector oracle calls required to find $\epsilon$-critical points.
• To avoid reliance on large batch sizes, which are common in prior second-order methods.
• To design an adaptive variant that automatically improves convergence when gradient noise is low.
• To validate the method empirically across diverse deep learning architectures and tasks.

Proposed method

• The algorithm introduces a Hessian-vector product to correct a bias term in the momentum of SGD, improving gradient estimates.
• It uses a Lyapunov function $\Phi_t$ to track progress and derive convergence bounds, incorporating both gradient and error terms.
• The method employs a decreasing learning rate schedule $\eta_t$ to balance exploration and convergence.
• A normalized variant removes the need for a Lipschitz gradient assumption, enhancing robustness.
• An adaptive version dynamically adjusts learning rates based on observed gradient variance, achieving $O(\epsilon^{-2})$ rate under low noise.
• Theoretical analysis relies on bounding the expected decrease in the Lyapunov function using properties of Hessian-vector products and gradient noise.

Experimental results

Research questions

• RQ1 Can Hessian-vector products be used to improve the convergence of SGD in non-convex optimization without large batch sizes?
• RQ2 Does incorporating second-order information via Hessian-vector products lead to a faster convergence rate than standard SGD?
• RQ3 Can the algorithm adaptively improve convergence when gradient noise is reduced?
• RQ4 Is the proposed method practically effective across diverse deep learning architectures and tasks?
• RQ5 Can the algorithm achieve the theoretical lower bound of $O(\epsilon^{-3})$ oracle calls for finding $\epsilon$-critical points?

Key findings

• The proposed SGDHess algorithm achieves the optimal $O(\epsilon^{-3})$ convergence rate for finding $\epsilon$-critical points in non-convex stochastic optimization.
• The method does not require large batch sizes, unlike many prior second-order methods, making it more practical for large-scale learning.
• An adaptive variant of the algorithm achieves a faster $O(\epsilon^{-2})$ convergence rate when gradient noise is negligible.
• The normalized version of the algorithm removes the need for a Lipschitz gradient assumption, broadening its applicability.
• Empirical results show consistent performance gains across multiple deep learning benchmarks and architectures, often matching or exceeding state-of-the-art methods.
• The tuning process is simple, with many hyperparameters transferable from standard SGD.