Hybrid Summary Statistics

Implicit (simulation-based) inference makes solving otherwise intractable inverse problems possible by employing neural compressors which can flexibly map big data vectors into informative summaries (Cranmer et al., 2020; Hoffmann & Onnela, 2023). These mappings can be made optimal (Charnock et al., 2018; Makinen et al., 2021, 2024; Lanzieri et al., 2024; Jeffrey et al., 2020), but often require large numbers of simulations to achieve convergence.

For many physics problems, such as large N-body and hydrodynamical solvers, forward simulations are expensive to generate, so networks tasked to compress these data for parameter inference must learn informative features using sparsely-sampled training sets, especially when large numbers of parameters are varied.

In inference problems, we aim to infer quantities of interest ${\theta}$ from data d. In the physical sciences we often have domain knowledge which allows us to design a summary statistic of the data $\textbf{t}(\textbf{d})$ that capture some, but not all, of the information in a dataset. To enhance the information extraction, we aim to learn additional summaries $\textbf{s}(\textbf{d})$ that are complementary to t. We can formalise the extra information beyond t captured by s as the conditional mutual information $I(\textbf{s};{\theta}|\textbf{t})$, which measures how much the uncertainty about ${\theta}$ is reduced by knowing s given a static t. Evaluating the conditional mutual information can be difficult. However, using the chain rule of MI we can write

where $I(\textbf{t};{\theta})$ is a constant and $[a,b]$ denotes concatenation. We can then maximise the additional information captured by s by maximising $I([\textbf{t},\textbf{s}];{\theta})$. There are various ways to maximize MI; here, we focus on two specific objectives, detailed in Appendix A. The first, referred to as the Posterior Entropy (EPE) objective, is given by:

where $\textbf{s}(\textbf{d})$ is a neural network and $q$ is a neural density estimator. Hoffmann & Onnela (2023) demonstrate that this loss unifies many information-theoretic losses into a stable objective. This mutual information objective can also be parameterised as a cross-entropy classification problem by employing the Jensen-Shannon divergence (Chen et al., 2021; Devon Hjelm et al., 2018; Nowozin et al., 2016):

where $s$ is the summarizer, $c$ is a classifier tasked with distinguishing between data from $p(\theta,x)$ and $p(\theta)p(x)$ and $\mathrm{sp}(z)=\log(1+e^{z})$ is the softplus function. We refer to this as the Cross Entropy (CE) loss.

To test the information capture in learned summaries we employ a two-step process, illustrated in Fig. 1. We first optimise a neural compression (embedding network) to a few additional numbers conditional on the specified (not learned) existing summary t, using either EPE or CE losses (denoted MI maximiser in Fig. 1). For the EPE loss we train a simple mixture density network (MDN) with two small hidden layers to approximate $q(\theta|[\textbf{s},\textbf{t}])$. For the CE loss we train a classifier fully-connected network with hidden sizes $[128,64,64]$ to one output. We then take the static learned and existing summaries and parameterise a separate posterior estimator using a masked autoregressive flow (MAF) to minimise $p(\theta|[\textbf{s},\textbf{t}])$ from the LtU-ILI package¹¹1https://github.com/maho3/ltu-ili (Ho et al., 2024). We feed comparison summaries (power spectrum and competing network schemes) into the same MAF architecture to obtain a consistent comparison of information capture. We detail exact configurations in Appendix B.

The reionization parameters we vary are (1) $\zeta$, the ionizing efficiency. It primarily determines the timing of the EoR, with higher values leading to earlier ionization of the Universe. We vary $\zeta$ as $10\leq\zeta\leq 250$, and (2) $T_{\mathrm{vir}}$, the minimum virial temperature of halos that host ionizing sources. $T_{\mathrm{vir}}$ controls the timing of astrophysical epochs and influences the scales of heated and ionized regions. We vary this parameter as $4\leq\log_{10}\left(T_{\mathrm{vir}}/\mathrm{K}\right)\leq 6$. More details can be found in Appendix B.

In this initial experiment, we compare tree types of summaries, assuming a sufficiently large training set (10,000 samples): 1) Power spectrum only (11 $k$-bins). 2) Hybrid method: Power spectrum + two learned supplementary features (EPE Loss). 3) Hybrid method: Power spectrum + two learned supplementary features (CE Loss). One of the resulting inferences is shown in Figure 2. The results show that both approaches capture non-Gaussian information and improve inference performance, yielding similar outcomes. This suggests we can learn two additional parameters instead of new summaries, maintaining power spectrum interpretability while achieving near-optimal inference. The agreement also confirms both loss functions are effective, converging to the same results.

We perform an ablation study to demonstrate the effectiveness of the hybrid statistics over neural-only methods and display results in Fig. 3. The embedding network in the hybrid setting is the lightweight CNN with symmetric kernels adapted from Makinen et al. (2024) to output 3 additional numbers. For comparison, we train a larger, more expressive CNN under the same EPE loss without access to the $C_{\ell}$ vector to output 3 numbers, and then for the density estimation step concatenate its outputs to the $C_{\ell}$s. We train the networks and density estimators first using all $5000$ simulations available (split into 70% train and 30% validation sets). We then reduce the total number of available simulations to $1000$ and further to $500$ and re-learn the embeddings and posteriors from scratch. When the simulation budget is reduced, all inferences suffer, but the hybrid statistic formalism encourages the smaller network to find non-Gaussian features in the dataset that are more robust to this change. We also note that the more expressive, separately-trained CNN inference degrades almost to the level of the $C_{\ell}$ contour in the lowest-data setting, even when concatenated to the $C_{\ell}$ vector in the density estimation step.

We detailed a method to learn compressed summary statistics that explicitly complement an existing summary of the data through mutual information maximisation. We show that these techniques can capture non-Gaussian information in two cosmological applications using two different loss criteria to significantly improve parameter information capture.

We additionally demonstrate that these hybridised summaries improve information capture when the training simulation budget is limited. This suggests that requiring a network to find patterns in the data that are explicitly complementary to a provided summary “tells it where to look" and improves the compression optimisation in smaller datasets over wide parameter space.

We note that MI can also be used as a static metric to quantify the information content in arbitrary summaries as in Sui et al. (2023). This technique could be extended towards exhaustive information studies to measure how much more information might be unlocked with multiple traditional or neural statistics.