Λ-conditional flow matching: scale-aware representation learning for cosmology with flow matching

Abstract In modern cosmology, high-resolution simulations of the cosmic web generate petabytes of data, creating a critical need for dimensionality reduction schemes that preserve scientific information for downstream analysis. In this work, we demonstrate that \textit{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, $\Lambda$-Conditional Flow Matching (\OURS), learns representations 32x smaller than the raw field data, usable for field level reconstruction and synthetic data generation. We show that an additional 64x compression is possible, yielding summary statistics can be used to estimate the cosmological parameters, $\Omega_m$ and $\sigma_8$ to within 5\% accuracy. Finally, we show that through the use of a time-dependent masking scheme, our model also learns \textit{interpretable, scale-aware} representations, in which different channels of the compressed representation correspond to features at different cosmological scales.