[Paper Review] Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks

Boris Ginsburg, Patrice Castonguay+8 more|arXiv (Cornell University)|May 27, 2019

Stochastic Gradient Optimization Techniques39 references88 citations

TL;DR

NovoGrad is an adaptive SGD method using layer-wise gradient normalization and decoupled weight decay, performing as well or better than SGD with momentum and Adam/AdamW across vision, speech, translation, and language modeling with robustness to LR and initialization and half the memory footprint of Adam.

ABSTRACT

We propose NovoGrad, an adaptive stochastic gradient descent method with layer-wise gradient normalization and decoupled weight decay. In our experiments on neural networks for image classification, speech recognition, machine translation, and language modeling, it performs on par or better than well tuned SGD with momentum and Adam or AdamW. Additionally, NovoGrad (1) is robust to the choice of learning rate and weight initialization, (2) works well in a large batch setting, and (3) has two times smaller memory footprint than Adam.

Motivation & Objective

• Develop a single optimizer that performs well across image classification, speech recognition, machine translation, and language modeling.
• Improve robustness to learning rate and weight initialization.
• Reduce memory footprint compared to Adam while maintaining performance.
• Enhance regularization and support large-batch training scenarios.

Proposed method

• Introduce layer-wise second moments and normalize gradients by per-layer norms.
• Compute the first moment using gradients divided by the layer-wise second moment and add decoupled weight decay to the normalized gradient before updating the first moment.
• Update weights using the first moment in a SGD-with-momentum style with a decoupled weight decay term.
• Discuss stability improvements and provide connections to AMS-Grad for convex-like settings.

Experimental results

Research questions

• RQ1 Can NovoGrad achieve parity or superiority over SGD with momentum and Adam/AdamW across diverse deep learning tasks?
• RQ2 Is layer-wise gradient normalization more robust to learning rate and initialization than element-wise normalization?
• RQ3 Does decoupled weight decay improve regularization and generalization in NovoGrad?
• RQ4 How does NovoGrad perform in large-batch training regimes compared to existing optimizers?
• RQ5 What memory and computational advantages does NovoGrad offer relative to Adam?

Key findings

• NovoGrad matches or surpasses SGD with momentum and Adam/AdamW across ImageNet, WMT’14 English–German translation, LibriSpeech, and WikiText-103 language modeling.
• NovoGrad is robust to learning rate and weight initialization and performs well without LR warm-up in many setups.
• In large-batch image and speech tasks, NovoGrad achieves competitive or superior accuracy or WER without additional tricks.
• NovoGrad has about half the memory footprint of Adam, aiding large-scale model training.