Segmenting proto-halos with vision transformers

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.

39 pages, 14 figures, 11 tables; updated to match the published version: JCAP11(2025)083

References in corpus (42)

• The cosmological simulation code GADGET-2
• Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization
• Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations
• Toward a halo mass function for precision cosmology: the limits of universality
• Ellipsoidal collapse and an improved model for the number and spatial distribution of dark matter haloes
• Ahf: Amiga's Halo Finder
• Multi-scale initial conditions for cosmological simulations
• Precision Determination of the Mass Function of Dark Matter Halos
• Multiscale Gaussian Random Fields for Cosmological Simulations
• The Excursion Set Theory of Halo Mass Functions, Halo Clustering, and Halo Growth
• Testing Tidal-Torque Theory: I. Spin Amplitude and Direction
• PINOCCHIO: pinpointing orbit-crossing collapsed hierarchical objects in a linear density field
• Cosmic Shear from Galaxy Spins
• Testing Tidal-Torque Theory: II. Alignment of Inertia and Shear and the Characteristics of Proto-haloes
• An accurate tool for the fast generation of dark matter halo catalogs
• The Halo Mass Function from Excursion Set Theory. II. The Diffusing Barrier
• Deep symbolic regression for physics guided by units constraints: toward the automated discovery of physical laws
• Collapse Barriers and Halo Abundance: Testing the Excursion Set Ansatz
• The Gaussian streaming model and Lagrangian effective field theory
• ZOMG-I. How the cosmic web inhibits halo growth and generates assembly bias
• The mass-Peak Patch algorithm for fast generation of deep all-sky dark matter halo catalogues and its N-Body validation
• Predicting the number, spatial distribution and merging history of dark matter haloes
• Peaks theory and the excursion set approach
• The Peaks Formalism and the Formation of Cold Dark Matter Haloes
• Halo bias in Lagrangian Space: Estimators and theoretical predictions
• Machine learning cosmological structure formation
• Two Ways of Biasing Galaxy Formation
• Ellipsoidal halo finders and implications for models of triaxial halo formation
• The spatial and velocity bias of linear density peaks and proto-haloes in the Lambda cold dark matter cosmology
• A volumetric deep Convolutional Neural Network for simulation of mock dark matter halo catalogues
• An interpretable machine learning framework for dark matter halo formation
• The formation of CDM haloes I: Collapse thresholds and the ellipsoidal collapse model
• Priors on Lagrangian bias parameters from galaxy formation modelling
• Predicting dark matter halo formation in N-body simulations with deep regression networks
• The formation of CDM haloes II: collapse time and tides
• The locations of halo formation and the peaks formalism
• Testing approximate predictions of displacements of cosmological dark matter halos
• Halo assembly bias from a deep learning model of halo formation
• Culminating the peak CUSP to descry the dark side of halos
• Classification algorithms applied to structure formation simulations
• Deep learning insights into cosmological structure formation
• Machine Learning and Structure Formation in Modified Gravity