Hierarchical Contrastive Graph Representation Learning for Cosmological Merger Trees and Parameter Inference

Abstract: Analyzing the complex hierarchical structures of dark matter halo merger trees is crucial for understanding the impact of cosmological parameters on structure formation, but efficiently discriminating between trees originating from different cosmologies poses a significant challenge. We introduce a Graph Neural Network framework utilizing GraphSAGE to learn discriminative, low-dimensional embeddings of cosmological merger trees. Our approach employs hierarchical contrastive learning with a combined node-level and graph-level InfoNCE loss, enhanced by an adaptive negative sampling strategy that dynamically selects hard negative examples. Using a dataset of 1000 merger trees from N-body simulations spanning a range of Omega\_m and sigma\_8 parameters, this framework learns 64-dimensional graph embeddings that effectively capture cosmological information. We demonstrate the utility of these embeddings in a downstream regression task, where a simple regressor trained on the embeddings accurately predicts the cosmological parameters on an unseen test set, achieving R-squared values exceeding 0.97 for Omega\_m and 0.79 for sigma\_8. Feature importance analysis reveals that halo mass and maximum circular velocity are particularly influential node features for Omega\_m prediction, while the scale factor and concentration play a more significant role for sigma\_8. Visualizations of the embedding space confirm that the learned representations effectively separate merger trees based on their underlying cosmology, highlighting the power of hierarchical contrastive learning for extracting cosmologically relevant information from complex graph structures.

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