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

TL;DR

This work develops HB Flow (HBF), a piece-wise continuous differential equation that serves as a high-precision continuous-time approximation to the discrete Heavy-Ball momentum method by explicitly canceling discretization error via counter terms. By solving a functional integral equation, the authors show how G_k and γ_k can be chosen to achieve discretization errors of arbitrary order in η, providing HB-like dynamics with controllable accuracy. They characterize the implicit bias and regularization effects of HB through HBF on diagonal linear networks, revealing an initialization-rescaling effect and directional-smoothness regularization distinct from GD and GF. Numerical experiments on simple 2D models and diagonal networks illustrate that higher-order HBF (α=3) more faithfully tracks HB than RGF or lower-order HBF. The framework lays a foundation for analyzing momentum methods in continuous time and offers insights into their implicit biases and potential benefits for deep learning contexts.

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.

Figures (6)

• Figure 1: Comparison of directional smoothness for HB and GD for MLP on CIFAR-10 with full-batch GD and HB ($\mu=0.9$) with $\eta = 0.1$.
• Figure 2: (a). Trajectories for learning dynamics of GD, HB, RGF, and HBF with discretization error $\mathcal{O}\left({\eta^2}\right)$ and $\mathcal{O}\left({\eta^3}\right)$ in a 2-d model. All dynamics start from the same point $(a_1 = 2.8, a_2=3.5)$. The convergence point of HB is denoted as a black star. The black dotted line denotes the set of all global minima. (b). Discretization errors for different continuous approximations of HB during training in (a).
• Figure 3: Trajectories of continuous approximations with different $\eta_{\text{Euler}}$.
• Figure 4: Discretization error of HBF in MNIST when training a three non-linear networks.
• Figure 5: Discretization errors $\|\mathbf{w}_k^{\mathop{\mathrm{HB}}\nolimits} - \mathbf{w}(t_k)\|^2_2$ for HBF (dotted lines) and RGF (solid lines), respectively, when training the 2-layer diagonal linear networks.

Theorems & Definitions (14)

• Theorem 2.1: HBF with the discretization error $\varepsilon_k = \mathcal{O}\left({\eta^{\alpha}}\right)$
• Theorem 3.1: Implicit bias of HBF for diagonal linear networks
• Corollary 3.2: Implicit bias of HBF for diagonal linear networks with $\varepsilon_k = \mathcal{O}\left({\eta^2}\right)$
• Lemma A.1
• Lemma A.2
• proof
• proof
• Lemma B.1
• proof
• Lemma B.2