The Universe is worth 64³ pixels: Convolution Neural Network and Vision Transformers for Cosmology

We present a novel approach for estimating cosmological parameters, $\Omega_m$, $\sigma_8$, $w_0$, and one derived parameter, $S_8$, from 3D lightcone data of dark matter halos in redshift space covering a sky area of $40^\circ \times 40^\circ$ and redshift range of $0.3 < z < 0.8$, binned to $64^3$ voxels. Using two deep learning algorithms, Convolutional Neural Network (CNN) and Vision Transformer (ViT), we compare their performance with the standard two-point correlation (2pcf) function. Our results indicate that CNN yields the best performance, while ViT also demonstrates significant potential in predicting cosmological parameters. By combining the outcomes of Vision Transformer, Convolution Neural Network, and 2pcf, we achieved a substantial reduction in error compared to the 2pcf alone. To better understand the inner workings of the machine learning algorithms, we employed the Grad-CAM method to investigate the sources of essential information in activation maps of the CNN and ViT. Our findings suggest that the algorithms focus on different parts of the density field and redshift depending on which parameter they are predicting. This proof-of-concept work paves the way for incorporating deep learning methods to estimate cosmological parameters from large-scale structures, potentially leading to tighter constraints and improved understanding of the Universe.