On Second Order Behaviour in Augmented Neural ODEs

TL;DR

This paper introduces Second Order Neural ODEs (SONODEs), extending NODEs to better model second order physical systems, and compares them with Augmented NODEs (ANODEs) in terms of efficiency and interpretability.

Contribution

It extends the adjoint sensitivity method to SONODEs, proves their computational efficiency, and analyzes the capabilities and limitations of ANODEs for higher order dynamics.

Findings

SONODEs enable faster training and better performance on dynamical systems.

ANODEs can learn higher order dynamics with minimal augmented dimensions.

Theoretical analysis shows equivalence and efficiency of first order coupled ODEs.

Abstract

Neural Ordinary Differential Equations (NODEs) are a new class of models that transform data continuously through infinite-depth architectures. The continuous nature of NODEs has made them particularly suitable for learning the dynamics of complex physical systems. While previous work has mostly been focused on first order ODEs, the dynamics of many systems, especially in classical physics, are governed by second order laws. In this work, we consider Second Order Neural ODEs (SONODEs). We show how the adjoint sensitivity method can be extended to SONODEs and prove that the optimisation of a first order coupled ODE is equivalent and computationally more efficient. Furthermore, we extend the theoretical understanding of the broader class of Augmented NODEs (ANODEs) by showing they can also learn higher order dynamics with a minimal number of augmented dimensions, but at the cost of interpretability. This indicates that the advantages of ANODEs go beyond the extra space offered by the augmented dimensions, as originally thought. Finally, we compare SONODEs and ANODEs on synthetic and real dynamical systems and demonstrate that the inductive biases of the former generally result in faster training and better performance.