Score Matching on Large Geometric Graphs for Cosmology Generation

TL;DR

This paper presents a scalable, physically consistent score-based generative model with an equivariant graph neural network for simulating large-scale galaxy clustering in cosmology, outperforming existing models.

Contribution

It introduces a novel topology-aware noise schedule and an equivariant graph neural network to generate realistic cosmological simulations of large galaxy distributions.

Findings

Generates full-scale cosmological point clouds of up to 600,000 halos.

Respects periodic boundary conditions and physical priors.

Outperforms existing diffusion models in clustering statistics.

Abstract

Generative models are a promising tool to produce cosmological simulations but face significant challenges in scalability, physical consistency, and adherence to domain symmetries, limiting their utility as alternatives to $N$-body simulations. To address these limitations, we introduce a score-based generative model with an equivariant graph neural network that simulates gravitational clustering of galaxies across cosmologies starting from an informed prior, respects periodic boundaries, and scales to full galaxy counts in simulations. A novel topology-aware noise schedule, crucial for large geometric graphs, is introduced. The proposed equivariant score-based model successfully generates full-scale cosmological point clouds of up to 600,000 halos, respects periodicity and a uniform prior, and outperforms existing diffusion models in capturing clustering statistics while offering significant computational advantages. This work advances cosmology by introducing a generative model designed to closely resemble the underlying gravitational clustering of structure formation, moving closer to physically realistic and efficient simulators for the evolution of large-scale structures in the universe.