Mechanistic Neural Networks

Adeel Pervez, Francesco Locatello, Efstratios Gavves

Primary Area: neurosymbolic & hybrid AI systems (physics-informed, logic & formal reasoning, etc.)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: differential equations, differentiable optimization, dynamics

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Abstract:

We present Mechanistic Neural Networks, a new neural module that represent the evolution of its input data in the form of differential explicit equations.
Similar to regular neural networks, Mechanistic Neural Networks $\mathcal{}F(x)$ receive as input system observations $x$, \emph{e.g.} $n$-body trajectories or fluid dynamics recordings.
However, unlike regular neural network modules that return vector-valued outputs, mechanistic neural networks output (the parameters of) a \emph{mechanism} $\mathcal{U}_x=F(x)$ in the form of an explicit symbolic ordinary differential equation $\mathcal{U}_x$ (and not the numerical solution of the differential equation), that can be solved in the forward pass to solve arbitrary tasks, supervised and unsupervised. Providing explicit equations as part of multi-layer architectures, they differ from Neural ODEs, UDEs and symbolic regression methods like SINDy.
To learn explicit differential equations as representations, Mechanistic Neural Networks employ a new parallel and differentiable ODE solver design that (i) is able to solve large batches of independent ODEs in parallel on GPU and (ii) do so for hundreds of steps at once (iii) with \emph{learnable} step sizes.
The new solver overcomes the limitations of traditional ODE solvers that proceed sequentially and do not scale for large numbers of independent ODEs.
Mechanistic Neural Networks can be employed in diverse settings including governing equation discovery, prediction for dynamical systems, PDE solving and yield competitive or state-of-the-art results.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6069