Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation

Technical Summary: Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation

Problem Statement
Reconstructing the primordial matter density field (cosmological initial conditions, or ICs) from late-time observations of the large-scale structure (LSS) is a fundamental challenge in cosmology. While the early Universe is well-described by linear Gaussian fluctuations, the present-day cosmic web is highly non-linear and non-Gaussian. Reversing this evolution is an ill-posed problem: the mapping from early to late times is non-injective (different ICs can yield the same late-time field), and structure formation on halo scales is chaotic, amplifying measurement noise.

Traditional probabilistic approaches, such as the Bayesian Origin Reconstruction from Galaxies (BORG), utilize Hamiltonian Monte Carlo (HMC) to sample the posterior distribution. However, these methods require differentiable simulators (often relying on approximation schemes) and are computationally prohibitive, requiring months of CPU time to converge on high-dimensional parameter spaces ($\gtrsim 10^6$ parameters). Conversely, many machine learning approaches in this domain rely on point estimators (e.g., U-Nets) that provide a single "best-fit" reconstruction without quantifying uncertainty, or on generative models (e.g., score-based diffusion) that remain computationally expensive even at inference time.

Methodology
The authors propose a Simulation-Based Inference (SBI) framework termed Gaussian Neural Posterior Estimation (NPE). The core innovation is modeling the posterior distribution of the initial conditions, $p(\mathbf{z}|\mathbf{x})$, as a Gaussian distribution in the inferred field $\mathbf{z}$, with a diagonal covariance matrix in Fourier space.

The methodology proceeds as follows:

1. Statistical Formulation:

   • Prior: The prior $p(\mathbf{z})$ is assumed to be a Gaussian random field defined by the linear matter power spectrum at high redshift ($z=127$).
   • Likelihood: The likelihood $p(\mathbf{x}|\mathbf{z})$ is modeled as a Gaussian with an observation-dependent mean $\hat{\mathbf{z}}_\theta(\mathbf{x})$ and a precision matrix $\mathbf{Q}^L_\theta$. Crucially, $\mathbf{Q}^L_\theta$ is constrained to be diagonal in Fourier space (and approximately rotationally invariant), reflecting the translational invariance of the physics.
   • Posterior: The combination of Gaussian prior and Gaussian likelihood yields a Gaussian posterior. The posterior precision matrix is the sum of the prior and likelihood precision matrices ($\mathbf{Q}_\theta = \mathbf{Q}_P + \mathbf{Q}^L_\theta$).

2. Neural Architecture:

   • MAP Estimator: The mean of the posterior, $\hat{\boldsymbol{\mu}}_\theta(\mathbf{x})$, is modeled using a 3D U-Net architecture. To improve performance, the authors introduce a custom activation layer that separates the reconstruction into two scales: a trainable scaling factor $\alpha_\theta(\mathbf{k})$ handles large-scale linear evolution, while the U-Net focuses on reconstructing non-linear small-scale residuals.
   • Precision Matrix: The diagonal elements of the likelihood precision matrix, $D^L_\theta(k)$, are learned. The authors tested two parametrizations: (1) treating values as free parameters for each Fourier mode, and (2) enforcing rotational symmetry via a dense neural network taking the wavenumber $k$ as input. The "free" matrix approach showed slightly better stability.

3. Training and Data:

   • The model is trained on 2,000 pairs of initial ($z=127$) and final ($z=0$) dark matter overdensity fields from the Quijote $N$-body simulation suite.
   • The forward model uses full, non-differentiable $N$-body simulations (Gadget-III), avoiding the need for differentiable approximations required by HMC methods.
   • Training minimizes the negative log-posterior (Eq. 7) over 30 epochs on a single Nvidia A100 GPU, taking approximately 1.5 hours.

Key Contributions

• Efficient Sampling: By assuming a Gaussian posterior with a diagonal Fourier-space precision matrix, the authors enable the generation of thousands of posterior samples in seconds (approx. 3 seconds for 1,000 samples on a single GPU), orders of magnitude faster than existing HMC or diffusion-based methods.
• Non-Differentiable Simulators: The method does not require the forward model to be differentiable, allowing the use of state-of-the-art, full $N$-body simulations without approximation schemes.
• Analytical Fit for Fast Sampling: The authors demonstrate that the learned precision matrix diagonal values $D_\theta(k)$ can be accurately fitted with a simple analytical function (a Gaussian plus a constant) up to $k \approx 0.2 \, h/\text{Mpc}$. This allows any trained point estimator (MAP) to be converted into a fast posterior sampler without retraining the full precision matrix.
• Uncertainty Quantification: Unlike standard point estimators, this method provides a full posterior distribution, enabling rigorous uncertainty quantification for field-level reconstruction.

Results
The authors validate the method using a fiducial simulation not present in the training set:

• Visual and Statistical Consistency: Generated samples visually resemble the ground truth. Summary statistics (power spectrum, transfer function, cross-correlation) show high agreement with the true initial conditions. The power spectrum matches the truth within 1–2% up to the Nyquist scale ($k_{Nyq} \approx 0.4 \, h/\text{Mpc}$).
• Cross-Correlation: The cross-correlation between generated samples and the truth remains high ($\gtrsim 70\%$) up to $k \approx 0.3 \, h/\text{Mpc}$.
• Bayesian Consistency: A scale-dependent consistency test confirms that the normalized differences between the true modes and the posterior mean follow a standard normal distribution ($N(0,1)$) across various $k$-bins. This verifies that the inferred uncertainties are well-calibrated.
• Robustness: The method maintains performance even when Gaussian noise is added to the observed final field.

Significance and Claims
The paper claims that this approach provides a simulation-efficient pathway for sequential SBI in cosmology. By reducing the complexity of the posterior to a Gaussian form with a diagonal precision matrix, the authors achieve a balance between physical accuracy and computational speed that was previously unattainable for high-dimensional field reconstruction.

The authors modestly note that the Gaussian Fourier-diagonal approximation is sufficient for the selected scale range ($k \lesssim 0.4 \, h/\text{Mpc}$) but may require more sophisticated precision matrices (e.g., convolutional or wavelet-based) for smaller scales or when observational effects break translational invariance. They emphasize that while the current work assumes fixed cosmological parameters and idealized dark matter observations, the framework is designed to be extended to varying cosmologies and realistic galaxy surveys in future work. The primary significance lies in enabling the generation of thousands of constrained realizations in seconds, paving the way for robust statistical analyses of the LSS that were previously computationally infeasible.