1 of 1

Semi-Implicit Neural ODEs for Learning Chaotic Systems�FASTMath partnership

Scientific Achievement

We have developed a semi-implicit neural ordinary differential equation (ODE) approach that can learn chaotic dynamical systems efficiently and robustly from short-term trajectory data.

Significance and Impact

Chaotic dynamical systems cannot be learned by classical machine learning time-series methods (short-term accuracy is obtained but long-term statistics are lost). We reconcile this.
The approach also provides significant speed-up in training and inference over standard variants of neural ODEs due to the use of higher-order integrators.

Technical Approach

We partition the neural ODE into a linear part treated implicitly for enhanced stability and a nonlinear part treated explicitly.
The linear and nonlinear decomposition pushes trajectories to have steady-state statistical accuracy.
Adjoint-capable Implicit-Explicit Runge-Kutta solvers in PetSC are used for training the partitioned ODE with reverse-accuracy and memory efficiency
IMEX-RK allows for much larger step sizes due to its excellent stability, resulting in significant training speedup.

PI(s)/Facility Lead(s): Person Name; Romit Maulik, Hong Zhang

Collaborating Institutions: Penn State, ANL

ASCR Program: ASCR Data-intensive scientific machine learning

ASCR PM: Steve Lee, Kalyan Perumalla

IMEX methods outperform the traditional methods by 47X for training time-to-solution for the Kuramoto–Sivashinsky equations

1. H. Zhang, Y. Liu, R. Maulik. Semi-Implicit Neural ODEs for Learning Chaotic Systems, NeurIPS workshop on Heavy Tails (2023)

Ground Truth

Prediction

TALKING POINTS:

Developed a novel neural ODE that can learn chaotic dynamical systems that are typically not predictable (in terms of long-term statistics) by classical ML methods.
Used advanced solvers with PetSC to accelerate training and inference significantly.
Chaotic systems like weather and climate are now being used for training larger variants of the proposed architecture.

METADATA:

Name of the associated awarded project: ASCR Data-intensive scientific machine learning

PI name(s):Qi Tang (LANL), Romit Maulik (ANL), Jinqiao Duan (IIT Chicago), Yannis Kevrekidis (Johns Hopkins)

Name of the program manager: Steve Lee, Kalyan Perumalla

Supporting technology supported by ASCR research: ALCF facilities

CITATIONS:

Zhang, Hong, Ying Liu, and Romit Maulik. "Semi-implicit neural ordinary differential equations for learning chaotic systems." In NeurIPS 2023 Workshop Heavy Tails in Machine Learning. 2023.