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

2024-09-05

arXiv.org

AI Summary

Existing dynamical system forecasting models suffer from significant limitations in long-horizon prediction accuracy, modeling of long-range dependencies, capture of chaotic evolution, and extrapolation under scarce-data conditions. This paper introduces the first Mamba-enhanced state-space model (SSM) for physics-informed machine learning (PIML), innovatively embedding Mamba’s structured state evolution mechanism into a neural operator framework while integrating parameter remapping and quantitative systems pharmacology priors. We design a rigorous extrapolation benchmark encompassing chaotic systems and multiscale dynamics to systematically address generalization bottlenecks under long-range dependencies, strong nonlinearity, and data scarcity. Experiments demonstrate state-of-the-art performance across diverse interpolation and extrapolation tasks with the lowest computational overhead. In real-world evaluation of anticancer drug efficacy, the method achieves highly robust predictions using only minimal clinical data—marking a critical advance in interpretable, sample-efficient PIML for complex biological dynamics.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.)

Problem

• Long-term Prediction
• Distant Correlation
• Complex Dynamics

Innovation

• Mamba
• Physical Information Machine Learning (PIML)
• Long-range Correlation Capturing