Generative Diffusion Priors for 3D Mapping of the Dark Universe

CVPR 2026 (Highlight)

Brandon Zhao¹   Diana Scognamiglio²   Olivier Doré^2,3   Katherine L. Bouman^1,3

¹Department of Computing and Mathematical Sciences, California Institute of Technology
²Jet Propulsion Laboratory, California Institute of Technology
³Cahill Center for Astronomy and Astrophysics, California Institute of Technology

Abstract

Reconstructing the three-dimensional distribution of dark matter from weak-lensing observations is a central but highly ill-posed inverse problem in cosmology. Unlike standard 3D reconstruction with multiple viewpoints, we observe the universe from a single line of sight, through noisy shape distortions of galaxies with uncertain distances, so meaningful recovery of the 3D matter field requires strong prior assumptions. Existing methods either produce point estimates with handcrafted priors or use neural ensembles for approximate Bayesian uncertainty, and struggle to capture the non-Gaussian, filamentary structure of the cosmic web. With the advent of new high-resolution cosmological simulations, we now have an alternative source of prior knowledge that captures the nonlinear statistics of structure formation with far greater fidelity than analytic prescriptions. We leverage these simulations to build a new dataset \( \texttt{Conicus3D} \), which enables us to learn a data-driven diffusion-model prior capturing the full 3D distribution of dark matter structure across cosmic time. Building on recent plug-and-play approaches, we modify a diffusion-based posterior sampling scheme to the 3D weak-lensing setting, combining the learned prior with a differentiable physical forward model. On realistic simulations targeting a modern weak lensing survey, our approach yields substantially improved 2D and 3D reconstruction accuracy over baseline methods. Moreover, it produces posterior samples whose statistics closely track the underlying simulations, while remaining robust to moderate shifts in cosmology.

arxiv [pdf] code [Inference][Dataset] dataset [Zenodo]

Video

Citation

@inproceedings{zhao2026generative,
    author    = {Zhao, Brandon and Scognamiglio, Diana and Dor{\'{e}}, Olivier and Bouman, Katherine L.},
    title     = {Generative Diffusion Priors for 3D Mapping of the Dark Universe},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2026},
    pages     = {23581-23590}
}

Acknowledgements