Heavy Ball Neural Ordinary Differential Equations

Abstract

We propose heavy ball neural ordinary differential equations (HBNODEs), leveraging the continuous limit of the classical momentum accelerated gradient descent, to improve neural ODEs (NODEs) training and inference. HBNODEs have two properties that imply practical advantages over NODEs: (i) The adjoint state of an HBNODE also satisfies an HBNODE, accelerating both forward and backward ODE solvers, thus significantly reducing the number of function evaluations (NFEs) and improving the utility of the trained models. (ii) The spectrum of HBNODEs is well structured, enabling effective learning of long-term dependencies from complex sequential data. We verify the advantages of HBNODEs over NODEs on benchmark tasks, including image classification, learning complex dynamics, and sequential modeling. Our method requires remarkably fewer forward and backward NFEs, is more accurate, and learns long-term dependencies more effectively than the other ODE-based neural network models. Code is available at \url{https://github.com/hedixia/HeavyBallNODE}.

arXiv comment23 pages, 9 figures, Accepted for publication at Advances in Neural Information Processing Systems (NeurIPS) 2021

References in corpus (15)

• Adam: A Method for Stochastic Optimization
• On the difficulty of training Recurrent Neural Networks
• Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation
• A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights
• Stochastic Gradient Hamiltonian Monte Carlo
• A Lyapunov Analysis of Momentum Methods in Optimization
• Full-Capacity Unitary Recurrent Neural Networks
• Tunable Efficient Unitary Neural Networks (EUNN) and their application to RNNs
• Graph Neural Ordinary Differential Equations
• Acceleration via Symplectic Discretization of High-Resolution Differential Equations
• Neural Stochastic Differential Equations: Deep Latent Gaussian Models in the Diffusion Limit
• SNODE: Spectral Discretization of Neural ODEs for System Identification
• Momentum Residual Neural Networks
• Optimization Algorithm Inspired Deep Neural Network Structure Design
• Understanding the Learned Iterative Soft Thresholding Algorithm with matrix factorization