C3NN-SBI: Learning Hierarchies of $N$-Point Statistics from Cosmological Fields with Physics-Informed Neural Networks

TL;DR

A simulation-based inference pipeline that not only benefits from the efficiency of machine learned summaries through optimization, but also holds on to the NPCF program, using the heavily constrained Cosmological Correlator Convolutional Neural Network.

Abstract

Cosmological analyses are moving past the well understood 2-point statistics to extract more information from cosmological fields. A natural step in extending inference pipelines to other summary statistics is to include higher order N-point correlation functions (NPCFs), which are computationally expensive and difficult to model. At the same time it is unclear how many NPCFs one would have to include to reasonably exhaust the cosmological information in the observable fields. An efficient alternative is given by learned and optimized summary statistics, largely driven by overparametrization through neural networks. This, however, largely abandons our physical intuition on the NPCF formalism and information extraction becomes opaque to the practitioner. We design a simulation-based inference pipeline, that not only benefits from the efficiency of machine learned summaries through optimization, but also holds on to the NPCF program. We employ the heavily constrained Cosmological Correlator Convolutional Neural Network (C3NN) which extracts summary statistics that can be directly linked to a given order NPCF. We present an application of our framework to simulated lensing convergence maps and study the information content of our learned summary at various orders in NPCFs for this idealized example. We view our approach as an exciting new avenue for physics-informed simulation-based inference.

Figures (6)

• Figure 1: Cartoon illustration of our summary extraction and neural posterior estimation (NPE) pipeline. The input convergence maps ($\kappa$, with $zs1$ to $zs3$ indicating source galaxy tomographic bins, i.e. different channels) are convolved once with a set of isotropic filters $\alpha$ with learnable weights, producing the first-order moment maps $C_\alpha^{(1)}(\pmb{x})$. Higher-order moment maps $C_\alpha^{(N)}(\pmb{x})$ are then computed recursively using the relation in Eq. \ref{['equ:recursion']}, avoiding the need for repeated convolutions. Spatial averaging of each moment map yields the corresponding moment estimates $c_\alpha^{(N)}$, which serve as our summary statistics. These summaries, together with the associated simulation parameters, are passed to a masked autoregressive flow (MAF)–based neural posterior estimator to infer the posterior distribution $p(\theta \mid \{c_\alpha^{(N)}\})$. The entire pipeline is trained end-to-end, with gradients propagating from the posterior estimator through the summary network.
• Figure 2: Mock inference on a withheld test simulation at the fiducial cosmology. In (a) we compare the constraining power of C3NNs at different orders. As we increase the order, i.e. up to the 4th order which is the highest order of correlation function considered, the constraining power increases due to the non-gaussian information in the convergence map. In (b) we compare a traditional CNN with our C3NN at 4th order. We find that at the same low amount of training simulations, the CNN model can extract more cosmological information possibly drawing from even higher orders in NPCFs. We also add the mock inference results from measured 2PCFs, and its constraining power is comparable to our C3NN at 4th order, due to the limited training simulation budget.
• Figure 3: Histograms of the improvement in the FoM in terms of fractional difference for 200 simulation realizations at the fiducial cosmology. We also show the median improvement in black, and the improvement seen in the posterior of the selected fiducial simulation inference shown in Figure \ref{['fig:C3NN_orders']} in red.
• Figure 4: (a) Comparison of a hybrid makinen_24makinen_25 C3NN at 4th order to the measured 2PCF with NPE. The hybrid model outperforms the 2PCF NPE posterior. (b) Comparison of our best C3NN 4th order NPE result with a corresponding C3NN 4th order NLE model using twice the training data for the inference network. By including more simulations in the training of the NLE inference network, the constraining power increases substantially.
• Figure 5: The constraining power of C3NN at different orders where the training maps are on the same cosmological nodes as used in Figure \ref{['fig:C3NN_orders']} but phase-randomized: The fields are first Fast Fourier Transformed (FFT), then their phases randomized following a uniform distribution within 0 and $2\pi$, and then with inverse FFT transformed back to configuration space.