CosmoMamba: State-Space Models for Cosmological Parameter Inference from 2D Field Maps

We present CosmoMamba, the first application of selective state-space models (SSMs) to cosmo- logical parameter inference from 2D field maps. Built on the Mamba architecture, CosmoMamba processes cosmological fields through multi-directional scans that capture spatial structure along row, column, and diagonal axes, achieving linear computational complexity O(n) in the number of spatial tokens. We benchmark CosmoMamba against convolutional neural networks (CNNs) and Vi- sion Transformers (ViTs) on the CAMELS Multifield Dataset (CMD), predicting the matter density parameter Ωm and the amplitude of matter fluctuations σ8 from dark matter density maps. Cosmo- Mamba achieves R2 = 0.936 for Ωm and R2 = 0.886 for σ8, outperforming the Vision Transformer (R2 = 0.938, 0.818) on σ8 with 2.7×fewer parameters and linear scaling. In cross-simulation- suite transfer (IllustrisTNG →SIMBA), CosmoMamba exhibits 30% less performance degradation than ViT, demonstrating superior generalization to unseen astrophysical models. While CNNs with strong spatial inductive biases retain an advantage on single-field local tasks, our results establish SSMs as a promising and efficient alternative to transformers for cosmological field analysis. Code and trained models are publicly available.