# Segmenting proto-halos with vision transformers

arXiv: https://arxiv.org/abs/2508.00049 · HTML: https://mlsanity.com/paper/2508.00049/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 76 references that are in the map.
- First posted on arXiv 31 Jul 2025.

## Builds on (most cited first, 14 of 14)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations](https://mlsanity.com/paper/1707.03237/) (2017)
- [A survey of loss functions for semantic segmentation](https://mlsanity.com/paper/2006.14822/) (2020)
- [Improved Texture Networks: Maximizing Quality and Diversity in Feed-forward Stylization and Texture Synthesis](https://mlsanity.com/paper/1701.02096/) (2017)
- [Deep symbolic regression for physics guided by units constraints: toward the automated discovery of physical laws](https://mlsanity.com/paper/2303.03192/) (2023)
- [A volumetric deep Convolutional Neural Network for simulation of mock dark matter halo catalogues](https://mlsanity.com/paper/1805.04537/) (2018)
- [Deep learning insights into cosmological structure formation](https://mlsanity.com/paper/2011.10577/) (2020)
- [Classification algorithms applied to structure formation simulations](https://mlsanity.com/paper/2106.06587/) (2021)

## Abstract (arXiv)

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 textscpinocchio, 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.