An Invariant Compiler for Neural ODEs in AI-Accelerated Scientific Simulation

Fangzhou Yu; Yiqi Su; Ray Lee; Shenfeng Cheng; Naren Ramakrishnan

cs.LGMar 25, 2026

An Invariant Compiler for Neural ODEs in AI-Accelerated Scientific Simulation

Authors: Fangzhou Yu, Yiqi Su, Ray Lee, Shenfeng Cheng, Naren Ramakrishnan

Organizations: Virginia Tech, VA, USA

Abstract

Neural ODEs are increasingly used as continuous-time models for scientific and sensor data, but unconstrained neural ODEs can drift and violate domain invariants (e.g., conservation laws), yielding physically implausible solutions. In turn, this can compound error in long-horizon prediction and surrogate simulation. Existing solutions typically aim to enforce invariance by soft penalties or other forms of regularization, which can reduce overall error but do not guarantee that trajectories will not leave the constraint manifold. We introduce the invariant compiler, a framework that enforces invariants by construction: it treats invariants as first-class types and uses an LLM-driven compilation workflow to translate a generic neural ODE specification into a structure-preserving architecture whose trajectories remain on the admissible manifold in continuous time (and up to numerical integration error in practice). This compiler view cleanly separates what must be preserved (scientific structure) from what is learned from data (dynamics within that structure). It provides a systematic design pattern for invariant-respecting neural surrogates across scientific domains.

Figures & tables

Figure 1: Invariant Compiler Pipeline for Neural ODEs. Epidemiological SIR dynamics are used as a representative example.

Invariant Type	Constraint	Geometric IR	Enforcement Mechanism	Applications	Key Refs.	App.
Simplex / Norm	$\sum_{i}x_{i}=1$ , $x_{i}\geq 0$	Unit sphere $S^{n-1}$ via $\sqrt{\cdot}$ map	Skew-sym. dynamics: $\dot{u}=A(u)u$ , $A^{\top}=-A$ ; $x_{i}=u_{i}^{2}$	Epidemiology, market share, pop. genetics	[ 40 , 41 ]	A
Lorentz Cone	$t\geq\\|x\\|_{2}$	Lorentz cone $\mathcal{L}^{n+1}$	Tangent-cone proj.: $\dot{z}=\pi_{T_{K}(z)}(\hat{f}_{\theta}(z))$	Relativity, robust opt., signal proc.	[ 45 , 42 , 49 ]	B
PSD Cone	$P\succeq 0$	Cholesky manifold	Cholesky param.: $P=LL^{\top}$ , evolve $L$	Covariance dynamics, Kalman filtering	[ 47 , 48 ]	C
Center of Mass	$\sum_{i}m_{i}\mathbf{r}_{i}=\mathbf{0}$ , $\sum_{i}m_{i}\mathbf{v}_{i}=\mathbf{0}$	Null space of mass-weighted sum	Mean subtraction: $\dot{\mathbf{r}}_{i}=\mathbf{v}_{i}-\bar{\mathbf{v}}$ , $\dot{\mathbf{v}}_{i}=\mathbf{a}_{i}-\bar{\mathbf{a}}$	$N$ -body, molecular dynamics	[ 50 , 51 ]	D
Stoichiometric	$\mathbf{M}c(t)=\text{const}$	Null space of molecular matrix $\mathbf{M}$	Null-space proj.: $\dot{c}=B\cdot r_{\theta}(c,t)$ , $\mathbf{M}B=0$	Chemical kinetics, metabolic networks	[ 21 , 20 ]	E
Hamiltonian	$\frac{dH}{dt}=0$ , Jacobi identity	Symplectic / Poisson manifold	Latent canonical coords via INN: $\dot{z}=J_{0}\nabla_{z}K(z)$	Pendulum, Lotka-Volterra, rigid body	[ 52 ]	F

Table 1: Catalog of invariant types supported by the Invariant Compiler. Each row shows one invariant type with its constraint, the geometric or algebraic mechanism used to enforce it, and representative application domains. Full derivations and proofs appear in the referenced appendix.

System	Dim.	Invariant type	Architecture	Key challenge
SIR epidemiological model	3	Simplex ( $S{+}I{+}R{=}1$ )	Skew-symmetric	Population conservation
NOx reaction network	5	Stoichiometric ( $\mathbf{S}\dot{\mathbf{c}}{=}\mathbf{0}$ )	Null-space	Nonlinear kinetics, product inhibition
Chemical reaction network	6	Stoichiometric ( $\mathbf{S}\dot{\mathbf{c}}{=}\mathbf{0}$ )	Null-space	Element conservation (C, H, O)
Lorentz cone spiral	3	Cone ( $t\geq\\|\mathbf{x}\\|_{2}$ )	Projection	Boundary dynamics
Coupled radial-angular dynamics	3	Cone ( $t\geq\\|\mathbf{x}\\|_{2}$ )	Projection	Nonlinear radial-angular coupling
Replicator-mutator ( $N{=}5$ )	5	Simplex ( $\sum x_{i}{=}1$ )	Skew-symmetric	Diverse dynamical regimes

Table 2: Catalog of dynamical systems used in the experiments. Each system is designated with its dimensionality, the invariant type, and the compiled architecture used to enforce the invariant.

Metric	Unconst.	Penalty	Ours	IF
MSE (Train)	$5.86\text{e-}7$ $\pm 6.10\text{e-}7$	$1.02\text{e-}2$ $\pm 1.02\text{e-}2$	$\mathbf{1.80\text{e-}7}$ $\pm 1.48\text{e-}7$	$3.3\times$
MSE (Extrap)	$1.23\text{e-}6$ $\pm 1.68\text{e-}6$	$1.22\text{e-}2$ $\pm 1.10\text{e-}2$	$\mathbf{9.18\text{e-}8}$ $\pm 8.39\text{e-}8$	$13\times$
MSE (Total)	$9.09\text{e-}7$ $\pm 1.04\text{e-}6$	$1.12\text{e-}2$ $\pm 1.06\text{e-}2$	$\mathbf{1.36\text{e-}7}$ $\pm 9.19\text{e-}8$	$6.7\times$
Elem. Dev.	$6.80\text{e-}3$ $\pm 4.15\text{e-}3$	$1.10\text{e-}4$ $\pm 7.42\text{e-}5$	$\mathbf{1.06\text{e-}6}$ $\pm 8.26\text{e-}7$	$\mathbf{104\times}$

Table 3: NOx reaction network ( Q1 : constraint satisfaction). Mean

\pm

std over 200 test trajectories. IF = improvement factor of our method over the best baseline.

Figure 2: NOx reaction network ( Q1 ). Trajectory comparison showing species concentrations over time under three different initial conditions: (a), (b), and (c). Each initial condition leads to distinct transient dynamics and converges to different equilibrium values. The vertical dashed line marks the end of the training window.

Metric	Unconst.	Penalty	Ours	IF
MSE (Train)	$3.04\text{e-}4$ $\pm 2.46\text{e-}3$	$2.45\text{e-}1$ $\pm 8.31\text{e-}1$	$\mathbf{1.02\text{e-}4}$ $\pm 4.81\text{e-}4$	$3.0\times$
MSE (Extrap)	$4.90\text{e-}4$ $\pm 4.03\text{e-}3$	$4.43\text{e-}1$ $\pm 1.40\text{e}{+}0$	$\mathbf{1.55\text{e-}4}$ $\pm 7.55\text{e-}4$	$3.2\times$
MSE (Total)	$3.74\text{e-}4$ $\pm 3.05\text{e-}3$	$3.20\text{e-}1$ $\pm 1.05\text{e}{+}0$	$\mathbf{1.22\text{e-}4}$ $\pm 5.84\text{e-}4$	$3.1\times$
Viol. (Mean)	$4.13\text{e-}4$ $\pm 4.37\text{e-}4$	$\mathbf{1.07\text{e-}11}$ $\pm 3.30\text{e-}11$	$3.75\text{e-}11$ $\pm 1.34\text{e-}10$	—

Table 4: Coupled radial-angular dynamics on the Lorentz cone ( Q2 : prediction accuracy). Mean

\pm

std over 200 test trajectories.

Figure 3: Coupled radial-angular dynamics on the Lorentz cone ( Q2 ). (a) Trajectories in 3D cone space. (b) Projection onto the

(x_{1},x_{2})

plane. (c) Time evolution of the temporal component

t

, radial norm

|x|

, and spatial coordinates

x_{1}

x_{2}

, with the vertical dashed line marking the training horizon.

Metric	Model	$\times 2$	$\times 5$	$\times 10$	IF ( $\times 10$ )
MSE (Train)	Unconst.	2.91e-06	2.91e-06	2.91e-06
	PINNs	2.56e-06	2.56e-06	2.56e-06
	Ours	1.16e-06	1.16e-06	1.16e-06	$2.2\times$
MSE (Extrap)	Unconst.	2.41e-05	1.40e-04	2.33e-04
	PINNs	2.16e-06	6.75e-06	1.12e-05
	Ours	7.21e-07	8.47e-07	9.77e-07	$\mathbf{11\times}$

Table 5: Thermomechanical system ( Q3 : long-horizon extrapolation). A single training trajectory is extrapolated to

2\times

5\times

, and

10\times

the training horizon. IF at

10\times

relative to best baseline.

Figure 4: 3D phase-space portraits. (a) Thermomechanical system ( Q3 ) (b) Extended pendulum ( Q6 ) Green dots indicate initial conditions.

Metric	Model	$\times 2$	$\times 5$	$\times 10$	IF ( $\times 10$ )
MSE (Train)	Unconst.	4.39e-06	4.39e-06	4.39e-06
	PINNs	7.54e-06	7.54e-06	7.54e-06
	Ours	7.61e-07	7.61e-07	7.61e-07	$5.8\times$
MSE (Extrap)	Unconst.	7.38e-05	1.08e-03	1.15e-02
	PINNs	6.03e-05	1.66e-04	3.47e-04
	Ours	2.25e-06	1.11e-05	4.76e-05	$\mathbf{7.3\times}$

Table 6: Lotka-Volterra system ( Q4 : conservation law fidelity). IF at

10\times

relative to best baseline.

Metric	Unconst.	Penalty	Ours	IF
MSE (Train)	$8.31\text{e-}5$ $\pm 1.83\text{e-}4$	$1.03\text{e-}2$ $\pm 1.33\text{e-}2$	$\mathbf{1.38\text{e-}5}$ $\pm 3.62\text{e-}5$	$6.0\times$
MSE (Extrap)	$7.46\text{e-}5$ $\pm 1.60\text{e-}4$	$7.81\text{e-}3$ $\pm 1.33\text{e-}2$	$\mathbf{1.19\text{e-}5}$ $\pm 2.75\text{e-}5$	$6.3\times$
MSE (Total)	$7.88\text{e-}5$ $\pm 1.71\text{e-}4$	$9.07\text{e-}3$ $\pm 1.18\text{e-}2$	$\mathbf{1.28\text{e-}5}$ $\pm 3.18\text{e-}5$	$6.2\times$
Viol. (Mean)	$3.08\text{e-}4$ $\pm 6.54\text{e-}5$	$3.11\text{e-}4$ $\pm 1.75\text{e-}4$	$\mathbf{2.79\text{e-}5}$ $\pm 6.97\text{e-}6$	$11\times$

Table 7: Replicator-Mutator system,

N{=}5

( Q5 : qualitative behaviour). Mean

\pm

std over 200 test initial conditions.

Figure 5: Replicator-Mutator system ( Q5 ). Trajectories for three held-out initial conditions exhibiting qualitatively different dynamics: (a) transient oscillation, (b) monotone convergence, and (c) sustained oscillations. The vertical dashed line marks the training horizon.

Metric	Unconst.	PINNs	Ours	IF
MSE (Total)	2.01e-05 $\pm$ 1.79e-05	6.75e-05 $\pm$ 9.91e-05	2.92e-06 $\pm$ 1.62e-06	$\mathbf{6.9\times}$
Energy Dev.	2.94e-04 $\pm$ 2.49e-04	4.84e-04 $\pm$ 3.04e-04	1.81e-04 $\pm$ 9.11e-05	$\mathbf{1.6\times}$
Momentum Dev.	1.77e-04 $\pm$ 6.89e-05	5.18e-04 $\pm$ 2.51e-04	2.42e-07 $\pm$ 1.16e-07	$\mathbf{731\times}$
Ang-Mom Dev.	4.86e-04 $\pm$ 4.61e-04	7.07e-04 $\pm$ 4.95e-04	2.75e-04 $\pm$ 1.78e-04	$\mathbf{1.8\times}$

Table 8: Two-body gravitational system ( Q7 : composability). Mean

\pm

std over 200 test trajectories. IF relative to best baseline.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Metric	Unconst	Penalty	Ours	IF
MSE (Train)	$3.62\text{e-}3$ $\pm 4.33\text{e-}3$	$1.20\text{e-}2$ $\pm 8.37\text{e-}3$	$\mathbf{1.37\text{e-}3}$ $\pm 3.83\text{e-}3$	$2.6\times$
MSE (Extrap)	$2.67\text{e}{+}2$ $\pm 7.46\text{e}{+}2$	$1.82\text{e-}2$ $\pm 1.73\text{e-}2$	$\mathbf{7.79\text{e-}3}$ $\pm 2.22\text{e-}2$	$2.3\times$
MSE (Total)	$1.33\text{e}{+}2$ $\pm 3.73\text{e}{+}2$	$1.51\text{e-}2$ $\pm 1.23\text{e-}2$	$\mathbf{4.58\text{e-}3}$ $\pm 1.29\text{e-}2$	$3.3\times$
Viol. (Mean)	$4.52\text{e}{+}0$ $\pm 7.98\text{e}{+}0$	$1.59\text{e-}2$ $\pm 1.13\text{e-}2$	$\mathbf{1.69\text{e-}5}$ $\pm 9.52\text{e-}6$	$941\times$

Appendix

Table A1: SIR model (Q1 supplementary). Mean

\pm

std over 20 test initial conditions.

Metric	Unconstrained	Penalty	Null-Space (Ours)	IF
MSE (Train)	$6.24\text{e-}6$ $\pm 6.33\text{e-}6$	$7.52\text{e-}3$ $\pm 7.94\text{e-}3$	$\mathbf{4.43\text{e-}6}$ $\pm 2.65\text{e-}6$	$1.4\times$
MSE (Extrap)	$2.89\text{e-}4$ $\pm 4.48\text{e-}4$	$3.82\text{e-}1$ $\pm 1.01\text{e-}1$	$\mathbf{1.26\text{e-}4}$ $\pm 1.40\text{e-}4$	$2.3\times$
MSE (Total)	$1.47\text{e-}4$ $\pm 2.26\text{e-}4$	$1.95\text{e-}1$ $\pm 5.09\text{e-}2$	$\mathbf{6.51\text{e-}5}$ $\pm 6.94\text{e-}5$	$2.3\times$
Viol. (Mean)	$2.50\text{e-}2$ $\pm 1.14\text{e-}2$	$5.53\text{e-}3$ $\pm 2.66\text{e-}3$	$\mathbf{3.16\text{e-}6}$ $\pm 1.09\text{e-}6$	$1750\times$

Appendix

Table A2: Chemical reaction network (Q1 supplementary). Mean

\pm

std over 20 test initial conditions.

Metric	Unconst.	Penalty	Ours	IF
MSE (Train)	$7.00\text{e-}8$ $\pm 6.64\text{e-}8$	$5.55\text{e-}5$ $\pm 5.67\text{e-}5$	$\mathbf{5.18\text{e-}8}$ $\pm 6.17\text{e-}8$	$1.4\times$
MSE (Extrap)	$\mathbf{3.03\text{e-}2}$ $\pm 3.83\text{e-}2$	$1.48\text{e}{+}0$ $\pm 1.92\text{e}{+}0$	$7.05\text{e-}2$ $\pm 9.98\text{e-}2$	—
MSE (Total)	$\mathbf{1.52\text{e-}2}$ $\pm 1.91\text{e-}2$	$7.40\text{e-}1$ $\pm 9.62\text{e-}1$	$3.53\text{e-}2$ $\pm 4.99\text{e-}2$	—
Viol. (Mean)	$3.79\text{e-}2$ $\pm 2.65\text{e-}2$	$9.27\text{e-}4$ $\pm 1.15\text{e-}3$	$\mathbf{2.85\text{e-}8}$ $\pm 2.22\text{e-}8$	$32\,500\times$

Appendix

Table A3: Lorentz cone spiral (Q2 supplementary). Mean

\pm

std over 20 test initial conditions.

Metric	Model	$\times 2$	$\times 5$	$\times 10$	IF ( $\times 10$ )
MSE (Train)	Unconst.	2.81e-06	2.81e-06	2.81e-06
	PINNs	8.18e-06	8.18e-06	8.18e-06
	Ours	3.12e-07	3.12e-07	3.12e-07	$9.0\times$
MSE (Extrap)	Unconst.	8.65e-07	2.52e-07	1.30e-07
	PINNs	6.97e-06	1.95e-06	9.35e-07
	Ours	2.50e-07	1.45e-07	1.25e-07	$1.04\times$

Appendix

Table A4: Damped harmonic oscillator (Q3 supplementary).

Metric	Model	$\times 2$	$\times 5$	$\times 10$	IF ( $\times 10$ )
MSE (Train)	Unconst.	1.68e-07	1.68e-07	1.68e-07
	PINNs	1.32e-05	1.32e-05	1.32e-05
	Ours	9.71e-08	9.71e-08	9.71e-08	$1.7\times$
MSE (Extrap)	Unconst.	2.91e-06	3.30e-05	1.29e-04
	PINNs	1.13e-04	2.19e-03	5.68e-02
	Ours	5.15e-07	1.62e-06	3.23e-06	$\mathbf{40\times}$

Appendix

Table A5: Extended pendulum ( Q6 : scalability with multiple simultaneous invariants). IF at

10\times

relative to best baseline.

Invariant Type	Prompt Specification	Compiler Rewrite
Simplex / Norm	State lives on a simplex (mass conservation + non-negativity). Require exact preservation without penalties, normalization, or clamping.	Compile to a norm-preserving latent representation (e.g., $\sqrt{\cdot}$ embedding) and enforce tangency via skew-symmetric generators; recover physical states by squaring.
Lorentz Cone	State must remain in a Lorentz/convex cone for all $t$ (forward invariance).	Enforce viability by projecting the vector field onto the tangent cone at the current state (closed-form cases for interior/boundary/apex).
PSD Cone	Matrix-valued state must remain positive semidefinite ( $\succeq 0$ ) throughout evolution.	Reparameterize using a Cholesky factorization $P=LL^{\top}$ and evolve $L$ ; PSD is guaranteed by construction of $LL^{\top}$ .
Center of Mass	Center-of-mass position and total momentum must remain zero (mass-weighted constraints).	Project out the center-of-mass components by mean subtraction (mass-weighted) in both velocity and acceleration updates.
Stoichiometric	Linear conservation laws must hold (e.g., elemental/mass balances encoded by a matrix).	Restrict dynamics to the null space of the conservation matrix via $\dot{c}=B\,r_{\theta}(c,t)$ with $\mathbf{M}B=0$ , guaranteeing conservation for any $r_{\theta}$ .
Hamiltonian	Preserve energy and Poisson structure (Jacobi identity).	Compile into canonical latent coordinates and impose structured Hamiltonian/Poisson dynamics (canonical $J$ ); induce the physical Poisson tensor via coordinate transformation.

Appendix

Table A6: Prompt Taxonomy for the Invariant Compiler.

Structure-Preserving Neural Networks

See all 14 papers

MPINeuralODE: Multiple-Initial-Condition Physics-Informed Neural ODEs for Globally Consistent Dynamical System Learning
May 13, 2026Lake Yang, Antonio Malpica-Morales, Frank Ioannis Papadakis Wood +1Neural Network Generalization System Identification
Department of Chemical Engineering, Imperial College London, London SW7 2AZ, United Kingdom
Exact Fixed-Point Constraints in Neural-ODEs with Provable Universality
May 11, 2026Feliciano Giuseppe Pacifico, Duccio Fanelli, Lorenzo Buffoni +3Constrained Optimization Neural Network Approximation Theory
Department of Informatics and Computer Science, University of Pisa, Italy · Department of Physics and Astronomy and INFN, University of Florence, Sesto Fiorentino, Italy
Local Truncation Error-Guided Neural ODEs for Large Scale Traffic Forecasting
May 5, 2026Xiao Zhang, Yafei Li, Ruixiang Wang +3Spatiotemporal Forecasting Neural ODEs
Zhengzhou University · Northwestern Polytechnical University

May 13, 2026J/K move · Enter open · S save

MPINeuralODE: Multiple-Initial-Condition Physics-Informed Neural ODEs for Globally Consistent Dynamical System Learning

Lake Yang, Antonio Malpica-Morales, Frank Ioannis Papadakis Wood, Serafim Kalliadasis

Department of Chemical Engineering, Imperial College London, London SW7 2AZ, United Kingdom

Neural Network Generalization System Identification Physics-Informed ML Neural ODEs Ordinary Differential Equations

Open paper AskPDF arXiv

Neural ordinary differential equations (Neural ODEs) often fit training trajectories while generalizing poorly to unseen initial conditions and long horizons. We propose MPINeuralODE, which combines a soft physics-informed residual with a Multiple-Initial-Condition (MIC) multiple-shooting curriculum whose ingredients are structurally complementary: the physics term anchors the vector-field magnitude on the support that MIC enlarges. We evaluate along three axes: out-of-sample error, long-horizon stability, and Hamiltonian drift, which together expose whether the learned dynamics recover the underlying vector field. On Lotka-Volterra, MPINeuralODE achieves the lowest out-of-sample and long-horizon MSE among data-driven methods, with a 24% reduction over the baseline Neural ODE, while essentially matching the PINN ablation on Hamiltonian drift.

Abstract

Structure-Preserving Neural Networks

Explore similar work

MPINeuralODE: Multiple-Initial-Condition Physics-Informed Neural ODEs for Globally Consistent Dynamical System Learning

Exact Fixed-Point Constraints in Neural-ODEs with Provable Universality

Local Truncation Error-Guided Neural ODEs for Large Scale Traffic Forecasting