Efficient Training of Physics-enhanced Neural ODEs via Direct Collocation and Nonlinear Programming

TL;DR

This paper introduces a new training method for Physics-enhanced Neural ODEs using direct collocation and nonlinear programming, improving stability, speed, and accuracy over traditional approaches.

Contribution

It formulates the training as a large-scale NLP solved by advanced solvers, generalizes to physics constraints, and provides an open-source implementation for efficient Neural ODE training.

Findings

Superior accuracy and speed on benchmark models

Smaller networks achieve comparable or better results

Enhanced generalization capabilities

Abstract

We propose a novel approach for training Physics-enhanced Neural ODEs (PeN-ODEs) by expressing the training process as a dynamic optimization problem. The full model, including neural components, is discretized using a high-order implicit Runge-Kutta method with flipped Legendre-Gauss-Radau points, resulting in a large-scale nonlinear program (NLP) efficiently solved by state-of-the-art NLP solvers such as Ipopt. This formulation enables simultaneous optimization of network parameters and state trajectories, addressing key limitations of ODE solver-based training in terms of stability, runtime, and accuracy. Extending on a recent direct collocation-based method for Neural ODEs, we generalize to PeN-ODEs, incorporate physical constraints, and present a custom, parallelized, open-source implementation. Benchmarks on a Quarter Vehicle Model and a Van-der-Pol oscillator demonstrate superior accuracy, speed, generalization with smaller networks compared to other training techniques. We also outline a planned integration into OpenModelica to enable accessible training of Neural DAEs.