Cosmo3DFlow: Wavelet Flow Matching for Spatial-to-Spectral Compression in Reconstructing the Early Universe

TL;DR

Cosmo3DFlow introduces a wavelet-based flow matching framework that significantly accelerates the reconstruction of the early Universe from large-scale cosmological simulations, reducing sampling time from minutes to seconds.

Contribution

The paper presents a novel generative model combining 3D wavelet transforms with flow matching to efficiently handle high-dimensional, sparse cosmological data for inference.

Findings

Achieves up to 50x faster sampling than diffusion models.

Reduces integration steps by 10x with lower computational cost.

Enables rapid sampling of initial conditions in seconds.

Abstract

Reconstructing the early Universe from the evolved present-day Universe is a challenging and computationally demanding problem in modern astrophysics. We devise a novel generative framework, Cosmo3DFlow, designed to address dimensionality and sparsity, the critical bottlenecks inherent in current state-of-the-art methods for cosmological inference. By integrating 3D Discrete Wavelet Transform (DWT) with flow matching, we effectively represent high-dimensional cosmological structures. The Wavelet Transform addresses the ``void problem'' by translating spatial emptiness into spectral sparsity. It decouples high-frequency details from low-frequency structures through spatial compression, and wavelet-space velocity fields facilitate stable ordinary differential equation (ODE) solvers with large step sizes. Using large-scale cosmological $N$-body simulations, at $128^3$ resolution, we achieve up to $50\times$ faster sampling than diffusion models, combining a $10\times$ reduction in integration steps with lower per-step computational cost from wavelet compression. Our results enable initial conditions to be sampled in seconds, compared to minutes for previous methods.