Latent Mamba Operator for Partial Differential Equations
Authors: Karn Tiwari
Niladri Dutta
N M Anoop Krishnan
Prathosh AP
Venue: ICML 2025 poster (2025)
Keywords: Neural Operator
Scientific Machine Learning
Partial Differential Equations (PDEs)
Primary area: applications->chemistry_physics_and_earth_sciences
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.

Latent Mamba Operator for Partial Differential Equations
Karn Tiwari' Niladri Dutta 2
N M Anoop Krishnan Prathosh A P
Abstract
approaches face significant computational challenges, are clusive to real-time forecasting, are limited in handling di- verse boundary conditions, and often require coarse grids for stable solutions (El-metwaly & Kamal, 2024). Addition- ally, many PDEs, including the Navier-Stokes equations, lack closed-form solutions, posing significant challenges for real-time weather modeling and robotics applications.
Neural operators have emerged as powerful data- driven frameworks for solving Partial Differential Equations (PDEs), offering significant speedups over numerical methods. However, existing neu- ral operators struggle with scalability in high- dimensional spaces, incur high computational costs, and face challenges in capturing continu- ous and long-range dependencies in PDE dynam- ies. To address these limitations, we introduce the Latent Mamba Operator (LaMO), which inte- grates the efficiency of state-space models (SSMs) in latent space with the expressive power of ker- nel integral formulations in neural operators. We also establish a theoretical connection between state-space models (SSMs) and the kernel inte- gral of neural operators. Extensive experiments across diverse PDE benchmarks regular grids, structured meshes, and point clouds covering solid and fluid physics datasets, LaMOs achieve consistent state-of-the-an (SOTA) performance, with a 32.3% improvement over existing base- lines in solution operator approximation, high- lighting its efficacy in modeling complex PDE solutions. Our code implementation is available at http:///gitbbbooo////GG..TDDD
In recent years, scientific machine learning (SciMI.) has in- troduced neural operators as promising alternative for solv- ing PDEs through data-driven approaches. Neural operators, an extension of neural networks, have emerged as a powerful tool for mapping between spaces, serving as universal functional approximations (Li et al., 2020). Unlike traditional methods, neural operators require no prior knowledge of the underlying PDEs, Instead, they rely data-driven training, enabling faster inference and making them compelling choice for real-time and complex applications. Neural operators have demonstrated remarkable success across diverse applications, including weather forecasting (Kurth et al., 2023; Bedi et al., 2025), biomedical surrogate modeling (Guan et al., 2021), acceler- ating sampling processes in diffusion models (Zheng et al., 2023), and foundation models for solving PDEs (Hao et al., 2024; Herde et al., 2024).
functional
Recent advances in neural operators for solving PDEs have been propelled by transformer based architectures, with models such as GNOT (Hao et al., 2023), ONO (Xiao et al., 2023), and Transolver (Wu et al., 2024) achieving SOTA performance across diverse PDE tasks. Despite their suc- cess, these methods face critical challenges in computational efficiency and scalability, primarily when applied to high- dimensional PDEs. In such cases, where they struggle to parameterize kernel integral transforms (Guihas et al., 2021), leading to overfitting, which limits their generalization. A fundamental bottleneck lies in the quadratic complexity of the self-attention mechanism inherent to conventional transformers, which incurs prohibitive computational costs, impedes the processing of continuous signals, and degrades inference speed as spatial resolution increases (Hao et al., 2023; Wu et al., 2024).
Continuum models describing physical systems are formu- lated as PDEs across various disciplines, including physics, chemistry, fluid mechanics, and roboties (Debnath & Deb- nath, 2005). Traditionally, these PDEs are solved by clas- sical numerical methods, such as finite element and spec- tral methods (Solin, 2005; Costa, 2004). However, these
Department of Electrical Communication Engineering, Indian Institute of Science Bangalore, India Department of Computer Science and Automation, Indian Institute Science, Bangalore, In- dia Yardi School of AI, Indian Institute of Technology, New Delhi, India 'Department of Civil Engineering, Indian Institute of Tech- nology, New Delhi, India. Conresoondence Karn Tiwari <karn- tiwari @ iisc.ac.in>,. Prathosh A P <prathosh @ ins..ac.in>.
Anoop Krishnan <krishhan @ aitd.ac.ins.
While recent efforts, such as Galerkin-type attention (Cao, 2021), address the issue of computational scaling by reduc- ing the complexity to linear, this compromise significantly
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Latent Mamba Operator for Partial Differential Equations
diminishes the expressive power of the self-attention mech- anism, limiting the ability capture intricate dynamical interactions in complex PDE systems. Similarly, global convolution-based paradigms like FNOs (Li et al., 2020) leverage spectral convolution to encode global information but encounter persistent challenges, including spectral de- cay, vanishing gradients, and computational bottlenecks when processing high-resolution inputs (Tran et al., 2021: Fanaskov & Oseledets, 2022). These limitations underscore pressing need for framework that balances expressive power, computational efficiency, and scalability in PDE learning, which can generalize better with limited data.
leverages frequency domain techniques. FNO employs Fourier kernel-based integral transformations facilitated by fast Fourier transform and projection blocks. An enhanced version, F-FNO (Tran et al., 2021), improves upon FNO by incorporating advanced spectral mixing and residual connections. Various kernel integral neural operators have been proposed based on frequency kernel methods. For instance, Fanaskov & Oseledets (2022) introduced spectral methods based Chebyshev and Fourier series to reduce aliasing errors and improve the clarity of FNO outputs. In contrast, CoNO (Tiwari et al., 2024) employs Fractional Fourier transform (FrFT)-based integral kernel. Interest- ingly, Kovachki et al. (2021) showed that the self-attention mechanism can be interpreted as a specific case of neural operators learning an integral kernel for solving PDEs.
Structured state-space models (SSMs) have emerged as a promising approach for sequence models, demonstrating ex- cellent performance in language processing. These models offer linear scaling with improved long-range dependency modeling, efficient memory usage, and superior handling of language and vision data (Gu & Dao, 2023; Zhu et al., 2024). However, the application of such framework toward PDEs remains poorly explored. Here, we propose SSMs in op- erator learning and introduce a unified framework, Latent Mamba Operator (LaMO). LaMO leverages latent space and bidirectional SSMs address the limitations of exist- ing Neural Operators, providing efficient, scalable, and effective solution for complex tasks in SciML. The main contributions of our work are briefly summarized below.
1.2. Transformer-Basse Neural Operators
Recent work has sought to enhance the efficiency of attention- based neural operators. Cao (2021) pioneered this direction by replacing traditional softmax layers with two novel self-attention operators, significantly reducing computational overhead while grounding attention mecha- nisms in kernel integral theory. Subsequent work by Hao et al. (2023) introduced the GNOT operator, which lever- ages a linear cross-attention block to improve geometric encoding for irregular domains. Despite these advances, transformer-based operators remain vulnerable to overfit- ting in data-scarce regimes, limiting their generalization. To mitigate this, Xiao et al. (2023) proposed orthogonal regular- ization, enhancing robustness by enforcing orthogonality in attention weights Meanwhile, the SOTA Transolver opera- tor (Wu et al., 2024) addresses scalability challenges in large meshes through latent attention, decoupling computational complexity from mesh resolution.
1. LaMO: We propose Latent Mamba Operator (LaMO). This framework leverages a latent physical space combined with bidirectional state-space models (SSMs) to learn underlying solution operators.
2. Theoretical Insights: We provide a theoretical analy- sis of LaMO from the operator's perspective, demon- strating its capability to act as a kernel integral for learning underlying solution operators in Theorem 3.4.
1.3. State Space Models (SSMs)
3. Superior Performance:
Extensive experiments are conducted across various regular and irregular grid datasets, showcasing an average improvement of 32.3% over the SOTA models as presented in Table 1.
Structured State Space (S4) models, introduced by (Gu et al., 2021a), provide novel alternative to transformers for mod- eling long-range sequences with linear scaling in sequence length, unlike the quadratic complexity of transformers. Gu et al. 2021b highlighted the potential of SSMs for captur- ing long-range dependencies. Subsequent advancements include complex diagonal structures, low-rank decomposi- tions, and selection mechanisms (Dao & Gu, 2024). Re- cently, Hu et al. 2024 proposed the integration of SSMs in solving ordinary differential equations (ODEs) for time series datasets, and MemNO (Buitrago et al.) introduced the benefits of incorporating SSMs for time-dependent PDEs. While SSMs have excelled in language (Gu & Dao, 2023; Qu et al., 2024), audio (Erol et al., 2024), and computer vision tasks (Zhu et al., 2024; Liu et al., 2024), their applica- bility to solving parametric high-dimensional PDEs remains largely unexplored in SciMI.
1. Related Work
1.1. Neural Operator (NO)
Neural operators have demonstrated significant success in solving parametric PDEs through data-driven approaches (Kovachki et al., 2021). Lu et al. (2021) introduced Deep- ONet, which established the universal functional approxi- mation capability of neural operators. DeepONet employs two networks: the branch network, responsible for learning the input function operator, and the trunk network, which projects this operator onto the target function space. An- other prominent neural operator, the FNO (Li et al., 2020),
2

---

Latent Mamba Operator for Partial Differential Equations
2. Preliminaries
and I represents the identity matrix. Subsequently, the pro- cess in Eq. can be realized in global convolution form:
2.1. Problem Statement
y x K == [CB, CAB. .... B, (6)
The main objective is to Iearn the nonlinear mapping be- tween two ciinite-dimesiional spaces through observed input-output pairs (Li et al., 2020; Lu al., 2021). Let D represent an open and bounded domain as D Rdd, with A - A(D; (R)) and U (D; R separable banach spaces of functions representing elements in R and R respectively. Let G A + U denotes a non-linear mapping arising from the underlying solution operator for parametric PDEs.
where K E RR represents the convolutional kernel.
Selective State Space Models: The Selective State Space (S6) mechanism, proposed in Mamba (Gu & Dao, 2023), introduces input-dependence for the parameters B, C, and A, significandly improving the performance of SSMs and parallel scanning to make the training performance com- parable with wansformer-based models. By making these parameters uppu-dependent, the global convolution kernel in Equation 6 can be reformulated as follows:
In operator learning, we aim to construct an approximation for gt via parametrie mapping G A O - U, equiv- alently, Go A + U,O E θ, within finitedimmensional parameter space θ via i..d. observations (ay, 1z));=1 where aj ~ μ., drawn from the underlying measure jl supported on A, and uj = (a,)) The aim is to learn Q E θ such that 0t, g = gt. It allows learning in infiniteddimensiinall spaces as the solution to the optimiza- tion problem in Eq. constructed using a loss function CUUXU R.
L-1 k= (C,BL, CLALIBLLLLLL CL II A.B.). O
The reformulation highlights that integrating selective state- space parameters (data-dependens parameters) enhances the expressiveness of the SSMs for complex tasks. Better elfi- ciency is attained through an 10-aware implementation of as- sociative scans, leveraging work-efficient parallel scanners to enable parallelization on modern hardware. In LaMO, Go is parameterized using the SSMs as an integral kernel. Remark 2.1. The matrix A of Mamba has diagonal structure i. A = Diag .p, with X E C Vi.
Bawn (t)) 0=8
(1)
The optimization problem is solved using a data-driven empirical risk minimization of the loss function, akin to the supervised learning approach using train-test observations.
3. Proposed Method
2.2. State Space Models
Transformers struggle to capture kemnll integral transforms efficiently in complex, high high-dimensional continuous PDEs (Guibas et al., 2021; Karniadakis et al., 2021). To address this, we introduce LaMO (Latent Mamba Operator), in- spired by structured state-space models like Mamba (Gu Dao, 2023). LaMO begins with Latent Encoder inspired by the Perceiver (Jaegle al., 2021), which condenses PDE's physical tokens into a compact latent representation. These tokens are then processed by an SSMs block (bidirectional), utilizing SSM's capability to learn data-dependent kernels, thereby effectively modeling complex PDE dynamics.
SSMs: Classical SSMs define continuous system (Linear (LTT) system) that maps an input sequence x(t) E R into hidden representation htt) E R which is then used to predict an output response (1) E RR2. For- mally, an SSMs can be formulated as ordinary differential equations (ODEs) as follows:
h' h'(1)= Ah(() Bx((),
(2)
yt) Ch(t),
(3)
where A E NXN, B E Nx1, and C E RIIN are learn-
able model parameters.
Discretization: To use continuous SSMs for integration into neural networks, it is essential to apply discretization oper- ations. By introducing a timescale parameter E IR and using the zero-order hold (ZOH) method as the discretiza- tion rule, the discrete counterparts of A and B (denoted as A and B, respectively). Consequently, Equation 3 can be rephrased into discretized form as:
3.1. Latent Mamba Operator (LaMO) Architecture
Overview: The LaMO operator takes the following form:
Go = @o= C oEo.
(8)
where denotes composition. The operators P R g and Q Rd - Rd correspond to the lifting and projection operations, encoding lower-dimensional spaces into higher- dimensional spaces or vice versa, which helps in convert- ing the non-linear dynamics into linear dynamics (Bevanda et al., 2021) using the ferd-forward neural network. The
hikk = Ahlk - 1] Bx[k], y[k] = Ch(k).
(4)
where,
A A B= ( 1JB AB,
(5) operations E and D correspond to the encoder and decoder,
3

---

Latent Mamba Operator for Partial Differential Equations
(z)
Latent
Lafent Block
D
E
Latent
(a)
Heads
Forward Convtd
Forward
SSM
Latent- SSM
Layynn LayerNorm
Linear
Linear
Linear
LayerNorm
Activation
Lay
Backward Convid
Backward
SSM
Figure 1. Overview. (1) The imput function a(x) is lifted to a nigher-dimensionn representation using the lifting operator P. (2) The encoder & maps the input from the physical domain the latent domain, where latent block (Rottom Left) performs the kernel integral via SSMs, applies channel mixing and decodes the latent tokens back the physical domain using the decoder 'D (3) A multi-headed bidirectional SSM (Bottom Right) is applied across the tokens within the latent SSMs block. (4) This process is repeated L times. Finally, the channel dimensions are reduced to the desired output size using the projection operator Q, yielding the final output function u(z).
transforming the input domain into latent tokens and vice versa. These operators constructed using layers of non-linear integral operators, denoted as L' consisting of an integral kernel XC, parametrized using the time-variant SSMs. The design of C' is inspired by the integral operators (Green's functions) commonly used for solving linear PDEs.
patch, the patchify operation can be approximated by set- ting Z w and optimizing latent weights W, such that W i.j = 1 and W i, Rs 0 for k + j. Thus, the patch operation is special case of latent tokens for regular grids.
Latent Block:
Inspired by the attention architecture, we describe the LaMO block as a combination of latent SSMs operation, which defines the kernel integral operator, fol- using T. layers;
The encoder module draws inspiration from the design principles of Perceiver (Jaegle al., 2021), which efficiently reduces the number of tokens for computational tractability. Unlike the Perceiver, where latent tokens corre- spond to fixed latent representation, this module derives latent tokens directly from the input physical tokens. Let X E ND denote the input tokens, where N is the number of tokens and is the embedding dimension. We project physical N tokens into M latent tokens where M << N which can be described follows:
Encoder (E):
lowed by channel mixing. Suppose the (-th layer of the LaMO block can be expressed as:
Ž' = 7'-1 *+ -Latent-SSM (LayerNorm (Z'-)),, (11)
z = ž²' + Linear (LayerNorm (Z))).
(12)
W Softmax (oneeaax(inear(X))),
(9)
where E (1,.... L}. Linear() represent a feed-forward neural network and LayerNorm() represent the layer nor- malization (Lei Ba et al., 2016). Here, Z E TR M denotes the output of the (-th layer, and z" E M represents the latent tokens embedded from the input physical tokens.
Z W
(10)
where Linear() projects the channels into M channels using a feed-forward neural network, W E N denotes the latent weights which are adaptive as its function of input tokens and Z denotes the latent tokens which correspond to latent space which in contrast a frequency -based method such as FNO, we learn the latent space depending the input tokens. The Softmax() operation, similar to that used in Transolver and Perceiver, ensures that the latent weights have a low entropy, resulting in an informative latent space. Remark 3.1 (VIT Patches as Special Cases of Latent Tokens (Alexey, 2020)). The latent encoder operation di- vides the grid into M equal-sized square patches for regu- lar computational grids. a point X, belongs to the j-th
Latent SSM Block: SSMs are primarily designed for 1- dimensional sequences, typically operating in a causal man- ner, where output at any step depends on the arrangement of preceding tokens. However, accounting for the entire se- quence in non-causal setting for tasks involving PDEs is essential. To address this, we use the bidirectional SSM block, which leverages both forward and backward kernels to process the sequence holistically.
The appendix Algorithm I provides the algorithm detailing the operations performed. Specifically, the latent token

---

Latent Mamba Operator for Partial Differential Equations
sequence Z undergoes following operation follows:
experiments demonstrate that the multi-head mechanism accelerates training convergence and reduces over-fitting.
x = Linear(Z), (Z),
Shared Weights: The transformation between latent and physical tokens allows LaMO to build deeper models but may limit its ability to fully capture physical space informa- tion due to increased computational complexity. We share encoder and decoder parameters between blocks to balance performance and efficiency (Jaegle et al., 2021).
? Linear (Z),
(A, B, C) + Discretization (X),
(13)
Y = Σ
(X)(Xx) @ Ac)
de Direcction
Z Linear(Y).
Computational Analysis: For simplicity, the above process is summarized as Z' = Latent-SSM((Z) with an overall computational complexity of ONN MD + MD). Since M is set as a constant and M < N, the computational complexity becomes linear with respect to the number of mesh points. Keeping the mesh size constant, it is also linear with respect the number of latent tokens.
where Linear() represents a feed-forward neural network, while direction refers to the SSM recurrence applied in the forward and backward directions. The activation function Act() is implemented as SiLU (Elfwing et al., 2018), and Discretization() denotes the process of converting continu- ous dynamics into discrete dynamics using the zero-order hold (ZOH). Finally, the block's 's output is obtained by ap- plying a gating mechanism followed by a linear projection. Remark 3.2 (Multidirectional Scan Regular Grid). For latent tokens regular grid, the spatial arrangement of tokens is crucial due to the inherent continuity in reg- ular grids. Unlike irregular grids, we utilize a multidirec- tional scan along four paths: (i) top-left to bottom-right, (ii) bottom-right to top-left, (iii) top-right to bottom-left, and (iv) bottom-leff to top-right. Each direction is processed in parallel, reducing the overall computational complexity to linear, and has proven effective for regular grids. The multidirectional scan is further analyzed and validated ex- perimentally in the appendix Section E.I.
3.2. Theoretical Analysis
In this subsection, we present a theoretical analysis of the components of the proposed method. Specifically, (i) we establish the equivalence between the SSM and the Euler method for solving PDEs, which is essential for understand- ing the performance of SSM (Proposition 3.3). (ii) Next, we prove that SSM serves as a Monte Carlo approximation of a learnable integral kernel (appendix Lemma 1.10). This result is then used to prove that the latent SSM is equivalent
to learnable kernel integral in Theorem 3.4. appendix Section B provides more detailed analysis.
Proposition 3.3. The of continuous parameters is expressed as follows:
Hold (ZOH) discretization
Decoder (D): Decoder can be defined converting the latent tokens back to physical tokens. We need to convert the latent token Z E R MDD into the physical tokens Y E
A exp(AA), B = (AA) (CS((A)) - -J)
R NDD where M < N. Formally, it can be described as:
(17)
W- Softmax(Linear(2 (Z)),
(14)
where A denotes the time step. i the identity matrix, and A, B are continuous system matrices. The discretization method aligns with the Euler method, approximating the matrix exponential by truncating its Taylor series expansion to the first-order term.
Y = WTZ,
(15)
where Linear() projects the D channels into N channels using feed-forward neural network, W E M denotes the latent weights and Softmax() operation, ensures that the latent weights have a low entropy similar to the encoder.
The complete proof is in the appendix Section B.4. The above proposition establishes the equivalence between the continuous and discrete SSM parameters and their relation to the Fuler method for solving the PDE. Understanding SSM's performance from the Euler method's perspective is crucial. The ZOH discretization method preserves higher- order terms, making more accurate for solving PDEs.
Multihead: To enhance the model's effectiveness, we em- ploy multihead architecture similar to the attention mecha- nism (Vaswani, 2017) that captures information from differ- ent representational subspaces. Specifically, a latent token embedding dimension D is split into H heads, each with a dimensionality of D/ H. The latent tokens are processed as:
SSM(Z)= Concat(SSM(())) SSM(())) (16)
Next, we demonstrate that SSM functions kernel, following prior works (Kovachki et al., 2021; Cao, 2021). The Lemma 1.10 shows that the canonical SSM is the Monte Carlo approximation of the integral kernel, used iteratively to solve the solution operator. However, we apply the kernel integral on latent tokens. For better understanding.
integral
To integrate features from different heads, the outputs of all heads are concatenated and linearly projected, yielding the final output. This architecture increases model diver- sity while maintaining a compact size. Furthermore, our
5

---

Latent Mamba Operator for Partial Differential Equations
Table main results across all benchmark datasets are presented using the mean relative & error (Equation 21) as the evaluation metric. A lower &2 error signifies better performance. The "INCREMENT % quantifies the relative error reduction achieved by our model compared to the second-hest performer on each benchmark. For clarity, in the color legend, orange represents the best performance, blue indicates the second-best performance, and violet signifies the third-best performance among the baselines.
POINT CLOUD
REGULAR GRID
STRUCTURED MESH
OPERATOR
ELASTICITY
HAVIER-STOKES
DARCY
PLASTICITY
AIRFOIL
PIPE
CLASSIC
UNET (2015)
0.0235
0.1982
0.0080
0.0051
0.0079
0.0065
RESNET (2016)
0.0262
0.2753
0.0587
0.0233
0.0391
0.0120
SWIN (2021)
0.0283
0,2248
0.0397
0.0170
0.0270
0.0109
DEEPONET (2021)
0.0965
0,2972
0.0588
0.0135
0.0385
0.0097
WMT (2021)
0.0359
0.1541
0.0082
0.0076
0.0075
0.0077
FREQUENCY
U-FNO (2022)
0,0239
0,2231
0.0183
0.0039
0.0269
0.0056
FNO (2020)
0,0229
0,1556
0.0108
0.0074
0.0138
0.0067
U-NO (2022)
0.0258
0.1713
0.0113
0.0034
0.0078
0.0100
F-FNO (2021)
0.0263
0,2322
0.0077
0.0047
0.0078
0.0070
LSM (2023)
0.0218
0.1535
0.0065
0.0025
0.0059
0,0050
GALERKIN (2021)
0.0240
0.1401
0.0084
0.0120
0.0118
0.0098
SSM TRANSFORMED
HT-NET (2022)
0,1847
0.0079
0.0333
0.0065
0.0059
OFORMER (2022B)
0.0183
0,1705
0.0124
0.0017
0.0183
0.0168
GNOT (2023)
0.0086
0.1380
0.0105
0,0336
0.0076
0.0047
FACTFORMER (2024)
0.1214
0.0109
0.0312
0.0071
0.0060
ONO (2023)
0.0118
0.1195
0.0076
0.0048
0.0061
0.0052
TRANSOLVER (2024)
0.0064
0.0957
0.0059
0.0013
0.0053
0.0046
LAMO (OuRS)
0.0050
0.0460
0.0039
0.0007
0.0041
0.0038
INCREMENT
21.8% %
51.9
33.9 %
46.1 %
22.6% %
17.4%
we show that the latent SSM in our framework is a learnable kernel integral in the latent domain in Theorem 3.4.
Connection with ONO (Xiao et al., 2023): ONO intro- duces orthogonal attention to model the kernel update using Cholesky Decomposition (Higham, 1990), ensuring the or- thogonality and positive definiteness the attention kernel. The orthogonal attention kernel was defined follows:
Theorem 3.4 (SSM as an equivalent integral kernel I). Let Q S. R" be bounded domain, a Ω + Rdd be given input function, and x E Q be a mesh point. An latent-SSM layer approximates the integral operator (2( R4) + 22( (D), defined as follows:
<(I, y) v'(z) (=)))aa(x))))))
(19)
where vb(-) Ω Rd maps the input embedding space. As shown in the appendix Section B, this structure resembles the inherent form of the SSM kernl::
G(a)(x)= = - k(x, v) aly) dy.
(18)
where k Ω x Ω + Raxd is the kernel function character- izing the operator G.
<(x, v) = Ciz)
B(y)-
(20)
By leveraging this inherent structure, LaMO avoids explicit decomposition for orthonormalization, thereby mitigating over-fitting and improving generalization.
The appendix Theorem B.12 provides the complete proof. In prior work, (Li et al., 2020) formalized the neural op- erator as an iterative process, where the key component is linear kernel integral. This formulation consists of an integral operator combined with a nonlinear activation function, which enables the learning of a nonlinear surro- gate mapping. Since the nonlinear activation function can be efficiently parametrized using feedforward neural net- work, the above theorem demonstrates that the SSM can be viewed as a linear kernel integral. This result implies that an 3SM-based operator can effectively learn the nonlinear surrogate mapping for the underlying PDE solution. Now, we show the connection the proposed method with the recent transformer-baat based baseline ONO (Xiao al., 2023).
4. Numerical Experiments
This section presents a comprehensive empirical evaluation of LaMO against several neural operator baselines. We con- duct extensive experiments on diverse standard benchmarks to validate the proposed method.
4.1. Implementation Details
Benchmark: We evaluate LaMO's performance on reg- ular grids using the Darcy and Navier-Stokes (Li al.,
6

---

Latent Mamba Operator for Partial Differential Equations
Relative L2 (x10 -3
20
15
5-
5
4
5
2
3
400
800
1,000
85
85"106"111
211
421a
8
12
16
20
32.
64
128
256
Number of Training Samples
Resolution
Number of Layers
Embedding Dimension
LaMO (Ours)
Transolver
Figure 2. Model performance on the scalability of the Darey Flow benchmark evaluated across various aspects: (Left) Data Efficiency, measuring performance with varying amounts of training data; (Middle Left) Resolution, assessing the impact of different input spatial resolutions; (Middle Right) Model Depth, analyzing performance with increasing layers; (Right) Embedding Dimension, examining the effect of varying latent space dimensionality. Lower relative I2 error (x10"* indicates better performance.
2020) benchmarks. We then extend our experiment to ir- regular 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.
4.2. Main Results
We benchmark LaMO against standard baselines and various latest SOTA neural operators (Wu et al., 2023; 2024; Hao et al., 2023), ranging from frequency-based to transformer-based neural operators. Table shows that LaMO outperforms these baselines across diverse physics domains, including solid and fluid dynamics on various geometries. On average. LaMO achieves 32.3% im- provement over the second-best baseline, with remark- able 49% gain time-dependent PDEs like Navier-Stokes (0.095 - 0.046) and Plasticity (0.0013 - 0.0007) bench- mark. It demonstrates LaMO's superior ability to han- dle complex dynamics compared to transformer- -based ap- proaches. While transformer-basee models such as GNOT and Oformer apply attention directly to mesh points, they struggle capture complex fluid interactions, evidenced by the Darcy and turbulent Navier- flows (Foias et al., 2001; Bungartz & Schäfer, 2006). In contrast, LaMO's 's use of latent tokens and SSM proves more effective for model- ing the intricate dynamics of regular fluid datasets, despite Transolver addressing some of these issues with physics- informed attention. These results highlight the effectiveness of framework in operator learning for learning accurate surrogate solution approximation.
Baselines: We evaluate LaMO against over 15+ baselines, ranging from frequency-based operators such FNO (Li al., 2020), U-FNO (Wen et al., 2022), WMT (Gupta et al., 2021), F-FNO (Tran et al., 2021), U-NO (Rahman et al., 2022) to transformer-basse operators such as GalerkinTrans- former (Cao, 2021), HT-Nee (Liu et al., 2022), GNOT (Hao et al., 2023), Factformer (Li et al., 2024), OFormer (Li et al., 2022b), ONO (Xiao et al., 2023), and Transolver (Wu al., 2024), representing SOTA neural operator.
Evaluation Metric (Relative 12) (Li et al., 2020): Mean Rel- ative E2 eтror is used as metric throughout the experiments.
N (Go(a)) (t))ll C
(21)
(a:(2l)
The regular mean-squared error (MSE) is enhanced with normalizer [[[t (a.)ll2 take account of discrepancies in absolute resolution scale across different benchmarks.
Implementation Details: We use the mean relative & error (Equation 21) as the training and evaluation metric. All mod- els are trained for 500 epochs with the AdamW (Loshchilov et al., 2017) optimizer and OneCycleLR scheduler (Smith & Topin, 2019). To ensure a fair comparison, we have kept our model parameters equal or fewer than those of the transformer-bas based baselines. Further details on implemen- tation and hyperparameters are provided in the appendix. Experiments are conducted on Linux machine with Ubuntu 20.04.3 LTS, an Intel(R) Core(TM) 19- 10900X processor, and a single NVIDIA A100 40 GB GPU. More comprehen- sive details are available in the appendix Section D.
4.3. Ablation Results
Ablations: In addition to the main result, we carry the ablations of the components in LaMO. Table 2 highlights the impact of SSM's varying numbers of latent tokens and state dimensions (DState) on model performance. Increasing the token number consistently yields lower errors, indicating that finer partitioning provides finer input representation. Increasing DState dimensions generally improves perfor- mance, particularly with larger latent token sizes, though it shows diminishing returns with decrease in latent to- kens. The best performance (0.0038) on Darcy is achieved

---

Latent Mamba Operator for Partial Differential Equations
with increasing the latent tokens and (DState 64) with expand dim 2, underscoring the importance of balancing fine-grained latent tokens and is crucial for achieving better performance with balanced computation.
Table 3. Ablation on the number of latent tokens the Airfoil benchmark of LaMO compared with Transolver. Lower relative l2 error indicates better model performance.
RELATIVE L2
ABLATIONS
TRANSOLVER LAMO (OURS)
Table 2. Ablation Study on a Number of Latent Tokens and SSM Parameters: We evaluate three variants on the Darcy, varying the latent token number, SSM dstate, and SSM expanding dimensions. Lower relative 2e error indicates better model performance.
0.0085
0.0079
8
0.0058
0.0050
16
0.0057
0.0048
RELATIVE L.2 Loss
32
0.0055
0,0045
NUMBER
64
0.0054
0,0041
LATENT TOKENS
STATE DIMENSIONS OF SSM
OF LATENT
96
0.0053
0.0040
TOKENS
128
0.0052
0.0039
16
32
64
256
0.0058
0.0043
1936
0.0045
0.0044
0.0041
0.0040
512
0.0059
0.0048
EXPAND DIM
484
0.0053
0.0052
0.0047
0.0056
121
0.0068
0.0065
0.0070
0.0071
olutions, depths, and embedding dimensions. As shown in Figure 2(a), LaMO achieves superior performance on the Darey flow benchmark even with 40% of the training data, demonstrating its remarkable efficiency in low-data regimes compared to second- best baseline models (Wu et al., 2024). This demonstrates LaMO's efficiency in learning meaningful representation with fewer data points, making it ideal for real-world applications with limited data. Addi- tionally, we evaluated LaMO under different benchmark res- olutions (Figure 2 (b)), depth (Figure (c)). and embedding dimensions (Figure (d)). The results indicate that the ini- tial LaMO configuration maintains consistent performance across varying resolutions, depths, and embedding dimen- sions compared with the second-best operator, Transolver. These findings suggest that LaMO exhibits robust scalability properties and can serve as a large-scale pre-trained PDE solver, paving the way for its use as a foundation model in SciML applications. The appendix Subsection E.5 presents more experiments across the benchmark.
1936
0.0041
0.0042
0.0039
0.0038
484
0.0049
0.0048
0.0051
0.0047
121
0.0065
0.0067
0.0066
0.0066
Ablation Study on Bidirectional SSM: To further assess the necessity of bidirectional SSM, conducted experiments with unidirectional SSM, which failed to perform consis- tently across benchmarks (on Plasticity, it becomes worse as 0.0007 + 0.0022). Thus confirming the significance of capturing dynamics in non-causal setting. More Ablation can be found in the appendix, Subsection E.1.
Effeec of Non-Periodic Boundary Condition: LaMO per- forms well challenging non-periodic boundary bench- marks, including Darcy flow and Navier-Stokes, as evi- denced by the results in Table 1. It demonstrates its ro- bustness in handling complex real-world scenarios where traditional periodic boundary assumptions fail.
Effect of the Number of Latent Tokens:
As shown in Table 3, we systematically increase the number of latent tokens in the Airfoil benchmark. The results indicate that higher number of latent tokens allows both Transolver and LaMO to capture finer details in the physical space, albeit at the cost of increased computation. Moreover, performance ini- tially improves with an increasing number of latent tokens. Eventually, it deteriorates beyond a certain threshold, sug- gesting the existence of an optimal number of latent tokens that balances accuracy and computational efficiency. LaMO consistently outperforms Transolver ver across all tested latent token configurations, demonstrating its robustness and supe- rior capability in capturing underlying physical dynamics using LaMO.
Visual Demonstrations: To visualize the performance of LaMO compared to Transolver, we plot heatmaps in Figure 3 for an intuitive performance comparison. It is observed that LaMO significantly outperforms Transolver in captur- ing medium boundaries in Darey flow problems (Figure 3())). Specifically, LaMO excels at resolving discontinuity boundaries and junctions with fewer artifacts than Tran- solver (Figure 3(b).. Furthermore, in time-dependent PDEs, such as the Navier-Stokes equations, LaMO demonstrates superior performance in accurately capturing turbulent fluid flow dynamics (Figure 3(c)). It handles boundaries and medium transitions more effectively, showeasing its robust- ness over Transolver in such scenarios. The appendix Sec- tion F presents more heat map plots.
Performance on Turbulent Fluid Dynamics: As shown in Table 1, LaMO demonstrates superior performance on Navier-Stoker (Reynold number 10*), a highly turbulent flow, achieving an average gain of 51.9% over the second- best baseline. For lower turbulence (Reynold number 10'),
4.4. Model Analysis
Scaling:
To further investigate the scalability of LaMO as a potential foundation model for PDE solvers, we analyzed its performance under varying training data samples, res-
8

---

Latent Mamba Operator for Partial Differential Equations
0,0025
0.0025
0.2
Erro
0.0004
0.0002
0.1
0.0000
0.0
-0.1
0.0004
Enror(uurs)
0.0002
0.0000
-0.0002
-0,0004
0.2
(b)
(c)
Figure Visual Comparison. The (Top Row) displays the ground truth, the (Middle Row) presents the heatmap of Transolver, and the (Bottom Row) presents the error hcatmap for LaMO on (a) Darcy Flow, (b) Plasticity, and (c) Navier-Stokes benchmark.
LaMO achieves relative l2 error improvement of 75% (0.0454 + 0.0117) compared to Transolver. Thus validat- ing it as consistent across the turbulent regions.
GNOT and ONO, LaMO exhibits even more significant improvement, requiring 5-7x fewer parameters and achiev- ing 3x speedup in running time. LaMO presents favorable efficiency considering the running time, GPU memory con- sumption, and model parameters compared to Transolver, and achieves better results. These results highlight the efficiency of LaMO in solving PDEs.
Using latent tokens formed on a regular grid, in Perceiver, we observed that while bidirectional SSM marginally improved performance 0.0059 > 0.0050, compared to patches (0.0039), it lacked scalability with data resolution, with increased computation. It highlights the significance of ViT-type latent tokens in scaling operators to higher resolutions, making them an optimal choice for foundational models, as demonstrated in transformer-bant foundation models (Poseidon (Herde et al., 2024)).
Latent Tokens:
Interpretability: To further interpret the performance, we analyze inter-cosine similarity across layers. As shown in the appendix Subsection E.6, LaMO demonstrates a reduc- tion inter-cosine similarity across layers compared to Transolver, indicating better representation learning. This improvement be attributed to the absence of the softmax operation, which in Transolver has been identified as caus- ing over-smoothing issues affecting its ability to capture the intrinsic relationship (Ali et al., 2023; Wang et al., 2022).
10
Parim Count
Params Court
ONO
3,0D
ONO
10.0MB
10,0BB
GPU Menory (GB)
GNOT
5. Conclusion and Future Work
Oues
OFormer
GNOT
In summary, introduce new approach to solving PDEs introducing the SSM-based Neural Operator. The pro- posed method leverages SSMs achieve superior perfor- mance while ensuring practical efficiency. The theoretical analysis reveals the equivalence of LaMO with kernel inte- gration. Through extensive experiments, we demonstrate that SSM-based operators outperform transformer-bee counterparts in performance and computational efficiency, While the current work establishes the potential of SSM- based neural operators, certain aspects remain unexplored. For instance, compatibility with pretraining paradigms us- ing SSM as a backbone has yet to be explored. Future work aims to extend the application SSM-based operators as foundation models for PDE solutions.
Transolver
Transolver
2
Galertan
Ours
Galerkin
10
15
20
25
30
35
40
20
30
40
50
GD
70
Training Time (s(epoch)
Training Time (skepoch)
Figure Efficiency comparison of top five baselines on (a) Darey and (b) Airfoil benchmark per epoch, respectively.
To further analyze the performance of the pro- posed model, we present its efficiency metrics in Figure 4 compared with transformer-bant based baselines. Specifically, on the Darcy flow and Navier-Stokes benchmark, LaMO uti- lizes 3x fewer parameters and is 1.8x faster than transformer- based baselines. Compared to other baselines, such as
Efficiency:

---

Latent Mamba Operator for Partial Differential Equations
Acknowledgement
Cao, S. Choose transformer: Fourier or galerkin. Ad- vances in neural information processing systems, 34: 24924-24940, 2021.
This work was partly supported by the Indian Institute of Science under a start-up grant to set up the GPU computing infrastructure. Prathosh was supported by the Infosys Foun- dation Young Investigator Award. Anoop acknowledges support from the Alexander von Humboldt Foundation and the Google Research Scholar Award. The Government of India supported Karn through the Prime Minister's Research Fellowship (PMRF).
Costa, B. Spectral methods for partial differential equations. CUBO, A Mathematical Journal, 6(4):1-32, 2004.
Dao, T. and Gu, A. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, arXiv preprint 2rXiv:2405.21060 2024.
Debnath, L. and Debnath, L. Nonlinear partial differential equations for scientists and engineers. Springer, 2005.
Impact Statement
El-metwaly, A. and Kamal, M. A brief review of numeri- cal methods for solving the boundary value problems of pde. In The International Conference on Mathematics and Engineering Physics, volume 10, pp- 1-7. Military Technical College, 2024.
This paper introduces a new neural operator to advance deep learning research in solving PDEs. The method leverages state-space models to capture correlations for long-range de- pendencies, addressing the limitations of transformer-base -based operators. Our model demonstrates superior performance in PDE tasks, showing potential for industrial applications and scalability for foundation models. This includes appli- cations for several practical large-scale problems, including, but not limited to, weather forecasting, pollution forecasting, and electronic structure, to name few. The focus remains on scientific challenges, with no identified ethical risks in this work. Thus, we believe there are no potential ethical risks our work in society.
Elfwing. S., Uchibe, E., and Doya, K. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018. Erol, M. H., Senocak, A., Feng, J., and Chung. J. S. Audio mamba: Bidirectional state space model for audio rep- resentation learning. arXiv preprint arXiv:2406.0334 03344,, 2024.
References
Fanaskov, V. and Oseledets, I. Spectral neural operators.
arXiv preprint arXiv:2205. 10573, 2022.
Alexey, D. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv: 2010. 2020.
Foias, C., Manley, O., Rosa, R., and Temam, R. Navier- Stokes equations and turbulence, volume 83. Cambridge University Press, 2001.
Ali, A., Galanti, T., and Wolf, L. Centered self-attention layers. arXiv preprint 100.:::3006.160100 2023.
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXi::2312.00752 2023.
Alonso, C. A., Sieber, J., and Zeilinger, M. N. State space models as foundation models: A control theoretic overview. arXiv preprint arXiv:2403. 16899, 2024.
Gu, A., Goel, K., and Ré, C. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021a
Bedi, S., Tiwari, K., P., P. A., Kota, S. H., and Krishnan, N. M. A. Conoair: A neural operator for forecasting carbon monoxide evolution in cities, 2025. URL https: //arxiv. , org/ abs 2501. 06007.
Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A., and Ré, C. Combining recurrent, convolutional, and continuous-tin time models with linear state space layers. Advances neural information processing systems, 34: 572-585, 2021b.
Bevanda, P., Sosnowski, S., and Hirche, S. Koopman oper- ator dynamical models: Learning, analysis and control. Annual Reviews in Control, 52:197-212, 2021.
Guan, S., Hsu, K.-T., and Chitnis, P. V. Fourier neural oper- ator networks: A fast and general solver for the photoa- coustic wave equation. arXiv preprint 108.v:2108.000 09374, 2021.
Buitrago, R., Marwah, T., Gu, A., and Risteski, A. On the benefits of memory for modeling denendepnddent pdes. In The Thineenth International Conference on Learning Representations.
Guibas, J., Mardani, M., Li, Z., Tao, A., Anandkumar, A., and Catanzaro, B. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint 1rXiv:2111.1335 2021.
Bungartz, H.-. and Schäfer, M. Fluid-structur interaction: modelling. simulation, optimisation, volume 53. Springer Science & Business Media, 2006.
10

---

Latent Mamba Operator for Partial Differential Equations
Gupta, G., Xiao, X., and Bogdan, P. Multiwavelet-based operator learning for differential equations. Advances in neural information processing systems, 34:24448-44062, 2021.
Lei Ba, J., Kiros, R., and Hinton, G. E. Layer normaliza- tion. ArXiv e-prints, pp. arXiv-1607, 2016.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:2010.08895 2020.
Hao, Z., Wang, Z., Su, H., Ying, C., Dong, Y., Liu, S., Cheng, Z.,, Song, J., and Zhu, J. Gnot: A general neu- ral operator transformer for operator learning. In Inter- national Conference on Machine Learning, pp. 12556- 12569. PMLR, 2023.
Li, Z., Huang, D. Z., Liu, B., and Anandkumar, A. Fourier neural operator with learned deformations for pdes on general geometries. arXiv preprint arXiv:2207.05209, 2022a.
Hao, Z., Su, C., Liu, S., Berner, J., Ying, C., Su, H., Anand- kumar, A., Song. J., and Zhu, J. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre- training, arXiv preprini 243::22403.03542, 2024.
Li, Z., Meidani, K., and Farimani, A. B. Transformer for partial differential equations operator learning. arXiv preprint arXiv:2205. 13677,, 2022b.
He, K., Zhang, X., Ren, S., and Sun, Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Li, Z., Shu, D., and Barati Farimani, A. Scalable trans- former for pde surrogate modeling. Advances in Neural Information Processing Systems, 36, 2024.
Herde, M. Raonić, B., Rohner, T., Käppeli, R., Moli- naro, R., de Bézenac, E., and Mishra, S. Poseidon: Efficient foundation models for pdes. arXiv preprint 2911::2405./9001, 2024.
Liu, X., Xu, B., and Zhang, L. Ht-net: Hierarchical trans- former based operator learning model for multiscale pdes. arXiv preprint arXiv:2210. 10090000. 2022.
Liu, Y., Tian, Y., Zhao, Y. Yu, H., Xie, L., Wang, Y., Ye, Q., and Liu, Y. Vmamba: Visual state space model, 2024. URL https:///// //t:://rraa //tp///avvvv///////////e.//aaav/ org/ abbs/241.. 10.01.10166 10166.
Higham, N. J. Analysis of the cholesky decomposition of semi-definite matrix. 1990.
Hu, Z., Daryakenari, N. A., Shen, Q., Kawaguchi, K., and Karniadakis, G. E. State-space models are accurate and efficient neural operators for dynamical systems. arXiv preprint ar arXiv:2409.032. 2024.
Liu, Z., Lin, Y., Cao. Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021.
Jaegle, A., Gimeno, F., Brock, A., Vinyals, O., Zisserman, A., and Carreira, Perceiver: General perception with it- erative attention. In International conference on machine learning, pp. 4651 1-4664 PMLR, 2021.
Loshchilov, I., Hutter, F., et al. Fixing weight decay regu- larization in adam. arXiv preprint arXiv:1711.05100 5. 2017.
Karniadakis, G. E., Kevrekidis, I. G., Lu, L., Perdikaris, P., Wang, S., and Yang, L. Physics-informe machine learning. Nature Reviews Physies, 3(6):422 440, 2021.
Lu, L., Jin, P., Pang. G., Zhang. Z., and Karniadakis, G. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, ((3):218-229) 2021.
Katsch, T. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:2311.0192 2023.
Qu, H., Ning, L., An, R., Fan, W., Derr, T., Liu, H., Xu, X., and Li, Q. A survey of mamba. arXiv preprint urXhv:2408.01222 2024.
Kovachki, N., Li, Z., Liu, B.. Azizzadenesheli, K., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Learning maps between function spaces. arXiv preprint arXiv:2108.0844 08481, 2021.
Rahman, M. A., Ross, Z. E., and Azizzadenesheli, K. U-no: U-shaped neural operators. arXiv preprint arXiv:2204. 11127, 2022.
Kurth, T., Subramanian, S., Harrington, P., Pathak, J., Mardani, M., Hall, D., Miele, A., Kashinath, K., and Anandkumar, A. Fourcastnet: Accelerating global high- resolution weather forecasting using adaptive fourier neu- ral operators. In Proceedings of the Platform for Ad- vanced Scientific Computing Conference, pp. 1-11, 2023.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention-MICCAA 2015: 18th International Confer- ence, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234 241. Springer, 2015.

---

Latent Mamba Operator for Partial Differential Equations
Smith, L. N. and Topin, N. Super-convergence: Very fast training of neural networks using large learning rates. In Artificial intelligence and machine learning for multi- domain operations applications, volume 1 11006, pp. 369- 386. SPIE, 2019.
Solín, P. Partial differential equations and the finite element method. John Wiley & Sons, 2005.
Tiwari, K., Krishnan, N., and Prathosh, A. Cono: Complex neural operator for continous dynamical physical systems. arXiv preprini arXiv:2406.02555 025977, 2024.
Tran, A., Mathews, A., Xie, L., and Ong, C. S. Factorized fourier neural operators. arXiv preprint arXiv:2111.13802. 2021.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wang, P. Zheng, W., Chen, T., and Wang, Z. Anti- oversmoothing in deep vision transformers via the fourier domain analysis: From theory to practice. arXiv preprint arXiv:2203.05962 2022.
Wen, G., Li, Z., Azizzadenesheli, K., Anandkumar, A., and Benson, S. M. U-fno- an enhanced fourier neural aperator-based deep-learning model for multiphase flow. Advances in Water Resources, 163: (41880, 2022.
Williams, R. L., Lawrence, D. A., et al. Linear state-space control systems. John Wiley & Sons, 2007.
Wu, H., Hu, T., Luo, H., Wang, J., and Long, M. Solving high-dlmensional pdes with latent spectral models. arXiv preprint ar arXiv:2301. 12664, 2023.
Wu, H., Luo, H., Wang, H., Wang, J., and Long, M. Tran- solver: fast transformer solver for pdes on general geometries. arXiv preprint 4rXIv:2402.023666 2024.
Xiao, Z., Hao, Z., Lin, B., Deng, Z., and Su, H. Improved operator learning by orthogonal attention. arXiv preprint arXiv:2310.1 12487, 2023.
Zheng, H., Nie, W., Vahdat, A., Azizzadenesheli, K., and Anandkumar, A. Fast sampling of diffusion models via operator learning. In International Conference on Ma- chine Learning, pp. 42390-42402. PMLR, 2023.
7hu, I.,, Liao, B., Zhang, Q. Wang, X., Liu, W., and Wang, X. Vision mamba: Efficient visual representation learn- ing with bidirectional state space model. arXiv preprint ADt.22401.09117. 2024.
12

---

Latent Mamba Operator for Partial Differential Equations
A. Table of Notations
Table 4 provides a comprehensive list of notations used throughout the main content for clarity and reference.
Table 4. The following table summarizes the notations and symbols used throughout the main text for consistency and clarity.
NOTATIONS
DESCRIPTIONS
LaMO
Latent Mamba Operator
PDEs
Partial Differential Equations
ODEs
Ordinary Differential Equations
SSMs
State Space Models
ZOH
Zero Order Hold
Linear D C Rd
Feedforward Neural Network
Spatial Domain for the PDE
Spatial Domain Points
E A (D; Rd,)
Input Functions Coefficient
de
Dimension of Input Function a(x)
uEU = (D;
Target Functions Solution
da
Dimension of Output Function u(x)
Dj
The Discretization of (a,, Uy)
gi, AA
Solution Operator
Go
Neural Operator
P
A Probability Measure on A
C
Loss Function
x(() E R
System Input Sequence
h(t) E RN
System State
v(t) E RM
System Output Sequence
2[k] E R
Discrete Input Sequence
[[]] E RR
Discrete State Discrete Output Sequence System Matrix in Continuous SSM
v[k] R
A E R
B E RN
Input Matrix in Continuous SSM
C E TR. M
Output Matrix in Continuous SSM
D E RMXH
Direct Transition Matrix in Continuous SSM
Á E R
System Matrix in Discrete SSM
E N
Input Matrix in Discrete SSM
C E M
Output Matrix Discrete SSM
D M
Direct Transition Matrix in Discrete SSM
A E RR+
Discrete Time Step in Discrete SSM
K
State Kernel in Convolutional SSM
Kernel Integral Operator
X
Physical Tokens
Latent Tokens
N
Number of Training Samples
B. Theoretical Insights
This section delves into rigorous mathematical analysis of the state-space model (SSM) from the perspective of a neural operator. We provide proof its core properties and present new insights. This comprehensive examination will highlight the theoretical foundations and unique advantages of leveraging SSMs neural operators. We abuse the notations when there is no misleading context.
Lemma B.1. (Williams et al., 2007) For any differential equation of the following form
h'(t) Ah(t) + Bx(t),
(22)
13

---

Latent Mamba Operator for Partial Differential Equations
Its general solution is given as follows:
h(t) = ((lt-ba)((to)) + [ edlt-s) Bx((s) ds.
(23)
Proof: Consider the n-dimensional dynamical system represented by the equation as follows:
(1()) Ah(t) + Bx(t),
(24)
where htt) E IR" is the state vector, A E R is the system matrix, B E Inxm is the input matrix, and x(t) E Rm
represents the input. Rearranging the terms, we get the following:
(t) Ah(t) = Bx(().
(25)
By multiplying both sides by the integrating factor e we obtain the following:
(t) c Ah(tA = - Br(().. d (e (E- = =
tA
(26)
(27)
di
Integrating both sides with respect to time over the interval [0, t). we have the following:
L-¹ d (e *A **\(s)) ds ( E - Br(s)ds.
(28)
ds
e e=()((t)= h(0) S e c- Bx()sss
(29)
htt) = c ((0) / Alt-s) Br(s)ds.
(30)
This can be rewritten in convolutional form follows:
M(I) = e ((0) + (u (1),
(31)
where u({) = Ae represents the impulse response function.
Thus, the above equation provides the analytical solution to the r-dimensional state-space equation.
Remark B.2. is important to note that the solution consists of two components:
Initial Response: (l)) which depends on the initial state.
Forced Response: f/r A((- Ba(s)ds, driven by the input st)).
Both components involve the exponential matrix function, which encapsulates the system's 's dynamics, and B denotes the kernel of the SSM.
Assumption B.3. We discretize the input signal to incorporate the State Space Model (SSM) into neural network to obtain discrete parameter representation equivalent to the continuous SSM. Specifically, the signal u(t) is sampled at uniform intervals of A, such that the time is represented as iA,, where k = 0, 1.... is non-negative integer. Utilizing the Zero-Order Hold (ZOH) method, we assume that the signal remains constant within each sampling interval, i.e.. u(f) = u(kA) for t [KA, (k + 1)A].
Proposition B.4, The Zero-Order Hold (ZOH) discretization of continuous-time system parameters can be represented as:
A = AA
(32)
B = (AAA -I)A
where A denotes the discretization time step, I the identity matrix, and A, B are the continuous- time system matrices.
14

---

Latent Mamba Operator for Partial Differential Equations
Proof: From the previous proposition, the general equation of the continuous state-space model (SSM) is given by:
h(t) h(tt)((tt) + f (x(ss))((ss) Br((s) ds.
(33)
Assuming a zero-order hold (ZOH) on the input 2(55 between & [K:A, (k + 1)4), the system evolves as:
(K+1) ((R+1)A-5)) Bookk)
M(K + 1)2) = (AAA(A)) +
(34)
By defining hi(k 1)4) hk41, h(kA) h2s and x(kA)( Tk. the equation simplifies to following:
h441 c (Fr
(k+t))
(((++1))-s)) B ds Fk-
(35)
This corresponds to the standard discrete state-space model with parameters as follows:
(k+1)4
Ā= A B
((K+1)A-s) B ds.
(36)
Now, assuming A is invertible, simplify B as follows:
B
Bds..
(37)
d
B =
A
(E
As)
ds B.
(38)
J0
ds
B= (e B)B
(39)
Thus, under the assumption of being invertible, the continuous SSM can be discretized with the parameters:
A A
B -6 I)B.
(40)
Remark B.5. Further B AA B)B can be further simplified by truncating higher order terms as follows:
B=
(41)
Corollary B.6. The discretization method described above aligns with the Euler method, which approximates the matrix
exponential by truncating its Taylor series expansion the first-order
Proof: By substituting the discrete parameters into the SSM equation, derive the following:
h(t At) Ah(t) + Bx(t),
(42)
-AAIh(t) + (t)
(43)
-I AAA + ((AP)) + AtBr(t),
(44)
-A(t) ALAN(!) + AIBa(I),
(45)
A((t) - At(AA()) + Br(t)),
(46)
= A(t) (1)
(47)
The above illustrates that the update step corresponds the Euler method's first-order approximation of the dynamics. Remark B.7. The above corollary establishes the equivalence between the Euler and Zero-Order Hold (ZOH) methods in solving PDEs. While the Euler method truncates higher-order terms from the Taylor series, the ZOH method retains them, making it a more generalized and accurate approach for solving PDEs. This distinction is crucial for understanding SSMss performance from the Fuler method's perspective.
15

---

Latent Mamba Operator for Partial Differential Equations
Lemma B.8. (Williams el al., 2007) Any continuous nonlinear differentiable dynamical system described by
h(t) = (t)))
(48)
can be approximated by its linear state-space model (SSM) representation
h(t) Ah(t) Bx(() ohh,
(49)
where the Jacobian matrices and B are given by:
df ((((),(()())) B dr
of (h(1), ((),
(50)
dh
and O(h,z) x) represents higher-order infinitesimal terms.
Proof: Using the Taylor series expansion around (())) we expand fhh(t), x(),, t) as follows:
05 (h)) -(t)) +
0s
=))
+ Öh
-((t) HOT,
(51)
Or
where HOT denotes higher-order terms.
By defining the Jacobian matrices as follows:
A(I) = of (h(), ah
B(I) of (h(1), dr
(52)
and rearranging terms, we obtain the following:
i()) Ah(() Bx((1) Ohh,)).
(53)
Thus, a linear SSM with higher-order corrections can approximate continuous nonlinear differentiable dynamics. Remark B.9. (Alonso et al., 2024) A nonlinear continuous differential system with the above dynamics has long-range memory, i.e., it captures information from past inputs if all eigenvalues of A are within the unit circle. Mathematically. it should satisfy the following conditions:
eii((A) < and leig(A))
Veig(A).
Lemma B.10. The SSM operator is Monte Carlo approximation of an integral operator.
Proof: Let (S, F,) be probability space, where Ω is measurable space equipped with probability measure Let
Q + RC be function in the Hilbert space (2, p1;;;;; R Let define, the integral operator G L μé; RC) + (2 p1 C) as follows:
(a))(a) J. n(x, y)u(y) p(dy),
(54)
where x Q Ω R is a measurable kernel function.
Now let's discretize the domain, i.c., consider partition of Ω into N distinct mesh points {u,) in particular order sequence, and approximate the measure y by uniform measure over these points. The integral operator G can then be approximated via Monte Carlo integration as follows:
Ω
((u)(x)
M(,,))(MM))
(55)
where [2] = =o2 du denotes the total measure of Ω.
In SSM, the integral kernel x is parameterized as follows:
<(y.,)) = Wo(l))
II
W Ww())
(56)
Relative Pocition
16

---

Latent Mamba Operator for Partial Differential Equations
the relative position denotes the matrix exponential of A depends the position, yk y/ denotes the relative distance between the tokens and W and W B denotes the SSM input dependent weights.
Substituting this parameterization into the Monte Carlo approximation, we obtain the following:
N
(z))())
W (w)
II Uk F.
W W
(57)
This expression demonstrates that the SSM operator can be interpreted as Monte Carlo approximation of the integral operator G, where the state-space dynamics parameterize the kernel
Lemma B.11. (Wu et al., 2024) St is countable domain, the latent domain Q2 is isomorphic to Ω
Proof: Let x; E 52 denote the i-th element in the input domain, and let 2j represent the j-th latent token in the latent domain Ωg The latent weight of x; with respect 2j is denoted by Wa,,25 R. Given constant K 2 1 and K E N, define projection f Ω - 94 as follows:
f(x:) arg max Way.yy
(58)
23
where the mapping is constrained such that:
()/Kj K <j S ((i +1) K,
(59)
and the j-th latent token zy has not been assigned projection previously.
This construction guarantees a bijective mapping between elements of the input domain Q and latent token in the latent domain Q2, ensuring that the cardinalities of the two domains are equivalent. Hence, f2 is isomorphic to Q, i.e. Ω ~ 2g
Using the above results, we will show that SSM an integral kernel on Ω similar to transformer-hased operator (Wu et al., 2024; Cao, 2021; Wu et al., 2023).
Theorem B.12 (SSM as an equivalent integral kernel on ?). Let Ω C R be bounded domain, a Ω: -> Re be a given function and x E Ω be mesh point. An SSM (Structured State-Space Model) layer approximates the linear integral operator g (22 Rd) > (12, R.d), defined as follows:
((a)(x) / <<(x,, y) aly) dy.
(60)
where K x Q Rdxd the kernel function characterizing the operator
Proof: According to Lemma A.3, we can obtain an isomorphic projection f between a countable input domain Ω and the latent domain Ωg Suppose that the latent weight w.,. Ω Ω2 - RR is smooth in both Ω and Ω2 where St and Ωg denote the continuation Ω and Ω2, respectively. Then, we can obtain f as diffeomorphism projection.
Then, we define the value function u, on the latent token domain Ω2 as follows:
u(z)) = / We u(z) dæ
(61)
which corresponds the latent token definition Eq. (2). Based on the above assumptions and definitions, we have:
17

---

Latent Mamba Operator for Partial Differential Equations
((u)(x)= =
k(a, v)u(v) dy
(62)
Kma (2, ya) (16) of (ya)
(.ma(:))): S2 - RCxC is kernel function)
(63)
Bu
Kizz (1, (1, (1b2) det V vn.f (
(64)
Jaa sa,
(o. Wz. (v': 1b) u.)dd/ us (16) ((=))()))))
(65)
(Jo, Wainttin (V': yo duy uz (Uz) det( v. (d)
(66)
Wz.8/ Decader
Fozz (Ues 3ka)
U2 (312)
VV.ffV..ft g (wa))d
(Lemma B.10)
(67)
St.
among latent
Latent token
M
M
N
W
(Wouz(Zas))
II
(i) (WBus(Ze.))
Σ
(pxp)(x))
(68)
(=1
Eq.
(15)
(13)
Fg. (10)
M
M
Wi
C
II 1St
Diag A)
(Br)
Zt,
(69)
where Niz denotes the kernel function between mesh points and latent tokens, and Kiz is the kernel defined among latent tokens. For simplicity, we take y points, the usage of Lemma B.10 here is based on the Monte-Carlo approximation in the latent domain.
1 for simplification. Different from the kernel integral among mesh
Remark B.13. The kernel function <(2, y) encapsulates the interactions between the evaluation point and the domain points y E Ω. The SSM layer implements numerical approximation of the integral operator using finit-ddimensional representation of the input function a over the discretized domain Ω. Specifically, the SSM architecture constructs learned representation of x and combines it with the discretized function to approximate the integral.
In the subsequent section, we comprehensively analyze the learning mechanism underlying the State-Space Model (SSM) kernel and explore its theoretical and empirical connections to transformer-base architectures.
Matrices Structures State Space models: Given an input sequence := [T,, ,II] E LX consisting of L feature vectors, the hidden state each vector is denoted as follows:
BII1,
(70)
h2 = A2h + B282 = A2 + B2
(71)
hg Azhi2 B333 A1 + Ag + B3g3,
(72)
h1 = A BLd - ... B.B A, A ABB2a2
+ BLØL.
(73)
The above recurrence can be rewritten matrix form as:
B
0
F1
A₂B
B2
0
T2
H [h,, 12, ha, ... Jr.]
A3A2B1
Bg
B3
0
T3
(74)
1| A,
11 A,
11 A,B3
Br.
[L
18

---

Latent Mamba Operator for Partial Differential Equations
For the output sequence Y W/, 1. vr]" each vector y; ,,.", expressed as follows:
y C,h
(75)
where and in matrix form, the above equation can be compactly rewritten as follows:
C
h1
C2
hv
Y
C3
0
ha
CH.
(76)
CL
h
Substituting // from the Eq, we obtain the following relation:
Bi
211
Ag
B2
0
X22
Y=C
A1A2B1
AgB2
B3
33
(77)
T A,B1
A,
BL_
[
In general, = -, we have following recurrence for hidden states and output.
h4 - Σ Σ II A,
(78)
j=lkijt
C. Σ
II
A B,aj.
(79)
j=1k=jt1
which can be further simplified and written as follows:
C,B
0 C₂B₂
0
CrArBI
0
0
22
Y
C1AgABBI
CzAgB2
CgB3
373
(80)
Cz 11 A,B,
CL
A,B2
Cr 11 A,Bs
Ct Br
This can be compactly written as:
= = SSM(A,
(81)
where M e Lx represents the matrix. Hence, the M can be seen as a controlled linear operator. The above Equation can be seen as a variant self-attention, specifically, the casual version attention. The element of M are as follows:
Mc C, II ABB, kmj+1
(82)
Remark B.14. Linear SSM is linear shift-invariant and thus is causal system.
Similarity with Attention: The input sequence X E Lx are linearly projected into Query Q E RLXD Key K E RLXD and Value E L
Q = = K= XW V= =
(83)
with Wo, WK E DDDD and Wv E DDDD as the corresponding weight matrices. More specifically, Q := qq, (q,...,, qr) = ki..... ......,,, kg] V [Vi,...,,Vu where qi k,, v; (for I,.... I.) represent the query, key, and
19

---

Latent Mamba Operator for Partial Differential Equations
value vectors, respectively, for input 24. Based Q and K, the attention matrix S IR contains the correlations between all query and key vectors, with softmax function applied each row of which reduce it to transition matrix (row stochastic matrix):
QK
S = softmax
(84)
VD
where the above matrix can be written as follows (for clarity, we have ignored the normalization constant V(d):
exp(qi ki)
explan-kr
explq kz)
exp(qrkk)
explg .k))
exp(q) kk)
S
(85)
exp(gz- ki)
kxp
exp k..)
exp(qr.k)
explar.k) kj)
exp(gr.ky)
Each element S (for i.j ,,.., L) represents the attention score between qi and k,.
For the output sequence Y := 1' :ur]"" is calculated based on Attention Matrix (S) follows:
==
(86)
which can be rewritten term of as follows:
Y= SXWT follows that each output vector y, (for i 1, 1..... L) can be written in vector form as:
(87)
y:
Vi
(88)
which implics the output yi is linear combination of vectors vy(i = I.... L) with coefficient S. The greater the attention score, the larger its influence on the output. Remark B.15. From both the matrix form of SSM and attention, we observe that both of them satisfy the pseudo-linear form, i. as follows:
Y Y= xXXXX
(89)
we can observe that both are linear equations, but the weight matrix is data-dependent.
Interpreting the Hidden Matrices of SSM: From Equation 82 of SSM, we have the element of the matrix as follows:
Mij C, II A B3
(90)
Discouut Factor
Let abbreviate Hag li Aş above equation be written as,
Mi, CjH.,
(91)
In Structured State-Space Models (SSMs), the softmax normalization typically used in transformers is removed. Additionally, SSMs incorporate masking mechanism that introduces input-dependent relative positional encodings (denoted H1,)), in contrast to transformers, where positional encodings typically fixed. This masking mechanism be interpreted as "discount factor" that modulates interactions based on the relative distance between positions and j. formulated as:
H II A k=j+1
(92)
which the context of attention mechanisms, this input-dependent positional mask plays a crucial role in encoding the "selectivity" of SSMs neural operator (dataddependent kernel), thereby influencing their ability to capture dependencies in modeling PDEs solution (Katsch, 2023; Dao & Gu, 2024).
20

---

Latent Mamba Operator for Partial Differential Equations
C. Details of Benchmark
This section provides a detailed overview the benchmark datasets and specific tasks associated with each benchmark.
FLUID PHYSICS
SOLID PHYSICS
Regular
Structured Mesh
Point Cloud
Darcy
Airfoil
Elasticity
OUTPUT: Fluid
INPUT: Airfoil Structure
OUTPUT: Airfoil velocity
Material
OUTPUT: Inner
Medium
Pressure
Structure
Structured Mesh
Navier-Stokes
Pipe
Plasticity
INPUT: Past Velocity
Pipe
INPUT: External
OUTPUT: Mesh Displacement
Velocity
Structure
Velocity
Force
Figure 5. The diagram presents an overview of the Neural Operator learning task benchmark PDEs classified as Fluid Physics and Solid Physics. (Top Row) showcases three specific PDEs: Darcy Flow, Airfoil, and Elasticity, (Bottom Row), an additional set three PDEs is shown: Navier Stokes, Pipe, and Plasticity.
Table benchmark details follow settings from (Li al., 2022a), with input-output resolutions presented as (temporal, spatial, variate). The "Y" denotes the ahsence of that dimension.
SOLID PHYSICS
FLUID PHYSICS
OVERVIEW
ELASTICITY
PLASTICITY
NAVIER-STOKES
PIPE
DARCY
PHYSICS
TASK
ESTIMATE STRESS
MODEL DEFORMATION MODEL VISCOUS FLOW
ESTIMATE VELOCITY
ESTIMATE PRESSURE
INPUT
MATERIILL STRUCTURE
EXTERNAL FORCE
PAST VELOCITY
STRECTURE
POROUS MEDRIM
OUTPUT
INNER
MESI DISPLACEMENT
FUTURE VELOCITY
MACH NUMBER FLUID VELOCITY
FLUID PRESSURE
GEOMETRY
POINT CLOUD
STRUCTURED MESH
REGULAR GRID
STRUCTURED MESH
REGULAR GRID
2D
TIME
DATA
TRAIN
1000
900
1000
1000
1000
1000
200
200
INPUT TENSOR
(1,972,)) 2)
(/, 101 x 31,2)
(10, x 6)
[/,221 x 51,2) (/, x 129, 2)
(/.85 85,
OUTPUT TENSOR
(/,972,1) 1)
(20, 101 4),
(10.64 64,1)
(/, x 51, 1) (/, 129, 1)
(/.85 6/.85. 85, 1)
Table 5 and Fig. 5 provide comprehensive overview of the benchmark details. The categorization of the generation details based on the governing PDEs, which are follows:
Elasticity (Li et al., 2022a): This benchmark aims to estimate the internal stress of an clastic material based on its structure, discretized into 972 points. For each sample, the input is a tensor with a shape of 972 x 2, representing the 2D position of each discretized point. The output 1000 samples with varying structures are generated for training, while 200 samples are used for testing.
stress each point, formatted as a tensor of shape 972 x In experiment,
Plasticity (Li et al., 2022a): This benchmark focuses predicting the future deformation of plastic material impacted from above by die with arbitrary shape. For each sample, the input is the die's shape, discretized into structured mesh and represented
tensor with shape of 101 x 31. The output is the deformation of each mesh point over 20 future timesteps, recorded as tensor of shape 20 x 101 x 31 x 4, capturing deformation in four directions. In the experiment, 900 samples with varying die shapes are used for training, while 80 additional samples are reserved for testing.
Navier-Stokes (Li et al., 2020): 2D Navier-Stokes equation mathematically describes the flow of viscous, incompressible
21

---

Latent Mamba Operator for Partial Differential Equations
fluid in vorticity form on the unit torus follows:
8,w(z, 1) + ua I) Vw(a, 1) = vAw( () +()) * (0,1))), (0,7)
(93)
V u(x, 1) = 0, x E (0, LE [0,
(94)
w(x, 0) = mo( a E (0,1)²
(95)
where, represents the velocity field, w V x is the vorticity, wn is the initial vorticity, v is the viscosity coefficient, and f is the forcing function. In this dataset, the viscosity (v) is fixed at 10-5, and the 2D field has a resolution of 64 x 64. Each sample within the dataset comprises 20 consecutive frames. The objective is to predict the subsequent ten frames based on the preceding ten. The experiment uses 1000 fluid samples with different initial conditions for training, while 200 additional samples are reserved for testing.
Pipe (Li et al., 2022a): This benchmark aims to estimate the horizontal fluid velocity based on the structure of the pipe. The governing equations are follows:
V U U 0,
(96)
ØU at VU
f- 1 VD VD+
U
(97)
Each sample represents the pipe as a structured mesh with dimensions 129 x 129. The input tensor, shaped as 129 129 x encodes the position of each discretized mesh point. The output tensor, with shape of 129 129 x 1, provides the velocity value at each point. For training, 1000 samples with varying pipe shapes are generated, while 200 additional samples, created by altering the pipe's 's centerline, are reserved for testing.
Airfoil (Li et al., 2022a): This benchmark pertains to transonic flow over an airfoil. Due to the negligible viscosity of air, the viscous term u omitted from the Navier-Stokes equation. Consequently, the governing equations for this scenario are expressed as follows:
Öpf
+ V .. (DfU) =0
(98)
01
H(fUU) Öt
+ V (pFUU + +pI)= =0
(99)
aE
Öt + v ((E + p)))= = 0.
(100)
where pʃ represents fluid density, and E denotes total energy. The input shape is discretized into a structured mesh with dimensions 221 x 51, and the output represents the Mach number at each mesh point. All shapes are derived from the NACA-0012 case provided by the National Advisory Committee for Aeronautics. In the training, 1000 samples from various airfoil designs are used for training, while the remaining 200 samples are reserved for testing.
Darey Flow (Li et al., 2020): This benchmark represents the flow through porous media. 2D Darcy flow over unit square is given as follows:
V (a(x)Vu(z)) = f(x), æ xE (0,
(101)
u(x) 00 x e 1)0
(102)
where a(x) is the viscosity, fxx is the forcing term, and u(x) is the solution. This dataset employs constant value of forcing term F(z)= = 3. Further, Equation 101 is modified in the form of temporal evolution as follows:
0,l(,, t) - (u(z) f)) = f(x), x E (0,
(103)
In this dataset, the input represented by the parameter a, and the corresponding output is the solution u. The process is discretized into 421 x 421 regular grid and then downsampled resolution of 85 85 for main experiments. For training, 1000 samples are used, 200 samples are generated for testing, and different cases contain different medium structures.
D. Implementation Details
This section presents a comprehensive overview of the experimental setup, covering benchmark datasets, evaluation metrics, and implementation details to ensure rigorous and reproducible analysis.
22

---

Latent Mamba Operator for Partial Differential Equations
Table Training and model configurations of LaMO. Training configurations directly from previous works without extra tuning. For the Darcy dataset, we adopt additional spatial gradient regularization term (aoll following ONO work.
BENCHMARKS
CONFIGURATION
DARCY
NAVIER-STOKES
ELASTICITY
PLASTICITY
AIRFOIL
PIPE
Loss FUNCTION
0.1
RELATIVE l2
TRAINING
EPOCHS
500
INITIAL LR
5 10
10-3
OPTIMIZER
ADAMW
BATCH SIZE
4
2
8
4
SCHEDULER
ONECYCLELR
LAYERS
8
ARCHITECTURE
EMBEDDING DIM
64
256
128
LATENT TOKENS
1936
1024
64
DSTATE
64
HEADS
D.1. Training Detail
As detailed in Table 5 and Table 6, the datasets were split into training and testing sets by configurations adopted from prior works. Each dataset corresponds to specific task and follows these established settings. We employed the Adam optimizer with the OneCycleLR scheduler to optimize model performance for training. The relative l2 error metric evaluated and reported results across all benchmarks. Further specifics on the dataset splits, tasks, and training configurations are available the respective tables for additional clarity for each benchmark dataset.
D.2. Hyperparameters Details
Table 6 shows that all baselines and benchmarks were trained using consistent configurations, ensuring that LaMO maintains fewer or comparable parameters relative to transformer-basee baselines. The table provides an overview of the training and model configurations for LaMO across multiple benchmark datasets. The training configurations include the use of relative l2 loss term across all datasets, with an additional spatial gradient regularization term (ga) incorporated for the Darcy dataset 12 + 0.1 following ONO (Xiao et al., 2023). The model is trained for 500 epochs using the AdamW optimizer, with OneCycleLR scheduler employed for learning rate adjustment.
On the architectural details of the LaMO, all benchmarks utilize an 8-layer latent architecture, with a fixed state dimension (DState) 64 and a convolutional kernel size of 3 across all datasets. SSM heads are configured as for Darcy and 4 for remaining benchmarks. These configurations align with prior works and ensure robust training without additional hyperparameter tuning. For regular grids, a single encoder-decoder pair is used. In contrast, the encoder and decoder share parameters for other cases for irregular mesh and point cloud and are applied each latent block.
D.3. Baselines Details
All baseline models have been extensively validated in prior work. For our experiments, we utilized the official codebase for each operator, ensuring all settings were implemented as prescribed prior work.
Typical Neural Operator:
We rely on widely recognized baselines, reporting performance metrics for most operators based on the results provided in their respective official papers for fair comparison. For models that cannot directly handle irregular datasets, we applied the transformation method outlined in the GeoFNO paper to evaluate and report their performance on the corresponding benchmarks.
Transformer Neural Operator: For GNOT, OFormer, and ONO, the performance metrics reported on the benchmark were
23

---

Latent Mamba Operator for Partial Differential Equations
Algorithm Latent SSM (Multidirectional Latent SSM)
1: Input: Latent tokens -111 BBXxD
2: - < LayerNorm 2(-
3: Z
1
4: X Tinear(Zi: 1
5: A, B, t Lincar(X)
6: A, B Discretization(A A, A, B)
7: Y 0
8: for in Direction-Scan do
9: Y += Mulhhhead-SSM
10: end for
11: + Y Activation(2)
Z t Linear(YY
13: Output: Latent tokens Z RBxMxD
directly taken from their respective official publications. In the case of Transolver, recognized as the state-of-the-art (SOTA) operator, we conducted additional evaluations by rumning the official codebase three times. This ensured consistency and accuracy in reporting the benchmark results. We conducted experiments for the remaining baseline models using their official codehases on GitHub. For reproducibility, we have provided our implementation codebase as supplementary.
D.4. Evaluation Metric
Our experimental evaluation focuses on standard PDE benchmarks, utilizing the mean relative lz error (Li et al., 2020) as primary metric to assess accuracy of the predicted physics fields. This metric consistently reported across all experiments and defined as follows:
[Ga(a;) (a,)ll2
L
(104)
[[:t(1;)]2
where N represents the number of samples, (a(a) denotes the predicted solution, and (ax) is the ground truth. The inclusion of the normalizing term (a:)ll ensures the metric accounts for variations in absolute resolution scales across different benchmarks, enhancing comparability.
D.5. Algorithm
In this subsection, we introduce the latent SSM algorithms. In Algorithm 1, the direction refers to the axis along which the SSM operates. Experimentally, we employ bidirectional scan for irregular mesh data. At the same time, for regular grids, we utilize multidirectional scan as shown in Figure along four paths: (i) top-left to bottom-right, (ii) bottom-right to top-left, (iii) top-right to bottom-left, and (iv) bottom-left to top-right. Each direction processed in parallel, reducing the overall computational complexity to linear. The multidirectional sean ensures non-causal kernel, enhancing the model's ability capture the underlying latent dynamics more effectively.
Figure 6. Visualization of multidirectional scan for regular grid benchmark dataset.
24

---

Latent Mamba Operator for Partial Differential Equations
E. Additional Experiments
This section presents additional experimental results to evaluate our proposed method comprehensively. This includes ablation studies assess the impact of key components, scaling experiments to analyze performance across different problem sizes, standard deviation analysis quantify robustness, and interpretability studies for deeper insights.
E.1. Ablations
Ablation on Bidirectional SSM and Weight Sharing: Table 7 demonstrates the impact of different scanning strategies on LaMO's performance across various datasets. The unidirectional scan consistently underperforms compared to the bidirectional setting, emphasizing the necessity of non-causal bidirectional sean. This is particularly evident in the Airfoil dataset, where the error worsens from 0.48 - 0.91. Furthermore, employing multidirectional scan leads to further performance gains for regular grid data, underscoring the advantages of incorporating directional diversity (for Navier Stokes, it improves from 0.10 -> 0.04). Additionally, sharing encoder and decoder parameters across blocks effectively reduces model complexity while maintaining competitive performance, as seen in the improved results for Airfoil and Elasticity and comparable accuracy in Plasticity and Pipe.
Table Comparison of different LaMO configurations with different scanning directions and weight sharing across all the benchmarks. w/ denotes the with, and w/o denotes the without settings. Lower relative l2 error indicates better model performance.
POINT CLOUD
REGULAR GRID
STRUCTURED MESH
OPERATOR (x10-2)
ELASTICITY NAVIER-STOKES DARCY PLASTICITY AIRFOIL
PIPE
LAMO (UNIDIRECTIONAL)
0.97+0044
16.24+0.320
1.33
0.22+0.022
0.91 0.88+0.041
LAMO (BIDIRECTIONAL W/O WEIGHTS SHARED)
0. 1..00.
10. 20.12
0.64
0.07
0.48 0.38 10.025
LAMO (BIDIRECTIONAL W/ WEIGHTS SHARED)
0.50
12.2.10108
0.77.0.015
0.08
0.41 0.4 10.027
LAMO (MULTIDIRECTIONAA)
0.62+0.030
4.60=0.108
0.39+0.015
0. 0.12+4.004
0.50+0.066
E.2. Standard Deviations
In this section, we report the standard deviation of all benchmarks while presenting the mean values in the main results table. We repeat all experiments across multiple independent runs and compute the standard deviation to ensure statistical rigor. Table 8 shows that LaMO consistently outperforms all haselines within a 95% confidence interval. We also run the second-best baselines three times for a fair comparison and report their standard deviations. Given the variability in their results across different benchmarks, achieving superior performance across all previous models is challenging. This consistency further validates the robustness and effectiveness of our proposed model.
Table 8. The standard deviations of LaMO are reported across all experiments. For comparison, the performance of the second-bess operator also included. Each experiment was conducted over five independent runs to calculate the standard deviation. Lower relative l2 error indicates better model performance.
POINT CLOUD
REGULAR GRID
STRUCTURED MESH
OPERATOR (x10-3)
ELASTICITY
SAVIER-STOKES
DARCY
PLASTICITY
AIRFOIL
PIPE
0.6410.02
9.57
0.59
0310000
0.53
0.461000
BECOND-BEST MODEL
(TRANSOLVER)
(TRANSOLVER)
(TRANSOLVER) (TRANSOLVER) (TRANSOLVER) (TRANSOLVER)
LAMO (OuRS)
0.5040.02
4.60
0.39
0.07
0.41
4.38
E.3. Efficiency
As shown in Figure 7, we compare the training time, inference time, and memory consumption of LaMO against Transolver across all datasets and in Table 9, we compare the model parameter count range across the benchmark compared with baselines. Our results indicate that LaMO demonstrates comparable performance in terms of efficiency metrics. Furthermore, LaMO exhibits significantly better performance across all efficiency metrics for ime-deeendent PDEs such as Slvier-Stooke and Darcy. This advantage makes it well-suited for scaling over regular grids, facilitating the development of foundation
25

---

Latent Mamba Operator for Partial Differential Equations
Table Comparison of models' parameters count of LaMO with baselines.
OPERATOR
FNO
U-FNO
LSM
GNOT
GALERKIN
TRANSOLVER
LAMO
PARAMETER(IN M)
0.9-18.9
1.0-19.4
4.8-13.9
9.0-14.0
2.2-2.5
2.8-11.2
1.1-4.0
models for PDEs.
Triining time (Slepech)
Inference time (slepoch)
15
200
Menory (GB)
10
10
100
5
Danyy
Airbil
Plosbill
Pipe
Dangy
NS
Airoll
Pleakity
Fluskity
NS
Plaskitt
Eleakilly
Plye
LaMO (Ours)
Transolver
Figure 7. Efficiency comparison between LaMO and Transolver (Left) Training Time, (Middle) Inference Time and (Right) Memory
consumption per epoch on all the benchmark datasct.
E.4. Time Complexity
Table 10 compares the computational complexity of different neural operator models. The models are characterized by their kernel range (Global or Global + Local), kernel type (Fixed or Data- metaa-Dependnt), and computational complexity for the number of sampling points in a continuous function. FNO achieves O(N log N) complexity with a fixed global kernel, Transformers incur ( complexity due to their global, data-dependent kernel, and LaMO balances efficiency and flexibility with ONN) complexity by leveraging both global and local, data-dependent kernels.
Table 10. Computational complexity of neural operators. N represents the number of sampling points of continuous function.
MODEL
KERNEL RANGE
KERNEL TYPE
COMPLEXITY
FNO
GLOBAL
FIXED DATA DEPENDENT
ONN N)
TRANSFORMER
GLOBAL
O(N3) O(N)
LAMO
GLOBAL + LOCAL
DATA DEPENDENT
E.5. Model Scalability
From Figure observe that LaMO consistently outperforms Transolver across all six benchmarks, with notable margin in most cases, except for Airfoil, where the improvement is marginal. LaMO achieves a lower relative l2 error than Transolver, with only 40% of the training data for the Darey and Navier-Stokes benchmarks. In the Plasticity benchmark, LaMO surpasses Transolver with just 20% of the training samples, exhibiting minimal performance degradation despite the reduced dataset size. This is particularly valuable in practical scenarios where collecting large-scale PDE datasets is costly, emphasizing the importance of models that maintain high accuracy with limited training samples.
Furthermore, as presented in Table 11, we evaluate the performance of neural operators in comparison with Transolver across varying network depths the Navier-Stokes (NS) dataset, which presents significant challenges due to its complex, nonlinear fluid dynamics. Our experiments focus assessing how the number of layers in the operator architecture influences predictive accuracy. We observe a consistent improvement in performance as the depth of the operator increases, indicating that deeper architectures are better able to capture the intricate temporal and spatial correlations inherent in turbulent flow regimes. Remarkably, even with a relatively shallow configuration of just four layers, LaMO outperforms Transolver, demonstrating superior efficiency and representational capacity in modeling complex physical systems.
26

---

Latent Mamba Operator for Partial Differential Equations
(a) Navicr-Stokes
(b) Darcy
Airfoil
Relative 1.2 (x100
250
25
25
Relative 1.2 (x10
Relative L2 (x 10
20
150
15
16
100
10
10
50
0
200
400
600
800
1,000
200
400
600
800
1,000
200
400
600
800
1,000
Number of training samples
Number of Training Samples
Number of Training Samples
(d) Pipe
(c) Plasticity
(f) Elasticity
Relative L.2 (x10 -)
30
Relative L2 (x
Relative L2 (x110
10
26
1.5
8
20
15
6
10
4
0.5
200
100
600
800
1,000
180
300
540
720
900
200
100
600
800
1,000
Number of Training Samples
Number of Training Samples
Number of Training Samples
LaMO (Ours)
Transolver
Figure 8. Data Efficiency for different benchmarks for LaMO with Transolver (Top) (a) Navier Stokes, (b) Darcy (c) Airfoil, (Buttom) (d)
Pipe, (c) Plasticity and (f) Elasticity benchmark respectively.
Table 11. Comparison of relative l2 error on Navier-Stokes henchmark of I.MM with second best baseline model Transolver. Lower relative la error indicates better model performance.
NUMBER OF LAYERS
2
4
TRANSOLVER
0.1601
0.1518
0.1241
0.0957
LAMO
0.1038
0.0608
0.0524
0.0460
E.6. Interpretability
In Figures 9 and 10, we present the hidden attention matrices of Transolver and LaMO across all layers. For LaMO, we plot the bidirectional attention maps for the Airfoil and Elasticity benchmarks to facilitate an intuitive comparison. As observed, Transolver's attention maps exhibit values concentrated on few tokens, whereas LaMO's s attention is more diffused across tokens, thus establishing better representation with more diverse latent tokens. The difference can likely be attributed to using softmax in Transolver, which sharpens the attention maps and is known cause over-smoothing issues.
To further analyze LaMO's performance compared Transolver, we examine the inter-cosine similarity across layers for both models as shown in Figure 11. LaMO significantly reduces inter-cosine similarity across layers compared to Transolver, indicating better representation learning through SSM. This improvement can be attributed to the absence of the softmax operation, which in Transolver has been identified as causing over-smoothing issues (Ali et al., 2023; Wang et al., 2022).
F. Visualization
This section presents comparative analysis of model predictions for Transolver and LaMO across all benchmark datasets. For each dataset, we visualize the predicted outputs alongside their corresponding ground truth values to assess the performance of both models. Additionally, we include heatmaps to highlight the spatial distribution of errors and deviations, providing deeper insights into the strengths and weaknesses of each approach. These visualizations demonstrate the effectiveness of LaMO in capturing intricate patterns within the data while benchmarking it against the existing state-of-the- art model, Transolver.
27

---

Latent Mamba Operator for Partial Differential Equations
Layer Layar2
Transolver
Durs
Figure Hidden Attention Matrices comparison of Transolver (Top) with LaMO (Bottom) the layers Airfoil benchmark.
Layer
Layor22
Layer
Layer
Layer
Layer
Lager
Layer'
Transolver
Dura
Figure 10. Hidden Attention Matrices comparison of Transolver (Top) with LaMO (Bottom) across the layers on Elasticity benchmark.
(a) Airfoil
Elasticity
Cosine Similarity
Cosine Similarity
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
5
6
2
3
8
Layer Index
Layer Index
LaMO (Ours)
Transolver
Figure Intercosine similarity of tokens across layers for LaMO Transolver on (a) Airfoil and (b) Elasticity benchmark, respectively.
G. Limitations and Future Work
LaMO demonstrates promising results in solving parametric PDEs, but its efficiency in unsupervised setting remains un- explored. Investigating unsupervised training methods and improved token scanning techniques could enhance performance. Additionally, evaluating LaMO's scalability and optimizing its training strategies are key areas for research. For future work, we plan to investigate the potential using SSM-based operators as foundation models for more efficient PDE solving, focusing on unsupervised learning to achieve better representations. However, the compatibility of existing pretraining methods with SSM-based architectures for operators and developing pretraining techniques specifically tailored for these models remains open of exploration. Additionally, aim explore improved training strategies for SSM-based operators, including enhanced scanning methods for regular grid PDEs, optimize their performance further.
28

---

Latent Mamba Operator for Partial Differential Equations
Pred(Transoiven Ground truth
0.0075
0.0075
Erron(Ours) Transsoinn Pred[Ours)
20075
0030
0034
0034
Figure 12. Model Prediction Comparison: The figure compares Transolver and LaMO on the Darcy dataset. The (Top) row represents the ground truth, the (Middle) row shows the predictions from Transolver and LaMO, and the (Bottom) row illustrates the error heatmap, capturing the differences between the ground truth and the predicted results.
Preet (ra(nsrrns Ground truth
EnrortTanzzaiee PrestOurs)
Error(Ours)
Figure 13. Model Prediction Comparison: The figure compares Transolver and LaMO the Navier Stokes dataset. The (Top) row represents the ground truth, (Middle) row shows the predictions from Transolver and LaMO, and the (Bottom) row illustrates the error heatmap, capturing the differences between the ground truth and the predicted results.
29


Reviews:
[{"review_id": "n8rZCVwNvQ", "reviewer": "Reviewer_SQnF", "summary": "This paper introduces a methodology to use Mamba based architecture to model the spatial dynamics of a PDE, in an operator based setting. The authors take inspiration from the Perceive model and use a latent space, in addition to directly modeling the input physical properties, however instead of a Transformer layer, the authors use a Mamba layer. \n\nThe key motivation here it to reduce the quadratic complexity of Transformer to that a linear complexity incase of Mamba. Furthermore, the hope is that Mamba based model is also able to effectively model long range spatial dynamics of the input physical process. \n\nThe authors also theoretically prove that in general an SSM based architecture should be able to model the kernel integral operator on a domain, and proving that parametrically at least that the Mamba layer should matches the required structure. \n\nThe authors back their claims with extensive empirical analysis, where across a series of benchmarks (like Darcy Flow and 2D Navier-Stokes on regular grids, as well as airfoil, plasticity etc on irregular meshes) Latent Mamba operator is able to outperform baselines such as Transolver. \n\nStrengths: \n\n- The experiments are extensive, and the authors show that their method outperforms the baselines consistently.\n- The authors perform ablation studies to explain the importance of different architectural design choices that they make in the design of their architecture.\n- They show that their method is efficient and often requires 5-7 times fewer parameters when compared to baselines.\n\nFew Limitations\n\n- They establish that an SSM is a Monte-Carlo Approximation of a Kernel integral operator. However, the theorem does not give any kind of a guarantee in the number of samples (given some assumption on the kernel $\\kappa$ (like boundedness, or smoothness assumptions) or the error). I guess the proof is based on the parametric form (Equation 57 of Appendix) of how the SSM approximates the kernel, but even then something that establishes the approximation capacity of the approximation help establish the true equivalence.\n    - In general the equivalence is made using $\\approx$ which is not formal.\n- The scaling results are on Darcy Flow, which potentially tells us nothing. I think similar scaling results should be shown on relatively complex PDEs (perhaps even Navier Stokes).\n    - In fact, Darcy Flow is quite a simple dataset that can potentially be learned with very few layers.\n\nFew Suggestions: \n\n- It will be useful to add the number of parameters (or FLOPs) of various methods used in Table 1.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments and for supporting the work. Please see the responses to your questions below.\n\n>Sample Complexity: However, the theorem does not...\n \n**A:** A Monte Carlo integral approximation typically has $O\\left(\\frac{1}{\\sqrt{n}}\\right)$ convergence rate which is empirically corroborated from Figure 2. We do have a working proof (outlined below) which is subjected to the assumptions stated on $\\kappa$. We will be adding it in the supplementary of the manuscript. A formal proof will be explored as part of future work. \n\n**Statement:** A Monte Carlo integral approximation of SSM with suitable assumptions on the kernel (e.g. boundedness and smoothness) provides an upper bound on the error as $\\epsilon \\sim O\\left( \\frac{1}{\\sqrt{n}} \\right)$.\n\n**Proof sketch:**\n\nBoundedness assumption: $\\kappa(x, y) \\le B$.\n\nIn an SSM layer, the approximation of the operator computed via a Monte Carlo sum is: $\\[T_n f\\](t) = \\frac{1}{n} \\sum_{i=1}^{n} \\varphi(t)^\\top A \\psi(s_i) f(s_i)$.\n\nDefine the error operator as $E_n = T_n - T$.\n\nWith probability of at least $1-\\delta$ (with $0<\\delta<1$), Berstein's inequality suggests that:\n\n$\\|T_n - T\\| \\le C \\ \\sqrt{\\sigma^2 \\ \\frac{\\log(1/\\delta)}{n}}$,\n\nwhere the variance parameter $\\sigma^2$ is given by $\\sigma^2 = \\left\\| \\sum_{i=1}^{n} \\mathbb{E}\\left[(Z_i - T)^2\\right] \\right\\|$, and $C$ is a constant.\n\nHence, to achieve an error of at most $\\epsilon$, substituting $\\epsilon$ on the LHS, we note that $n \\sim O\\left( \\frac{M}{\\epsilon^2} \\right)$ where $M$ is constant. \n\nTherefore, $\\epsilon \\sim O\\left( \\frac{1}{\\sqrt{n}} \\right)$.\n \n>Scaling results on relatively complex PDEs (perhaps even Navier Stokes).\n\n **A:** Following your suggestion, we performed additional experiments to show the scaling results on the Navier-Stokes (NS) dataset in Tables 1 and 2. We note that the operator's performance improves as the number of layers increases on the NS dataset. Notably, even with only 4 layers, LaMO achieves superior performance compared to Transolver (see Tab. 1). Furthermore, Tab. 2 demonstrates that LaMO trained on 400 NS samples delivers a performance comparable to that of Transolver trained on 1000 samples, highlighting LaMO's efficiency and effectiveness with fewer training samples. For a more comprehensive examination of scalability across different sample sizes, please refer to Appendix Section E.5, which presents the results for all benchmark datasets.\n\n| # Layers | 2    | 4    | 6    | 8    |\n| ---------------- | ---  | ---  | ---  | ---- |\n| Transolver       |0.1601|0.1518|0.1241|0.0957|\n| LaMO (Ours)       |0.1038|0.0608|0.0524|0.0460|\n\n***Table 1: Relative L2 error vs. layer scaling on NS (ν = 1e-5).***\n\n\n| # Training samples | 200    | 400    |  600 | 800    | 1000   |\n| -------------------------- | ------ | ------ | ---  | ------ | ------ |\n| Transolver                 | 0.2330 | 0.1874 |0.1552| 0.1120 | 0.0957 |\n| LaMO (Ours)                 | 0.1490 | 0.0972 |0.0648| 0.0570 | 0.0460 |\n \n ***Table 2: Relative L2 error vs. training samples on NS (ν = 1e-5).***\n \n >Number of parameters (or FLOPs).\n\n**A:** In Table 3, we present the number of parameters for a range of neural operators across all benchmarks. The exact number of parameters used for each dataset is available and will be included in the Appendix. We will incorporate these changes into the manuscript and ensure they are reflected in the final version.\n\n| Operator | FNO      | U-FNO    | LSM      | GNOT    | Galerkin | Transolver | LaMO    |\n| --------------- | -------- | -------- | -------- | --- | -------- | ---------- | ------- |\n| Parameters (M)  | 0.9-18.9 | 1.0-19.4 | 4.8-13.9 | 9-14    | 2.2-2.5  | 2.8-11.2   | 1.1-4.0 |\n\n***Table 3: Baselines parameter range (in M)***\n\n>Performance on 3D. \n\n**A:** As shown in Table 4 of the Appendix, the benchmark datasets used, such as Navier-Stokes (regular) and Plasticity (point cloud), consist of 2D spatial and 1D temporal dimensions. However, it would be interesting to explore the performance of the methods on datasets with higher physical dimensions (e.g., 3D spatial and 1D temporal) as part of future work.\n\n>Recent literature.\n\n**A:** Thank you for highlighting the relevant references we missed. We will include these recent works on SSMs for PDEs in the related work section of the final version.\n\nWe have addressed all your concerns regarding these. Please let us know if you have any further comments or questions.", "claims_and_evidence": "The claim regarding the empirical benefits of the method are appropriately validated through experiments. However, the fact that the SSM layers can approximate the kernel integral operator are lacking given the lack of either number of samples or error bounds.", "methods_and_evaluation_criteria": "The benchmarks used are on 2D datasets, though it will be interesting to see how the methods perform on 3D. Other than that, the authors also show results on irregular meshes with is good.", "theoretical_claims": "There are theoretical claims, though the statement can be made more formal and precise. I think its somewhat straightforward to show though. Please see the main review for more details.", "experimental_designs_or_analyses": "Yes, the empirical results are extensive. The authors show the scaling results on relatively simple datasets on Darcy flow, it will be interesting to see the scaling behaviour on more complex dynamics (such as Navier Stokes).", "relation_to_broader_scientific_literature": "Very relevant, given the key point that the quadratic complexity of transformers may be prohibitive for higher dimensions.", "essential_references_not_discussed": "There are a few recent works that try to use SSM based architectures that the authors may have missed: \n\nHu, Zheyuan, et al. \"Deepomamba: State-Space Model for Spatio-Temporal Pde Neural Operator Learning.\" *Available at SSRN 5149007*.\n\nRuiz, Ricardo Buitrago, et al. \"On the Benefits of Memory for Modeling Time-Dependent PDEs.\" *arXiv preprint arXiv:2409.02313* (2024).", "other_strengths_and_weaknesses": "Mostly discussed in the main review.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r68rAx8Cp4", "reviewer": "Reviewer_9Sjn", "summary": "This paper introduces a new approach to solving PDEs by introducing the SSM-based Neural Operator on the latent space. The proposed method achieves a good balance on performance and efficiency. The authors provide theoretical analysis that reveals the equivalence of LaMO with kernel integration. With extensive experiments, the propsed method outperforms SOTA transformer-based model and achieving comparable computational efficiency.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments and for supporting the work. Please see the responses to your questions below.\n\n>Results in Supplementary E.3 show that the efficiency of your model is only comparable to TRANSOLVER, but the main text presents it as if your model is superior in all aspects/datasets, which is misleading.\n\n**A:** Thank you for pointing this out. We agree that LaMO's efficiency is indeed comparable to that of TRANSOLVER. However, in the main text, our reference to efficiency pertains specifically to comparisons with transformer-based baselines such as GNOT and ONO, where LaMO demonstrates superior efficiency. We will revise the final version's text to clarify this distinction and avoid any potential misinterpretation.\n\n>The performance of TRANSOLVER on the Pipe dataset differs significantly from the results reported in the original paper. The original result is 0.0033, which is better than yours. Please explain this discrepancy.\n\n**A:** All the baseline errors were sourced from existing literature, except Transolver. Since Transolver represents the SOTA, we ran its official codebase (keeping the hyperparameters, data splits, and other settings the same as the original paper) and reported the best errors from multiple runs in our manuscript. However, despite our best efforts, we could not replicate the result of 0.0033 reported in the original paper. This discrepancy may be due to differences in the computing environment or random initialization settings. To demonstrate this further, we have made the log file for multiple runs on the pipe dataset available on our anonymous GitHub repository. We shall clarify this in the final version of the paper.\n\n>The effect of the latent encoder/decoder and the SSM has not been quantified, making it hard to distinguish their individual contributions.\n\n**A:** Following your suggestion, in Table 4 (below), we present an ablation study on the latent encoder/decoder, where we observe that ViT patches outperform the Perceiver encoder on regular grids. Additionally, the effects of the SSM and the encoder are detailed in Table 2 of the main text and further elaborated in Appendix Section E.1.\n\n\n| Encoder/Decoder type | Relative L2 |\n| -------------------- | ----------- |\n| ViT                  | 0.0039      |\n| Perceiver (Unshared) | 0.0051      |\n| Perceiver (Shared)   | 0.0056      |\n\n***Table 4: Ablation on effect of latent encoder/decoder on Darcy.***\n\n>In the paragraphs \"Latent Tokens\" and \"Efficiency\" on Page 8, the authors do not indicate which exact table/figure they are referring to.\n\n**A:** In the \"Efficiency\" paragraph, we are referring to Figure 4 in the main text. However, in the \"Latent Tokens\" paragraph, we are not referencing any specific table or figure, as it presents an additional ablation study on using latent tokens on a regular grid for the foundation model.\n\n>In the analysis of computational complexity (Page 5 Computational Analysis and Supplementary E.4.), while M is theoreticallt constant, but in practice it is a hyperparameter and will be set due to the size of N and other factors. You can compare the size of M and logN to give readers a better sense of the complexity of your model.\n\n**A:** We acknowledge that while $M$ is theoretically constant, it functions as a hyperparameter in practice, influenced by the size of $N$ and other factors. We will enhance the discussion by comparing $M$ to $log N$ better to illustrate its impact on the model's computational complexity.\n\nWe have addressed all your concerns. Please let us know if you have any further comments or questions.", "claims_and_evidence": "Results in Supplementary E.3 show that the efficiency of your model is only comparable to TRANSOLVER, but the main text presents it as if your model is superior in all aspects/datasets, which is misleading.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for the application.", "theoretical_claims": "I have checked all the proofs and found no obvious issues.", "experimental_designs_or_analyses": "1. The performance of TRANSOLVER on the Pipe dataset differs significantly from the results reported in the original paper. The original result is 0.0033, which is better than yours. Please explain this discrepancy.\n\n2. The effect of the latent encoder/decoder and the SSM has not been quantified, making it hard to distinguish their individual contributions.", "relation_to_broader_scientific_literature": "They used Mamba (Gu & Dao, 2023) as the key component of the model and their Latent Encoder is inspired by the Perceiver (Jaegle et al., 2021).", "essential_references_not_discussed": "Essential references are included.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper is well written.\n2. Comprehensive experiments, outperforming SOTA\n3. Establish a connection between SSMs and kernel integral operators", "other_comments_or_suggestions": "1. In the paragraphs \"Latent Tokens\" and \"Efficiency\" on Page 8, the authors do not indicate which exact table/figure they are referring to.\n\n2. In the analysis of computational complexity (Page 5 Computational Analysis and Supplementary E.4.), while M is theoreticallt constant, but in practice it is a hyperparameter and will be set due to the size of N and other factors. You can compare the size of M and logN to give readers a better sense of the complexity of your model.", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oal46yyOIX", "reviewer": "Reviewer_ZQ3y", "summary": "- The Latent Mamba Operator (LaMO) is a scalable state-space model integrated with a kernel integral formulation. \n\n- The authors also provide a theoretical foundation for their approach. \n\n- LaMO demonstrates state-of-the-art performance across various problems.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments. Please see the responses to your questions below.\n\n>They also conduct scaling experiments, particularly focusing on the amount of training data. While Figure 2 implicitly provides scaling behavior with model size, a more detailed analysis of how the model scales with parameter count—whether it follows a power law or logarithmic trend—would be valuable.\n\n**A:** We provide the parameter count scaling for the Darcy and NS datasets in Table 5. To better illustrate the scaling trend, we plot the parameter count on a log scale against the number of layers. The plot shows that the parameter growth follows a consistent linear trend, confirming the proportional relationship between the layer count and the model complexity.\n\n| # Layers           | 2    | 4    | 8     | 12    | 24    |\n| -------------------------- | ---- | ---- | ----- | ----- | ----- |\n| Navier-Stokes (Transolver) | 2.93 | 5.70 | 11.23 | 16.76 | 33.35 |\n| Navier-Stokes (LaMO)       | 2.72 | 5.17 | 10.06 | 14.95 | 29.62 |\n| Darcy (Transolver)         | 0.74 | 1.43 | 2.82  | 4.21  | 8.38  |\n| Darcy (LaMO)               | 0.38 | 0.63 | 1.14  | 1.65  | 3.18  |\n\n***Table 5: Parameter count (in M) vs. layer count for NS (ν = 1e-5) and Darcy.***\n\n>Additionally, a clearer examination of model scaling with respect to parameter count would enhance the discussion of efficiency.\n\n**A:** Please refer to Table 2, Table 4, and Figure 2 in the main text for details on model scaling concerning parameters and efficiency on the Darcy and Navier-Stokes datasets. As shown, the model performance improves with an increase in the number of parameters (layers).\n\n>However, a deeper exploration of the theoretical implications, particularly in comparison to existing SSM-based PDE solvers, would further solidify the claims.\n\n**A:**  We acknowledge the importance of further exploring the theoretical implications, particularly in comparison with existing SSM-based PDE solvers. We will consider this as part of our future work.\n\n>Several recent studies explore SSM-based neural operators:\n\n**A:**  We have cited [3] under SSMs for PDEs in the related work section. We will ensure that the recent works on SSMs for PDEs [1, 2] are included in the related work section of the final version and a relevant discussion.\n\n>Were the baseline errors sourced from existing literature, or did you train the baseline models yourselves? If you conducted the training, how were the architectures chosen, and what training procedures were followed?\n\n**A:** The baseline errors were sourced from existing literature, except Transolver. Since Transolver represents the SOTA, we ran its official codebase (keeping the hyperparameters, data splits, and other settings the same as the original paper) and reported the best errors from multiple runs. We adhered to the standard training procedures described in Appendix Section C for all experiments.\n\n>How does your approach compare to previous works on SSMs for PDEs?\n\n**A:** Our approach differs from previous works on SSMs for PDEs in several key aspects. While prior studies such as [3] employed unidirectional Mamba models for ODEs (1D baselines), we utilize a latent bidirectional Mamba architecture designed explicitly for PDEs (2D baselines). Furthermore, recent works [1, 2] also adopt unidirectional Mamba models (combined with convolution) but are limited to regular grid PDEs only. In contrast, our method leverages a latent bidirectional Mamba, making it applicable to regular and irregular benchmark datasets. We will include additional discussion in the manuscript to address these points. \n\nWith these, we hope to have addressed all your comments. Please let us know if you have any further concerns or questions. If not, we request you support the manuscript by raising the score.", "claims_and_evidence": "The authors present a novel approach that achieves state-of-the-art performance across different PDEs, conducting extensive experiments on both regular and irregular domains, including turbulence-related problems. They compare their model against multiple baselines, providing a solid empirical foundation. However, additional evaluations on PDEs such as the Poisson equation, Wave equation, and problems involving temporal dynamics would further strengthen their claims.\n\nThey also conduct scaling experiments, particularly focusing on the amount of training data. While Figure 2 implicitly provides scaling behavior with model size, a more detailed analysis of how the model scales with parameter count—whether it follows a power law or logarithmic trend—would be valuable.", "methods_and_evaluation_criteria": "The study includes several important ablation experiments, such as the comparison between unidirectional and bidirectional approaches. This is a critical aspect of the analysis. Moreover, the authors investigate resolution invariance, a fundamental property for neural operators, and provide supporting evidence for it.\n\nAdditionally, a clearer examination of model scaling with respect to parameter count would enhance the discussion of efficiency.", "theoretical_claims": "The paper includes an accessible theoretical study of the proposed model, demonstrating that its SSM formulation approximates a class of integral linear operators.\n\nHowever, a deeper exploration of the theoretical implications, particularly in comparison to existing SSM-based PDE solvers, would further solidify the claims.", "experimental_designs_or_analyses": "See above", "relation_to_broader_scientific_literature": "See my comments above.", "essential_references_not_discussed": "The paper does not mention prior works on SSMs for PDEs, despite their relevance. \n\nSeveral recent studies explore SSM-based neural operators:\n\n[1] Cheng, C. W., Huang, J., Zhang, Y., Yang, G., Schönlieb, C. B., & Aviles-Rivero, A. I. (2024). Mamba neural operator: Who wins? transformers vs. state-space models for pdes. arXiv preprint arXiv:2410.02113.\n\n[2] Zheng, J., Li, W., Xu, N., Zhu, J., & Zhang, X. (2025). Alias-Free Mamba Neural Operator. Advances in Neural Information Processing Systems, 37, 52962-52995.\n\n[3] Hu, Z., Daryakenari, N. A., Shen, Q., Kawaguchi, K., & Karniadakis, G. E. (2024). State-space models are accurate and efficient neural operators for dynamical systems. arXiv preprint arXiv:2409.03231.\n\n\nIncluding discussions on these works would provide a more comprehensive positioning of the proposed approach within the broader landscape of SSM-based PDE solvers.", "other_strengths_and_weaknesses": "/", "other_comments_or_suggestions": "I recommend accepting the paper, provided the authors incorporate prior work on SSMs for PDEs. The study is robust and well-executed, with no significant gaps.", "questions_for_authors": "- Were the baseline errors sourced from existing literature, or did you train the baseline models yourselves? If you conducted the training, how were the architectures chosen, and what training procedures were followed?\n\n- How does your approach compare to previous works on SSMs for PDEs?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VmQuyaWeG1", "reviewer": "Reviewer_hGk2", "summary": "This paper introduces LaMO, which is an SSM-based neural operator designed to overcome the computational limitations of traditional neural operators for solving PDEs. It establishes a kernel integral interpretation of the SSM framework, proving its equivalence to integral kernel neural operators. It achieves an average 32.3% improvement over existing neural operators across multiple PDE benchmarks, including Navier–Stokes, Darcy flow, and elasticity problems.\n\n## update after rebuttal\nThank you for addressing the concerns. I am keeping my original score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive comments. Please see the responses to your questions below.\n\n>Computational complexity analysis: While LaMO is claimed to be more efficient, a more precise breakdown of runtime (e.g., FLOPs, GPU memory usage per operator) would improve the claims.\n\n**A:** The above response in Table 3 presents the range of parameters per operator. Additionally, Figure 4 in the main text and Appendix Section E.3 offer a detailed analysis of the training, inference, and memory consumption, providing a precise breakdown of the computational requirements.\n\n>Ablation studies on kernel choices: The study does not fully explore different kernel configurations for SSM parameterization.\n\n**A:** In our experiments, we primarily investigate the effect of directionality in kernel parameterization, as presented in Table 6 of the Appendix. The kernel formulation is inspired by Mamba, which employs time-variant parameterization for matrices B and C. Additionally, matrix A follows a diagonal structure, demonstrating superior performance compared to time-invariant parameterization in computer vision and LLM tasks.\n\n>Weaknesses: While the benchmarks are diverse, an application to real-world turbulent fluid dynamics (e.g., weather modeling, aerodynamics) would strengthen the evaluation.\n\n**A:** In response to your suggestion, we have conducted additional experiments on the ERA dataset, which is derived from the fifth generation of ECMWF reanalysis data. Following standard weather forecasting practices, we selected the geopotential height at 500 hPa (Z500) with a resolution of 2.5° and a time interval of 3 hours as our data. The objective is to predict the geopotential height for 10 timesteps given 2 observations. Due to time constraints and the computational limitation of the experiments, we adopted a data split of 1000/200, consistent with the NS benchmark. We utilized 4 layers for the architecture for both Transolver and LaMO models. Our results indicate that LaMO outperforms Transolver by 25%, demonstrating its superior accuracy in this setting. We plan to explore this direction further as part of our future work.\n\n\n| Operator |  Train error | Test error |\n| --------------- | --------     | --------   |\n| Transolver      | 0.04197      | 0.31665    |\n| LaMO            | 0.04072      | 0.23746    |\n\n***Table 6: Relative L2 error on ERAZ500 dataset.***\n\n>Limited discussion on hyperparameter selection: The role of latent dimension choices, SSM state sizes, and discretization steps could be better explained.\n\n**A:**  We have employed a learning rate selected from the set {1e-3, 5e-4, 1e-4}, while the remaining hyperparameters are consistent with those used in Transolver. For the SSM, we utilized a state dim (DState) of 64 and varied the number of heads in the range {1, 4}.\n\nThe influence of the latent dim and SSM state size, which follow the same configuration as Mamba, is presented in Table 2 of the main text and Table 7 in the Appendix. Our ablation indicates an optimal latent dimension, beyond which performance decreases before increasing again. Regarding discretization, we employed the ZOH scheme, which is similar to Mamba.\n\n>More runtime benchmarks needed: While Figure 4 suggests efficiency, a breakdown of training vs. inference time cost would be helpful.\n\n**A:** We have included the efficiency analysis, covering training time, inference time, and memory consumption, in Appendix Section E.3, where we compare our method with Transolver.\n\n>Sparse kernel methods (e.g., Gaussian Processes for PDEs) could be referenced for comparison.\n\n**A:** We highlight that extending sparse kernel methods, such as Gaussian Processes for PDEs, is computationally demanding and may become impractical for high-dimensional or large-scale PDE problems. Additionally, these methods often struggle with capturing complex, non-stationary dynamics, which neural operators are better equipped to handle. We will ensure that the relevant work is appropriately cited in the final version of the manuscript.\n\n>Impressive work! Besides the theoretical contributions, this work adds a novel type of kernel operators (Mamba) to the family of neural operators for PDEs. Its experiments also cover many SOTA baselines, providing a good benchmarking framework for future work on neural operators and Mamba-like kernels. The community would appreciate it if the authors make their code publicly available. (I have gone through the Supplementary Material.) I believe the authors will, yeah^^?\n\n**A:** Yes, we plan to open-source our codebase, which will contribute to benchmarking future operators and fostering further research and development in this area. Please note that our code is already a part of the supplementary material. \n\nWith these, we hope to have addressed all the comments. Please let us know if you have any further concerns or questions. If not, we request you support the manuscript by raising the score.", "claims_and_evidence": "- Weaknesses:\n  - Computational complexity analysis: While LaMO is claimed to be more efficient, a more precise breakdown of runtime (e.g., FLOPs, GPU memory usage per operator) would improve the claims.\n  - Ablation studies on kernel choices: The study does not fully explore different kernel configurations for SSM parameterization.", "methods_and_evaluation_criteria": "- Strengths: Multiple benchmark PDEs and baselines are tested. The evaluation metric (relative L2) also complies with standard practice. \n- Weaknesses: While the benchmarks are diverse, an application to real-world turbulent fluid dynamics (e.g., weather modeling, aerodynamics) would strengthen the evaluation.", "theoretical_claims": "Claims look good in general. Details of math are not checked line by line.", "experimental_designs_or_analyses": "- Weaknesses:\n  - Limited discussion on hyperparameter selection: The role of latent dimension choices, SSM state sizes, and discretization steps could be better explained.\n  - More runtime benchmarks needed: While Figure 4 suggests efficiency, a breakdown of training vs. inference time cost would be helpful.", "relation_to_broader_scientific_literature": "Sparse kernel methods (e.g., Gaussian Processes for PDEs) could be referenced for comparison.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "Impressive work! Besides the theoretical contributions, this work adds a novel type of kernel operators (Mamba) to the family of neural operators for PDEs. Its experiments also cover many SOTA baselines, providing a good benchmarking framework for future work on neural operators and Mamba-like kernels. The community would appreciate it if the authors make their code publicly available. (I have gone through the Supplementary Material.) I believe the authors will, yeah^^?", "questions_for_authors": "See the \"Weaknesses\" items above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: There is a good consensus for acceptance, despite some concerns, including missed references and small issues in the experimental setting. The reviewers all appreciated the methodological novelty and generalization of MAMBA to learning on function spaces.  recommend the paper for acceptance. 

Please incorporate all the additional discussions, missed references and results from the rebuttal and discussions with the reviewers into the final version of the manuscript. 

A few references were missed that should be added, and some of the comparisons seemed a little cherry picked, in particular, the discrepancy on the results of the Transolver should be clearly addressed in the final manuscript. Ideally I would recommend the authors reach out to the Transolver authors to try and resolve any potential issue. When comparing the methods in terms of number of parameters, please incorporate the added scaling against the number of layers. It would also be good to look at, or at least discuss Neural Operators methods that are geared towards parameter efficiency, such as the FFNO that is in the table but not the plot, and the Tensorized Fourier Neural Operator (TFNO), that compresses the FNO also presented in the table. The reviewers also asked for more details on the hyper parameters selection, this should be added to the final version. 

Similarly, some of the reviewers asked for more complex experiments, which the authors addressed in their rebuttal with an experiment on era. The results in their current state aren’t interpretable: 2.5degrees corresponds to a very small resolution. Why did the authors choose a 3 hour prediction window? Era5 is typically use for 1 or 6hours prediction. Did the authors perform the downsampling themselves to 2.5 degree? If so was that with conservative regrinding? The choice of 1000/200 samples seems arbitrary: how were these chosen? I would expect to use e.g. 5 years for training and an out-of-sample year for validation + 1 for testing. Lastly, please also report the RMSE as is standard for z500 prediction. 
These results need further details when added to the paper, including rigorous description of the setting, parameters used, metrics, etc, to ensure reproducibility. 

Finally, the authors promised to release the code: please make sure the trained model and scripts to reproduce the new experiments are added to the final version.
Decision: Accept (poster)