Robust cosmological inference with Graph Neural Networks

Current and future bis astronomical surveys will provide huge datasets with billions of objects such as galaxies and their distribution in space. These can be used to better understand the formateion and evolution of our Universe i.e. infer cosmological paramerters.

Deep learning models have been shown to outperform methods that rely on summary statistics (which often cannot account for non-Gaussianity, which we know is present), like the power spectrum, in extracting information from complex cosmological data sets. Deep learning models are usually trained on simualtions where we can vary cosmologucal and astrophysical parameters and match those with observable galaxy distributions. Ideally, we hope to then apply these trained models to precisely infer cosmology from observations. However, due to differences in the subgrid physics implementation and numerical approximations across different simulation suites, models trained on data from one cosmological simulation show a drop in performance when tested on another. Similarly, models trained on any of the simulations would also likely experience a drop in performance when applied to observational data.

Training on data from two different suites of the CAMELS hydrodynamic cosmo-logical simulations, we examine the generalization capabilities of Domain Adaptive Graph Neural Networks (DA-GNNs). By utilizing GNNs, we capitalize on their capacity to capture structured scale-free cosmological information from galaxy dis- tributions. Moreover, by including unsupervised domain adaptation via Maximum Mean Discrepancy (MMD), we enable our models to extract domain-invariant features. We demonstrate that DA-GNN achieves higher accuracy and robustness on cross-dataset tasks (up to 28% better relative error and up to almost an order of magnitude better χ2). Using data visualizations, we show the effects of domain adaptation on proper latent space data alignment. This shows that DA-GNNs are a promising method for extracting domain-independent cosmological information, a vital step toward robust deep learning for real cosmic survey data!