Mechanistic Neural Networks for Scientific Machine Learning

TL;DR

Mechanistic Neural Networks (MNNs) introduce a physics-grounded neural design that learns explicit differential equation representations via Mechanistic Blocks and solves them with a novel NeuRLP, a differentiable, GPU-friendly quadratic programming-based ODE solver. The framework enables scalable, interpretable discovery of governing dynamics, equation learning, and parameter inference across diverse scientific domains, including ODE and PDE settings, N-body dynamics, and time-series forecasting. Empirical results show competitive or superior performance relative to SINDy, Neural ODEs, and Fourier/Lie neural PDE solvers, with notable speedups and the ability to learn adaptive time steps. Overall, the work provides a principled, scalable pathway for integrating mechanistic equations into neural models for robust scientific machine learning and dynamic system analysis.

Abstract

This paper presents Mechanistic Neural Networks, a neural network design for machine learning applications in the sciences. It incorporates a new Mechanistic Block in standard architectures to explicitly learn governing differential equations as representations, revealing the underlying dynamics of data and enhancing interpretability and efficiency in data modeling. Central to our approach is a novel Relaxed Linear Programming Solver (NeuRLP) inspired by a technique that reduces solving linear ODEs to solving linear programs. This integrates well with neural networks and surpasses the limitations of traditional ODE solvers enabling scalable GPU parallel processing. Overall, Mechanistic Neural Networks demonstrate their versatility for scientific machine learning applications, adeptly managing tasks from equation discovery to dynamic systems modeling. We prove their comprehensive capabilities in analyzing and interpreting complex scientific data across various applications, showing significant performance against specialized state-of-the-art methods.

Figures (19)

• Figure 1: Mechanistic Neural Networks are a new neural network design that learn explicit ODE representations. Mechanistic Blocks can used as bottlenecks in other neural networks to approximate dynamical systems and discover governing equations underlying data. Additional encoders and decoders are optional and depend on the application.
• Figure 2: Learned ODE vector fields for Mechanistic NN (left) and SINDy (right) for ODEs with non-linear tanh and rational functions of polynomial basis.
• Figure 3: Solving 1d KdVbrandstetter2022lie (N train samples) (left). Comparison of ground truth and MNN prediction for the 1d KdV equation for 100 seconds (right).
• Figure 4: Ephemerides experiment predictions for orbits of Earth, Mars (left) for 1000 days (2000 steps) and eval loss (right). Showing x,y coordinates with time for visualization.
• Figure 5: Normalized true and learned force vectors during 550 steps for 2-body parameter discovery and comparison.