Predicting Halo Assembly Bias from Merger Trees using Graph Neural Networks with Formation Time Regularization

Abstract: Halo assembly bias, where halo clustering depends on formation history beyond just mass, poses a challenge for accurate cosmological modeling. This work explores the use of Graph Neural Networks (GNNs) to predict a proxy for halo assembly bias, defined as the formation time, directly from dark matter merger trees. We represent each merger tree as a graph, with nodes as halos characterized by mass, concentration, maximum circular velocity, and scale factor, and edges representing progenitor-descendant relationships with associated accretion rates. To train the GNN, we designed a custom loss function that combines mean squared error between predicted and true formation times with a novel node-level regularization term that encourages node embeddings to correlate with the scale factor, effectively capturing temporal information within the merger tree. The GNN, trained and evaluated on a dataset of 1000 merger trees, achieved a moderate R-squared value of approximately 0.48 on the test set. Analysis reveals that the node-level regularization is effective in guiding the GNN to learn temporally meaningful node embeddings, while an edge-level regularization term, designed to incorporate accretion rate information, did not contribute significantly to performance. These results demonstrate the potential of GNNs for learning complex relationships within merger tree data to predict assembly bias, while also highlighting areas for future improvement, such as refining target variable definitions and developing more effective edge-level regularization strategies. \

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