GeoMaNO: Geometric Mamba Neural Operator for Partial Differential Equations

TL;DR

GeoMaNO presents a geometry-aware neural operator framework for PDEs on regular grids, integrating a geometry-aware encoder/decoder with the GeoMamba-SSM module to achieve linear-time complexity and reduce duplicate hidden states via a geometric correction. It extends Mamba-SSM to 2D via a 2DGeoMamba-SSM and uses a four-way cross-scan kernel to preserve geometric locality while maintaining efficiency. Empirically, GeoMaNO delivers state-of-the-art accuracy and speed on standard PDE benchmarks (Darcy flow and Navier–Stokes), with improvements up to 58.9% in accuracy and significant inference-time gains over the prior SOTA, and shows competitive generalization capabilities to other domains like ImageNet-100. Limitations include its current restriction to regular grids, with future work aiming to extend the geometric Mamba framework to unstructured meshes, point clouds, and 3D temporal settings.

Abstract

The neural operator (NO) framework has emerged as a powerful tool for solving partial differential equations (PDEs). Recent NOs are dominated by the Transformer architecture, which offers NOs the capability to capture long-range dependencies in PDE dynamics. However, existing Transformer-based NOs suffer from quadratic complexity, lack geometric rigor, and thus suffer from sub-optimal performance on regular grids. As a remedy, we propose the Geometric Mamba Neural Operator (GeoMaNO) framework, which empowers NOs with Mamba's modeling capability, linear complexity, plus geometric rigor. We evaluate GeoMaNO's performance on multiple standard and popularly employed PDE benchmarks, spanning from Darcy flow problems to Navier-Stokes problems. GeoMaNO improves existing baselines in solution operator approximation by as much as 58.9%.

Figures (8)

• Figure 1: Overview of the novel GeoMaNO architecture. (1) The input function $a(x)$ is lifted and patchified to a higher-dimensional latent representation by the geometry-aware encoder layer $\mathcal{E}$; (2) $T$ stacked GeoMaNO layers $M_1 \cdots M_T$ (middle) performs the kernel integral via GeoMamba kernels (bottom); (3) Each GeoMamba kernel performs geometry-aware Mamba scans, combined with skip-connections and non-linear activations; and (4) The last GeoMaNO layer's output is transformed back to the physical domain with the decoder layer $\mathcal{D}$. This yields the output function $u(x)$.
• Figure 2: The four-way cross-scan pattern. Left: In the data preparation stage, the input grid is traversed in four orders, and the traversal results are stacked at an additional order dimension. Right: For each slice along the order dimension, we process it with our GeoMamba-SSM module (shown in blue color). As a special case, for 2D PDEs, in GeoMamba-SSM, we employ a 2D SSM representation with geometric rigor (shown in green color). The results for all four cross-scan directions are merged.
• Figure 3: Illustration of our GeoMamba-SSM variants. Left: A sample $4 \times 4$ input grid. Middle: 1DGeoMamba-SSM flattens the input into a 1D sequence and regresses the hidden states over it. However, after flattening, vertically-adjacent patches $x_{33}$ (highlighted with red margins) and $x_{43}$ (highlighted with blue margins) are no longer adjacent. This leads to geometric inconsistency in the hidden states. Right: 2DGeoMamba-SSM regresses over the input grid in a 2D manner. Adjacent patches remain adjacent in the hidden states, properly preserving geometric information.
• Figure 4: GeoMaNO's scalability. We ablate model depth and embedding dimensions, and report the relative L2 errors ($\%$).
• Figure 5: Illustration of 1D v.s. 2D scanning paths and spatial discrepancy. Left: The two-dimensional input grid. Middle: 1D Mamba-SSM flattens the input into a 1D sequence, but adjacent patches (the red and blue) are far away in this sequence, leading to the geometric inconsistency in the hidden states. Right: 2DMamba-SSM the input grid in a 2D manner, and maintains spatial continuity.

Theorems & Definitions (3)

• Remark 4.1: Mamba's D Coefficients
• Remark 4.2: Design Choices for Fixed Geometric Correction Coefficients
• Remark A.1: GeoMaNO vs. Zheng et al. zheng24-nips-aliasfreemambaneuraloperator