Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adaptive higher order reversible integrators for memory efficient deep learning

Sofya Maslovskaya, Sina Ober-Blöbaum, Christian Offen, Pranav Singh, Boris Wembe

TL;DR

This work presents an approach for constructing high-order reversible methods that allow adaptive time-stepping and shows the advantages in computational speed when applied to the task of learning dynamical systems.

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.

Adaptive higher order reversible integrators for memory efficient deep learning

TL;DR

This work presents an approach for constructing high-order reversible methods that allow adaptive time-stepping and shows the advantages in computational speed when applied to the task of learning dynamical systems.

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.

Paper Structure

This paper contains 33 sections, 4 theorems, 49 equations, 10 figures, 5 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Neural ODE
  4. Methods of gradient computations
  5. Backpropagation
  6. Adjoint method
  7. Checkpointing
  8. Reversible neural network
  9. Asynchronous Leapfrog (ALF) method
  10. New reversible architectures
  11. Adaptive stepping
  12. Gradient computations
  13. Costs comparison
  14. Computational costs
  15. Memory costs
  16. ...and 18 more sections

Key Result

Theorem 3.1

Composition of two steps of ALF methods, i.e. $\Psi^{ALF}_{h/2}\circ \Psi^{ALF}_{h/2}$, applied to $\dot z = f(z,t)$ provides second order accurate approximations of position $z$ and velocity $v=\dot z$.

Figures (10)

  • Figure 1: Log-log plot of the global error of trajectories $(z(t),v(t))$ of \ref{['eq:example']} defined on time interval $[0, 1.0]$ and obtained by ALF with $h$ ranging from $0.5$ to $10^{-3}$.
  • Figure 2: Error of the learned parameter with respect to the ground truth $\alpha$ as a function of time.
  • Figure 3: Error of the learned parameter with respect to the ground truth $\alpha$ as a function of epochs.
  • Figure 4: Error landscape of ALF and Y4 methods for Kepler problem showing the loss computed for the parameters in a neighbourhood of the true value of $\alpha$ displayed by a vertical line.
  • Figure 5: Error in learned parameters as a function of time.
  • ...and 5 more figures

Theorems & Definitions (9)

  • Theorem 3.1
  • Theorem 3.2: Yoshida1990
  • Remark 3.3
  • Remark 3.4
  • Example A.1
  • Theorem B.1
  • proof
  • Theorem B.2
  • proof