Mamba Neural Operator: Who Wins? Transformers vs. State-Space Models for PDEs

## Abstract

Abstract Partial differential equations (PDEs) are widely used to model complex physical systems, but solving them efficiently remains a significant challenge. Recently, Transformers have emerged as the preferred architecture for PDEs due to their ability to capture intricate dependencies. However, they struggle with representing continuous dynamics and long-range interactions. To overcome these limitations, we introduce the Mamba Neural Operator (MNO), a novel framework that enhances neural operator-based techniques for solving PDEs. MNO establishes a formal theoretical connection between structured state-space models (SSMs) and neural operators, offering a unified structure that can adapt to diverse architectures, including Transformer-based models. By leveraging the structured design of SSMs, MNO captures long-range dependencies and continuous dynamics more effectively than traditional Transformers. Through extensive analysis, we show that MNO significantly boosts the expressive power and accuracy of neural operators, making it not just a complement but a superior framework for PDE-related tasks, bridging the gap between efficient representation and accurate solution approximation.

## 1 Introduction

Partial differential equations (PDEs) describe various real-world phenomena, such as heat transfer (Heat Equation), fluid dynamics (Navier-Stokes), and biological systems (Reaction-Diffusion). While analytical solutions are sought, many PDEs—like the Navier-Stokes equations—lack closed-form solutions, making them computationally intensive to solve. Numerical methods, such as finite element, finite difference (Mehra et al., ), and spectral methods, discretise these equations but involve trade-offs between computational cost and accuracy. Coarser grids reduce computational load but sacrifice precision, while finer grids increase both accuracy and computational expense.
Recent advancements in deep learning have changed techniques for solving PDEs. Physics-Informed Neural Networks (PINNs) (Raissi et al., ; Mattey and Ghosh, ) integrate governing equations and boundary conditions into the loss function, but often struggle with generalisation and require retraining for changes in coefficients. Neural operators (Bhattacharya et al., ; Kovachki et al., ), on the other hand, learn mappings between function spaces, offering a mesh-free, data-driven approach that generalises better across different PDE instances.

Operator learning has gained traction with models like DeepONet (Lu et al., ) and the Fourier Neural Operator (FNO) (Li et al., ), which achieved state-of-the-art performance. These models learn input-output mappings to approximate complex operators, similar to sequence-to-sequence problems.
Transformers (Vaswani, ) have become a go-to architecture for PDEs (Cao, ; Li et al., ; Bryutkin et al., ) due to their ability to capture long-range dependencies. However, their quadratic complexity limits efficiency for tasks such as long-time integration. To overcome this, efficient variants like Galerkin attention (Cao, ) reduce computational cost to linear scaling.
While these models improve efficiency, they trade off model capacity by approximating the self-attention mechanism, potentially reducing accuracy for tasks that need precise attention. Moreover, Transformers face challenges with PDEs due to limited context windows, inefficiency with continuous data, and high memory usage, making them less effective for capturing dependencies over continuous domains and high-resolution grids.

While Transformers are popular for PDE modelling, they have limitations in handling continuous data and high-resolution grids. An emerging alternative is State-Space Models (SSMs) (Gu et al., ; ), which offer better scalability, reduced memory usage, and improved handling of long-range dependencies in continuous domains compared to Transformers. In particular, Mamba (Gu and Dao, ) is a novel way designed to effectively capture long-range dependencies, handle continuous data efficiently, and reduce memory consumption in sequence-to-sequence problems. Although Transformers dominate applications like foundational models and computer vision, the use of SSMs—especially Mamba—for neural operators in PDEs remains underexplored, and their theoretical connections and potential advantages are yet to be fully understood.

Contributions. We introduce the concept of Mamba Neural Operator (MNO), which provides a novel perspective applicable to Transformer-based techniques for PDEs. Unlike closely related works, we offer a formal theoretical connection between Mamba and Neural Operators, demonstrating its advantages for PDEs. MNO addresses key challenges in PDE modelling by leveraging its structured state-space design to capture long-range dependencies and continuous dynamics more effectively than Transformers. Our particular contributions are as follows.

★ We introduce the concept of the Mamba Neural Operator (MNO), where we underline:

- –
Mamba Neural Operator expands the SSM framework into a unified neural operator approach, making it adaptable to diverse architectures, including any Transformer-based model.
- –
Unlike existing related works, we provide a theoretical understanding that shows how neural operator layers share a comparable structural framework with time-varying SSMs, offering a new perspective on their underlying principles.

★ We evaluate MNO on various architectures and PDEs, showing through systematic analysis that Mamba enhances the expressive power and accuracy of neural operators. This indicates that Mamba is not just a complement to Transformers, but a superior framework for PDE-related tasks, bridging the gap between efficient representation and accurate solutions.

## 2 Related Work

Data-Driven PDEs.
Recent advances in fluid dynamics and solving PDEs have led to architectures modelling continuous-time solutions and multiparticle dynamics (Kochkov et al., ; Lusch et al., ). Physics-informed models now offer solutions in unsupervised and semi-supervised settings (Raissi et al., ; Li et al., ). These models typically encode spatial data and evolve over time, utilising methods like convolutional layers (Ronneberger et al., ; Wiewel et al., ) symbolic neural networks (Udrescu and Tegmark, ), and residual networks (He et al., ). Finite element methods (FEM), including Galerkin and Ritz, are also integrated into learning frameworks (Chen et al., ).

Neural Operators.
Neural operators, such as the Graph Neural Operator (Li et al., ) and Fourier Neural Operator (Li et al., ), excel at learning mappings in infinite-dimensional spaces, particularly by leveraging techniques like graph structures or transformations in Fourier space. The Fourier Neural Operator (FNO) and its variations, including the incremental, factorised, adaptive FNO, and FNO+ (Zhao et al., ; Tran et al., ; Guibas et al., ) have shown exceptional performance in both speed and accuracy. Their key advantage lies in their ability to maintain discretisation invariance, which sets them apart in many applications. DeepONet (Lu et al., ) pioneered the nonlinear operator approximation using separate networks for inputs and query points, while extensions like MIONet handle multiple inputs (Jin et al., ). Challenges like irregular grids are being addressed through grid mapping and subdomain partitioning (Li et al., ; Wen et al., ) though scalability for diverse inputs remains a key focus.

Transfomers for PDEs.
The Transformer model (Vaswani, ) stands out due to its distinctive features, primarily its use of attention mechanisms to model the relationships among input elements.
Initially, it was developed for NLP, and attention mechanisms have been adapted to PDEs, providing flexible and efficient mappings between function spaces.
Galerkin attention (Cao, ) introduced linear complexity to reduce computational costs, inspiring further developments like GNOT (Hao et al., ) and OFormer (Li et al., ), which achieve state-of-the-art results. Additionally, graph-based Transformers have also been explored to capture complex interactions in irregular domains (Bryutkin et al., ).

State-Space Models for PDEs & Comparison to Ours.
Initial studies on SSMs for PDEs, like MemNO (Buitrago Ruiz et al., ), explored combining FNO with S4 but were restricted to low-resolution or noisy inputs. In contrast, we introduce the Mamba Neural Operator, which generalises the SSM framework to neural operators, making it compatible with any architecture, including Transformers. Our approach extends the theoretical foundations for broad applicability to any PDE family, highlighting Mamba’s effectiveness in diverse scenarios.
At the time of our submission, the work of that (Hu et al., ) proposed integrating state-space models into neural operators for dynamical systems. While related, our work differs significantly, their approach focuses on dynamical systems and tests only on ordinary differential equations (ODEs), whereas we target parametric partial differential equations (PDEs). Additionally, we provide a theoretical understanding showing that neural operator layers share a comparable structural framework with time-varying SSMs, demonstrating alignment between hidden space updates and the iterative process in neural operators.

## 3 Mamba Neural Operator

This section details the theoretical underpinning and practicalities of the Mamba Neural Operator. We outline its design, key components, and operational mechanisms, explaining how it efficiently models partial differential equations by leveraging structured state-space models (SSMs).

### 3.1 Problem Statement

We consider parametric partial differential equations (PDEs) defined on a domain $\Omega\subset\mathbb{R}^{n}$, parameterised by $\theta\in S\subset\mathbb{R}^{p}$, where $\theta$ is sampled from a distribution $w$. The general form of the PDE is:

$$ $\displaystyle P:\mathcal{P}\times\Omega\times\mathcal{W}\times\mathbb{R}^{m} \times\ldots\times\mathbb{R}^{m}\to\mathbb{R}^{\ell},\quad\Omega\subset\mathbb {R}^{n},W\subset\mathbb{R}^{m},$ (1) $\displaystyle P(\theta,x,u,\partial_{x_{1}}u,\ldots,\partial_{x_{n}}u,\ldots, \partial^{\beta_{1}}_{x_{1}}\cdots\partial^{\beta_{n}}_{x_{n}}u)=0,$ $$

where the unknown function $u:\Omega\rightarrow V$ solves $P$. The multi-index $\beta=(\beta_{1},\ldots,\beta_{n})$, with $|\beta|=\sum_{i=1}^{n}\beta_{i}$, determines the differentiation orders. If time is involved, $\Omega$ reduces to $\mathcal{T}\subset\mathbb{R}_{\geq 0}$ and $\Omega\subset\mathbb{R}^{n-1}$. To ensure well-posedness, initial and boundary conditions must hold:

$$ $\displaystyle u(x,T_{0})=u_{0}(x),\;x\in\Omega_{\theta},\quad\quad u(x,t)=u_{b }(x),\;x\in\partial\Omega_{\theta},\;t\in\mathcal{T},$ (2) $$

for $x\in\Omega_{\theta}$ and $t\in\mathcal{T}$, where $u_{0}$ and $u_{b}$ are the initial and boundary conditions, respectively. Assume $\Omega,\mathcal{P},V$ are Banach spaces, and there exists an analytic solution operator: $O:\mathcal{P}\times\Omega\times\mathbb{R}^{m}\times\ldots\times\mathbb{R}^{m}
\times\mathbb{R}^{\ell}\times\mathbb{R}^{\ell}\rightarrow V.$
Our aim is to design a neural network $\tilde{S}\mu:(\theta,u_{0},u_{b})\mapsto u$ that approximates this operator, with $\mu$ as the network’s parameters. Given a dataset ${(\theta^{(n)},u^{(n)})}{n=1}^{N}$, where $\theta^{(n)}$ and $u^{(n)}$ correspond to the system’s discretised parameters, we simplify the notation as $\theta^{(n)}=\theta(x^{(n)})$ and $u^{(n)}=u(x^{(n)})$.

### 3.2 Preliminaries: Mamba

Transformers are the leading architecture for many state-of-the-art techniques for PDEs, with preliminaries introduced in Appendix A. In this section, we outline the background on State Space Sequence models (SSM).

Structured State Space (S4) models introduce a new approach in deep learning sequence modelling, incorporating elements from Recurrent Neural Networks (RNNs), Convolutional Neural Networks (CNNs), and classical state space models. These models are inspired by control theory, where the process involves mapping an input sequence $u(t)\in\mathbb{R}^{L}$ to an output sequence $y(t)\in\mathbb{R}^{L}$ through a hidden latent state $h(t)\in\mathbb{R}^{N}$. The core mechanism of State Space Models (SSMs) is formulated using linear first-order ordinary differential equations, enabling efficient handling of temporal data, which reads:

$$ $\displaystyle h^{\prime}(t)=Ah(t)+Bu(t),\quad y(t)=Ch(t)+Du(t),$ (3) $$

where $A\in\mathbb{C}^{N\times N}$ and $B,C\in\mathbb{C}^{N}$. Mamba, a more advanced variant of SSMs, refines this formulation by incorporating efficient state space parameterisation and selection mechanisms. Unlike earlier models such as S4, which uses bilinear method, Mamba adopts zero-order holds, allowing it to handle larger hidden states and longer sequences more effectively. This makes Mamba particularly well-suited for complex sequence modeling tasks, such as natural language processing and time-series analysis.

### 3.3 State Space Models Discretisation for PDEs

State Space Models (SSMs) have emerged as a strong alternative to Transformers in deep learning. While Transformers dominate in areas like foundational models and computer vision, the application of SSMs, particularly the Mamba architecture, to neural operators for PDEs is still underexplored.

We start by demonstrating that the discretisation of S6 (Mamba) is equivalent to the well-known Eular method when the Taylor series expansion is applied. Mamba utilises zero-order holds, resulting in the following discretisation, which reads:

$$ $\displaystyle A=\exp(\Delta A),\quad B=(\Delta A)^{-1}\left(\exp(\Delta A)-I \right)\cdot\Delta B.$ (4) $$

Discretisation of SSM. To incorporate the SSM into deep learning frameworks, we need to transform the continuous-time SSM into a discrete formulation. This is done by expressing the continuous-time system as an ordinary differential equation (ODE) and then solving it numerically. As discussed in (Liu et al., ), the discrete SSM reads:

$$ $\begin{split}h_{a+1}&=e^{A\Delta a}(h_{a}+B_{a}U_{a}e^{-A\Delta a\Delta a})\\ &=e^{A\Delta a}h_{a}+B_{a}\Delta_{a}u_{a}=\bar{A}_{a}h_{a}+\bar{B}_{a}u_{a}, \end{split}$ (5) $$

where $\Delta$ is the time step size, and $\bar{A}_{a}=e^{A\Delta a}$ and $\bar{B}=B_{a}\Delta_{a}$ are the discretised system matrices.
In the S6 model, we define $\tilde{A}=e^{\Delta A}$ and $\tilde{B}=(\Delta A)^{-1}(e^{\Delta A}-I)\cdot\Delta B$. By applying sampling in $\tilde{A}$, we have $\bar{A}=\tilde{A}$.
For $\tilde{B}$, applying the sampling process yields to:

$$ $\begin{split}\tilde{B}&=(\Delta A)^{-1}(e^{\Delta A}-I)\cdot\Delta B=(\Delta A )^{-1}(I+\Delta A+O(\Delta^{2})-I)\cdot\Delta B\\ &=(\Delta A)^{-1}(\Delta A+O(\Delta^{2}))\cdot\Delta B=\Delta B(Drop~{}O( \Delta^{2}))=\bar{B}\end{split}$ (6) $$

Thus, we have shown that our discretisation method is equivalent to the zero-order hold method, where $\bar{A}=\tilde{A}$ and $\bar{B}=\tilde{B}$.

###### Proposition 1 .

The zero-order hold discretisation method, as in (4), is equivalent to the Euler method in SSM when the Taylor series expansion of the exponential function is truncated to its first-order term.

###### Proof.

In SSM, if we define the matrices as
$\hat{A}=I+\Delta A$ and $\hat{B}=\Delta B$. Then the discretised form of the state update can be written as:

$$ $\begin{split}h(t+\Delta t)&=\hat{A}h(t)+\hat{B}u(t)=(I+\Delta A)h(t)+\Delta Bu (t)\\ &=h(t)+\Delta(Ah(t)+Bu(t))=h(t)+\Delta^{\prime}(t).\end{split}$ (7) $$

which implies it is a first order eular method.
It is straightforward to show that
$\tilde{A}=\hat{A}$ since $\tilde{A}=e^{A\Delta}=I+A\Delta+O(\Delta^{2})=I+A\Delta=I+A\Delta=\hat{A}$. Similarly, we observe that $\tilde{B}=\bar{B}=\hat{B}$.
Therefore, the discretisation used in the SSM method can be replaced with the zero-order hold method by substituting $\hat{A}=\tilde{A}$ and $\hat{B}=\tilde{B}$, we get:

$$ $\begin{split}h(t+\Delta t)&=\hat{A}h(t)+\hat{B}u(t)=\tilde{A}h(t)+\tilde{B}u(t )\\ &=(e^{\Delta A})h(t)+((\Delta A)^{-1}(e^{\Delta A}-I)\cdot\Delta B)u(t)=(I+ \Delta A+O(\Delta^{2})h(t)+(\Delta B)u(t)\\ &=(I+\Delta A)h(t)+(\Delta B)u(t)=h(t)+\Delta Ah(t)+\Delta Bu(t)\\ &=h(t)+\Delta(Ah(t)+Bu(t))=h(t)+\Delta^{\prime}(t).\end{split}$ (8) $$

This shows that the zero-order hold discretisation method is equivalent to the Euler method, as both yield the same discrete update formula.

∎

Why is Proposition 1 important for PDEs? Proposition 1, which establishes the equivalence between the Zero-Order Hold (ZOH) method and the Euler method, is crucial for understanding Mamba’s performance in solving partial differential equations (PDEs). This equivalence demonstrates that ZOH can be viewed as a more generalised and accurate variant of the Euler method. Specifically, while the Euler method is derived by applying the Taylor series expansion and truncating it at the first order, the ZOH method retains additional higher-order terms from the Taylor series, making it inherently more accurate.
This distinction has significant implications when solving PDEs. Higher-order methods like ZOH provide better approximations of a system’s behaviour without requiring excessively small step sizes $\Delta$, which are often necessary for the Euler method to achieve a similar level of accuracy. Smaller step sizes can result in increased computational cost and potential numerical instability. By utilising ZOH’s higher-order accuracy, the Mamba architecture can handle a wide range of step sizes, ensuring both stability and convergence without compromising precision.

###### Proposition 1 .

###### Proof.

### 3.4 Mamba for Neural Operators

Neural Operators (Li et al., ) aim to learn mappings between function spaces, providing a framework for solving partial differential equations (PDEs) and other problems involving continuous functions. It updates the value by an iterative method: $i_{0}\rightarrow i_{1}\rightarrow\ldots\rightarrow i_{T},$ where each $i_{j}$ (for $j=0,1,...,T-1$) maps to $\mathbb{R}^{d_{v}}$. Let the input be $a(x)$ and the output be $u(x)$ . The input $a$, drawn from set $A$, is initially lifted to a higher-dimensional representation: $v_{0}(x)=P(a(x))$
where $P$ is a local transformation, typically parameterised by a fully-connected neural network. We then apply iterations to update $i_{t}\rightarrow i_{t+1}$ as defined in Definition 1. The final output: $u(x)=Q(v_{T}(x))$ is the result of projecting $v_{T}$ via the transformation: $Q:\mathbb{R}^{d_{v}}\to\mathbb{R}^{d_{u}}.$ Each update from $i_{t}$ to $i_{t+1}$ involves the integration of a non-local integral operator $K$ and a local nonlinear activation function $\sigma$. One of the main results of this work is establishing the equivalence between neural operators and the Mamba framework. Therefore, we first introduce fundamental definitions stated in  (Li et al., ) that are essential for demonstrating this relationship.

Figure: Figure 1: (A) Illustration of Mamba Neural Operator. Input image patches are processed by following two distinct scanning paths (referred to as Bidirectional -Scan). Each sequence generated from these paths is passed through separate S6 blocks/ Cross S6 Blocks for independent processing. Afterwards, the outputs from the S6 blocks / Cross S6 Blocks are combined to form a feature map, resulting in the final output (Bidirectional-Merge). (B) and (C) are the detailed block of the S6 Block and Cross S6 Block respectively. The detail network architecture and definition of Cross S6 Block can be found in Appendix A.
Refer to caption: extracted/5896769/image/teaserpde.jpg

###### Definition 1 .

(Iterative updates): The update from $i_{t}\rightarrow i_{t+1}$ is defined as follows:

$$ $\displaystyle i_{t+1}(x):=\sigma\left(Wi_{t}(x)+K_{\phi}(a)i_{t}(x)\right), \quad\forall x\in D,$ (9) $$

where $K:A\times\Theta_{K}\rightarrow L(U(D;\mathbb{R}^{d_{v}}),U(D;\mathbb{R}^{d_{v}
}))$ represents a mapping to bounded linear operators on $U(D;\mathbb{R}^{d_{v}})$, parameterised by $\phi\in\Theta_{K}$. The function $W:\mathbb{R}^{d_{v}}\rightarrow\mathbb{R}^{d_{v}}$ is a linear transformation, and $\sigma:\mathbb{R}\rightarrow\mathbb{R}$ is a nonlinear activation function applied component-wise.

###### Definition 2 .

(Kernel integral operator $K$): Define the kernel integral operator mapping in 1 by

$$ $\displaystyle K_{\phi}(a)i_{t}(x):=\int_{D}\kappa_{\phi}(x,y,a(x),a(y))i_{t}(y )\,dy,\quad\forall x,$ (10) $$

where $\kappa_{\phi}:\mathbb{R}^{2(d+d_{a})}\rightarrow\mathbb{R}^{d_{v}\times d_{v}}$ is a neural network parameterised by $\phi\in\Theta_{K}$.

As mentioned in the previous section, we can be discrete SSM into the form of (5). This representation can be rewrite as (Liu et al., ): $\mathbf{h}_{b}=\mathbf{w}_{T}\odot\mathbf{h}_{a}+\sum_{i=1}^{T}\frac{\mathbf{w
}_{T}}{\mathbf{w}_{i}}\odot(\mathbf{K}_{i}^{\top}\mathbf{V}_{i}).$
We define $\mathbf{V}=[\mathbf{V}_{1};\ldots;\mathbf{V}T]\in\mathbb{R}^{T\times D_{v}}$, where $\mathbf{V}_{i}=u_{a+i-1}\Delta_{a+i-1}\in\mathbb{R}^{1\times D_{v}}$, $\mathbf{K}=[\mathbf{K}_{1};\ldots;\mathbf{K}_{T}]\in\mathbb{R}^{T\times D_{k}}$, where $\mathbf{K}_{i}=\mathbf{B}_{a+i-1}\in\mathbb{R}^{1\times D_{k}}$ , and $\mathbf{Q}=[\mathbf{Q}_{1};\ldots;\mathbf{Q}_{T}]\in\mathbb{R}^{T\times D_{k}}$, where $\mathbf{Q}_{i}=\mathbf{C}_{a+i-1}\in\mathbb{R}^{1\times D_{k}}$. We further define $\mathbf{w}=[\mathbf{w}_{1};\ldots;\mathbf{w}_{T}]\in\mathbb{R}^{T\times D_{k}
\times D_{v}}$, where $\mathbf{w}_{i}=\prod_{j=1}^{i}e^{\mathbf{A}\Delta_{a-1+j}}\in\mathbb{R}^{D_{k}
\times D_{v}}$, and $\mathbf{H}=[\mathbf{h}_{a};\ldots;\mathbf{h}_{b}]\in\mathbb{R}^{T\times D_{k}
\times D_{v}}$, where $\mathbf{h}_{i}\in\mathbb{R}^{D_{k}\times D_{v}}$. Finally, we set $\mathbf{Y}=[\mathbf{y}_{a};\ldots;\mathbf{y}_{b}]\in\mathbb{R}^{T\times D_{v}}$, where $\mathbf{y}_{i}\in\mathbb{R}^{D_{v}}$

This formulation indicates that Gated Linear Attention (Yang et al., ) is actually a specific variant of Mamba.
We next present our main result is how neural operator layers share a comparable structural framework with time-varying SSMs, which, to the best of our knowledge, is established here for the first time.

###### Proposition 2 .

The hidden space in time-varying state-space models demonstrates a structural similarity to neural operator layers.

###### Proof.

We first rewrite the time-varying SSMs (5) as:

$$ $\displaystyle\mathbf{h}_{b}=\mathbf{w}_{T}\odot\mathbf{h}_{a}+\sum_{i=1}^{T} \frac{\mathbf{w}_{T}}{\mathbf{w}_{i}}\odot(\mathbf{K}_{i}^{\top}\mathbf{V}_{i}),$ (11) $$

where $\mathbf{w}_{T},\mathbf{w}_{i},\mathbf{K}_{i},\mathbf{V}_{i}$ are as defined previously.

To demonstrate that the hidden space update in our Mamba Operator has a similar structural framework to neural operator layers, we assume the shapes of $\mathbf{w}$ and $\mathbf{h}$ are $(T,D_{k})$, represented as vectors. Our goal is to show that the iterative process in (5) aligns with that of Definition 1. Consider the first part of Definition 1, represented by $W_{i}(x)$. We set $W=\mathbf{W}_{T}$ and $i_{t}(x)=\mathbf{h}_{a}$, where $\mathbf{h}_{a}$ is the hidden state from the previous iteration. We then verify that $\mathbf{W}_{T}$ satisfies the properties of a linear transformation, ensuring consistency with the neural operator framework.

We can proved this as follows: without loss of generality, let us assume that $W_{T}=W_{1}=e^{\mathbf{A}\Delta_{a}}$. Next, we apply this transformation to a vector $x$ and check the conditions for linearity: $T(x+y)=T(x)+T(y)$ and $T(ax)=aT(x)$. By applying the Taylor expansion to $e^{\mathbf{A}\Delta_{a}}$ , then we get $I+A\Delta_{a}+O(\Delta_{a}^{2})$. To show that $T(x+y)=T(x)+T(y)$, it suffices to demonstrate: $e^{A\Delta_{a}}(x_{1}+x_{1})=e^{A\Delta_{a}}(x_{1})+e^{A\Delta_{a}}(x_{2})$, we have:

$$ $\displaystyle e^{A\Delta_{a}}(x_{1}+x_{1})$ $\displaystyle=(I+A\Delta_{a}+O(\Delta_{a}^{2})(x_{1}+x_{2})=I(x_{1}+x_{2})+A \Delta_{a}(x_{1}+x_{2})+O(\Delta_{a}^{2})(x_{1}+x_{2})$ (12) $\displaystyle=Ix_{1}+A\Delta_{a}(x_{1})+O(\Delta_{a}^{2})(x_{1})Ix_{2}+A\Delta _{a}(x_{2})+O(\Delta_{a}^{2})(x_{2})$ $\displaystyle=(I+A\Delta_{a}+O(\Delta_{a}^{2}))(x_{1})+(I+A\Delta_{a}+O(\Delta _{a}^{2}))(x_{2})$ $\displaystyle=e^{A\Delta_{a}}(x_{1})+e^{A\Delta_{a}}(x_{2})$ $$

This shows $T(x+y)=T(x)+T(y)$. For the second condition, we want to show $T(\alpha x)=\alpha T(x)$, which is equivalent to demonstrating that $e^{A\Delta_{a}}(\alpha x)=\alpha e^{A\Delta_{a}}x$, we get:

$$ $\displaystyle e^{A\Delta_{a}}(\alpha x)$ $\displaystyle=(I+A\Delta_{a}+O(\Delta_{a}^{2}))(\alpha x)=I\alpha x+A\Delta_{a }\alpha x+O(\Delta_{a}^{2})\alpha x$ (13) $\displaystyle=\alpha(Ix+A\Delta_{a}x+O(\Delta_{a}^{2})x)=\alpha(e^{A\Delta_{a} }x)$ $$

Thus, we have shown that $e^{A\Delta_{a}}$ satisfies the two conditions, and hence it is a linear transformation. This shows the update in hidden space is the same as neural operator.

Secondly, we need to check the second part of Definition 1, which involves showing that:

$$ $\displaystyle K_{\phi}(a)i_{t}(x):=\sum_{i=1}^{T}\frac{\mathbf{w}_{T}}{\mathbf {w}_{i}}\odot(\mathbf{K}_{i}^{\top}\mathbf{V}_{i})$ (14) $$

has a similar structure.

According to Definition 2, it suffices to demonstrate that:
$\int_{D}\kappa_{\phi}(x,y,a(x),a(y))i_{t}(y)\,dy=\sum_{i=1}^{T}\frac{\mathbf{w
}_{T}}{\mathbf{w}_{i}}\odot(\mathbf{K}_{i}^{\top}\mathbf{V}_{i})$. We assume the kernel $\kappa_{\phi}$ can be decomposed into a finite sum of separable basis functions: $\kappa_{\phi}(x,y,a(x),a(y))=\sum_{i=1}^{T}\omega_{i}\varphi_{i}(x)\psi_{i}(y)$ such that $\omega_{i}$ is learnable weights for each basis function.
and Basis functions capturing interactions between $x$ and $y$. Then we substitute it into the integral such that : $\int_{D}\sum_{i=1}^{T}\omega_{i}\varphi_{i}(x)\psi_{i}(y)v_{t}(y)\,dy=\sum_{j=
1}^{T}\omega_{i}\int_{D}\varphi_{i}(x)\psi_{i}(y)v_{t}(y)\,dy$. We further discretise the domain
$D$ into $T$ points $\left\{\mathbf{y}_{i}\right\}_{i=1}^{T}$ with corresponding weights $\Delta y$. The integral becomes $K_{\phi}(a)i_{t}(\mathbf{x})\approx\sum_{i=1}^{T}\omega_{i}\sum_{j=1}^{T}
\varphi_{i}(x)\psi_{i}(y_{j})i_{t}(\mathbf{y}_{j})\Delta y$. We represent
$\varphi_{i}(\mathbf{x})$ as vector $\mathbf{K}_{i}$ and the input $i_{t}(\mathbf{y}_{i})$
as vector $\mathbf{V}_{i}$: $\mathbf{K}_{i}=\left[\varphi_{i}(\mathbf{x}),\varphi_{i}(\mathbf{x}),\dots,
\varphi_{i}(\mathbf{x})\right]^{\top}\in\mathbb{R}^{1\times D_{k}},\quad
\mathbf{V}_{i}=\left[\psi_{i}(y_{1})i_{t}(y_{1})\Delta y,\dots,\psi_{i}(y_{T})
i_{t}(y_{T})\Delta y\right]^{\top}\in\mathbb{R}^{1\times D_{k}}$. If we further factorise $\omega_{i}$ as $\frac{w_{T}}{w_{i}}$, where $w_{T}$ is a hyperparameter and $w_{i}$ represents a set of parameters to be learned, we obtain the update:
$K_{\phi}(a)i_{t}(x):=\sum_{i=1}^{T}\frac{\mathbf{w}_{T}}{\mathbf{w}_{i}}\odot(
\mathbf{K}_{i}^{\top}\mathbf{V}_{i})$. Consequently, the neural operator layer shares a comparable structural framework with time-varying SSMs, demonstrating that the hidden space update in these models aligns with the iterative process in neural operator layers.

∎

###### Definition 1 .

###### Definition 2 .

###### Proposition 2 .

###### Proof.

## 4 Experiments and Discussion

In this section, we thoroughly describe the implementation setup and present experimental results to validate Mamba Neural Operators along with Transformers.

### 4.1 Dataset Description & Implementation Protocol

★ PDEs Selection.
We utilise datasets from PDEBench (Takamoto et al., ), a publicly available benchmark for partial differential equations (PDEs). We focus on three PDEs representing both stationary and time-dependent problems: Darcy Flow, Shallow Water 2D (SW2D), and Diffusion Reaction 2D (DR2D). All simulations are performed on a uniform grid. Detailed information about the datasets is provided in the Appendix.

★ Implementation & Evaluation Protocol.
As Transformers have become the go-to architecture for PDE modelling and serve as the primary counterpart to SSM models, we selected three state-of-the-art Transformers as our baselines: GNOT (Hao et al., ), Galerkin Transformer (G.T.) (Cao, ), and OFormer (Li et al., ).
To achieve a fair comparison between Transformers and Mamba, we integrated the S6 block and Cross S6 block to replace self-attention and cross-attention in each model, creating modified versions of the original architectures.
All three experimental methods initially adopt a linear attention mechanism as described in their original publications, while we evaluated two configurations for each of them: an implementation with standard softmax attention mechanism (w/S.A.) and a Mamba-enhanced implementation (our Mamba Neural Operator principle) (w/Mamba).
All experiments were conducted on a single NVIDIA RTX 4090 GPU with 24GB of memory to ensure consistent and fair comparison conditions.
Three metrics including Root Mean Squared Error (RMSE), Normalised RMSE (nRMSE), and Relative L2 Norm (RL2) were utilised for evaluation.

**Table 1: Quantitative comparison on Darcy Flow ($\beta=1$) across three methods with linear attention (original version), softmax attention and Mamba. The performance is measured in terms of Root Mean Squared Error (RMSE), Normalised RMSE (nRMSE), and Relative L2 Norm (RL2), with the best-performing results highlighted.**
|  |  | DarcyFlow |  |  |
| --- | --- | --- | --- | --- |
| Method | Type | RMSE$\downarrow$ | nRMSE$\downarrow$ | RL2$\downarrow$ |
| GNOT (Hao et al., ) | Galerkin | 0.0070 | 0.0485 | 0.0370 |
| w/S.A. | Softmax | 0.0061 | 0.0394 | 0.0299 |
| w/Mamba (MNO) | Mamba | 0.0061 | 0.0367 | 0.0297 |
| G.T. (Cao, ) | Galerkin | 0.0188 | 0.2027 | 0.1261 |
| w/S.A. | Softmax | 0.0103 | 0.1050 | 0.0648 |
| w/Mamba (MNO) | Mamba | 0.0061 | 0.0382 | 0.0286 |
| OFormer (Li et al., ) | Normalised | 0.0054 | 0.0253 | 0.0242 |
| w/S.A. | Softmax | 0.0066 | 0.0324 | 0.0323 |
| w/Mamba (MNO) | Mamba | 0.0054 | 0.0244 | 0.0241 |

**Table 2: Quantitative comparisons on Shallow Water 2D (SW2D) and Diffusion Reaction 2D (DR2D) across three methods with linear attention (original version) and Mamba. The performance is measured in terms of Root Mean Squared Error (RMSE), Normalised RMSE (nRMSE), and Relative L2 Norm (RL2), with the best-performing results highlighted in green.**
|  |  | SW2D | DR2D |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Method | Type | RMSE$\downarrow$ | nRMSE$\downarrow$ | RL2$\downarrow$ | RMSE$\downarrow$ | nRMSE$\downarrow$ | RL2$\downarrow$ |
| GNOT (Hao et al., ) | Galerkin | 0.0026 | 0.0025 | 0.0027 | 0.0567 | 0.6953 | 0.7233 |
| w/Mamba (MNO) | Mamba | 0.0023 | 0.0022 | 0.0024 | 0.0060 | 0.0811 | 0.0570 |
| G.T. (Cao, ) | Galerkin | 0.0037 | 0.0035 | 0.0038 | 0.0083 | 0.1259 | 0.0723 |
| w/Mamba (MNO) | Mamba | 0.0013 | 0.0013 | 0.0014 | 0.0012 | 0.0183 | 0.0099 |
| OFormer (Li et al., ) | Normalised | 0.0020 | 0.0020 | 0.0021 | 0.0177 | 0.2681 | 0.1559 |
| w/Mamba (MNO) | Mamba | 0.0021 | 0.0021 | 0.0022 | 0.0123 | 0.1712 | 0.1134 |

### 4.2 Chose Your Winner: Transformer vs. Mamba for PDEs

We begin by evaluating the performance of Transformers, their variants, and Mamba on the Darcy Flow dataset, as presented in Table 1. The results demonstrate that incorporating Mamba consistently improves performance across all metrics and models. For GNOT, while the RMSE remains close, the nRMSE and RL2 values are reduced, indicating that Mamba effectively refines predictions. The G.T. sees the most significant enhancement, with the RMSE dropping by 40% when Mamba is used. This suggests that Mamba’s design addresses the shortcomings of traditional Galerkin-type attention in capturing complex PDE dynamics.
For OFormer, Mamba not only retains the strong baseline performance but also achieves improvements across all metrics. The reduction in RL2 indicates that Mamba’s mechanism is better at mapping the solution space of PDEs with higher precision. Mamba also demonstrates an enhanced ability to capture the complex spatial correlations inherent to Darcy Flow more effectively.

Figure: Figure 2: Results of prediction map and error map of the GNOT across three versions: Galerkin attention, Softmax attention, and Mamba.
Refer to caption: x1.png

On the SW2D dataset, Mamba consistently outperforms the original Transformer models across all metrics. GNOT with Mamba achieves a lower RMSE and RL2, demonstrating Mamba’s ability to capture complex flow dynamics. The G.T. shows the most significant improvement, with RMSE of 65% reduction—highlighting Mamba’s superior capability in accurately representing the system’s behaviour. For OFormer, Mamba maintains comparable values but increases the RL2.
On the DR2D dataset, the Mamba-enhanced models exhibit even more substantial gains. The G.T. sees a dramatic reduction in RMSE and RL2, showing Mamba’s strength in handling the complex dynamics.

The results across all datasets demonstrate a clear advantage of the Mamba Neural Operator overr architectures for PDEs. While Transformers are effective at capturing dependencies and patterns, Mamba’s specialised attention mechanisms provide a more understanding of the complex dynamics involved. By leveraging its unique cross-attention and self-attention blocks, Mamba not only achieves lower error rates but also enhances the stability and precision of predictions, particularly in highly nonlinear systems. These results suggest that Mamba enhances the expressive power and accuracy of neural operators, indicating that it is not just a complement to Transformers but a superior framework for PDE-related tasks, bridging the gap between efficient representation and accurate solution approximation.

We further validate Mamba’s potential through visualisations, as shown in Figure 2. The prediction and error maps reveal that Mamba consistently outperforms all Transformer variants, delivering more accurate solutions with lower error across challenging regions. Mamba handles fine details, particularly in capturing sharp gradients and subtle variations that standard attention mechanisms often miss. Compared to the Galerkin and Softmax attention Transformer models, Mamba reduces error propagation and improves spatial coherence. More Visualization result can be seen in Appendix C.

Figure: Figure 3: Fourier analysis comparing three GNOT versions: Galerkin attention, Softmax attention, and Mamba. The $\Delta$ log amplitude shows how each model handles frequency components. We calculate the change by comparing the log amplitude at the center (0.0 $\pi$) and boundary frequencies (1.0 $\pi$). For clarity, only half-diagonal components of the 2D Fourier-transformed feature maps are shown.
Refer to caption: x2.png

Figure: Figure 4: $\Delta$ log amplitudes for Galerkin attention, Softmax attention, and Mamba. Gray regions indicate the operator, and white regions show MLP. Mamba shows a more stable response across frequencies.
Refer to caption: x3.png

### 4.3 Why the Winner Wins: Breaking Down Mamba’s Win

We aim to explore why Mamba outperforms Transformers by examining the frequency response of feature maps. This analysis helps us understand how each model handles high-frequency signals and evaluate its ability to maintain stability and robustness.
The results in Figure 3 compare the frequency response of three GNOT variants: Galerkin attention, Softmax attention, and Mamba. The Galerkin version shows a sharp decline in high-frequency components, indicating underfitting and loss of fine details. The Softmax version retains more high frequencies but risks instability and noise sensitivity. Mamba, on the other hand, demonstrates a balanced suppression of high-frequency signals, maintaining stability and robustness. The change in log amplitude across the frequency range is more uniform, indicating that Mamba effectively balances between capturing necessary high-frequency information and filtering out noise.
This controlled response across the spectrum highlights why Mamba is better suited for PDEs.

Figure 4 shows the $\Delta$ log magnitudes across the normalised depth for the Galerkin attention, Softmax attention, and Mamba versions of GNOT. The Galerkin and Softmax versions exhibit sharp fluctuations, indicating instability and inconsistent feature extraction at different depths. In contrast, Mamba maintains a steady and flat profile, reflecting robust and stable feature extraction. The gray and white bands indicate the alternating roles of the operator and NLP components, further emphasising Mamba’s balanced performance across layers, making it ideal for handling complex PDEs.

**Table 3: Comparisons with different query positions using nRMSE.**
| Method | Query Positions |  |
| --- | --- | --- |
|  | Identical | Diagonal |
| OFormer | 0.0253 | 0.0318 |
| w/S.A. | 0.0324 | 0.0382 |
| w/Mamba | 0.0244 | 0.0314 |

### 4.4 Ablation Study: Final Battles, Winner Takes All

#### Mamba vs. Transformers in Misalignment: A Battle for Query Positioning.

Table 4 compares nRMSE performance across two query scenarios: Identical positions (input and query points are the same) and Diagonal positions (shifted inputs creating a mismatch). Experiments are done using Darcy Flow. While prior work shows that Transformers handle inconsistent input-query positions well (Li et al., ), our results demonstrate a clear advantage of Mamba in both configurations. For Identical query positions, Mamba version achieves the lowest error, outperforming OFormer and its softmax variant, demonstrating Mamba’s superior ability to capture relationships when input and query points are perfectly aligned.
For Diagonal query positions, where inputs and queries are misaligned, Mamba achieves the best performance compared to OFormer and its variant, demonstrating its superior ability to generalise under spatial shifts.

#### Scaling Down Without Sacrifice: A Battle for Resilience with Limited Data.

Table 4 compares nRMSE performance across different dataset sizes for GNOT, GNOT with softmax attention (w/ S.A.), and GNOT with Mamba. Experiments are carried out using Darcy Flow. As dataset size decreases, Mamba consistently achieves the lowest error, demonstrating superior performance and robustness in data-scarce scenarios. For instance, with the smallest dataset (1K), Mamba achieves an nRMSE of 0.0617, significantly lower than GNOT’s and GNOT w/ S.A.’s, showcasing its resilience and generalisation capability even with limited data. This highlights Mamba’s efficiency in learning meaningful representations with fewer data points, making it a powerful choice for real-world applications where data availability is a constraint.

## 5 Conclusion

We have introduced the concept of the Mamba Neural Operator (MNO), a framework that redefines how neural operators approach PDEs by integrating structured state-space models. Unlike closely related works, we formalise this connection by providing a theoretical understanding that demonstrates how neural operator layers share a comparable structural framework with time-varying SSMs, offering a fresh perspective on their underlying principles. Experimental results show that MNO significantly enhances the expressive power and accuracy of neural operators across various architectures and PDEs. This indicates that MNO is not merely a complement to Transformers, but a superior framework for PDE-related tasks, bridging the gap between efficient representation and precise solution approximation.

## Acknowledgments

CWC is supported by the Swiss National Science Foundation (SNSF) under grant number 20HW-1_220785. It also acknowledge CMI, University of Cambridge.
JH was supported in part by the Imperial College Bioengineering Department PhD Scholarship and the UKRI Future Leaders Fellowship (MR/V023799/1).
GY were supported in part by the ERC IMI (101005122), the H2020 (952172), the MRC (MC/ PC/21013), the Royal Society (IEC NSFC211235), the NVIDIA Academic Hardware Grant Program, the SABER project supported by Boehringer Ingelheim Ltd, Wellcome Leap Dynamic Resilience, and the UKRI Future Leaders Fellowship (MR/V023799/1).
CBS acknowledges support from the Philip Leverhulme Prize, the Royal Society Wolfson Fellowship, the EPSRC advanced career fellowship EP/V029428/1, EPSRC grants EP/S026045/1 and EP/T003553/1, EP/N014588/1, EP/T017961/1,
the Wellcome Innovator Awards 215733/Z/19/Z and 221633/Z/20/Z, CCMI and the Alan Turing Institute.
AAR gratefully acknowledges funding from the Cambridge Centre for Data-Driven Discovery and Accelerate Programme for Scientific Discovery, made possible by a donation from Schmidt Futures, ESPRC Digital Core Capability Award, the EPSRC programme grant in ‘The Mathematics of Deep Learning’ EP/V026259/1, and CMIH and CCIMI, University of Cambridge.

## Supplementary Material

In this supplementary material, we provide additional details regarding our methodology and a more comprehensive description of the dataset used in our experiments.

## Appendix A Supplementary Information

### A.1 Preliminaries: Transformer.

In each Transformer layer, an attention mechanism enables interaction between inputs at varying positions, followed by a position-wise fully connected network applied independently to each position. Specifically, the attention mechanism involves projecting an intermediate representation into three components—query $Q\in\mathbb{R}^{N\times d_{k}}$, key $K\in\mathbb{R}^{N\times d_{k}}$, and value $V\in\mathbb{R}^{N\times d_{v}}$—using three separate position-wise linear layers. These representations are then used to calculate the output as:

$$ $\displaystyle\operatorname{Attention}(Q,K,V)=\mathrm{softmax}\bigg{(}\frac{QK^ {T}}{\sqrt{d_{k}}}\bigg{)}V,$ (S.1) $$

of which the memory complexity is $O(n^{2})$. To reduce the computational inefficiency, Galerkin-type attention was proposed by [Cao, ] to remove Softmax attention with linear complexity. It defines as follows:

$$ $\displaystyle\operatorname{Attention}_{g}(Q,K,V)=\frac{Q(\widetilde{K}^{T} \widetilde{V})}{d},$ (S.2) $$

in which $\widetilde{\cdot}$ denotes layer normalisation, as described in [Lei Ba et al., ]. The Galerkin-type attention mechanism involves two matrix product operations, resulting in a computational complexity of $O(nd^{2})$. This reduces the sequence length dependency to only $O(n)$.

### A.2 Network Architecture.

As depicted in Figure 1 from the main paper, the data processing pipeline in our Mamba Neural Operator (MNO) is composed of three key stages: Bi-Directional Scan Expand, S6/Cross S6 Block, and Bi-Directional Scan Merge. When solving PDEs over a fixed grid, the input data can be structured as grid-based data, similar to an image. In the first stage, Bi-Directional Scan Expand, the MNO unfolds the input data into sequences by traversing the grid along two distinct paths. These sequences, representing input patches, are processed independently in the next step. The second stage, S6/Cross S6 Block, involves processing each patch sequence using either an S6 or Cross S6 block, depending on the model variation being employed. For instance, in the enhanced version of Mamba, the GNOT model utilises a Cross S6 block followed by an S6 block for further refinement. Finally, in the Bi-Directional Scan Merge stage, the processed sequences are reshaped and merged back together to generate the output map, completing the data forwarding process. This structured approach allows the MNO to efficiently handle grid-based input data, enabling scalable solutions for PDEs.

### A.3 Definition of Cross S6.

Let $x$ and $x^{\prime}$ be two independent input vectors. Each input is processed through two independent linear transformation, resulting in corresponding parameter sets $(B,C,\Delta)$ for $x$ and $(B^{\prime},C^{\prime},\Delta^{\prime})$ for $x^{\prime}$. Specifically, these transformations are defined as:

$$ $\displaystyle B,C,\Delta$ $\displaystyle=\text{Linear}_{x}(x),$ (S.3) $\displaystyle B^{\prime},C^{\prime},\Delta^{\prime}$ $\displaystyle=\text{Linear}_{x^{\prime}}(x^{\prime}),$ $$

where $\text{Linear}_{x}$ and $\text{Linear}_{x^{\prime}}$ are the respective linear layers applied to $x$ and $x^{\prime}$.

Next, the parameters $(\tilde{B},\tilde{C},\tilde{\Delta})$ are computed by combining the updated values from both inputs according to the following equations:

$$ $\displaystyle\tilde{B}$ $\displaystyle=B+qB^{\prime},$ (S.4) $\displaystyle\tilde{C}$ $\displaystyle=C+qC^{\prime},$ $\displaystyle\tilde{\Delta}$ $\displaystyle=\Delta+q\Delta^{\prime},$ $$

where $q$ is a scalar ratio controlling the contribution of the second input $x^{\prime}$ to the combined output. Once we have these updated parameters, we apply the State Space Model (SSM) to compute the final output y.

## Appendix B Further Implementation Details

#### Darcy Flow.

The two-dimensional Darcy Flow equation defines as follows:

$$ $\begin{cases}-\nabla\cdot\left(a(x,y)\nabla u(x,y)\right)=f(x,y),&\text{for }( x,y)\in\Omega,\\ u(x,y)=0,&\text{for }(x,y)\in\partial\Omega,\end{cases}$ (S.5) $$

where $a(x,y)$ is the diffusion coefficient, $u(x,y)$ is the solution respectively, and $\Omega=(0,1)^{2}$ is a square domain. In Darcy Flow, the force term $f(x,y)$ is set to be a hyperparameter $\beta$, which influences the scale of the solution $u(x,y)$.
Experiments were performed on the steady-state solution of the 2D Darcy Flow over a uniform square domain. The goal is to approximate the solution operator $S$ defined by:

$$ $\displaystyle S:a\mapsto u,\quad\text{for }(x,y)\in\Omega,$ (S.6) $$

with $a(x,y)$ and $u(x,y)$ as previously defined. Similar as PDEBench [Takamoto et al., ] protocol, we used only $\beta=1.0$ and we divided the training and testing ratio into 9:1 which contains 9,000 samples for training and 1,000 samples for testing.

#### Shallow Water.

We conducted experiments on the two-dimensional Shallow Water equations, which are effective for modeling free-surface flow problems. The equations are formulated as follows:

$$ $\displaystyle\partial_{t}h+\partial_{x}(hu)+\partial_{y}(hv)$ $\displaystyle=0,$ (S.7) $\displaystyle\partial_{t}(hu)+\partial_{x}\left(u^{2}h+\tfrac{1}{2}g_{r}h^{2}\right)$ $\displaystyle=-g_{r}h\,\partial_{x}b,$ $\displaystyle\partial_{t}(hv)+\partial_{y}\left(v^{2}h+\tfrac{1}{2}g_{r}h^{2}\right)$ $\displaystyle=-g_{r}h\,\partial_{y}b,$ $$

where $u=u(x,y,t)$ and $v=v(x,y,t)$ represent the velocities in the horizontal and vertical directions, respectively, and $h=h(x,y,t)$ denotes the water depth. The term $b=b(x,y)$ stands for the spatially varying bathymetry, and $g_{r}$ is the gravitational acceleration.

The dataset simulates a 2D radial dam-break scenario within a square domain $\Omega=[-2.5,2.5]^{2}$ over the time interval $t\in[0,1]$. The initial condition is defined by:

$$ $h(t=0,x,y)=\begin{cases}2.0,&\text{if }\sqrt{x^{2}+y^{2}}<r,\\ 1.0,&\text{if }\sqrt{x^{2}+y^{2}}\geq r,\end{cases}$ (S.8) $$

where the radius $r$ is randomly drawn from a distribution $D(0.3,0.7)$.

Our objective is to approximate the solution operator $S$, defined as:

$$ $\displaystyle S:h|_{t\in[0,t^{\prime}]}\mapsto h|_{t\in(t^{\prime},T]},\quad(x ,y)\in\Omega,$ (S.9) $$

with $t^{\prime}=0.009\,\text{s}$ and $T=1.000\,\text{s}$. Here, $h=h(x,y,t)$ represents the water depth over time.

Each sample in the dataset is discretized on a spatial grid of $128^{2}$ points and a temporal grid of 101 time steps. The first 10 time steps are used as input to the model, while the remaining 91 time steps serve as the target output. Following the protocol established by PDEBench [Takamoto et al., ], the dataset consists of 900 samples for training and 100 samples for testing.

#### Diffusion Reaction.

The Diffusion Reaction equations are expressed as:

$$ $\displaystyle\partial_{t}u$ $\displaystyle=D_{u}\partial_{xx}u+D_{u}\partial_{yy}u+R_{u},$ (S.10) $\displaystyle\partial_{t}v$ $\displaystyle=D_{v}\partial_{xx}v+D_{v}\partial_{yy}v+R_{v},$ $$

where the activator and inhibitor are represented by the functions $u=u(x,y,t)$ and $v=v(x,y,t)$. In addition, these two variables are non-linearly coupled variables. These functions describe the interaction between the activator and inhibitor in the system.
$D_{u}=1\times 10^{-3}$ and $D_{v}=5\times 10^{-3}$ are the diffusion coefficients for the activator and inhibitor, respectively.

The reaction terms for the activator and inhibitor are then defined as follows:

$$ $\displaystyle R_{u}(u,v)=u-u^{3}-k-v,\quad R_{v}(u,v)=u-v,$ (S.11) $$

with $k=5\times 10^{-3}$.
The simulation is performed over the domain $\Omega=[-1,1]^{2}$ with the time interval $t\in[0,5]$. The solution operator $S$ is defined as:

$$ $\displaystyle S:\{u,v\}_{t\in[0,t^{\prime}]}\mapsto\{u,v\}_{t\in(t^{\prime},T] },\quad(x,y)\in\Omega,$ (S.12) $$

where $t^{\prime}=0.045\,\text{s}$ and $T=5.000\,\text{s}$, and the spatial domain is $\Omega=[-1,1]^{2}$. Here, $u=u(x,y,t)$ and $v=v(x,y,t)$ represent the activator and inhibitor, respectively. In this dataset, we follow the same discretization scheme similar to the Shallow Water equation, where each sample is downsampled to a spatial resolution of $128^{2}$ and a temporal resolution of 101 time steps (with 10 for input and rest of the 91 for target). Similar as the PDEBench protocol [Takamoto et al., ], the dataset includes 900 samples for training and 100 samples for testing.

## Appendix C Extended Visual Results

This section extends the visual results presented in the main paper, providing a deeper comparison of the predictive performance and error distribution across different models. The supplementary figures illustrate the impact of incorporating Mamba on prediction accuracy and spatial coherence, further validating its advantages over traditional Transformer-based approaches.

Figure S.1 presents the prediction and error maps for the Galerkin Transformer (G.T.) and OFormer across three configurations: Galerkin attention, standard softmax attention, and Mamba (MNO). The results show that Mamba consistently achieves lower prediction errors, especially in regions with high variability, highlighting its ability to capture complex dynamics with greater precision compared to other configurations.

Figure S.2 and Figure S.3 provide visualised predictions over time for the Shallow Water and Diffusion Reaction datasets, respectively, using the original Galerkin Transformer and its Mamba-enhanced version. For the Shallow Water dataset, the Mamba-integrated model better preserves fine details and the circular wavefronts as time progresses, reflecting its superior capability to maintain spatial coherence. Similarly, in the Diffusion Reaction dataset, Mamba reduces the spread of error and better approximates the reference solution, demonstrating improved stability and generalisation in long-term simulations.

Overall, these visualisations clearly indicate that incorporating Mamba significantly enhances predictive accuracy and robustness, making it a superior choice for capturing intricate spatial-temporal patterns in complex PDE systems.

Figure: Figure S.1: Results of prediction map and error map of the Galerkin Transformer and OFormer across three versions: Galerkin attention, Softmax attention, and Mamba.
Refer to caption: x4.png

Figure: Figure S.2: Visualised prediction on Shallow Water dataset using Galerkin Transformer (G.T.) across the original and Mamba version.
Refer to caption: extracted/5896769/image/SWfinal.jpg

Figure: Figure S.3: Visualised prediction on Diffusion Reaction dataset using Galerkin Transformer (G.T.) across the original and Mamba version.
Refer to caption: extracted/5896769/image/DRvisu.jpg

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see the weaknesses part.
soundness: 2
strengths: 1. The topic explored in this paper is interesting and engaging.

2. The paper offers some theoretical insights that enhance our understanding of operator learning.

3. The experiments conducted are thorough.
summary: This paper presents the Mamba Neural Operator (MNO), a new approach designed to address the challenges in modeling and solving partial differential equations (PDEs) using neural networks. The MNO integrates structured state-space models (SSMs) into neural operators, allowing it to capture long-range dependencies and continuous dynamics more effectively than existing Transformer-based methods.
weaknesses: 1.  The paper feels somewhat rushed, which has resulted in several typos that hinder understanding of the content. Here are some specific issues I've noticed:

   a. In Section 3.1, there is a mistake in one of the formulas ($n=1^N$), which could confuse readers trying to grasp the main concepts.

   b. In Section 3.2, it’s unclear how $B$, which is an n-dimensional complex vector, and $u(t)$, which is an L-dimensional real vector, are multiplied in Equation 3. This lack of clarity raises questions about the mathematical operations being performed.

   c. In Equation 4, the variables $A$ and $B$ are introduced, but it's unclear whether they are the same as what’s defined in Equation 3.

   d. In Equation 5, there seems to be inconsistency with the notation: are $U_a$ and $u_a$ meant to represent the same entity? Additionally, is $\Delta_a$ the same as $\Delta a$?

   e. In Equation 6, the term "Drop" is used without adequate explanation, leaving readers uncertain about its intended meaning.

2. In Section 3.3 (Page 5), the authors should specify which order of terms they are referring to when they mention "higher-order terms." Providing this detail is essential for understanding the implications of their claims. Furthermore, if the authors intend to use this aspect of the ZOH to demonstrate the superiority of Mamba, they need to establish that traditional transformers cannot align with the Taylor expansion. 

3. In the experimental results, the authors should report the differences in training and inference times between Mamba and other methods, particularly focusing on inference time, as this is crucial for neural operators. Additionally, I recommend that the authors include comparisons with traditional numerical methods. 

4. The authors primarily focus their experiments on a two-dimensional setting. However, it raises an important question: when extending to a three-dimensional setting, which approach would be more suitable—Mamba or transformers? It’s worth mentioning that there are already existing works that utilize transformers for 3D neural operators [1].

[1] Li Z, Liu T, Peng W, et al. A transformer-based neural operator for large-eddy simulation of turbulence[J]. arXiv preprint arXiv:2403.16026, 2024.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 1
questions: 1) MNO uses bidirectional scans, which are from two directions. As continuity is essential for PDEs, are there other directions you can also consider, like ZigZag Mamba? How is it different from vision mamba? 

2) I don't get why we need to consider the cross-scan in MNO. Have you done ablation removing the cross scan, as you have already considered the bidirectional scan? Or does it help boost performance? Have you tried running naive vision Mamba as a baseline?

3) I found the paper difficult to read on Page 6 after the Kernel Integral Operator definition and spent quite a lot of time trying to understand it better. Notation could be more precise. Presentation could be improved as in the main script, MNO architecture can be adequately described. Use consisent notation as in Equ 5 (Captial U and small u) and in proposition 2.

4) How is Proposition 2 Kernel Integral when making specific points out of domain? Isn't it more like an approximation? Also, part 2 proof is confusing.

5) Could you please provide a more comprehensive description of the MNO architecture? If you are adding nonlinearity such as FNO, hyperparameters, etc. How did you finetune the hyperparamters?

6) Does MNO hold a universal approximation of an operator like FNO?

7) See Weakness
soundness: 3
strengths: 1) Theoretical Connection between Mamba and Neural Operator.

2) MNO addresses long-range dependencies and continuous dynamics better than a transformer.
summary: Introduces Mamba Neural Operator as an alternative to transformer.
weaknesses: 1) The presentation could be improved by addressing notational inconsistencies. 

2) Lack of Novelty as compared to Vision Mamba.

3) Weak baselines. SOTA transformers such as Transolver, Latent Neural Operator, DPOT, LSM, etc.

4) Missing model parameters, training time, inference time comparison and MNO Hyperparamters.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 3
questions: Personally, I find some problems with the description of Transformer-based PDE solving in the introduction part. First, Transformer architectures are not mainstream in Neural Operators, so many models such as GNO, FNO, KNO, CNO, and DeepONet, which tend to prefer more lightweight structures instead of Transformer. This raises questions about the framing in the introduction. The authors mention, "While Transformers are popular for PDE modelling, they have limitations in handling continuous data and high-resolution grids." However, in PDE problems, Transformer architectures often perform poorly at lower resolutions (e.g., 64x64 or 128x128) and achieve better results at higher resolutions (such as 720x1440), as evidenced by AFNO and FourCastNet.

In summary, there are two questions.
1. Please explain why Transformers were chosen as the comparison baseline while other operator methods were overlooked.
2. Please clarify the reasoning behind the claim that Transformers perform poorly at high resolutions?
soundness: 1
strengths: Mathematical derivations are complete, coherent, and consistent. Combines the mathematical frameworks of Mamba (Gu, 2023) and Neural Operator (Kovachki, 2023).
summary: This paper integrates Mamba and designs the MNO model to solve PDE equations. It compares Transformer-based models and draws some inconsequential conclusions.
weaknesses: (1) This paper sets the Transformer as a target for comparison in studying neural operators, disregarding all other neural operators. However, the Transformer is neither mainstream nor representative in the neural operator field (see Questions part for details). This setting limits the value of the conclusions drawn in the paper.

(2) The paper uses RMSE as a metric for evaluating PDE solutions but fails to address critical properties of PDE solving. This undermines the persuasiveness of the proposed model, suggesting that MNO may be a better tool for approximation series rather than an effective PDE solver. The paper overlooks the generalization capability of Neural Operators in Banach spaces—an essential feature of neural operators. It also fails to discuss resolution invariance or any specific physical or mathematical consistency with respect to the target PDE system. This makes it difficult to determine if the model is locally fitting a few PDE solutions or genuinely learning the target operator.

In summary, the significance and value of this paper are unclear, resulting in limited research impact.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: See Weaknesses.
soundness: 2
strengths: 1. The summary of related work is detailed.
2. The paper establishes a theoretical connection between the Mamba architecture and neural operators.
3. Although not comprehensive, the paper includes experiments on several datasets and compares the proposed method with some baseline models.
summary: This paper introduces Mamba in the PDE-solving area. It establishes a formal theoretical connection between structured state-space models (SSMs) and neural operators. The proposed method performs relatively well on the considered datasets.
weaknesses: 1. **Presentation Issues:** The paper lacks line numbers, which are essential for precise referencing during reviews. Additionally, the appendix is included in the supplementary materials rather than in the main document, which hampers the ease of accessing critical information.
2. **Writing Clarity:** The manuscript’s writing requires enhancement, as some sentences lack logical coherence. For instance, the statement: *“While Transformers dominate in areas like foundational models and computer vision, the application of SSMs, particularly the Mamba architecture, to neural operators for PDEs is still underexplored.”* The connection between Transformers’ dominance in certain fields and the underexplored application of SSMs in PDEs is unclear.
3. **Figure 1 Interpretation:** The purpose and contribution of Figure 1 are not clear. The main paper contains excessive mathematical details that could be better placed in the appendix. And a detailed description of the main architecture of MNO is necessary. Specifically, does MNO include a patchification layer as shown in Figure 1? Furthermore, since your baseline models operate on mesh points directly without patchifying, how does this influence your comparisons?
4. **Novelty Concern:** The paper’s novelty appears limited. Introducing the Mamba architecture into PDE solving is not new. The approach seems to involve simply replacing the Transformer block with a Mamba block without significant modifications. It would be helpful to explain why PDE-specific adaptations are not needed. Additionally,  while you theoretically demonstrate that the MNO is a type of neural operator, this does not enhance the understanding of the benefits of employing Mamba in PDE solving.
5. **Efficiency Evaluation:** It is well-known that Mamba offers efficiency advantages over Transformers in long sequence data. However, your experiments do not include efficiency comparisons with baseline models, which is a missed opportunity to highlight Mamba’s benefits.
6. **Dataset and Baseline Limitations:** The datasets and baselines used are not comprehensive enough. The datasets selected are relatively simple. Incorporating more complex datasets, such as the 2D Navier-Stokes dataset from FNO [1], would strengthen your evaluation. Additionally, there are other Transformer-based solvers beyond those utilizing linear attention, such as Transolver [2], which computes attention based on learned physical states.
7. **Implementation Details:** The implementation details provided are insufficient. Did you keep all hyperparameters consistent across the three variants? What is the total parameter count for each model? It is hard to believe that merely replacing the Transformer block with a Mamba block results in significant performance improvements, especially since such gains have not been observed in NLP and CV domains. Providing more detailed implementation information or sharing inference code would greatly enhance the paper’s credibility.
8. **Clarity of Figure 3:** The representation in Figure 3 is not clear. What do the different colored lines mean, and do these lines correspond to any ground truth? The sentence *“The change in log amplitude across the frequency range is more uniform, indicating that Mamba effectively balances between capturing necessary high-frequency information and filtering out noise.”* is difficult to understand. It would be helpful if you could provide more explanations or clarify this statement.
9. **Irregular data:** The paper focuses solely on grid data. I wonder what distinctions do you see between grid-based PDE data and images. Futhermore, how would you extend the MNO to handle irregular data, considering that linear Transformers can readily address such types?

[1] Fourier Neural Operator for Parametric Partial Differential Equations

[2] Transolver: A Fast Transformer Solver for PDEs on General Geometries