Vision Mamba for Permeability Prediction of Porous Media

TL;DR

This work addresses the computational challenge of predicting permeability in 3D porous media from voxel data. It introduces Vision Mamba as a backbone, leveraging linear ($O(N)$) scaling with token size to outperform ViTs and CNNs in both accuracy and memory efficiency for this task. The proposed architecture patches 3D volumes, applies axis-wise selective state-space scans with bidirectional updates, and aggregates features via global pooling to a permeability scalar, achieving an $R^2$ around $0.997$ with far fewer parameters than CNN baselines. The results suggest Vision Mamba as a strong, scalable building block for large vision models in porous media, with open-source code to enable reproducibility and broader adoption.

Abstract

Vision Mamba has recently received attention as an alternative to Vision Transformers (ViTs) for image classification. The network size of Vision Mamba scales linearly with input image resolution, whereas ViTs scale quadratically, a feature that improves computational and memory efficiency. Moreover, Vision Mamba requires a significantly smaller number of trainable parameters than traditional convolutional neural networks (CNNs), and thus, they can be more memory efficient. Because of these features, we introduce, for the first time, a neural network that uses Vision Mamba as its backbone for predicting the permeability of three-dimensional porous media. We compare the performance of Vision Mamba with ViT and CNN models across multiple aspects of permeability prediction and perform an ablation study to assess the effects of its components on accuracy. We demonstrate in practice the aforementioned advantages of Vision Mamba over ViTs and CNNs in the permeability prediction of three-dimensional porous media. We make the source code publicly available to facilitate reproducibility and to enable other researchers to build on and extend this work. We believe the proposed framework has the potential to be integrated into large vision models in which Vision Mamba is used instead of ViTs.

Figures (6)

• Figure 1: Three representative examples of the synthetically generated three-dimensional digital porous media (with $n=64$) used to train Vision Mamba are shown; phases are color-coded with blue indicating the solid grain matrix and red indicating the pore space.
• Figure 2: Schematic architecture of the proposed network based on Vision Mamba for deep learning of the permeability of three-dimensional porous media. If the input is a $64 \times 64 \times 64$ porous medium (i.e., $n = 64$) with a patch size of 32, there are eight subcubes, labeled 1 through 8. By scanning along the $y$-axis, the subcubes are arranged in the order 1, 5, 2, 6, 3, 7, 4, 8. Scanning along the $x$ and $z$ directions is defined similarly.
• Figure 3: Comparison of Vision Mamba and CNN performance using the $R^2$ score (left: Vision Mamba; right: CNN)
• Figure 4: GPU memory (per epoch) scaling versus relative patch size (i.e., relative token size) for Vision Mamba (ViM) and ViT on log--log axes. Least-squares fits reveal linear scaling for Vision Mamba and quadratic scaling for ViT.
• Figure 5: Comparison of Vision Mamba (ViM) and ViT performance based on the $R^2$ score. The first row corresponds to Vision Mamba (ViM), and the second row corresponds to ViT.