Contrastive Learning of Merger Tree Embeddings for Likelihood-Free Cosmological Inference

Abstract: Cosmological inference from dark matter halo merger trees is challenging due to the intricate relationships between tree structure, assembly bias, and underlying cosmological parameters. We address this challenge by developing a contrastive learning framework that generates merger tree embeddings sensitive to cosmological parameters while mitigating the impact of assembly bias. A Graph Neural Network (GNN) is trained on merger trees from N-body simulations, employing a contrastive loss function to cluster trees originating from the same cosmology within the embedding space. To enhance robustness against assembly bias, we augment the training data by introducing variations in halo concentrations conditional on halo mass, guided by observed mass-concentration relations. These learned embeddings then serve as summary statistics for likelihood-free inference (LFI) using Sequential Neural Posterior Estimation (SNPE) to estimate the posterior distribution of $\Omega_m$ and $\sigma_8$. Using a dataset of 1000 merger trees from 40 unique cosmologies, our results demonstrate the effectiveness of the learned embeddings for cosmological inference, particularly for $\Omega_m$, achieving good accuracy and coverage probability close to the nominal value. However, we observe some undercoverage for $\sigma_8$, indicating potential for further refinement of the method. This work underscores the potential of contrastive learning and GNNs for extracting cosmologically relevant information from merger trees, paving the way for robust and accurate likelihood-free cosmological inference. \

Subjects:	astro-ph.CO; cs.LG
Cite as:	PX:2508.00007