Developing a faster AI method for reconstructing the early universe from 3D maps of today’s cosmos

CosmicAI Researchers Md. Khairul Islam, Zeyu Xia, Ryan Goudjil, and Judy Fox (University of Virginia), together with Jialu Wang and Arya Farahi (University of Texas at Austin), developed a faster AI method for reconstructing the early universe from 3D maps of today’s cosmos. 

What the team did

Cosmologists use 3D maps of matter in the universe to study how galaxies, clusters, and the cosmic web formed over time. One major goal is to work backward from these present-day structures and estimate the universe’s initial conditions: the tiny variations in matter that later grew into the large-scale structures we see today.

This is difficult because the data are both massive and uneven. A 3D map divides space into millions of small regions, but much of that volume is relatively empty. Existing generative AI methods, including diffusion models, reconstruct these maps through many small steps and still spend substantial computing time on regions that contain little useful information.

The team’s method, Cosmo3DFlow, tackles both issues. It first represents the 3D data with wavelets, a mathematical tool that separates broad cosmic structures from fine-scale details and compresses the information into a more efficient form. It then uses flow matching, an AI technique that can generate a reconstruction in a small number of direct steps rather than the many iterative steps used by diffusion models. The researchers tested Cosmo3DFlow on large 3D N-body simulations from the Quijote suite, which model how matter evolves across the universe under many possible physical conditions.

What researchers found

Cosmo3DFlow reconstructed 3D initial conditions up to 46 times faster than a leading diffusion-based baseline. It produced a reconstruction in about five seconds, compared with roughly four minutes for the earlier method.

The speedup came from two sources. Cosmo3DFlow required about 10 times fewer generation steps, and each step cost about five times less to run. It also matched or outperformed the diffusion-based method across every accuracy measure the researchers evaluated.

Why the work matters

Reconstructing 3D initial conditions is an important part of testing models of how the universe evolved. Faster reconstructions could let researchers compare more possible models, run larger studies, and make better use of increasingly detailed simulations and galaxy surveys.

The work also addresses a broader challenge in scientific AI. Many physical systems are represented on large grids even though the information of interest is concentrated in a small fraction of the space. By combining a more compact representation with a faster generative method, Cosmo3DFlow reduces that computational burden.

The paper was published in the Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD 2026), August 9–13, 2026, Jeju Island, South Korea.