Losing momentum in continuous-time stochastic optimisation

TL;DR

This paper introduces a continuous-time stochastic optimization model with momentum, analyzes its convergence properties, and proposes a discretisation scheme that performs competitively in machine learning tasks.

Contribution

It develops a novel continuous-time model for stochastic gradient descent with momentum and provides convergence analysis and a stable discretisation scheme.

Findings

The model converges to the global minimiser under certain conditions.

The discretisation scheme is stable and effective in practice.

The algorithm performs competitively on neural network training tasks.

Abstract

The training of modern machine learning models often consists in solving high-dimensional non-convex optimisation problems that are subject to large-scale data. In this context, momentum-based stochastic optimisation algorithms have become particularly widespread. The stochasticity arises from data subsampling which reduces computational cost. Both, momentum and stochasticity help the algorithm to converge globally. In this work, we propose and analyse a continuous-time model for stochastic gradient descent with momentum. This model is a piecewise-deterministic Markov process that represents the optimiser by an underdamped dynamical system and the data subsampling through a stochastic switching. We investigate longtime limits, the subsampling-to-no-subsampling limit, and the momentum-to-no-momentum limit. We are particularly interested in the case of reducing the momentum over time. Under convexity assumptions, we show convergence of our dynamical system to the global minimiser when reducing momentum over time and letting the subsampling rate go to infinity. We then propose a stable, symplectic discretisation scheme to construct an algorithm from our continuous-time dynamical system. In experiments, we study our scheme in convex and non-convex test problems. Additionally, we train a convolutional neural network in an image classification problem. Our algorithm {attains} competitive results compared to stochastic gradient descent with momentum.