Jacobian Regularization Stabilizes Long-Term Integration of Neural Differential Equations

TL;DR

This paper introduces Jacobian regularization techniques for Neural Differential Equations, significantly enhancing their stability during long-term integration without the high computational costs of unrolled training.

Contribution

It proposes novel Jacobian regularization methods for NDEs that improve long-term stability, applicable to both known and unknown dynamic systems, with lower training costs.

Findings

Regularization stabilizes long-term NDE integration.

Methods are effective for ordinary and partial differential equations.

Training costs are significantly reduced compared to unrolled trajectories.

Abstract

Hybrid models and Neural Differential Equations (NDE) are getting increasingly important for the modeling of physical systems, however they often encounter stability and accuracy issues during long-term integration. Training on unrolled trajectories is known to limit these divergences but quickly becomes too expensive due to the need for computing gradients over an iterative process. In this paper, we demonstrate that regularizing the Jacobian of the NDE model via its directional derivatives during training stabilizes long-term integration in the challenging context of short training rollouts. We design two regularizations, one for the case of known dynamics where we can directly derive the directional derivatives of the dynamic and one for the case of unknown dynamics where they are approximated using finite differences. Both methods, while having a far lower cost compared to long rollouts during training, are successful in improving the stability of long-term simulations for several ordinary and partial differential equations, opening up the door to training NDE methods for long-term integration of large scale systems.