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

Abstract

We propose a novel approach for training Physics-enhancedNeural ODEs (PeN-ODEs) by expressing the training processas a dynamic optimization problem. The full model,including neural components, is discretized using ahigh-order implicit Runge-Kutta method with flippedLegendre-Gauss-Radau points, resulting in a large-scalenonlinear program (NLP) efficiently solved bystate-of-the-art NLP solvers such as Ipopt. Thisformulation enables simultaneous optimization of networkparameters and state trajectories, addressing keylimitations of ODE solver-based training in terms ofstability, runtime, and accuracy. Extending on a recentdirect collocation-based method for Neural ODEs, wegeneralize to PeN-ODEs, incorporate physical constraints,and present a custom, parallelized, open-sourceimplementation. Benchmarks on a Quarter Vehicle Model and aVan-der-Pol oscillator demonstrate superior accuracy,speed, generalization with smaller networks compared toother training techniques. We also outline a plannedintegration into OpenModelica to enable accessible trainingof Neural DAEs.