[Paper Review] Online Batch Selection for Faster Training of Neural Networks

Ilya Loshchilov, Frank Hutter|arXiv (Cornell University)|Nov 19, 2015

Stochastic Gradient Optimization Techniques32 references171 citations

TL;DR

This paper proposes an online batch selection strategy that prioritizes training samples with higher recent losses by ranking them and sampling with exponentially decaying probability. Experiments on MNIST show this method accelerates convergence by approximately 5x for both Adam and AdaDelta optimizers, significantly improving training speed without modifying the underlying optimization algorithms.

ABSTRACT

Deep neural networks are commonly trained using stochastic non-convex optimization procedures, which are driven by gradient information estimated on fractions (batches) of the dataset. While it is commonly accepted that batch size is an important parameter for offline tuning, the benefits of online selection of batches remain poorly understood. We investigate online batch selection strategies for two state-of-the-art methods of stochastic gradient-based optimization, AdaDelta and Adam. As the loss function to be minimized for the whole dataset is an aggregation of loss functions of individual datapoints, intuitively, datapoints with the greatest loss should be considered (selected in a batch) more frequently. However, the limitations of this intuition and the proper control of the selection pressure over time are open questions. We propose a simple strategy where all datapoints are ranked w.r.t. their latest known loss value and the probability to be selected decays exponentially as a function of rank. Our experimental results on the MNIST dataset suggest that selecting batches speeds up both AdaDelta and Adam by a factor of about 5.

Motivation & Objective

• To investigate whether online selection of training batches can accelerate deep neural network training beyond fixed-batch-size stochastic optimization.
• To address the underexplored potential of adaptive batch selection in improving convergence speed of state-of-the-art optimizers like Adam and AdaDelta.
• To develop and evaluate a simple, effective strategy for dynamically selecting batches based on real-time loss values of individual datapoints.
• To demonstrate that focusing on high-loss samples reduces training time without compromising model performance.

Proposed method

• Each training datapoint is ranked by its most recently computed loss value, with higher loss values receiving higher ranks.
• The probability of selecting a datapoint in a batch is set to decay exponentially with its rank, favoring samples with the largest losses.
• The selection process is updated periodically, with a configurable frequency (r_freq), to balance computational cost and adaptivity.
• The method is applied to both Adam and AdaDelta optimizers, maintaining their original update rules but modifying the batch sampling procedure.
• The approach is invariant to rank-preserving transformations of the loss function, ensuring robustness to scaling.
• The method avoids expensive re-sorting by using a bisection-based approach when needed, keeping overhead low.

Experimental results

Research questions

• RQ1 Can online batch selection based on real-time loss values accelerate training of deep neural networks compared to uniform random sampling?
• RQ2 How does the performance of Adam and AdaDelta change when batches are selected non-uniformly based on loss magnitude?
• RQ3 What is the optimal decay schedule for selection probability as a function of loss rank to maximize convergence speed?
• RQ4 How does the computational overhead of dynamic batch selection compare to the training speedup gained?
• RQ5 Does the proposed method generalize beyond MNIST to other datasets like CIFAR-10?

Key findings

• Online batch selection based on loss ranking reduces training time by approximately 5x for both Adam and AdaDelta on the MNIST dataset.
• The method achieves significant speedups without altering the underlying optimization algorithms or requiring additional hyperparameter tuning.
• The performance gain is consistent across both optimizers, indicating broad applicability to modern adaptive stochastic gradient methods.
• The computational overhead introduced by sorting and ranking is negligible compared to the overall training cost, especially for large models.
• The method remains effective even when the loss ranking is only partially stable over time, suggesting robustness to noise in loss estimates.
• Preliminary results on CIFAR-10 show benefits over random selection but not yet over shuffling, indicating potential dataset-dependent behavior.