Probing Cosmology with Diffusion Models: Memorization, Generalization, and Parameter Estimation
Abstract
Diffusion models are increasingly used as surrogates for expensive cosmological simulations, but simple metrics like reconstruction fidelity do not guarantee a surrogate can be trusted for scientific inference. To explore this, we train diffusion models on HI fields from the CAMELS IllustrisTNG simulations and test how three properties change with training-set size: memorization of individual training examples, agreement with one- and two-point statistics (the pixel-value distribution and the power spectrum), and estimation of cosmological parameters from the generated fields. At fixed model size, generative models trained with smaller, less diverse training data produced memorized samples. While this is a well known phenomenon, we show that this leads to miscalibrated parameter densities for cosmological inference tasks. At fixed size, there is a training set size with enough information content to enable a model to generalize, at which point posterior density estimation with neural surrogates can be well-calibrated. We demonstrate that a neural surrogate can therefore pass visual, one-point and power-spectrum convergence tests while still resulting in overconfident parameter estimation in a cosmological inference pipeline. Thus, tests of model generalization should be an additional check for inference pipelines adopting neural surrogates.
Related
- A Six-Parameter Teleparallel Cosmology Beyond ΛCDM: Sign-Changing Torsional Dark Energy and Implications for H0
- Cosmology-dependent covariance in galaxy cluster number counts: consequences for parameter inference
- Parameter Estimation for Eccentric Supermassive Black Hole Binaries with Pulsar Timing Arrays
- Parameter Estimation of Ringdown Quasinormal Modes with Autoencoder
- Crop biophysical parameter estimation via AI-enhanced hyperspectral remote sensing
BibTeX
@misc{pan2026probing,
title = {{Probing Cosmology with Diffusion Models: Memorization, Generalization, and Parameter Estimation}},
author = {Jiaming Pan and Nicholas Kern},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.37825},
eprint = {2610.37825},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.37825},
}
RIS
TY - UNPB
TI - Probing Cosmology with Diffusion Models: Memorization, Generalization, and Parameter Estimation
AU - Jiaming Pan
AU - Nicholas Kern
PY - 2026
DA - 2026/10/27
PB - arXiv
DO - 10.48550/arxiv.2610.37825
UR - https://arxiv.org/abs/2610.37825
AB - Diffusion models are increasingly used as surrogates for expensive cosmological simulations, but simple metrics like reconstruction fidelity do not guarantee a surrogate can be trusted for scientific inference. To explore this, we train diffusion models on HI fields from the CAMELS IllustrisTNG simulations and test how three properties change with training-set size: memorization of individual training examples, agreement with one- and two-point statistics (the pixel-value distribution and the power spectrum), and estimation of cosmological parameters from the generated fields. At fixed model size, generative models trained with smaller, less diverse training data produced memorized samples. While this is a well known phenomenon, we show that this leads to miscalibrated parameter densities for cosmological inference tasks. At fixed size, there is a training set size with enough information content to enable a model to generalize, at which point posterior density estimation with neural surrogates can be well-calibrated. We demonstrate that a neural surrogate can therefore pass visual, one-point and power-spectrum convergence tests while still resulting in overconfident parameter estimation in a cosmological inference pipeline. Thus, tests of model generalization should be an additional check for inference pipelines adopting neural surrogates.
ER -