[Paper Review] Optimal Mini-Batch Size Selection for Fast Gradient Descent

Michael Perrone, Haidar Khan+4 more|arXiv (Cornell University)|Nov 15, 2019

Stochastic Gradient Optimization Techniques30 references4 citations

TL;DR

This paper proposes a closed-form model for optimal mini-batch size selection in stochastic gradient descent (SGD) by identifying a robust empirical inverse relationship between mini-batch size and the number of updates required for convergence. The model decouples algorithmic behavior from hardware, enabling principled optimization of training time through the noise sensitivity parameter α and convergence floor N∞, with validation across image recognition and machine translation tasks.

ABSTRACT

This paper presents a methodology for selecting the mini-batch size that minimizes Stochastic Gradient Descent (SGD) learning time for single and multiple learner problems. By decoupling algorithmic analysis issues from hardware and software implementation details, we reveal a robust empirical inverse law between mini-batch size and the average number of SGD updates required to converge to a specified error threshold. Combining this empirical inverse law with measured system performance, we create an accurate, closed-form model of average training time and show how this model can be used to identify quantifiable implications for both algorithmic and hardware aspects of machine learning. We demonstrate the inverse law empirically, on both image recognition (MNIST, CIFAR10 and CIFAR100) and machine translation (Europarl) tasks, and provide a theoretic justification via proving a novel bound on mini-batch SGD training.

Motivation & Objective

• To develop a principled methodology for selecting mini-batch size that minimizes training time in SGD, independent of hardware and software implementation.
• To decouple algorithmic convergence behavior from system-level performance factors, enabling clearer analysis of learning dynamics.
• To identify and quantify the impact of mini-batch size on convergence iterations and training time across diverse machine learning tasks.
• To challenge the assumption that weak scaling (fixed mini-batch per worker) minimizes training time, showing it can be suboptimal.
• To introduce the concept of 'noise sensitivity' (α) as a key algorithmic property influencing optimal batch size selection.

Proposed method

• Proposes an empirical inverse relationship: $ N_{\text{Update}} = N_\infty + \frac{\alpha}{M} $, where $ N_{\text{Update}} $ is the number of SGD updates to convergence, $ M $ is the mini-batch size, and $ \alpha $ is the noise sensitivity.
• Decomposes total training time as $ T_C = N_{\text{Update}} \cdot T_{\text{Update}} $, separating algorithmic complexity from per-update computation time.
• Uses measured system performance to model $ T_{\text{Update}} $ as a function of mini-batch size $ M $ and number of parallel learners $ P $, enabling closed-form training time prediction.
• Derives a theoretical bound on SGD convergence that supports the inverse $ M $-dependence, reinforcing the empirical law via the Central Limit Theorem.
• Validates the model across multiple benchmarks: MNIST, CIFAR10, CIFAR100 (image recognition), and Europarl (machine translation).
• Introduces $ N_\infty $ as the theoretical lower bound on updates, representing the minimum number of steps needed even with full-batch gradients.

Experimental results

Research questions

• RQ1 What is the fundamental relationship between mini-batch size and the number of SGD updates required to converge?
• RQ2 How can algorithmic convergence behavior be decoupled from hardware and software implementation to enable principled optimization?
• RQ3 Does weak scaling (fixed mini-batch per worker) always minimize training time, or can it be suboptimal?
• RQ4 What role does 'noise sensitivity' (α) play in determining optimal mini-batch size and training efficiency?
• RQ5 Can the inverse relationship between batch size and convergence iterations be theoretically justified and generalized beyond SGD?

Key findings

• An empirical inverse relationship $ N_{\text{Update}} = N_\infty + \frac{\alpha}{M} $ holds robustly across image recognition (MNIST, CIFAR10, CIFAR100) and machine translation (Europarl) tasks.
• The noise sensitivity parameter $ \alpha $ increases with model and data complexity, pushing the convergence floor $ N_\infty $ to higher batch sizes for more complex problems.
• The theoretical bound derived for SGD convergence exhibits the same inverse $ M $-dependence as the empirical model, validating the relationship.
• The model shows that weak scaling can be suboptimal because it ignores the dependence of convergence time on mini-batch size, leading to longer training times.
• The parameter $ N_\infty $ represents a fundamental lower bound on updates, limiting the benefits of parallelization regardless of hardware scale.
• The framework enables closed-form prediction of training time and provides a principled basis for optimizing both algorithmic design and hardware system architecture.