Adaptive higher order reversible integrators for memory efficient deep learning

TL;DR

This paper introduces high-order reversible integrators with adaptive time-stepping for neural ODEs, enabling memory-efficient training and handling irregular time-series data in deep learning.

Contribution

It develops the first high-order reversible methods with adaptive time-stepping, improving accuracy and efficiency in neural ODE-based deep learning.

Findings

Enhanced computational speed in learning dynamical systems

Memory requirement independent of network depth due to reversibility

Effective handling of irregular time-series data

Abstract

The depth of networks plays a crucial role in the effectiveness of deep learning. However, the memory requirement for backpropagation scales linearly with the number of layers, which leads to memory bottlenecks during training. Moreover, deep networks are often unable to handle time-series data appearing at irregular intervals. These issues can be resolved by considering continuous-depth networks based on the neural ODE framework in combination with reversible integration methods that allow for variable time-steps. Reversibility of the method ensures that the memory requirement for training is independent of network depth, while variable time-steps are required for assimilating time-series data on irregular intervals. However, at present, there are no known higher-order reversible methods with this property. High-order methods are especially important when a high level of accuracy in learning is required or when small time-steps are necessary due to large errors in time integration of neural ODEs, for instance in context of complex dynamical systems such as Kepler systems and molecular dynamics. The requirement of small time-steps when using a low-order method can significantly increase the computational cost of training as well as inference. In this work, we present an approach for constructing high-order reversible methods that allow adaptive time-stepping. Our numerical tests show the advantages in computational speed when applied to the task of learning dynamical systems.