[Paper Review] Stagewise Enlargement of Batch Size for SGD-based Learning

Shen-Yi Zhao, Yin-Peng Xie+1 more|arXiv (Cornell University)|Feb 26, 2020

Machine Learning and Algorithms27 references4 citations

TL;DR

This paper proposes Stagewise Enlargement of Batch Size (SEBS), a method that geometrically increases the batch size across stages in SGD-based learning. SEBS reduces parameter updates by up to 64% without sacrificing test accuracy, outperforming classical stagewise SGD and large-batch baselines on CIFAR-10 and ImageNet.

ABSTRACT

Existing research shows that the batch size can seriously affect the performance of stochastic gradient descent~(SGD) based learning, including training speed and generalization ability. A larger batch size typically results in less parameter updates. In distributed training, a larger batch size also results in less frequent communication. However, a larger batch size can make a generalization gap more easily. Hence, how to set a proper batch size for SGD has recently attracted much attention. Although some methods about setting batch size have been proposed, the batch size problem has still not been well solved. In this paper, we first provide theory to show that a proper batch size is related to the gap between initialization and optimum of the model parameter. Then based on this theory, we propose a novel method, called \underline{s}tagewise \underline{e}nlargement of \underline{b}atch \underline{s}ize~(\mbox{SEBS}), to set proper batch size for SGD. More specifically, \mbox{SEBS} adopts a multi-stage scheme, and enlarges the batch size geometrically by stage. We theoretically prove that, compared to classical stagewise SGD which decreases learning rate by stage, \mbox{SEBS} can reduce the number of parameter updates without increasing generalization error. SEBS is suitable for \mbox{SGD}, momentum \mbox{SGD} and AdaGrad. Empirical results on real data successfully verify the theories of \mbox{SEBS}. Furthermore, empirical results also show that SEBS can outperform other baselines.

Motivation & Objective

• Address the challenge of choosing an optimal batch size in SGD to balance training speed and generalization.
• Overcome the limitations of fixed-batch-size methods and dynamic batch-size strategies that lack theoretical guarantees.
• Develop a method that reduces parameter updates while preserving generalization performance in deep learning training.
• Provide theoretical justification for batch size scaling based on the initialization-to-optimum gap in model parameters.
• Design a universal framework compatible with SGD, momentum SGD, and AdaGrad for efficient training.

Proposed method

• Propose a multi-stage training scheme where the batch size is geometrically increased by a ratio ρ at each stage.
• Theoretically link batch size scaling to the gap between model initialization and optimal parameters, ensuring convergence without increased generalization error.
• Maintain a constant learning rate per stage, contrasting with classical stagewise SGD that reduces learning rates.
• Derive a condition η_s / b_s = O(ε_s) to ensure convergence and generalization stability during batch size enlargement.
• Apply the method to standard optimizers: SGD, momentum SGD (mSGD), and AdaGrad, demonstrating broad applicability.
• Implement SEBS with stage transitions at predefined epochs (e.g., 80, 120) to balance computational efficiency and model performance.

Experimental results

Research questions

• RQ1 Can a geometrically increasing batch size strategy reduce the number of parameter updates without degrading generalization performance?
• RQ2 How does the batch size scaling ratio affect convergence and test accuracy in deep neural network training?
• RQ3 Can a stagewise batch size enlargement strategy outperform classical stagewise SGD with decreasing learning rates?
• RQ4 Does SEBS maintain generalization performance when compared to large-batch training methods that require extensive hyperparameter tuning?
• RQ5 Is the proposed method universally applicable across different optimization algorithms like SGD, momentum SGD, and AdaGrad?

Key findings

• SEBS reduces parameter updates by 64% on ImageNet (ResNet50) compared to classical mSGD, while maintaining test accuracy of 75.87%.
• On CIFAR-10 with ResNet20, SEBS with ρ=12 reduces parameter updates from 62.5k (classical stagewise) to 32.6k, achieving comparable test accuracy.
• SEBS outperforms DB-SGD, which increases batch size every epoch and suffers from poor generalization, achieving lower test accuracy.
• mSEBS achieves 91.74% test accuracy on CIFAR-10 with 160k parameter updates, outperforming LARS (90.97%) despite using fewer updates.
• The method in Smith et al. (2018), which scales batch size rapidly, drops to 73.44% accuracy on ResNet50, while mSEBS maintains 75.87%.
• Empirical results confirm that SEBS maintains generalization performance even with large batch size scaling ratios (e.g., ρ=12), unlike prior dynamic batch-size methods.