Latent selective state-space models for partial differential equations via progressive learning

Abstract

Partial Differential Equations (PDEs) play a fundamental role in modeling scientific and natural systems, such as fluid dynamics and weather forecasting. The integration of artificial intelligence into scientific computing has garnered significant attention in recent years. While many researchers focus on enhancing precision, the equally critical aspects of computational speed and training costs—particularly in long-term forecasting and control tasks—are often overlooked. This paper proposes MambaPDE, a novel approach that leverages latent vectors and state-space models to accelerate computations. To further improve accuracy, a progressive learning strategy is employed during the training phase. By comparing MambaPDE with multiple baseline methods across various datasets, this study demonstrates its effectiveness. Furthermore, our analysis shows that MambaPDE can effectively improve the condition number of the latent space, which is beneficial for the training stage. MambaPDE achieves competitive parameter efficiency, making it particularly well suited for applications requiring long-term forecasting, PDE control, and multi-modality simulation tasks.