Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stochastic Normalized Gradient Descent with Momentum for Large-Batch Training

Shen-Yi Zhao, Chang-Wei Shi, Yin-Peng Xie, Wu-Jun Li

TL;DR

The paper tackles the generalization gap observed with large-batch SGD by introducing stochastic normalized gradient descent with momentum (SNGM). SNGM updates a momentum term on the normalized gradient, enabling larger batch sizes without sacrificing convergence to an $\epsilon$-stationary point and achieving an ${\mathcal O}(1/\epsilon^4)$ complexity. The authors provide rigorous convergence analysis for both $L$-smooth and relaxed $(L,\lambda)$-smooth objectives and show that block-wise updates do not improve performance. Empirically, SNGM demonstrates superior test accuracy and efficiency over MSGD and several large-batch methods across image classification, NLP, and CTR tasks, confirming its practical value for large-scale training.

Abstract

Stochastic gradient descent~(SGD) and its variants have been the dominating optimization methods in machine learning. Compared to SGD with small-batch training, SGD with large-batch training can better utilize the computational power of current multi-core systems such as graphics processing units~(GPUs) and can reduce the number of communication rounds in distributed training settings. Thus, SGD with large-batch training has attracted considerable attention. However, existing empirical results showed that large-batch training typically leads to a drop in generalization accuracy. Hence, how to guarantee the generalization ability in large-batch training becomes a challenging task. In this paper, we propose a simple yet effective method, called stochastic normalized gradient descent with momentum~(SNGM), for large-batch training. We prove that with the same number of gradient computations, SNGM can adopt a larger batch size than momentum SGD~(MSGD), which is one of the most widely used variants of SGD, to converge to an $ε$-stationary point. Empirical results on deep learning verify that when adopting the same large batch size, SNGM can achieve better test accuracy than MSGD and other state-of-the-art large-batch training methods.

Stochastic Normalized Gradient Descent with Momentum for Large-Batch Training

TL;DR

The paper tackles the generalization gap observed with large-batch SGD by introducing stochastic normalized gradient descent with momentum (SNGM). SNGM updates a momentum term on the normalized gradient, enabling larger batch sizes without sacrificing convergence to an -stationary point and achieving an complexity. The authors provide rigorous convergence analysis for both -smooth and relaxed -smooth objectives and show that block-wise updates do not improve performance. Empirically, SNGM demonstrates superior test accuracy and efficiency over MSGD and several large-batch methods across image classification, NLP, and CTR tasks, confirming its practical value for large-scale training.

Abstract

Stochastic gradient descent~(SGD) and its variants have been the dominating optimization methods in machine learning. Compared to SGD with small-batch training, SGD with large-batch training can better utilize the computational power of current multi-core systems such as graphics processing units~(GPUs) and can reduce the number of communication rounds in distributed training settings. Thus, SGD with large-batch training has attracted considerable attention. However, existing empirical results showed that large-batch training typically leads to a drop in generalization accuracy. Hence, how to guarantee the generalization ability in large-batch training becomes a challenging task. In this paper, we propose a simple yet effective method, called stochastic normalized gradient descent with momentum~(SNGM), for large-batch training. We prove that with the same number of gradient computations, SNGM can adopt a larger batch size than momentum SGD~(MSGD), which is one of the most widely used variants of SGD, to converge to an -stationary point. Empirical results on deep learning verify that when adopting the same large batch size, SNGM can achieve better test accuracy than MSGD and other state-of-the-art large-batch training methods.

Paper Structure

This paper contains 21 sections, 8 theorems, 60 equations, 4 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Relationship between Smoothness Constant and Batch Size in MSGD
  4. Stochastic Normalized Gradient Descent with Momentum
  5. Comparison with MSGD
  6. Comparison with LARS
  7. Convergence Analysis
  8. Smooth Objective Function
  9. Relaxed Smooth Objective Function
  10. Block-wise Update Slows Down Convergence Rate
  11. Experiments
  12. Image Classification
  13. NLP
  14. Click-Through Rate Prediction
  15. Conclusion
  16. ...and 6 more sections

Key Result

Lemma 1

If $h(\cdot)$ is $(L,\lambda)$-smooth, then for any ${\bf u},{\bf w}, \alpha$ such that $\|{\bf u}-{\bf w}\|\leqslant \alpha$, we have

Figures (4)

  • Figure 1: Training loss and test accuracy of training a non-convex model (a network with two convolutional layers) on the CIFAR-10 dataset.
  • Figure 2: Training loss and test accuracy on CIFAR-10.
  • Figure 3: Validation perplexity on Wikitext-2.
  • Figure 4: Validation AUC on Criteo.

Theorems & Definitions (13)

  • Definition 1
  • Definition 2
  • Lemma 1
  • Lemma 2
  • Lemma 3
  • Theorem 1
  • Corollary 1
  • Remark 1
  • Theorem 2
  • Corollary 2
  • ...and 3 more