Benchmarking AI-evolved cosmological structure formation

TL;DR

This work assesses the feasibility of using a 3-D U-Net to emulate cosmological structure formation by learning a forward map from initial Lagrangian displacement fields to evolved fields, trained on ZA and PM datasets. It defines a comprehensive, physics-motivated benchmark suite—including density PDFs, power spectra, bispectra, percolation, and cross-power tests—to quantify fidelity beyond traditional statistics. A standard mean-squared error loss underperforms on small scales, while a density-weighted loss significantly improves accuracy for dense, nonlinear structures and key clustering statistics, illustrating how physics-informed objectives can boost scientific emulation. The study shows potential for fast covariance generation and parameter inference, while also outlining challenges in scalability, boundary effects, and the need for integrating physics constraints and multi-resolution approaches in future AI-based cosmological emulators.

Abstract

The potential of deep learning-based image-to-image translations has recently attracted significant attention. One possible application of such a framework is as a fast, approximate alternative to cosmological simulations, which would be particularly useful in various contexts, including covariance studies, investigations of systematics, and cosmological parameter inference. To investigate different aspects of learning-based cosmological mappings, we choose two approaches for generating suitable cosmological matter fields as datasets: a simple analytical prescription provided by the Zel'dovich approximation, and a numerical N-body method using the Particle-Mesh approach. The evolution of structure formation is modeled using U-Net, a widely employed convolutional image translation framework. Because of the lack of a controlled methodology, validation of these learned mappings requires multiple benchmarks beyond simple visual comparisons and summary statistics. A comprehensive list of metrics is considered, including higher-order correlation functions, conservation laws, topological indicators, and statistical independence of density fields. We find that the U-Net approach performs well only for some of these physical metrics, and accuracy is worse at increasingly smaller scales, where the dynamic range in density is large. By introducing a custom density-weighted loss function during training, we demonstrate a significant improvement in the U-Net results at smaller scales. This study provides an example of how a family of physically motivated benchmarks can, in turn, be used to fine-tune optimization schemes -- such as the density-weighted loss used here -- to significantly enhance the accuracy of scientific machine learning approaches by focusing attention on relevant features.

Figures (21)

• Figure 1: Architecture for the 3-d U-Net showing the contracting and expanding parts. The input volume is progressively downsampled through multiple convolutional blocks (blue bars) with stride 2 (orange arrows), halving the spatial dimensions at each stage while increasing the number of feature channels. After reaching the bottleneck, the decoder path upsamples the feature maps via transpose convolutions (green arrows), concatenating them (dashed lines) with the corresponding feature maps from the encoder at each resolution level. This skip-connection strategy preserves high-resolution details lost during downsampling. The final output volume matches the original spatial dimensions of the input.
• Figure 2: Illustration of the cross-power test (Section \ref{['cross']}), showing the density fields and cross-correlation power spectra computed between them.
• Figure 3: Convergence investigations with sample size and epoch (Section \ref{['subsec:mse-conv']}). Top panel: Training and validation loss curves as a function of epoch number, for different-sized training datasets. Bottom panel: The first five epochs during the training, validation and training losses are plotted as a function of the number of samples in the training set.
• Figure 4: Comparison of projected densities between the predicted cosmic web and ZA ground truth. Densities are derived from the displacement field with a Cloud-In-Cell (CIC) method and summed over one axis. The 'input' panel (first panel) presents the density field projection of the early ($z=z_0$) snapshot. The color bar in each panel shows the magnitude of the matter density, $\rho(\textbf{x})$. In the above case, the two scale factors are $a_0=0.0465$ and $a_1=0.215$, and the box size is 64 $h^{-1}$Mpc, with $64^3$ particles.
• Figure 5: Comparison of projected densities between the predicted cosmic web and PM-evolved ground truth. $\sim 600$ samples are used for the training, each of them containing $128^3$ particles in 50 $h^{-1}$Mpc boxes. Densities are derived from the displacement field with a CIC method and summed over one axis. The 'input' panel (first panel) presents the density field projection of the early ($z=z_0$) snapshot. The color bar in each panel shows the magnitude of the matter density $\rho(\textbf{x}$). The snapshots used are from $a_0=0.05$ and $a_1=0.1$.