Segmenting proto-halos with vision transformers

TL;DR

This paper demonstrates that transformer-based neural networks significantly improve the segmentation and classification of proto-halo regions in the early universe's density field, outperforming traditional CNNs and perturbation models.

Contribution

It introduces a transformer-based architecture for proto-halo segmentation, showing superior accuracy and providing insights into feature importance via Grad-CAM.

Findings

Transformer networks outperform CNNs in proto-halo segmentation.

Deep learning models surpass perturbation-theory-based models like PINOCCHIO.

Combining density and tidal shear features improves model performance.

Abstract

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 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 \textsc{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.