# A memory-efficient neural ODE framework based on high-level adjoint differentiation

arXiv: https://arxiv.org/abs/2206.01298 · HTML: https://mlsanity.com/paper/2206.01298/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 44 references that are in the map.
- First posted on arXiv 2 Jun 2022.

## Builds on (most cited first, 20 of 20)

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Hamiltonian Neural Networks](https://mlsanity.com/paper/1906.01563/) (2019)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Augmented Neural ODEs](https://mlsanity.com/paper/1904.01681/) (2019)
- [Universal Differential Equations for Scientific Machine Learning](https://mlsanity.com/paper/2001.04385/) (2020)
- [Stable Architectures for Deep Neural Networks](https://mlsanity.com/paper/1705.03341/) (2017)
- [How to train your neural ODE: the world of Jacobian and kinetic regularization](https://mlsanity.com/paper/2002.02798/) (2020)
- [Beyond Finite Layer Neural Networks: Bridging Deep Architectures and Numerical Differential Equations](https://mlsanity.com/paper/1710.10121/) (2017)
- [Deep Neural Networks Motivated by Partial Differential Equations](https://mlsanity.com/paper/1804.04272/) (2018)
- [Neural Jump Stochastic Differential Equations](https://mlsanity.com/paper/1905.10403/) (2019)
- [Implicit Deep Learning](https://mlsanity.com/paper/1908.06315/) (2019)
- [ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs](https://mlsanity.com/paper/1902.10298/) (2019)
- [Stiff Neural Ordinary Differential Equations](https://mlsanity.com/paper/2103.15341/) (2021)
- [Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE](https://mlsanity.com/paper/2006.02493/) (2020)
- [SqueezeNext: Hardware-Aware Neural Network Design](https://mlsanity.com/paper/1803.10615/) (2018)
- [STEER: Simple Temporal Regularization For Neural ODEs](https://mlsanity.com/paper/2006.10711/) (2020)
- [MALI: A memory efficient and reverse accurate integrator for Neural ODEs](https://mlsanity.com/paper/2102.04668/) (2021)
- [Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows](https://mlsanity.com/paper/2005.13420/) (2020)
- [Learning Dynamical Systems from Partial Observations](https://mlsanity.com/paper/1902.11136/) (2019)
- [Toward Performance-Portable PETSc for GPU-based Exascale Systems](https://mlsanity.com/paper/2011.00715/) (2020)

## Built on by (most cited first, 3 of 3)

- [Generalized Neural Closure Models with Interpretability](https://mlsanity.com/paper/2301.06198/) (2023)
- [Semi-Implicit Neural Ordinary Differential Equations](https://mlsanity.com/paper/2412.11301/) (2024)
- [OS-net: Orbitally Stable Neural Networks](https://mlsanity.com/paper/2309.14822/) (2023)

## Abstract (arXiv)

Neural ordinary differential equations (neural ODEs) have emerged as a novel network architecture that bridges dynamical systems and deep learning. However, the gradient obtained with the continuous adjoint method in the vanilla neural ODE is not reverse-accurate. Other approaches suffer either from an excessive memory requirement due to deep computational graphs or from limited choices for the time integration scheme, hampering their application to large-scale complex dynamical systems. To achieve accurate gradients without compromising memory efficiency and flexibility, we present a new neural ODE framework, PNODE, based on high-level discrete adjoint algorithmic differentiation. By leveraging discrete adjoint time integrators and advanced checkpointing strategies tailored for these integrators, PNODE can provide a balance between memory and computational costs, while computing the gradients consistently and accurately. We provide an open-source implementation based on PyTorch and PETSc, one of the most commonly used portable, scalable scientific computing libraries. We demonstrate the performance through extensive numerical experiments on image classification and continuous normalizing flow problems. We show that PNODE achieves the highest memory efficiency when compared with other reverse-accurate methods. On the image classification problems, PNODE is up to two times faster than the vanilla neural ODE and up to 2.3 times faster than the best existing reverse-accurate method. We also show that PNODE enables the use of the implicit time integration methods that are needed for stiff dynamical systems.