GRADSTOP: Early Stopping of Gradient Descent via Posterior Sampling

TL;DR

GRADSTOP introduces a gradient-based early stopping method that estimates the Bayesian posterior from gradients, avoiding the need for a hold-out validation set and improving performance in data-limited scenarios.

Contribution

It proposes a novel stochastic early stopping technique using gradient information to estimate the Bayesian posterior, eliminating the need for a validation set.

Findings

Achieves small test loss comparable to validation-based methods

Performs well in data-limited settings like transfer learning

Can be integrated into gradient descent libraries with minimal overhead

Abstract

Machine learning models are often learned by minimising a loss function on the training data using a gradient descent algorithm. These models often suffer from overfitting, leading to a decline in predictive performance on unseen data. A standard solution is early stopping using a hold-out validation set, which halts the minimisation when the validation loss stops decreasing. However, this hold-out set reduces the data available for training. This paper presents GRADSTOP, a novel stochastic early stopping method that only uses information in the gradients, which are produced by the gradient descent algorithm ``for free.'' Our main contributions are that we estimate the Bayesian posterior by the gradient information, define the early stopping problem as drawing sample from this posterior, and use the approximated posterior to obtain a stopping criterion. Our empirical evaluation shows that GRADSTOP achieves a small loss on test data and compares favourably to a validation-set-based stopping criterion. By leveraging the entire dataset for training, our method is particularly advantageous in data-limited settings, such as transfer learning. It can be incorporated as an optional feature in gradient descent libraries with only a small computational overhead. The source code is available at https://github.com/edahelsinki/gradstop.