A simultaneous approach for training neural differential-algebraic systems of equations

Neural differential-algebraic systems of equations (DAEs) are a modeling paradigm where some unknown relationships within a DAE are modeled with a neural network and learned from data. Training neural DAEs is more challenging than training neural ordinary-differential equations (ODEs), particularly for higher-index systems. Existing approaches utilize differentiable pipelines, usually comprising integration, projection, operator splitting, or penalty terms for algebraic constraints. The parameters are then updated in a sequential manner using gradient descent. In this work, we employ the simultaneous approach for DAE-constrained parameter estimation instead. This defines a fully discretized nonlinear programing problem (NLP), whose solution simultaneously obtains the neural network parameters and the trajectories of the corresponding DAE, while enforcing constraint satisfaction at the discretization points. We show that with careful initialization and handling of the neural network terms, this approach can be efficient for smaller-scale problems, including higher-index DAEs. As the number of parameters or the amount of data increase, decomposition strategies are necessary to make this method scalable. We present a bi-level approach, where a tractable, discretized NLP is solved at every iteration of an outer gradient descent method, and the sensitivity of its solution with respect to the neural network parameters is evaluated in an efficient manner. We demonstrate the scalability of this decomposition with respect to the number of parameters and the size of the training data set.