Skip to main content
Illustris Simulation
August 23, 2023

Through machine learning maps, cosmic history comes into focus

Written By: Jason Daley

For millennia, humans have used optical telescopes, radio telescopes and space telescopes to get a better view of the heavens.

Today, however, one of the most powerful tools for understanding the cosmos is the computer chip: Cosmologists rely on processing power to analyze astronomical data and create detailed simulations of cosmic evolution, galaxy formation and other far-out phenomena. These powerful simulations are starting to answer fundamental questions of how the universe began, what it is made of and where it’s likely headed.

Kangwook Lee

“It is extremely expensive to run these simulations and basically takes forever,” says Kangwook Lee, an assistant professor of electrical and computer engineering at the University of Wisconsin-Madison. “So they cannot run them for large-scale simulations or for high-resolution at that same time. There are a lot of issues coming from that.”

Instead, machine learning expert Lee and physics colleagues Moritz Münchmeyer and Gary Shiu are using emerging artificial intelligence techniques to speed up the process and get a clearer view of the cosmos.

Münchmeyer began by running large-scale, low-resolution simulations, then attempted to upgrade them using a machine-learning technique called normalizing flow. The results, however, weren’t as promising as he’d hoped.

That’s why he reached out to Lee, who specializes in a type of machine learning called diffusion—a model capable of generating data similar to a data set on which it is trained. In the last year, this type of artificial intelligence has gained wide popularity for generating text-to-image artwork through applications like Stable Diffusion.

With diffusion, the team first focused on dark matter, an unknown type of matter that could account for about 85% of all matter in the universe. The researchers trained the diffusion model on dark matter simulations of various resolutions. By applying the model to low-resolution datasets of dark matter distributions, the team could quickly upgrade the data into high-resolution images. The researchers also applied the technique to gas distribution models derived from the dark matter distributions.

Working across one small patch of sky at a time, then using an algorithm to stitch the data together, they made high-resolution images of the cosmic web of galaxies and galaxy clusters that forms the universe. The implementation of the project was led by the CS student Michael Gira and physics grad student Adam Rouhiainen, with help from Münchmeyer’s postdoctoral researcher Utkarsh Giri, Shiu’s grad student Jacky Yip, and Lee’s students Joseph Shenouda, Ying Fan and Yicong Chen.

And those machine-learning-enhanced images were accurate: The researchers’ images proved extremely similar to those produced by running full, high-resolution simulations of the dark matter and gas distributions.

Now the researchers are working on a similar technique that will allow them to apply the same diffusion model to high-resolution cubes of data that they will stitch together to produce 3D maps of the cosmos.

Creating these high-resolution simulations isn’t just about making a map of the universe, says Lee. There’s a much loftier goal: addressing the “inverse problem” in cosmology. To understand the origins of galaxies, matter and the Big Bang, researchers must work backward from what they can currently see—so they benefit from a more accurate window into today’s universe.

“We can observe the universe today, but we cannot observe the same part of the universe at different times,” he says. “But from the arrangement of matter and gas today, we are getting better and better at estimating what the universe looked like in the past, as far back as the time of the big bang. That’s what this project is about.”

Although his own understanding of cosmology is still fairly limited, Lee is happy his expertise can play a part in unraveling the secrets of the universe. “I feel like this type of interdisciplinary collaboration is one of the best things you can have at a university,” he says. “I found this very cool. My work has been very theoretical so far, so it feels rewarding to work on a real scientific application.”

A UW-Madison Research Forward grant funded portions of the research.

Top photo caption: Using machine learning techniques, Kangwook Lee and his collaborators are able to produce high-resolution images from low-resolution simulations. These types of techniques could help improve large scale models, like the Illustris Simulation, shown here. In this simulation, dark matter density is overlaid with the gas velocity field. Credit: Illustris Collaboration