State-space models are accurate and efficient neural operators for dynamical systems

TL;DR

This paper introduces Mamba, a state-space model-based neural operator that effectively captures long-range dependencies, offering superior accuracy and efficiency in dynamical systems prediction, especially in challenging extrapolation scenarios.

Contribution

The paper presents Mamba, a novel state-space neural operator that improves long-term dependency modeling and computational efficiency for dynamical systems, outperforming existing methods.

Findings

Mamba achieves top performance in both interpolation and extrapolation tasks.

Mamba maintains low computational cost compared to other models.

Mamba performs well in real-world pharmacology applications with limited data.

Abstract

Physics-informed machine learning (PIML) has emerged as a promising alternative to classical methods for predicting dynamical systems, offering faster and more generalizable solutions. However, existing models, including recurrent neural networks (RNNs), transformers, and neural operators, face challenges such as long-time integration, long-range dependencies, chaotic dynamics, and extrapolation, to name a few. To this end, this paper introduces state-space models implemented in Mamba for accurate and efficient dynamical system operator learning. Mamba addresses the limitations of existing architectures by dynamically capturing long-range dependencies and enhancing computational efficiency through reparameterization techniques. To extensively test Mamba and compare against another 11 baselines, we introduce several strict extrapolation testbeds that go beyond the standard interpolation benchmarks. We demonstrate Mamba's superior performance in both interpolation and challenging extrapolation tasks. Mamba consistently ranks among the top models while maintaining the lowest computational cost and exceptional extrapolation capabilities. Moreover, we demonstrate the good performance of Mamba for a real-world application in quantitative systems pharmacology for assessing the efficacy of drugs in tumor growth under limited data scenarios. Taken together, our findings highlight Mamba's potential as a powerful tool for advancing scientific machine learning in dynamical systems modeling. (The code will be available at https://github.com/zheyuanhu01/State_Space_Model_Neural_Operator upon acceptance.)