# CosmoBench: A Multiscale, Multiview, Multitask Cosmology Benchmark for Geometric Deep Learning

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 94 references that are in the map.
- First posted on arXiv 4 Jul 2025.

## Builds on (most cited first, 30 of 31)

- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [ShapeNet: An Information-Rich 3D Model Repository](https://mlsanity.com/paper/1512.03012/) (2015)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [TUDataset: A collection of benchmark datasets for learning with graphs](https://mlsanity.com/paper/2007.08663/) (2020)
- [Benchmarking Graph Neural Networks](https://mlsanity.com/paper/2003.00982/) (2020)
- [Long Range Graph Benchmark](https://mlsanity.com/paper/2206.08164/) (2022)
- [Tree-Structured Parzen Estimator: Understanding Its Algorithm Components and Their Roles for Better Empirical Performance](https://mlsanity.com/paper/2304.11127/) (2023)
- [Topological Deep Learning: Going Beyond Graph Data](https://mlsanity.com/paper/2206.00606/) (2022)
- [ATOM3D: Tasks On Molecules in Three Dimensions](https://mlsanity.com/paper/2012.04035/) (2020)
- [Position: Graph Learning Will Lose Relevance Due To Poor Benchmarks](https://mlsanity.com/paper/2502.14546/) (2025)
- [Fine-grained Expressivity of Graph Neural Networks](https://mlsanity.com/paper/2306.03698/) (2023)
- [TopoX: A Suite of Python Packages for Machine Learning on Topological Domains](https://mlsanity.com/paper/2402.02441/) (2024)
- [Learning cosmology and clustering with cosmic graphs](https://mlsanity.com/paper/2204.13713/) (2022)
- [The Cosmic Graph: Optimal Information Extraction from Large-Scale Structure using Catalogues](https://mlsanity.com/paper/2207.05202/) (2022)
- [Inferring halo masses with Graph Neural Networks](https://mlsanity.com/paper/2111.08683/) (2021)
- [The continuous categorical: a novel simplex-valued exponential family](https://mlsanity.com/paper/2002.08563/) (2020)
- [Estimating Cosmological Parameters from the Dark Matter Distribution](https://mlsanity.com/paper/1711.02033/) (2017)
- [Robust Field-level Likelihood-free Inference with Galaxies](https://mlsanity.com/paper/2302.14101/) (2023)
- [Unsupervised Resource Allocation with Graph Neural Networks](https://mlsanity.com/paper/2106.09761/) (2021)
- [Domain Adaptive Graph Neural Networks for Constraining Cosmological Parameters Across Multiple Data Sets](https://mlsanity.com/paper/2311.01588/) (2023)
- [Field-Level Inference with Microcanonical Langevin Monte Carlo](https://mlsanity.com/paper/2307.09504/) (2023)
- [How DREAMS are made: Emulating Satellite Galaxy and Subhalo Populations with Diffusion Models and Point Clouds](https://mlsanity.com/paper/2409.02980/) (2024)
- [A Cosmic-Scale Benchmark for Symmetry-Preserving Data Processing](https://mlsanity.com/paper/2410.20516/) (2024)
- [textttMangrove: Learning Galaxy Properties from Merger Trees](https://mlsanity.com/paper/2210.13473/) (2022)
- [A point cloud approach to generative modeling for galaxy surveys at the field level](https://mlsanity.com/paper/2311.17141/) (2023)
- [Graph Neural Network-based Resource Allocation Strategies for Multi-Object Spectroscopy](https://mlsanity.com/paper/2109.13361/) (2021)

## Built on by (most cited first, 3 of 3)

- [gwBenchmarks: Stress-Testing LLM Agents on High-Precision Gravitational Wave Astronomy](https://mlsanity.com/paper/2605.11269/) (2026)
- [Multi-View Graph Learning with Graph-Tuple](https://mlsanity.com/paper/2510.10341/) (2025)
- [Velocityformer: Broken-Symmetry-Matched Equivariant Graph Transformers for Cosmological Velocity Reconstruction](https://mlsanity.com/paper/2605.21483/) (2026)

## Abstract (arXiv)

Cosmological simulations provide a wealth of data in the form of point clouds and directed trees. A crucial goal is to extract insights from this data that shed light on the nature and composition of the Universe. In this paper we introduce CosmoBench, a benchmark dataset curated from state-of-the-art cosmological simulations whose runs required more than 41 million core-hours and generated over two petabytes of data. CosmoBench is the largest dataset of its kind: it contains 34 thousand point clouds from simulations of dark matter halos and galaxies at three different length scales, as well as 25 thousand directed trees that record the formation history of halos on two different time scales. The data in CosmoBench can be used for multiple tasks -- to predict cosmological parameters from point clouds and merger trees, to predict the velocities of individual halos and galaxies from their collective positions, and to reconstruct merger trees on finer time scales from those on coarser time scales. We provide several baselines on these tasks, some based on established approaches from cosmological modeling and others rooted in machine learning. For the latter, we study different approaches -- from simple linear models that are minimally constrained by symmetries to much larger and more computationally-demanding models in deep learning, such as graph neural networks. We find that least-squares fits with a handful of invariant features sometimes outperform deep architectures with many more parameters and far longer training time. Still there remains tremendous potential to improve these baselines by combining machine learning and cosmology to fully exploit the data. CosmoBench sets the stage for bridging cosmology and geometric deep learning at scale. We invite the community to push the frontier of scientific discovery by engaging with this dataset, available at https://cosmobench.streamlit.app