CosmoFlow: Scale-Aware Representation Learning for Cosmology with Flow Matching

TL;DR

CosmoFlow is a flow matching generative model that learns compact, interpretable, and semantically rich latent representations of cosmological simulation data, enabling reconstruction, synthetic data generation, and parameter inference.

Contribution

It introduces a flow matching-based approach to learn compact, interpretable representations of cosmological data without supervision, significantly reducing data size.

Findings

Learned representations are 32x smaller than raw data.

Representations support reconstruction, synthetic data generation, and parameter inference.

Latent channels correspond to features at different cosmological scales.

Abstract

Generative machine learning models have been demonstrated to be able to learn low dimensional representations of data that preserve information required for downstream tasks. In this work, we demonstrate that flow matching based generative models can learn compact, semantically rich latent representations of field level cold dark matter (CDM) simulation data without supervision. Our model, CosmoFlow, learns representations 32x smaller than the raw field data, usable for field level reconstruction, synthetic data generation, and parameter inference. Our model also learns interpretable representations, in which different latent channels correspond to features at different cosmological scales.