Vision Transformers for Cosmological Fields: Application to Weak Lensing Mass Maps

TL;DR

This paper evaluates the use of vision transformer architectures, including Swin transformers, for extracting cosmological parameters from weak lensing mass maps, comparing their performance to CNNs within a simulation-based inference framework.

Contribution

It introduces transformer-based models for cosmological parameter inference from weak lensing data and compares their effectiveness and data efficiency to traditional CNNs.

Findings

Swin transformers outperform vanilla ViTs with limited data

Transformers achieve comparable cosmological constraints to CNNs under realistic noise

Pre-training enhances transformer performance in cosmological inference

Abstract

Weak gravitational lensing is a powerful probe of the universe's growth history. While traditional two-point statistics capture only the Gaussian features of the convergence field, deep learning methods such as convolutional neural networks (CNNs) have shown promise in extracting non-Gaussian information from small-scale, nonlinear structures. In this work, we evaluate the effectiveness of attention-based architectures, including variants of vision transformers (ViTs) and shifted window (Swin) transformers, in constraining the cosmological parameters $\Omega_m$ and $S_8$ from weak lensing mass maps. Using a simulation-based inference (SBI) framework, we compare transformer-based methods to CNNs. We also examine performance scaling with the number of available $N$-body simulations, highlighting the importance of pre-training for transformer architectures. We find that the Swin transformer performs significantly better than vanilla ViTs, especially with limited training data. Despite their higher representational capacity, the Figure of Merit for cosmology achieved by transformers is comparable to that of CNNs under realistic noise conditions.