Latent Mamba Operator for Partial Differential Equations

Karn Tiwari, Niladri Dutta, N M Anoop Krishnan, Prathosh AP

TL;DR: Introduces a new neural operator that leverages state-space models (SSMs) to parameterize the integral kernel in latent space.

Abstract: Neural operators have emerged as powerful data-driven frameworks for solving Partial Differential Equations (PDEs), offering significant speedups over numerical methods. However, existing neural operators struggle with scalability in high-dimensional spaces, incur high computational costs, and face challenges in capturing continuous and long-range dependencies in PDE dynamics. To address these limitations, we introduce the Latent Mamba Operator (LaMO), which integrates the efficiency of state-space models (SSMs) in latent space with the expressive power of kernel integral formulations in neural operators. We also establish a theoretical connection between state-space models (SSMs) and the kernel integral of neural operators. Extensive experiments across diverse PDE benchmarks on regular grids, structured meshes, and point clouds covering solid and fluid physics datasets, LaMOs achieve consistent state-of-the-art (SOTA) performance, with a 32.3\% improvement over existing baselines in solution operator approximation, highlighting its efficacy in modeling complex PDEs solution.

Lay Summary: Introducing the Latent Mamba Operator (LaMO), a new AI-driven approach that solves the complex math behind physical processes, like airflow, heat flow, or material deformation, much faster and more accurately without knowing the underlying dynamics. Traditional simulations bog down when problems grow large or when distant interactions matter. Still, LaMO sidesteps this by compressing all the fine details into a compact “representation” and then using the sequence-learning method to predict how that representation changes. In tests across a wide range of physics problems—grids, meshes, and even scattered data—LaMO cut errors by about 32.3% and ran faster than the numerical solvers, opening the door to real-time forecasting, rapid engineering design, and other applications that once took days to compute.

Link To Code: https://github.com/M3RG-IITD/LaMO

Primary Area: Applications->Chemistry, Physics, and Earth Sciences

Keywords: Neural Operator, Scientific Machine Learning, Partial Differential Equations (PDEs)

Submission Number: 4640