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

Aug 2026·
Md Khairul Islam
Zeyu Xia
Zeyu Xia
,
Ryan Goudjil
,
Jialu Wang
,
Arya Farahi
,
Judy Fox
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, 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 \(46\times\) faster sampling than diffusion models. Our results enable initial conditions to be sampled in seconds, compared to minutes for previous methods.
Type
Publication
Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2 (KDD 2026), 11153-11164. Association for Computing Machinery
publications
Generative AI Efficient Sampling Flow Matching Generative Models Wavelet Representation Spatial Sparsity Spectral Sparsity Astronomy
Authors
Md Khairul Islam
Zeyu Xia
Authors
Zeyu Xia (he/him)
Computer Science PhD Student

I am a Computer Science PhD student at the University of Virginia, advised by Prof. Geoffrey C. Fox. I develop generative AI (GenAI) and scientific machine learning methods, with a focus on flow matching and GPU-accelerated differentiable simulation. My work connects efficient sampling, benchmark design, and model evaluation with inverse problems in the physical sciences. I contribute to DeepVELA, a Genesis Mission project exploring agentic benchmarks for scientific applications in collaboration with MLCommons. My work spans time, energy, space, cost, and precision evaluation and agent harness design with structured, verifiable outputs. I have also prototyped and trained JEPA world models.

I am seeking 2027 research internships in generative AI and ML systems. Get in touch.