Photographer: Samuel Avraham, Sören Andersson, Gunilla Häggström.

Link to the article: Learning the Universe: learning to optimize cosmic initial conditions with non-differentiable structure formation models

A central question in cosmology is: how did the Universe we see today come about? To answer this, researchers have for decades built physics-based simulations that evolve the cosmos forward in time. However, before we can run a simulation that looks like our observable Universe, we need to estimate the Universe’s early state. Reconstructing this starting point from today’s observations is a daunting task, as it means working backward across 13.8 billion years of cosmic history. 

The digital twin idea is that once we can recover this early state and run it forward with sophisticated physical modelling, we have a complete simulated counterpart of the Universe, galaxy by galaxy. This makes it possible to study the evolution of individual structures in our Universe and explore how their properties change if we alter the physical laws or cosmological parameters.

The challenge of numerical simulations

State-of-the-art cosmological simulations incorporate detailed physics on small scales, including the formation of dark matter halos and galaxies. These complex processes are, however, difficult to integrate into methods that infer the Universe’s early state. In technical terms, such simulations are often ‘non-differentiable’, meaning they cannot easily be adapted to reconstruction approaches that rely on gradient information. As a result, some of the most sophisticated physics models have remained out of reach for this line of research.

“This problem is highly non-trivial because the relationship between initial state and final observed structures is strongly non-linear, and corresponds to the evolution of structures in our Universe over roughly 13.8 billion years. If we identify a discrepancy in the observed data, we need to determine the corresponding update in the initial conditions needed to reduce the data discrepancy”, says Ludvig Doeser, Doctoral student at the OKC and lead author of the study.

Training AI to guide the search

This project was carried out within the cosmo-AI team at Oskar Klein Centre under the guidance of Associate Prof. Jens Jasche. Led by author Ludvig Doeser, together with co-author Metin Ata, the team developed an approach that enables using any simulation in the reconstruction. They trained a neural optimizer – an AI model that learns how to update the early state of the Universe so that, when simulated forward, it matches the structures we can observe in the night sky today.

Infographic: A loop showing how AI could be used to create a digital twin of the Universe: by iteratively updating the initial conditions until the simulation output matches observed galaxies. Image: Samuel Avraham

“One of the biggest challenges was training a model that could learn how changes in the data space (the outputs of a simulation or from real observations) correspond to changes in the initial conditions (the inputs to the simulation), including the generation of suitable training data and neural network architecture”, says Ludvig Doeser.

To make this work possible, the researchers ran 6,000 structure formation simulations on the Sunrise supercluster at Stockholm University in mid-2024. From these runs, they built a training set of ‘difference fields’ by subtracting the inputs and outputs of one simulation from another. The neural optimizer was trained on this set to learn how changes in the starting conditions affect the output. The training was enabled by multi-GPU training on the Berzelius supercomputer at Linköping University using the deep learning framework PyTorch. 

To test the model, the team used the output of one simulation to reconstruct its input. Running a new simulation from this reconstructed state faithfully reproduced the original, matching it down to the smallest details. This success hints at the exciting possibility of applying the method to real observations of the Universe.

A scalable framework for digital twins of the cosmos

The method called Learning the Universe by Learning to Optimize (LULO) is designed to scale. It can be applied to any physics simulator as long as example pairs of early and late states can be generated. This flexibility makes it possible to include complex, realistic physics models directly in the reconstruction process without needing years of code modifications by large teams of researchers.

"This type of framework demonstrates the potential to leverage state-of-the-art simulations directly, allowing them to be fitted to data without the need for further development, offering a more efficient approach to cosmological inference. Rather than spending valuable time making simulation codes differentiable or crafting approximate alternatives, researchers can directly focus on improving the accuracy, speed, and robustness of the physics models themselves", says Ludvig Doeser.

By automating the process of fitting simulations to data, the framework removes the need for large teams to manually adapt codes. This marks an important step toward applying these methods at the scale needed for upcoming galaxy surveys. 

This work shows how AI can help us build a digital twin of the Universe. Such a computer model captures our Universe’s early state, runs it forward with physics simulators, and compares the result to the galaxies we see in the sky today. In doing so, researchers gain new ways to explore how cosmic structures took shape and to deepen our understanding of the cosmos.