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

arXiv: https://arxiv.org/abs/2602.10172 · HTML: https://mlsanity.com/paper/2602.10172/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Inverse problems and image restoration, one of 14 directions in Diffusion models.
- Builds on 12 of its 35 references that are in the map.
- First posted on arXiv 10 Feb 2026.

## Placed in

- Diffusion models / [Inverse problems and image restoration](https://mlsanity.com/diffusion/inverse-problems-and-image-restoration/)

## Builds on (most cited first, 12 of 12)

- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains](https://mlsanity.com/paper/2006.10739/) (2020)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [The Well: a Large-Scale Collection of Diverse Physics Simulations for Machine Learning](https://mlsanity.com/paper/2412.00568/) (2024)
- [Wavelet Diffusion Models are fast and scalable Image Generators](https://mlsanity.com/paper/2211.16152/) (2022)
- [WDM: 3D Wavelet Diffusion Models for High-Resolution Medical Image Synthesis](https://mlsanity.com/paper/2402.19043/) (2024)
- [Latent Wavelet Diffusion For Ultra-High-Resolution Image Synthesis](https://mlsanity.com/paper/2506.00433/) (2025)
- [Conditional Diffusion-Flow models for generating 3D cosmic density fields: applications to f(R) cosmologies](https://mlsanity.com/paper/2502.17087/) (2025)
- [Initial Conditions from Galaxies: Machine-Learning Subgrid Correction to Standard Reconstruction](https://mlsanity.com/paper/2504.01092/) (2025)
- [Wavelet Flow For Extragalactic Foreground Simulations](https://mlsanity.com/paper/2505.21220/) (2025)
- [FlowLensing: Simulating Gravitational Lensing with Flow Matching](https://mlsanity.com/paper/2510.07878/) (2025)
- [Cosmo-FOLD: Fast generation and upscaling of field-level cosmological maps with overlap latent diffusion](https://mlsanity.com/paper/2601.14377/) (2026)

## Built on by (most cited first, 1 of 1)

- [ScatterPrism: convergence for generative simulation and inverse problems in particle and nuclear physics](https://mlsanity.com/paper/2604.01313/) (2026)

## Abstract (arXiv)

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, and wavelet-space velocity fields facilitate stable ordinary differential equation (ODE) solvers with large step sizes. Using large-scale cosmological N-body simulations at 1283 resolution, we achieve up to 46× faster sampling than diffusion models. Our results enable initial conditions to be sampled in seconds, compared to minutes for previous methods.