A simultaneous approach for training neural differential-algebraic systems of equations
Laurens R. Lueg, Victor Alves, Daniel Schicksnus, John R. Kitchin, Carl D. Laird, Lorenz T. Biegler, “A simultaneous approach for training neural differential-algebraic systems of equations”, Computational Optimization and Applications (2026). https://doi.org/10.1007/s10589-026-00823-y
This summary was written by a language model from the paper and has not been checked line by line. Rely on the published paper for details.
Note on the source text: the extracted text is the arXiv preprint (arXiv:2504.04665v2, April 2026) rather than the typeset journal version. The prose is readable; equations and figure axes are partly garbled, so results are described here from the text and the readable table values.
Question
Neural differential-algebraic equations (neural DAEs) are hybrid models in which some unknown terms of a mechanistic DAE are represented by a neural network learned from data. Existing training methods are “sequential”: they push a loss through a differentiable pipeline of integration plus projection, operator splitting or penalty terms for the algebraic constraints, then update the weights by gradient descent. These methods do not strictly enforce algebraic constraints and struggle with higher-index systems. The paper asks whether the simultaneous approach long used for DAE-constrained parameter estimation in process systems engineering (full discretization plus a nonlinear programming solver) can train neural DAEs, and how to make it scale (Lueg et al. 2026).
Approach
The training problem is posed as a DAE-constrained optimization over all observed trajectories: a least-squares fit to noisy differential-state observations plus weight regularization, subject to differential equations, algebraic equalities, optional inequality (path) constraints, and the constraint that the unknown terms equal the network output. The DAE is discretized by orthogonal collocation on finite elements (Radau points), so state trajectories, unknown terms, static parameters and network weights all become variables of one large NLP solved with IPOPT. Smooth activations are required because the interior point method needs second derivatives.
Three devices make this NLP tractable. First, an initialization problem drops the network and instead treats the unknown terms at each collocation point as free variables, fitting the data with a smoothness penalty; it is separable by trajectory, and its solution initializes the trajectories, fixes the input/output normalization layers, and gives targets for pretraining the weights by stochastic gradient descent. Second, the network constraints are implemented as Pyomo external gray-box models whose values, Jacobians and Hessians are computed in vectorized form with JAX, avoiding dense algebraic expressions. Third, an L-BFGS Hessian approximation replaces the dense, nonconvex exact Hessian terms in the weights.
For larger problems the authors propose a bi-level decomposition. The outer level updates the weights with Adam; the inner level solves, in parallel, one NLP per trajectory with the weights fixed and slack variables (with an absolute-value penalty) on the network constraints so a feasible solution always exists. The gradient of the outer loss is obtained by implicit differentiation of the inner KKT conditions, and the needed products reduce to vector-Jacobian products of the network, again evaluated with JAX. Sensitivities are well defined when LICQ and strong second-order sufficiency hold, which can be checked at each inner solution.
Key findings
On an index-2 four-tank manifold DAE (learning the pump flow and one outflow as functions of the tank heights), an ablation with a two-hidden-layer, 30-neuron tanh network over three trajectories showed that the fully simultaneous NLP with exact Hessians and no initialization gave the lowest training error (MSE 0.067) but took about 711 s, while L-BFGS cut this to about 32 s but failed on one of three problems. With trajectory and weight initialization plus L-BFGS, total time was about 37 s and all problems solved, at a somewhat higher error (MSE 0.084), reflecting a bias toward smooth learned terms. The bi-level decomposition kept the time and IPM iterations per Adam step roughly constant from 2 to 128 trajectories (one per core) and from one to eight hidden layers of 50 neurons, though the number of Adam steps grew with problem size. The fully simultaneous method could not reach those sizes because of memory, and its cost grew with trajectories and with network depth (denser linear systems), yet on small problems it produced markedly more accurate models than the decomposition.
In a predator-prey model, adding a known Lyapunov function as a path constraint (its time derivative nonpositive) gave a learned model that generalized much better beyond the data than the unconstrained one. In a fed-batch bioreactor, learning the Monod growth rate from two trajectories and then using the simultaneous formulation for inference (with slacks on the network output) kept the substrate concentration nonnegative where a standard Runge-Kutta integration of the same learned model violated it, giving trajectories closer to ground truth.
Significance
The work brings the collocation-plus-NLP machinery of dynamic optimization into scientific machine learning, showing that it can train hybrid models with higher-index DAEs and inequality constraints while enforcing those constraints exactly at the discretization points. The decomposition connects this to differentiable programming: NLP solutions become differentiable layers whose sensitivities are computed by implicit differentiation and AD. The authors note open issues: hyperparameter and stopping-criterion tuning for the outer loop, closer coupling of the solver and sensitivity computation, partial rather than full Hessian approximation, GPU solution, and tests on experimental data. Code is released as SiNDAE on GitHub.