Cosmo3DFlow: Wavelet Flow Matching for Spatial-to-Spectral Compression in Reconstructing the Early Universe
This work addresses the computational challenges in cosmological inference for astrophysics, representing an incremental improvement by combining existing techniques like wavelet transforms and flow matching to enhance efficiency.
The paper tackles the computationally demanding problem of reconstructing the early Universe from present-day data by developing Cosmo3DFlow, a generative framework that integrates 3D Discrete Wavelet Transform with flow matching to address dimensionality and sparsity bottlenecks, achieving up to 50× faster sampling than diffusion models and enabling initial conditions to be sampled in seconds instead of minutes.
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.