Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

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.

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

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.
Paper Structure (121 sections, 49 equations, 5 figures, 14 tables)

This paper contains 121 sections, 49 equations, 5 figures, 14 tables.

Table of Contents

  1. Introduction
  2. Invariant Compiler Design
  3. Stage 1: Input ODEs.
  4. Stage 2: Invariant Typing (Discovery or Specification).
  5. Stage 3: Prompt (Program Specification).
  6. Stage 4: Invariant Compiler.
  7. Stage 5: Generate Neural ODE Code.
  8. Stage 6: Simulation & Learning.
  9. Catalog of Invariant Types
  10. Simplex Preservation via Skew-Symmetric Dynamics
  11. Cone Invariants via Tangent-Space Projection
  12. Lorentz Cone.
  13. PSD Cone.
  14. Center of Mass Frame Constraints
  15. Stoichiometric Constraints via Null-Space Projection
  16. ...and 106 more sections

Figures (5)

  • Figure 1: Invariant Compiler Pipeline for Neural ODEs. Epidemiological SIR dynamics are used as a representative example.
  • 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.
  • 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.
  • Figure 4: 3D phase-space portraits. (a) Thermomechanical system (Q3) (b) Extended pendulum (Q6) Green dots indicate 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.