Convolutional Vision Transformer for Cosmology Parameter Inference

Parameter inference is a crucial task in modern cosmology, which requires accurate and fast computational methods to keep up with the high precision and volume of observational datasets. In this study, we experiment with a hybrid vision transformer, the Convolution vision Transformer (CvT), which simultaneously benefits from the advantages of vision transformers and convolutional neural networks, and use it to infer the $\Omega_m$ and $\sigma_8$ cosmological parameters from dark matter (pretraining) and halo distribution (fine-tuning) fields. Our experiments suggest that the CvT constraints on $\Omega_m$ and, more prominently, $\sigma_8$ are better than a simple vision transformer on both dark matter and halo fields. Pretraining on DM data proves advantageous for improving constraints on the two parameters using halo fields rather than training a model from the beginning. The CvT is more efficient than the traditional ViT since, despite more parameters, CvT requires a similar training time to the traditional ViT and is thus scalable to large datasets.