Our group has developed BORG (Bayesian Origin Reconstruction from Galaxies), a field-level Bayesian inference framework that reconstructs the Universe's initial conditions and three-dimensional matter distribution directly from galaxy survey data. BORG has enabled high-resolution reconstructions of the cosmic web, including the Manticore programme of large-scale structure reconstructions. BORG relies on an explicit physical forward model combined with an explicit model of galaxy bias, the connection between the underlying dark matter field and the observed galaxies. Bias modelling is a dominant source of modelling uncertainty in current field-level analyses. Existing bias models are typically perturbative or phenomenological, which limits the scales and accuracy to which galaxy data can be exploited.
In this project, we will go beyond this approach by developing a generative machine-learning model that directly learns the conditional distribution of cosmological initial conditions given an observed galaxy field. The model is trained on paired volumetric data from the FLAMINGO suite of large-volume cosmological hydrodynamical simulations, linking initial conditions to the corresponding simulated galaxy fields. Because the training data already encode the full non-linear, non-local and stochastic physics of galaxy formation as modelled in FLAMINGO, the approach bypasses the need for an explicit parametric bias model. Once trained, the model is designed to replace the BORG inference pipeline, generating samples of initial conditions consistent with galaxy data. We will validate it against existing BORG reconstructions and assess its robustness to the galaxy-formation physics assumed during training. The long-term goal is more accurate field-level inference from current and upcoming surveys such as Euclid and DESI.
The training data set has been fully prepared, and the model architecture has been designed and implemented. The purpose of this allocation is to benchmark training and inference performance on modern GPU hardware, in particular to:
1) establish memory requirements and throughput for three-dimensional volumetric inputs at the resolutions needed for scientific use;
2) assess scaling across multiple GPUs using data- and model-parallel strategies;
3) identify the configuration required for a subsequent larger production allocation.
The workload is dominated by large three-dimensional tensors. Individual training samples and intermediate activations are volumetric fields, so the project requires high-memory GPUs (80 GB class or equivalent for optimal efficiency) and efficient multi-GPU communication. Smaller-memory devices would force a reduction in volume or resolution that would compromise the physical fidelity the model is intended to achieve.
The results of this benchmarking phase will determine the computational strategy for scaling the approach to production-level training. The project connects state-of-the-art hydrodynamical simulations with field-level Bayesian inference. It aims to improve the accuracy with which galaxy surveys can constrain cosmological models and study the physics of galaxy formation.