A graph-based Neural Network surrogate model for accelerating semi-analytical model of galaxy formation and evolution

Abstract: Understanding how galaxy populations emerge and evolve from the growth of dark matter structure is a central challenge in galaxy formation theory. Semi-analytic models (SAMs) provide an efficient framework to address this problem, but exploring large ensembles of merger trees across broad parameter spaces remains computationally demanding. We develop a conditional graph neural network surrogate model that combines merger tree information with SAM parameters to predict galaxy properties across cosmic time. Using merger trees of dark matter halos from the Uchuu simulation and the Galacticus SAM, the model predicts stellar mass, luminosity, angular momentum, gas metal mass, and specific star formation rate across the wide redshift range of 0 <= z <= 5. For instance, the model can predict stellar mass at 0 <= z <= 3 with a scatter of 0.19-0.28 dex and coefficient of determination R^2 of 0.946-0.973 (R^2 close to 1 indicates prediction closely matching the truth). The results show that a single graph based model can reproduce these galaxy properties with good accuracy over multiple SAM realizations, merger trees and redshifts. This catalog-level model provides a practical route for accelerating SAM based studies of galaxy formation to enable a more detailed investigation of the model parameter space. The inference code, trained models, and example data products are publicly available at https://github.com/MutongCat/sam2galaxy-gnn.

• The paper introduces a graph-based neural network surrogate that leverages merger tree graphs and SAM parameter conditioning to predict key galaxy properties.
• It employs a GraphSAGE backbone with heteroscedastic regression and mixture-of-experts heads, achieving stellar mass predictions with scatter as low as 0.225 dex and R² up to 0.973.
• The model offers a substantial computational speedup for SAM evaluations, paving the way for scalable galaxy simulations and enhanced parameter exploration in cosmological surveys.

Graph Neural Network Surrogate Modeling for Semi-Analytical Galaxy Formation

Motivation and Context

Semi-analytical models (SAMs) are a cornerstone of theoretical galaxy formation, offering flexibility and efficiency for modeling baryonic evolution within dark matter halos under broad parameter regimes. However, calibrating and exploring SAMs over large ensembles of merger trees and high-dimensional parameter spaces remains computationally prohibitive. This paper introduces a graph-based neural network (GNN) surrogate model, leveraging merger tree graphs from cosmological N-body simulations and explicit SAM parameter conditioning to predict diverse galaxy properties across cosmic time.

Model Architecture and Data Representation

The proposed approach encodes each dark matter halo and its merger/history as a node within a graph, with edges for both progenitor-descendant and host-satellite relations. Halo features include mass, position, velocity, angular momentum, scale radius, spin parameter, and cosmic scale factor. Each Galacticus SAM realization is parameterized by a 17-dimensional vector regulating processes such as star formation, feedback, outflow, cooling, and AGN activity. The model aims to predict five key galaxy properties—stellar mass, z-band luminosity, angular momentum, specific star formation rate (sSFR), and gas metal mass—at nine redshift outputs across $0 \leq z \lesssim 5$.

Figure 1: Joint distribution and correlations of the five key galaxy properties at $z_1=0$ in a single Galacticus catalog.

The GNN backbone (GraphSAGE architecture) performs message passing, aggregating local and historical halo information to produce latent node embeddings. A conditioning module fuses each node's latent state with the corresponding SAM parameter vector, enabling generalization across both merger histories and baryonic model variations. Prediction heads are designed for both heteroscedastic regression (smooth targets) and mixture-of-experts (MoE) heads (branch/population-structured targets), enabling flexible modeling of complex conditional distributions.

Figure 2: Model pipeline: shared GNN backbone encodes merger trees, SAM parameters condition intermediate representations, and property-specific heads output predictions at multiple redshifts.

Quantitative Performance and Evaluation

Extensive evaluation was performed on a sample of 201,986 merger trees and 7,800 Galacticus catalogs. The main experiment splits merger trees and catalogs for training/testing, ensuring simultaneous generalization to unseen histories and parameter settings.

Stellar Mass Prediction: The GNN surrogate yields a scatter of $0.225$ dex, bias $0.009$, correlation coefficient $0.978$, and $R^2=0.957$ for $\log(M_\star)$ at $z=0$. These metrics demonstrate highly accurate recovery of stellar mass relative to full SAM evaluations. The GNN outperforms an MLP baseline, which uses only instantaneous halo features, verifying the critical role of explicit assembly history in predicting integrated galaxy properties.

Figure 3: One-to-one correspondence between predicted and true stellar masses for $z=0$ galaxies, including marginal distributions and metrics.

Across the nine redshifts, stellar mass prediction remains robust: scatter stays within $0.19-0.28$ dex, and $z_1=0$0 within $z_1=0$1, except at the highest redshift where sample sparsity diminishes performance. Population-level diagnostics with the stellar mass function indicate excellent agreement over the majority of the mass range, with systematic underprediction in the highest-mass bins, particularly at high redshift—a limitation imposed by both model and simulation volume.

Figure 4: Predicted vs. true stellar mass functions and residuals across redshift; median agreement and percentiles displayed.

Clustering statistics are also well reproduced by the surrogate over intermediate scales and masses, with largest deviations in sparse, high-mass/high-redshift regimes.

Figure 5: Two-point correlation functions across redshifts and mass bins, comparing surrogate to true catalogs.

Other Target Properties: Luminosity ($z_1=0$2)—scatter $z_1=0$3 dex, $z_1=0$4; angular momentum ($z_1=0$5)—scatter $z_1=0$6 dex, $z_1=0$7. sSFR ($z_1=0$8) and gas metal mass ($z_1=0$9) exhibit non-Gaussian/multimodal conditional distributions; MoE heads are necessary for capturing their branch/floor structure. The GNN reduces scatter compared to baselines (e.g., gas metal mass scatter from $0.225$0 to $0.225$1 dex) and substantially increases $0.225$2 (from $0.225$3 to $0.225$4).

Figure 6: Predicted versus true values for $0.225$5, $0.225$6, sSFR, and gas metal mass at $0.225$7, demonstrating both tight relations and branch structure.

Classification performance (F1 scores) for identifying quenched and floor states is improved using state-aware heads and balanced training, yet these discrete populations remain challenging, especially when state occupancy varies across both trees and SAM catalogs.

Analysis of Parameter and Tree Variations

The paper rigorously dissects sources of variance in target properties:

• Tree-to-Tree Variation: Differences in halo assembly history at fixed SAM parameters.
• Catalog-to-Catalog Variation: Changes in predicted galaxies stemming from baryonic model variations at fixed merger history.

Properties like stellar mass and luminosity are influenced by both axes, while quantities like angular momentum are primarily assembly-history dependent. sSFR and gas metal mass exhibit substantial catalog-driven dispersion, with state occupation (quenched/floor) varying across trees and parameterizations.

Figure 7: Distribution of quenched and floor-fraction across catalogs as a function of redshift; clear variation and label imbalance.

Figure 8: Tree-level state variability analysis demonstrating that most merger trees do not remain permanently in one state over all SAM catalogs.

Ensemble averaging modestly increases regression performance but leaves discrete state classification largely unchanged, highlighting the inherent difficulty of state prediction under imbalanced or variable occupancy conditions.

Architectural and Methodological Implications

The results validate that GNNs are capable surrogates for SAM evaluations when provided with detailed merger history and explicit baryonic parameter conditioning. The MoE head design is essential for branch-structured outputs, and label balancing during training enhances discrete classification.

Limitations include:

• Model applicability restricted to one simulation and SAM framework.
• Lack of explicit edge features and occupancy prediction.
• Reduced accuracy in rare object regimes and discrete state classification.
• Environmental effects and full merger tree utilization are not captured.

  Figure 9: F1 scores for quenched/floor classification under different data settings, illustrating performance dependence on head design and balanced training.

Practical Acceleration and Future Prospects

The surrogate achieves a notable speedup: processing the full UchuuMicro catalog with Galacticus takes $0.225$8 hours on CPU, while the GNN surrogate requires $0.225$9 minute (hardware and output details considered). This strategy can substantially accelerate parameter exploration, calibration, and mock catalog generation, particularly for large-scale cosmological surveys (e.g., Euclid, Roman Space Telescope).

Figure 10: Multi-property, multi-redshift overview showing one-to-one relations and prediction robustness of the Multi-SAM GNN across the galaxy property spectrum.

Future work should focus on expanding domain generalization (e.g., cross-SAM, cross-simulation), integrating richer merger/environmental information, learning galaxy occupancy alongside properties, and developing more flexible generative surrogates capable of modeling state/branch distributions natively.

Conclusion

A conditional GNN surrogate can accurately emulate multi-property outputs of a SAM across diverse merger trees and parameter choices. The model’s performance is strongest on integrated and smoothly-varying quantities; target-specific head design is necessary for branch/floor-structured properties. The approach provides substantial computational advantages for repeated SAM evaluations and paves the way for scalable galaxy modeling in forthcoming large-scale datasets. The limitations identified point directly to active research directions in surrogate modeling, graph-based representations, and robust state classification for galaxy evolution theory.