Latent Mamba Operator for Partial Differential Equations
Authors: Karn Tiwari, Niladri Dutta, N M Anoop Krishnan, Prathosh Ap
ICML 2025 | Venue PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | Extensive experiments across diverse PDE benchmarks on regular grids, structured meshes, and point clouds covering solid and fluid physics datasets, La MOs 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 PDE solutions. Our code implementation is available at https://github.com/M3RG-IITD/La MO. |
| Researcher Affiliation | Academia | 1Department of Electrical Communication Engineering, Indian Institute of Science, Bangalore, India 2Department of Computer Science and Automation, Indian Institute of Science, Bangalore, India 3Yardi School of AI, Indian Institute of Technology, New Delhi, India 4Department of Civil Engineering, Indian Institute of Technology, New Delhi, India. Correspondence to: Karn Tiwari <EMAIL>, N M Anoop Krishnan <EMAIL>, Prathosh A P <EMAIL>. |
| Pseudocode | Yes | Algorithm 1 Latent SSM (Multidirectional Latent SSM) 1: Input: Latent tokens Zl 1 RB M D 2: Zl 1 Layer Norm(Zl 1) 3: ˆZ Linear(Zl 1) 4: ˆX Linear(Zl 1) 5: , B, C Linear( ˆX) 6: A, B Discretization( , A, B) 7: Y 0 8: for d in Direction-Scan do 9: Y += Multihead-SSM( ˆXd) 10: end for 11: Y Y Activation(ˆZ) 12: Zl Linear(Y) 13: Output: Latent tokens Zl RB M D |
| Open Source Code | Yes | Our code implementation is available at https://github.com/M3RG-IITD/La MO. |
| Open Datasets | Yes | Benchmark: We evaluate La MO s performance on regular grids using the Darcy and Navier-Stokes (Li et al., 2020) benchmarks. We then extend our experiment to irregular geometries, including Airfoil, Plasticity, and Pipe (Li et al., 2022a) with structured meshes and Elasticity (Li et al., 2022a) represented as point clouds. Further details are provided in the appendix Section C. |
| Dataset Splits | Yes | Table 5. The benchmark details follow the settings from (Li et al., 2022a), with input-output resolutions presented as (temporal, spatial, variate). The / denotes the absence of that dimension. [...] TRAIN SET SIZE 1000 TEST SET SIZE 200 (for Elasticity) |
| Hardware Specification | Yes | Experiments are conducted on a Linux machine with Ubuntu 20.04.3 LTS, an Intel(R) Core(TM) i9-10900X processor, and a single NVIDIA A100 40 GB GPU. |
| Software Dependencies | No | The paper mentions the Adam W (Loshchilov et al., 2017) optimizer and One Cycle LR scheduler (Smith & Topin, 2019), but it does not provide specific version numbers for any key software components or libraries (e.g., Python, PyTorch, CUDA versions). |
| Experiment Setup | Yes | Table 6. Training and model configurations of La MO. Training configurations are directly from previous works without extra tuning. For the Darcy dataset, we adopt an additional spatial gradient regularization term lgdl following ONO work. CONFIGURATION BENCHMARKS DARCY NAVIER STOKES ELASTICITY PLASTICITY AIRFOIL PIPE LOSS FUNCTION l2 + 0.1lgdl RELATIVE l2 INITIAL LR 5 10 4 10 3 OPTIMIZER ADAMW BATCH SIZE 4 2 1 8 4 4 SCHEDULER ONECYCLELR ARCHITECTURE EMBEDDING DIM 64 256 128 LATENT TOKENS 1936 1024 64 |