Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation

TL;DR

The paper addresses the challenge of reconstructing cosmological initial conditions (ICs) from late-time large-scale structure data, a high-dimensional, ill-posed inverse problem. It introduces Gaussian Neural Posterior Estimation (NPE) within simulation-based inference (SBI), modeling the IC posterior as Gaussian in the initial field with a diagonal Fourier-space precision, which enables fast, amortised sampling even with full $N$-body forward models. A 3D U-Net-based MAP estimator coupled with a learnable diagonal likelihood precision in Fourier space yields an efficiently trainable model; an analytic Gaussian fit to the precision as a function of wavenumber further enables turning any IC point estimator into a fast sampler. Validation on Quijote simulations shows power spectrum accuracy at the ~1–2% level up to Nyquist and strong cross-correlations, with per-$k$-bin Bayesian coverage indicating well-calibrated uncertainties. The approach is differentiable-free and scalable, offering a practical route to sequential SBI and integration of observational effects in future work.

Abstract

Knowledge of the primordial matter density field from which the large-scale structure of the Universe emerged over cosmic time is of fundamental importance for cosmology. However, reconstructing these cosmological initial conditions from late-time observations is a notoriously difficult task, which requires advanced cosmological simulators and sophisticated statistical methods to explore a multi-million-dimensional parameter space. We show how simulation-based inference (SBI) can be used to tackle this problem and to obtain data-constrained realisations of the primordial dark matter density field in a simulation-efficient way with general non-differentiable simulators. Our method is applicable to full high-resolution dark matter $N$-body simulations and is based on modelling the posterior distribution of the constrained initial conditions to be Gaussian with a diagonal covariance matrix in Fourier space. As a result, we can generate thousands of posterior samples within seconds on a single GPU, orders of magnitude faster than existing methods, paving the way for sequential SBI for cosmological fields. Furthermore, we perform an analytical fit of the estimated dependence of the covariance on the wavenumber, effectively transforming any point-estimator of initial conditions into a fast sampler. We test the validity of our obtained samples by comparing them to the true values with summary statistics and performing a Bayesian consistency test.

Figures (7)

• Figure 1: Schematic description of the method presented in this work. The posterior probability distribution of initial conditions ${\boldsymbol{z}}$ for a given observation of the final conditions ${\boldsymbol{x}}_{\rm obs}$ is modelled to be Gaussian in ${\boldsymbol{z}}$, with the likelihood precision matrix $\bm Q_{\boldsymbol{\theta}}^L$ being diagonal in Fourier space. This leads to the ansatz for the posterior (\ref{['eq:posterior']}). The MAP estimator $\mathbf{\hat{\boldsymbol{\mu}}}_{\boldsymbol{\theta}}$ and the likelihood precision matrix $\bm Q_{\boldsymbol{\theta}}^L$ are then simultaneously trained to maximise the posterior probability for the training set under the assumed probability distribution. After training, which takes $\sim 1.5 \, \mathrm{h}$ on a single GPU at $128^3$ resolution, we can obtain samples of the initial conditions in an almost instantaneous way given the Gaussian approximation ($\lesssim 3 \, \mathrm{s}$ for $1000$ samples).
• Figure 2: Slices of the results of our Gaussian neural posterior estimation for 3D reconstruction of cosmological ICs. Top left: ground truth initial overdensity fields ${\boldsymbol{z}}_{\mathrm{truth}}$. Top centre: MAP estimate of the initial condition given the observed final conditions ${\boldsymbol{x}}_{\mathrm{obs}}$ computed with the estimator $\mathbf{\hat{\boldsymbol{\mu}}}_{\boldsymbol{\theta}}({\boldsymbol{x}}_{\mathrm{obs}})$. Bottom left and centre: two examples of the generated posterior ICs samples. Top right: given late-time density field, i.e. the observation ${\boldsymbol{x}}_{\mathrm{obs}}$. All the shown slices are averaged over the depth of $100 \ \text{Mpc}/h$ in the third axis direction. Bottom right: the posterior precision matrix is modelled to be diagonal in Fourier space: $\bm Q_{\boldsymbol{\theta}} = \bm Q^P + \bm Q_{\boldsymbol{\theta}}^L = \bm U^{\dagger} \bm D_{{\boldsymbol{\theta}}} \bm U$. Here we show the diagonal component $\bm D_{{\boldsymbol{\theta}}} = \bm D^P + \bm D_{{\boldsymbol{\theta}}}^L$ trained values as a function of the wavenumber $k$ ( green dots), as well as a simple analytical Gaussian fit of the $\bm D_{{\boldsymbol{\theta}}}(k)$ dependence accurate up to $k \simeq 0.2 \, h/\text{Mpc}$ ( blue line). Knowing the form of the $\bm D_{{\boldsymbol{\theta}}}(k)$ function allows one to turn any ICs point estimator into a fast sampler, as we discuss in Section \ref{['subsec:fit']}. The precision is indicative of the confidence in the reconstructed ICs at a given scale and therefore peaks at large scales, which are easier to infer.
• Figure 3: The trainable parts of the algorithm are the MAP estimator $\mathbf{\hat{\boldsymbol{\mu}}}_{\boldsymbol{\theta}}$ and the likelihood precision matrix $\bm Q_{\boldsymbol{\theta}}^L$. Orange vertical lines depict the mean of the posterior predictions for some randomly selected individual Hartley modes, equal to $\bm H \{\mathbf{\hat{\boldsymbol{\mu}}}_{\boldsymbol{\theta}}({\boldsymbol{x}}_{\mathrm{obs}})\}({\boldsymbol{k}})$ (with $\bm H$ denoting the Hartley transform, see App. \ref{['appendix_hartley']}), together with their $1\sigma$ errors, equal to $(\bm D^P + \bm D_{\boldsymbol{\theta}}^L)^{-\frac{1}{2}}$. Horizontal blue lines show the true value of a given mode.
• Figure 4: Comparison of the density field histograms (1-point functions) for the true field ${\boldsymbol{z}}_{\mathrm{truth}}$ and one of the generated samples.
• Figure 5: Comparison of the summary statistics at $z=127$ between the ground truth simulation and the generated samples together with the MAP estimate. Top: power spectra $P(k)$. Middle: transfer function $T(k)$. Bottom: cross-power spectra $C(k)$. For summaries of the generated samples, solid lines represent the mean, while shaded regions correspond to $1\sigma$ and $2\sigma$ uncertainties computed from a 1000 generated samples using the Pylians3 package 2018ascl.soft11008V. The plots demonstrate that generated samples reproduce the true field with a high level of precision.