Abstract

Contributed Talk - Splinter EScience   (MW-2235)

Segmenting proto-halos with vision transformers

Toka Alokda, Cristiano Porciani
Argelander Institute for Astronomy, University of Bonn / Excellence cluster "Our Dynamic Universe"

The formation of dark-matter halos from small cosmological perturbations generated in the early universe is a highly non-linear process typically modeled through N-body simulations. In this work, we explore the use of deep learning to segment and classify proto-halo regions in the initial density field according to their final halo mass at redshift z = 0. We compare two architectures: a fully convolutional neural network (CNN) based on the V-Net design and a U-Net vision transformer. We find that the transformer-based network significantly outperforms the CNN across all metrics, achieving sub-percent error in the total segmented mass per halo class. Both networks deliver much higher accuracy than the perturbation-theory-based model pinocchio, especially at low halo masses and in the detailed reconstruction of proto-halo boundaries. We also investigate the impact of different input features by training models on the density field, the tidal shear, and their combination. Finally, we use Grad-CAM to generate class-activation heatmaps for the CNN, providing preliminary yet suggestive insights into how the network exploits the input fields. Our results show both the promise of vision transformers and deep-learning models in general as a basis for scientific infrastructure that replaces expensive simulation codes, as well as the importance of the interpretability of these methods to guarantee accurate and physically meaningful results.