Λ-CFM: Representation Learning for Cosmology

Exploring the Dark Side of the Universe

Our world is made of things we can see: laptop screens, people, clouds, stars, and even tiny things like atoms. But it turns out that all of this familiar matter makes up only about 5% of the universe. The rest? Mostly two mysterious ingredients we call dark matter and dark energy. We call them dark because they do not interact with light in the familiar ways ordinary matter does, so we infer their existence mostly through their effects on the universe.

The leading hypothesis about this dark sector of the universe is Λ-CDM. The Λ represents the cosmological constant, the simplest model of dark energy. CDM stands for cold dark matter: dark matter whose particles move slowly enough to clump gravitationally, eventually helping form the vast network of clusters, filaments, and voids known as the cosmic web.

Scientists study the nature of dark matter and dark energy by running large-scale cosmological simulations based on Λ-CDM and comparing the simulated universe with observations (e.g., DESI, the Dark Energy Spectroscopic Instrument). In these simulations, cosmologists vary parameters such as $\Omega_m$, which describes how much matter the universe contains, and $\sigma_8$, which describes how strongly matter is clustered. The challenge is that both simulations and observations now contain enormous amounts of information. To test different hypotheses about dark matter and dark energy (e.g., Λ-CDM or other alternatives), we need ways to compress these data without throwing away the subtle features where new physics might be hiding.

Learning the Compressed Representation of the Universe

To deal with the enormous size of simulation and observational data, cosmologists traditionally compress these fields into summary statistics. One commonly used statistic is the power spectrum, which roughly quantifies how much structure exists at different spatial scales. While the power spectrum is sufficient for a Gaussian field, it cannot capture all the information in today’s highly non-Gaussian universe, which contains filaments, halos, clusters, and voids. Rather than hand-designing increasingly complex summary statistics, researchers began asking: can machine learning discover salient summaries that preserve the cosmological information we care about?

Existing machine-learning approaches$^{1-3}$, however, have two key limitations:

1. They are often supervised. A model trained to predict Λ-CDM parameters learns to preserve information useful for those parameters, but may discard unexpected features that could reveal where Λ-CDM fails.
2. They are not scale-aware. Unlike the power spectrum, a generic neural representation does not tell us whether a signal comes from large-scale cosmic structure or small-scale features.

We therefore wanted a representation that was both: unsupervised, so that it preserves information without being told in advance which cosmological parameters matter, and scale-aware, so that different pieces of the representation correspond to different physical scales.

Λ-CFM: Flow Matching Models Learn by Generating the Universe

We propose Λ-Conditional Flow Matching (Λ-CFM), a generative approach for scale-aware representation learning. Instead of asking the model, “What is $\Omega_m$?” we ask it to compress a simulated dark-matter field and then reconstruct the field itself. That changes what the representation has to learn. To reconstruct the cosmic web, the model must retain information about its filaments, clusters, voids, and structure across spatial scales, including information that may have nothing to do with the particular Λ-CDM parameters we chose beforehand.

The decoder is based on flow matching$^{4,5}$. It begins from noise and progressively transforms that noise into a cosmological density field. Interestingly, this generation process naturally proceeds from large scales to small ones: broad cosmic structure appears first, followed by increasingly fine detail. We exploit that behavior with a time-dependent masking scheme. Different latent channels are forced to become useful at different stages of generation, encouraging them to specialize in different spatial scales.

What Does Λ-CFM Learn?

A few results were particularly striking.

First, the compression can still reconstruct surprisingly detailed cosmic webs. Compared with a variational autoencoder using the same latent-code size, Λ-CFM preserves substantially more fine-scale structure and high-frequency power.

Second, the latent space behaves like a physically meaningful space rather than simply an arbitrary compression. Interpolating between two points in latent space produces realistic intermediate universes and smoothly changes the inferred cosmological parameter $\sigma_8$. Directly blending the two images instead produces blurred, unphysical structures.

Third, different pieces of the latent representation actually control different physical scales. Swapping one group of channels changes predominantly fine-scale structure; swapping another modifies large-scale structure. The separation disappears substantially when the temporal masking mechanism is removed.

And despite this being learned without directly supervising the representation on cosmological parameters, a tiny summary vector still predicts $\Omega_m$ and $\sigma_8$ to roughly 5% relative error. That is the broader idea behind Λ-CFM: rather than deciding beforehand what information from the universe deserves to survive compression, we can train a generative model to preserve the field itself, while organizing what it learns according to physical scale.

What's Next?

Λ-CFM learns its representation without being told the cosmological parameters, but there is still an important limitation: all of its training data come from ΛCDM simulations. As a result, when the model encounters physics outside that distribution, it can tend to map it back toward the ΛCDM universe it already knows. A key next step is therefore to build representations that are more sensitive to unexpected physics (e.g., warm dark matter).

We are also interested in understanding how well these representations transfer across different simulations and physical fields, extending the approach beyond field-level maps, and going from representation learning to actual compression. In this work, we reduce the dimensionality of the data, but we do not yet perform true bit-level compression for storage or transmission. Incorporating neural compression techniques could eventually turn Λ-CFM into a practical tool for storing and communicating increasingly large cosmological datasets.

1. Akhmetzhanova, A., S. Mishra-Sharma, and C. Dvorkin. 2023. “Data Compression and Inference in Cosmology with Self-Supervised Machine Learning.” Monthly Notices of the Royal Astronomical Society 527 (3): 7459–7481.
2. Andrianomena, S., and S. Hassan. 2023. “Latent Space Representations of Cosmological Fields.” arXiv preprint arXiv:2311.00799.
3. Charnock, T., G. Lavaux, and B. D. Wandelt. 2018. “Automatic Physical Inference with Information Maximizing Neural Networks.” Physical Review D 97 (8): 083004.
4. Albergo, Michael S., Nicholas M. Boffi, and Eric Vanden-Eijnden. 2023. “Stochastic Interpolants: A Unifying Framework for Flows and Diffusions.” arXiv:2303.08797.
5. Lipman, Yaron, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. 2023. “Flow Matching for Generative Modeling.” International Conference on Learning Representations (ICLR).