ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning

Zhewei Yao; Amir Gholami; Sheng Shen; Mustafa Mustafa; Kurt Keutzer; Michael W. Mahoney

Source-linked AI summary

ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning

Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, Michael W. Mahoney

arXiv:2006.00719v3Jun 1, 2020cs.LGmath.NAstat.ML

TL;DR

Traditional second-order optimization can be costly, while first-order neural-network training relies on task-specific optimizer and hyperparameter choices. ADAHESSIAN adaptively preconditions gradients using an approximate Hessian, achieving strong results across computer vision, NLP, and recommendation tasks while remaining robust to hyperparameters.

Problem
Neural-network optimization is difficult because high-dimensional, non-convex problems make classical methods inefficient, while first-order methods require ad-hoc optimizer and hyperparameter choices.
Method
ADAHESSIAN approximately computes an exponential moving average of the Hessian and uses it to adaptively precondition gradients, with diagonal Hessian approximation, spatial averaging, and Hessian momentum.
Results
ADAHESSIAN significantly outperforms adaptive baselines across computer vision, NLP, and recommendation tasks, including 0.13/0.33 BLEU gains over AdamW on IWSLT14/WMT14.
Takeaways & Limitations
The results support ADAHESSIAN as a broadly effective adaptive optimizer across ResNets, transformers, and DLRM models, with robustness to several hyperparameters.
Takeaways & Limitations
Naive second-order implementations can require cubic computational complexity and quadratic memory, motivating matrix-free or approximate approaches.

Abstract

from arXiv · show

We introduce ADAHESSIAN, a second order stochastic optimization algorithm which dynamically incorporates the curvature of the loss function via ADAptive estimates of the HESSIAN. Second order algorithms are among the most powerful optimization algorithms with superior convergence properties as compared to first order methods such as SGD and Adam. The main disadvantage of traditional second order methods is their heavier per iteration computation and poor accuracy as compared to first order methods. To address these, we incorporate several novel approaches in ADAHESSIAN, including: (i) a fast Hutchinson based method to approximate the curvature matrix with low computational overhead; (ii) a root-mean-square exponential moving average to smooth out variations of the Hessian diagonal across different iterations; and (iii) a block diagonal averaging to reduce the variance of Hessian diagonal elements. We show that ADAHESSIAN achieves new state-of-the-art results by a large margin as compared to other adaptive optimization methods, including variants of Adam. In particular, we perform extensive tests on CV, NLP, and recommendation system tasks and find that ADAHESSIAN: (i) achieves 1.80%/1.45% higher accuracy on ResNets20/32 on Cifar10, and 5.55% higher accuracy on ImageNet as compared to Adam; (ii) outperforms AdamW for transformers by 0.13/0.33 BLEU score on IWSLT14/WMT14 and 2.7/1.0 PPL on PTB/Wikitext-103; (iii) outperforms AdamW for SqueezeBert by 0.41 points on GLUE; and (iv) achieves 0.032% better score than Adagrad for DLRM on the Criteo Ad Kaggle dataset. Importantly, we show that the cost per iteration of ADAHESSIAN is comparable to first order methods, and that it exhibits robustness towards its hyperparameters.

Introduction

Machine-learning optimization is difficult because first-order methods require task-specific choices and tuning while ignoring loss curvature. ADAHESSIAN addresses this with an efficient adaptive Hessian approximation and reports strong results across vision, language, and recommendation tasks.

Motivation: First-order methods require ad-hoc optimizer and hyperparameter choices that can substantially affect convergence and performance across CV, NLP, and recommendation tasks.Typical choices are SGD with momentum for CV, Adam for transformers, and Adagrad for recommendation systems.
Motivation: Industrial-scale tuning can produce near-ideal SGD behavior for some tasks but is computationally and financially impractical to repeat routinely.The paper identifies this as a challenge for many large-scale learning problems.
Motivation: First-order methods use gradient information without loss curvature, whereas second-order methods incorporate gradient and Hessian information through gradient preconditioning.Curvature differences across parameter dimensions motivate using different effective step sizes.
Method: AdaHessian approximates the Hessian diagonally with Hutchinson’s method, applies root-mean-square exponential moving averages, and uses block diagonal averaging to reduce variance with O(d) memory complexity.These components target the computational and variability costs of adaptive second-order optimization.
Results: AdaHessian improves reported performance over Adam, AdamW, and Adagrad across computer vision, NLP, and recommendation-system benchmarks.Examples include 93.08% versus 91.63% for ResNet32 on Cifar10, 70.08% versus 64.53% for ResNet18 on ImageNet, 0.13/0.33 BLEU gains on IWSLT14/WMT14, 2.7/1.0 PPL gains on PTB/WikiText-103, 0.41 GLUE points, and 0.032% on Criteo.
Results: AdaHessian shows robustness to learning rate, spatial averaging size, and delayed Hessian computation, while using the same major training schedules and moment coefficients as baseline optimizers.The authors note that more extensive hyperparameter optimization could yield additional gains.

Problem Formulation and Related work

The section formulates stochastic non-convex learning optimization and reviews why first- and second-order methods face practical trade-offs. It motivates AdaHessian as a curvature-aware approach intended to reduce stochastic and computational difficulties.

Problem Formulation and Related work: Supervised learning is framed as solving a non-convex stochastic optimization problem over model parameters, losses, paired data, and mini-batches.The formulation defines θ as model parameters, l_i as the loss, (x_i, y_i) as paired data and labels, and N as the dataset size.
Problem Formulation and Related work: First-order methods dominate deep learning but require task-specific optimizer choices and carefully tuned learning rates, schedules, momentum, and warmup settings.Using an unsuitable SGD variant can significantly degrade performance, and training may require continuous monitoring without guaranteeing convergence within a fixed iteration budget.
Adaptive First Order Methods: Popular first-order optimizers differ in their moment estimates, with AdamW decoupling weight decay and becoming preferred for transformer-based NLP training.Adagrad can degrade on dense settings because accumulated gradients cause rapid learning-rate decay, while exponentially weighted moments limit the influence of early gradients.
Problem Formulation and Related work: First-order methods use gradient information without explicitly modeling loss curvature, whereas second-order methods incorporate curvature through gradient preconditioning.Second-order preconditioning can rotate and rescale gradients to account for different curvature across parameter directions and layers.
Second Order Methods: Second-order approaches can be computationally expensive: naïve Newton methods require solving systems involving the Hessian, with cubic computation and quadratic memory costs.Quasi-Newton methods such as LBFGS reduce complexity but may perform poorly for machine learning because stochastic gradients can badly distort their Hessian approximations.
Second Order Methods: Stochastic noise can make local Hessian estimates erroneous, producing suboptimal descent directions in complex neural-network architectures.The section identifies noisy curvature as a central reason second-order methods have struggled in machine learning compared with simpler models or other domains.
Second Order Methods: AdaHessian approximately tracks an exponential moving average of the Hessian and uses it to adaptively precondition gradients.The method addresses the difficulty of averaging an explicitly unavailable Hessian by approximating it and using a diagonal operator.
Second Order Methods: Hutchinson’s method and spatial averaging reduce the computational overhead and stochastic noise of Hessian-diagonal estimation.The Hessian diagonal can be computed through Hessian-vector products, while spatially averaged diagonals preserve the convergence rate for the stated strongly convex and strictly smooth setting.

Methodological Approach

AdaHessian approximates and smooths the Hessian diagonal to precondition stochastic gradient updates while avoiding the cost of applying the full inverse Hessian. Its design combines Hutchinson estimation, spatial averaging, and temporal Hessian momentum to reduce noisy curvature effects.

Hessian momentum: In a noisy 1D curvature example, moving-average Hessian information enabled convergence in 7 iterations, whereas omitting it failed to converge after 1000 iterations.The example illustrates why local curvature can be misleading and why temporal smoothing is central to the method.
Hessian diagonal approximation: AdaHessian uses the Hessian diagonal as a computationally efficient approximation to the full Hessian preconditioner.For simple strongly convex and strictly smooth functions, the diagonal update has the same convergence rate as the full-Hessian formulation.
Hessian diagonal approximation: Hutchinson’s method estimates the Hessian diagonal from Hessian-vector products without explicitly forming the Hessian.The Hessian-vector product is computed by backpropagating g^Tz at the cost of ordinary gradient backpropagation, and E[z ⊙(Hz)] equals the Hessian diagonal.
Hessian momentum: Hessian momentum uses an exponential moving average of the spatially averaged diagonal to replace noisy local curvature with smoother estimates.Because the diagonal is a vector rather than a quadratically large matrix, momentum can be applied efficiently across iterations.
Spatial averaging: AdaHessian applies spatial averaging to the Hessian diagonal to smooth parameter-wise curvature variations.For convolutional layers, it averages the diagonal over each 3×3 convolution kernel; the block size b controls the averaging groups.
AdaHessian: The training procedure estimates the diagonal Hessian, spatially averages it, updates its momentum estimate, and then uses the resulting quantities in the parameter update.The algorithm exposes learning rate, decay rates, block size, and Hessian power as configurable parameters.

Results

AdaHessian is evaluated across computer vision, NLP, and recommendation tasks, where it generally outperforms adaptive baselines including Adam, AdamW, and Adagrad.

Experimental scope: AdaHessian is extensively tested on image classification, neural machine translation, language modeling, natural language understanding, and recommendation tasks.The comparisons include SGD, Adam, AdamW, and Adagrad under shared optimizer settings where specified.
Computer Vision: 5.53% higher accuracy than Adam and 2.67% higher accuracy than AdamW are achieved on ResNet18 ImageNet classification.AdaHessian reaches similar performance to SGD on ImageNet.
Computer Vision: 1.80%/1.45% higher accuracy than Adam is achieved by AdaHessian for ResNet20/32 on Cifar10.AdaHessian is also comparable to SGD, with 0.05% higher accuracy for ResNet20 and 0.06% lower accuracy for ResNet32.
Natural Language Processing: 0.13/0.33 BLEU higher than AdamW is achieved on IWSLT14/WMT14 transformer translation.The comparison is against AdamW, described as the default and heavily tuned optimizer for these tasks.
Recommendation System: 0.032% higher testing accuracy than Adagrad is achieved for DLRM on the Criteo Ad Kaggle dataset.The reported AdaHessian testing accuracy is 79.167%.

Discussion

AdaHessian is reported to achieve state-of-the-art performance across a wide range of tasks and to be robust to learning-rate and block-size choices.

Robustness: AdaHessian achieves state-of-the-art performance on a wide range of learning tasks.The discussion identifies learning rate and block size as two hyperparameters whose sensitivity is examined.
Learning rate: 32.48 BLEU is achieved by AdaHessian with 10× learning-rate scaling on IWSLT14, while AdamW diverges at 6× scaling.This comparison supports greater robustness to large learning-rate changes for AdaHessian.
Block size: Small performance variation is observed when the Hessian spatial-averaging block size varies from 1 to 128.The best performance occurs at block size 32, and all tested settings are no worse than AdamW.

AdaHessian Overhead

AdaHessian adds Hessian-diagonal computation while remaining substantially less costly than a full second-order method, and its overhead can be reduced by lowering Hutchinson frequency.

Theoretical cost: 2× the flops of SGD are required by AdaHessian because one Hutchinson Hessian-matrix-vector product adds another gradient backpropagation per iteration.The stated theoretical overhead comes from computing the Hessian diagonal.
Measured runtime: 2.42× slower than SGD and 2.27× slower than Adam is AdaHessian for ResNet20 on a single RTX Titan GPU.The measured runtime is described as not orders of magnitude slower than first-order methods.
Overhead reduction: 1.5× and 1.2× theoretical cost are obtained by performing Hutchinson calculations every 2 and 5 iterations, respectively.Reducing calculation frequency lowers the cost from the standard 2× theoretical overhead.
Overhead reduction: Small performance variation accompanies reduced Hutchinson frequencies from 1 to 5 on Cifar10 ResNet20/32.The reduced-frequency approach significantly decreases overhead relative to SGD and Adam.

Conclusions

AdaHessian uses spatially averaged and temporally smoothed approximate Hessian diagonals to precondition gradients, improving generalization across diverse tasks while incurring a stated speed limitation.

Conclusions: AdaHessian combines an approximate Hessian diagonal with spatial averaging and momentum to rescale gradients toward better descent directions.Spatial averaging and exponential moving averages smooth noisy local Hessian information.
Conclusions: AdaHessian achieves comparable or higher generalization performance than highly tuned default optimizers across machine translation, language modeling, language understanding, vision, and recommendation tasks.The evaluated tasks include IWSLT14, WMT14, PTB, Wikitext-103, GLUE, Cifar10, ImageNet, and Criteo Ad Kaggle.
Conclusions: AdaHessian is reported as 2−3× slower than first-order methods such as SGD and Adam.The authors state that reducing this overhead requires more work.
Conclusions: Increasing Hessian computation frequency improves measured ResNet32 runtime relative to SGD from 3.23× to 1.45×.The speed comparison was measured on one RTX Titan GPU.
Conclusions: 2.7 better PPL is reported for AdaHessian than AdamW on a language-modeling task.The passage describes this improvement as significant for the task.
Conclusions: The paper positions AdaHessian as a step toward further progress in second-order methods for machine learning.The authors relate this direction to model compression, adversarial attacks, and loss-landscape studies.

Descending Property of Eq. 6

Under strong convexity and strict smoothness, the analyzed AdaHessian-related update formulations retain convergence-rate guarantees, with diagonal preconditioner elements bounded by the Hessian bounds.

Descending Property of Eq. 6: Under αI ≤∇2f(w) ≤βI for all w, the update formulation of Eq. 6 is shown to be convergent with a proper learning rate.The assumption applies to a strongly convex and strictly smooth function in R^d.
Descending Property of Eq. 6: For k = 0 or 1, the convergence rate is the same as gradient descent or Newton method, respectively.The passage identifies these cases as limiting comparisons for the analyzed update.
Descending Property of Eq. 6: Under the same strong-convexity and smoothness bounds, Eq. 7 has the same convergence rate as Eq. 6.The result is stated for αI ≤∇2f(w) ≤βI over all w.
Descending Property of Eq. 6: The diagonal elements in D are bounded within [α, β].The derivation uses positivity of the diagonal elements under strong convexity and basis-vector representations.
Descending Property of Eq. 6: Under αI ≤∇2f(w) ≤βI, Eq. 10 is stated to have the same convergence rate as Eq. 6.The proof extends the diagonal bounds to averages of subsets of diagonal elements.
Experimental context: The experimental setup includes ResNet20/32 Cifar10 evaluations, transformer translation and language-modeling experiments, and SqueezeBERT fine-tuning.The listed configurations include optimizer hyperparameters, datasets, and training settings for these evaluations.
Experimental context: For ResNet20/32 on Cifar10, SGD and AdaHessian consistently achieve better accuracy than Adam and AdamW in the reported curves.The final accuracy results are reported in Table 2.
Experimental context: AdaHessian’s Hessian diagonal is computed every iteration during a five-epoch warmup, then every one to five iterations.This delayed-update schedule is specified for ResNets on Cifar10.