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

TL;DR

Cosmo3DFlow tackles the high-dimensional cosmological initial-condition inference problem by marrying a 3D Discrete Wavelet Transform with flow matching in a multi-resolution 3D U-Net. By performing velocity-field learning in wavelet space and optionally enforcing a physically informed power spectrum, the method achieves deterministic, fast sampling with substantially reduced computation compared to diffusion baselines. It demonstrates up to 50x faster sampling at $128^3$ resolution while maintaining or improving fidelity across VRMSE, cross-correlation, and power-spectrum metrics on Quijote-derived datasets. The approach concentrates computation on information-dense structures (filaments and halos) through wavelet sparsity, enabling scalable inference for larger cosmological volumes.

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.

Figures (7)

• Figure 1: Schematic comparison of data representations for the Cosmic Web. Left: The voxel grid allocates memory uniformly, forcing high-resolution processing on empty void regions. Right: The wavelet representation adapts to structure; voids are represented by a few coarse coefficients (large blocks), while computational density is concentrated solely on the filaments and halos (dense nodes).
• Figure 2: Cosmo3DFlow targets a high-dimensional inverse problem, while achieving 50$\times$ faster sampling. Given observations, it reconstructs samples in 5 seconds vs. $\approx$240 seconds for diffusion models, while preserving fine-scale structure through wavelet sparsity. Based on variance-normalized root mean squared error (VRMSE) and correlation metrics, Cosmo3DFlow demonstrates superior performance compared to diffusion.
• Figure 3: Qualitative comparison of reconstruction quality. 2D slices through 3D density fields at $z=127$ for three test samples from the Standard Latin Hypercube. Columns show: observation $\mathbf{y}$ at $z=0$, ground truth initial conditions $\mathbf{x}$, Diffusion, and Cosmo3DFlow (ours) reconstruction, and absolute error maps. Dashed boxes highlight small-scale structure differences---the baseline exhibits smoothing artifacts (blurry), while Cosmo3DFlow preserves sharp features. Error maps demonstrate that Cosmo3DFlow achieves consistently lower reconstruction error across all samples, particularly in high-density regions. Colorbars indicate dark matter density contrast $\delta$ and absolute error magnitude.
• Figure 4: Computational efficiency vs. reconstruction quality: Cosmo3DFlow achieves significantly faster sampling ($4.4\times$ at the same number of steps) and maintains accuracy despite using $10\times$ fewer steps.
• Figure 5: The generation process for a test sample. Showing how the models reconstruct the target simulation from Gaussian noise. The VRMSE loss at the last row (lower is better) shows significantly faster convergence of our proposed model.