MBNO: Mamba-based neural operators for solving partial differential equations

초록

The recently released Mamba model leverages structured state space models (SSMs), incorporating hardware-efficient designs and selection mechanisms. The Mamba architecture demonstrates strong potential as a replacement for Transformer-based models across various tasks. In this work, we employ Mamba to train neural operators on infinite-dimensional spaces derived from partial differential equations. Using well-established theory on the Rough Path and Reproducing Kernel Hilbert Space (RKHS), we theoretically demonstrate that the SSM-based models can replace Transformer-based models for approximating operators. Our empirical findings further show that Mamba consistently outperforms Transformer models across various tasks while achieving faster inference, highlighting the potential of the Mamba architecture to outperform Transformer-based models in various operator learning tasks.

키워드