Jacobian Regularization Stabilizes Long-Term Integration of Neural Differential Equations
This addresses stability problems for researchers and practitioners using hybrid models in physical system modeling, though it is incremental as it builds on existing regularization techniques.
The paper tackles stability and accuracy issues in Neural Differential Equations during long-term integration by introducing Jacobian regularization via directional derivatives, achieving improved stability for several ordinary and partial differential equations with lower training costs compared to long rollouts.
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.