Heavy-Ball Momentum Method in Continuous Time and Discretization Error Analysis

TL;DR

This paper develops a continuous time differential equation model for the Heavy-Ball momentum method, explicitly accounting for discretization error, and explores its implications for deep learning and implicit regularization.

Contribution

It introduces a first-order piece-wise continuous differential equation for HB that explicitly controls discretization error, bridging the gap between discrete and continuous analyses.

Findings

Provides a new continuous time model for HB with explicit error control

Reveals implicit regularization effects in deep learning models

Supports theoretical results with numerical experiments

Abstract

This paper establishes a continuous time approximation, a piece-wise continuous differential equation, for the discrete Heavy-Ball (HB) momentum method with explicit discretization error. Investigating continuous differential equations has been a promising approach for studying the discrete optimization methods. Despite the crucial role of momentum in gradient-based optimization methods, the gap between the original discrete dynamics and the continuous time approximations due to the discretization error has not been comprehensively bridged yet. In this work, we study the HB momentum method in continuous time while putting more focus on the discretization error to provide additional theoretical tools to this area. In particular, we design a first-order piece-wise continuous differential equation, where we add a number of counter terms to account for the discretization error explicitly. As a result, we provide a continuous time model for the HB momentum method that allows the control of discretization error to arbitrary order of the step size. As an application, we leverage it to find a new implicit regularization of the directional smoothness and investigate the implicit bias of HB for diagonal linear networks, indicating how our results can be used in deep learning. Our theoretical findings are further supported by numerical experiments.