MALI: A memory efficient and reverse accurate integrator for Neural ODEs

Abstract

Neural ordinary differential equations (Neural ODEs) are a new family of deep-learning models with continuous depth. However, the numerical estimation of the gradient in the continuous case is not well solved: existing implementations of the adjoint method suffer from inaccuracy in reverse-time trajectory, while the naive method and the adaptive checkpoint adjoint method (ACA) have a memory cost that grows with integration time. In this project, based on the asynchronous leapfrog (ALF) solver, we propose the Memory-efficient ALF Integrator (MALI), which has a constant memory cost \textit{w.r.t} number of solver steps in integration similar to the adjoint method, and guarantees accuracy in reverse-time trajectory (hence accuracy in gradient estimation). We validate MALI in various tasks: on image recognition tasks, to our knowledge, MALI is the first to enable feasible training of a Neural ODE on ImageNet and outperform a well-tuned ResNet, while existing methods fail due to either heavy memory burden or inaccuracy; for time series modeling, MALI significantly outperforms the adjoint method; and for continuous generative models, MALI achieves new state-of-the-art performance. We provide a pypi package at \url{https://jzkay12.github.io/TorchDiffEqPack/}

https://openreview.net/forum?id=blfSjHeFM_e

References in corpus (15)

• On the difficulty of training Recurrent Neural Networks
• Density estimation using Real NVP
• DeepMind Control Suite
• A Machine Learning Framework for Solving High-Dimensional Mean Field Game and Mean Field Control Problems
• Residual Flows for Invertible Generative Modeling
• Symplectic ODE-Net: Learning Hamiltonian Dynamics with Control
• Scalable Gradients for Stochastic Differential Equations
• FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models
• Dissecting Neural ODEs
• Hamiltonian Graph Networks with ODE Integrators
• On Robustness of Neural Ordinary Differential Equations
• Continuous-in-Depth Neural Networks
• SNODE: Spectral Discretization of Neural ODEs for System Identification
• Universal Approximation Power of Deep Residual Neural Networks via Nonlinear Control Theory
• An asynchronous leapfrog method II

Cited by in corpus (6)

• Learning Neural Event Functions for Ordinary Differential Equations
• Machine learning structure preserving brackets for forecasting irreversible processes
• Prediction-based One-shot Dynamic Parking Pricing
• Opening the Blackbox: Accelerating Neural Differential Equations by Regularizing Internal Solver Heuristics
• Second-Order Neural ODE Optimizer
• Efficient and Accurate Gradients for Neural SDEs