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

2026-02-10

AI Summary

Efficient reconstruction of initial conditions in the early universe is hindered by high-dimensional sparsity and prohibitive computational costs. This work proposes the first method integrating three-dimensional discrete wavelet transform (DWT) into a flow matching framework, effectively alleviating the “emptiness problem” by converting spatial sparsity into spectral sparsity and enabling large-step ordinary differential equation (ODE) solvers. The approach achieves a 50-fold acceleration over diffusion models at a resolution of \(128^3\), generating initial conditions in seconds—a full order-of-magnitude speedup—while significantly enhancing both generation efficiency and numerical stability.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Web Mining and Content Analysis: Web data generation and simulation
• Systems and Infrastructure for Web, Mobile and WoT: Data management and stream processing for Web, mobile and wireless applications

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.

Problem

• early Universe reconstruction
• dimensionality
• sparsity
• cosmological inference
• void problem

Innovation

• Wavelet Flow Matching
• Spatial-to-Spectral Compression
• Cosmological Reconstruction
• 3D Discrete Wavelet Transform
• Generative Modeling