Marginal Bayesian Statistics Using Masked Autoregressive Flows and Kernel Density Estimators with Examples in Cosmology ^†

Abstract

Cosmological experiments often employ Bayesian workflows to derive constraints on cosmological and astrophysical parameters from their data. It has been shown that these constraints can be combined across different probes, such as Planck and the Dark Energy Survey, and that this can be a valuable exercise to improve our understanding of the universe and quantify tension between multiple experiments. However, these experiments are typically plagued by differing systematics, instrumental effects, and contaminating signals, which we collectively refer to as ‘nuisance’ components, which have to be modelled alongside target signals of interest. This leads to high dimensional parameter spaces, especially when combining data sets, with ≳20 dimensions of which only ∼5 correspond to key physical quantities. We present a means by which to combine constraints from different data sets in a computationally efficient manner by generating rapid, reusable, and reliable marginal probability density estimators, giving us access to nuisance-free likelihoods. This is possible through the unique combination of nested sampling, which gives us access to Bayesian evidence, and the marginal Bayesian statistics code margarine. Our method is lossless in the signal parameters, resulting in the same posterior distributions as would be found from a full nested sampling run over all nuisance parameters, and typically quicker than evaluating full likelihoods. We demonstrate our approach by applying it to the combination of posteriors from the Dark Energy Survey and Planck.

1. Introduction

Bayesian inference is a cornerstone of modern cosmology and astrophysics. It is frequently employed to derive parameter constraints on key signal parameters from data sets such as the Dark Energy Survey (DES, [1,2]), Planck [3], REACH [4], and SARAS2 [5], among others. Often, experiments are sensitive to different aspects of the same physics, and by combining constraints across probes we can improve our understanding of the Universe or reveal tensions between different experiments.

However, this can become a computationally expensive task as many experiments feature systematic, instrumental effects, as well as contamination from other physical signals that need to be modelled alongside the signal or parameters of interest. For individual experiments, this can lead to high dimensional problems with ≳20 parameters of which the majority can be considered ‘nuisance’ parameters. The problem is compounded when combining different data sets with different models for common nuisance components and different systematics or instrumental effects that have to be modelled.

In this work, we demonstrate that we can use density estimators, such as kernel density estimators [6,7] and masked autoregressive flows [8], to rapidly calculate reliable and reusable representations of marginal probability densities and marginal Bayesian summary statistics for key signal or cosmological parameters. This gives us access to the nuisance-free likelihood functions and allows us to combine parameter constraints from different data sets in a computationally efficient manner given marginal posterior samples from the different experiments. We use the publicly available code (https://github.com/htjb/margarine (accessed on 26 October 2022)) margarine [9] to generate density estimators.

In Section 2, we mathematically demonstrate that the application of margarine to the problem of combining the marginal posteriors from two data sets is equivalent to running a full nested sampling run including all ‘nuisance’ parameters. We define in this section the nuisance-free likelihood. Section 3 briefly discusses the methodology behind margarine with reference to a previously published work [9]. Finally, we show the results of combining samples from DES and Planck in Section 4 and conclude in Section 5.

2. Theory

2.1. Notation

Given a likelihood $\mathcal{L}\left(\mathsf{\Theta }\right)\equiv P\left(D\right|\mathsf{\Theta },\mathcal{M})$ representing the probability of data D given some model $\mathcal{M}$ with parameters $\mathsf{\Theta }$, Bayesian inference proceeds by defining a prior $\pi \left(\mathsf{\Theta }\right)\equiv P\left(\mathsf{\Theta }\right|\mathcal{M})$, and then through Bayes theorem computing a posterior distribution $\mathcal{P}\left(\mathsf{\Theta }\right)\equiv P\left(\mathsf{\Theta }\right|D,\mathcal{M})$ for the purposes of parameter estimation and an evidence $\mathcal{Z}\equiv P\left(D\right|\mathcal{M})$ in order to perform model comparison. In our notation, we suppress model dependence, but where we wish to refer to the likelihoods derived from different datasets, we denote this with a subscript so for example ${\mathcal{L}}_A\left(\mathsf{\Theta }\right)\equiv P\left(D_A\right|\mathsf{\Theta },\mathcal{M})$, and ${\mathcal{L}}_B\left(\mathsf{\Theta }\right)\equiv P\left(D_B\right|\mathsf{\Theta },\mathcal{M})$.

In our setting, the parameter vector is split into two sub-vectors $\mathsf{\Theta }=(\theta ,\alpha )$, where $\theta$ are parameters of scientific interest, and $\alpha$ are nuisance parameters, including for the purposes of data analysis. Such situations are common in astrophysics, where for example $\theta$ might be parameters governing the Universe’s evolution, whilst $\alpha$ might be associated with instrument calibration and foreground removal [3,4]. The $\alpha$ parameters are generally “marginalised out” and not considered further in final or future analyses.

2.2. Definitions

With this notation established, the Bayes theorem version (including nuisance parameters) takes the form

$$\mathcal{L}(\theta ,\alpha )\times \pi (\theta ,\alpha )=\mathcal{P}(\theta ,\alpha )\times \mathcal{Z},$$

(1)

where we have placed the inputs of inference (likelihood and prior) on the left-hand side, and the outputs (posterior and evidence) on the right. The evidence is as usual equivalent to the (fully-)marginalised likelihood $\mathcal{Z}=\int \mathcal{L}(\theta ,\alpha )\pi (\theta ,\alpha )d\theta d\alpha$.

We may marginalise any probability distribution so can straightforwardly define the nuisance marginalised posterior and prior by integrating over $\alpha$

$$\mathcal{P}(\theta )=\int \mathcal{P}(\theta ,\alpha )d\alpha ,\pi (\theta )=\int \pi (\theta ,\alpha )d\alpha .$$

(2)

The nuisance marginalised version of Bayes theorem Equation (1) takes the form

$$\mathcal{L}\left(\theta \right)\times \pi \left(\theta \right)=\mathcal{P}\left(\theta \right)\times \mathcal{Z}.$$

(3)

Here $\mathcal{Z}$ is the original evidence, whilst $\mathcal{L}\left(\theta \right)$ is non-trivially the nuisance-free likelihood

$$\mathcal{L}\left(\theta \right)\equiv \frac{\int \mathcal{L}(\theta ,\alpha )\pi (\theta ,\alpha )d\alpha }{\int \pi (\theta ,\alpha )d\alpha }=\frac{\mathcal{P}\left(\theta \right)\mathcal{Z}}{\pi \left(\theta \right)},$$

(4)

where the above is motivated by marginalising over $\alpha$ of the full Bayes theorem Equation (1) and substituting the definitions in Equations (2) and (4) recovers the marginalised Bayes theorem Equation (3).

The nuisance-free likelihood Equation (4) is straightforward to compute in our framework since we (uniquely) have NS-computed evidence $\mathcal{Z}$ combined with margarine trained distributions $\mathcal{P}\left(\theta \right)$ and $\pi \left(\theta \right)$ [9] (see Section 3).

We now explain why Equation (4) is a useful definition

Theorem 1.

Let ${\mathcal{L}}_A(\theta ,{\alpha }_A)$ and ${\mathcal{L}}_B(\theta ,{\alpha }_B)$ be two likelihoods with distinct datasets, each with their own nuisance parameters. The nuisance-free likelihoods ${\mathcal{L}}_A\left(\theta \right)$, ${\mathcal{L}}_B\left(\theta \right)$ form a lossless compression in θ. This means that we can recover the same (marginal) inference in combination that we would have made when performing a combined analysis with all nuisance parameters:

$$\begin{array}{cc}\hfill {\mathcal{L}}_A(\theta ,{\alpha }_A){\mathcal{L}}_B(\theta ,{\alpha }_B){\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)& ={\mathcal{P}}_{AB}(\theta ,{\alpha }_A,{\alpha }_B){\mathcal{Z}}_{AB},\hfill \end{array}$$

(5)

$$\begin{array}{cc}\hfill \Rightarrow {\mathcal{L}}_A\left(\theta \right){\mathcal{L}}_B\left(\theta \right)\pi \left(\theta \right)& ={\mathcal{P}}_{AB}\left(\theta \right){\mathcal{Z}}_{AB},\hfill \end{array}$$

(6)

if their respective priors ${\pi }_A(\theta ,{\alpha }_A)$ and ${\pi }_B(\theta ,{\alpha }_B)$ satisfy the marginal consistency relations:

$$\begin{array}{c}\pi \left(\theta \right)=\int {\pi }_A(\theta ,{\alpha }_A)d{\alpha }_A=\int {\pi }_B(\theta ,{\alpha }_B)d{\alpha }_B,\end{array}$$

(7)

$$\begin{array}{c}\int {\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)d{\alpha }_A={\pi }_B(\theta ,{\alpha }_B),\int {\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)d{\alpha }_B={\pi }_A(\theta ,{\alpha }_A).\end{array}$$

(8)

This process is represented graphically inFigure 1 and Figure 2.

Proof. 

Integrating the combined Bayes theorem Equation (5) with respect to ${\alpha }_B$, applying the definition of the marginal posterior Equation (2) on the right-hand side, and drawing out terms independent of ${\alpha }_B$ on the left, yields

$${\mathcal{L}}_A(\theta ,{\alpha }_A)\int {\mathcal{L}}_B(\theta ,{\alpha }_B){\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)d{\alpha }_B={\mathcal{P}}_{AB}(\theta ,{\alpha }_A){\mathcal{Z}}_{AB}.$$

(9)

From the definition of a nuisance-free likelihood Equation (4), and the marginal consistency Equation (8), we can say that the integral on the left-hand side becomes:

$$\begin{array}{cc}\hfill & \int {\mathcal{L}}_B(\theta ,{\alpha }_B){\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)d{\alpha }_B\hfill \\ \hfill & =\int {\mathcal{L}}_B(\theta ,{\alpha }_A,{\alpha }_B){\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)d{\alpha }_B\hfill & \hfill \left[{\mathcal{L}}_B(\theta ,{\alpha }_B)\equiv {\mathcal{L}}_B(\theta ,{\alpha }_A,{\alpha }_B)\sin \mathrm{ce}{\mathcal{L}}_B\mathrm{indep}\mathrm{of}{\alpha }_A\right]\\ \hfill & ={\mathcal{L}}_B(\theta ,{\alpha }_A)\int {\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)d{\alpha }_B\hfill & \hfill \left[\mathrm{Using}\mathrm{Equation}\left(4\right)\mathrm{for}{\mathcal{L}}_B\right]\\ \hfill & ={\mathcal{L}}_B\left(\theta \right)\int {\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)d{\alpha }_B\hfill & \hfill \left[{\mathcal{L}}_B(\theta ,{\alpha }_A)\equiv {\mathcal{L}}_B\left(\theta \right)\sin \mathrm{ce}{\mathcal{L}}_B\mathrm{indep}\mathrm{of}{\alpha }_A\right]\\ \hfill & ={\mathcal{L}}_B\left(\theta \right){\pi }_A(\theta ,{\alpha }_A).\hfill & \hfill \left[\mathrm{Using}\mathrm{marginal}\mathrm{consistency}\mathrm{Equation}\left(8\right)\right]\end{array}$$

(10)

Substituting Equation (10) back into Equation (9) we find

$${\mathcal{L}}_A(\theta ,{\alpha }_A){\mathcal{L}}_B\left(\theta \right){\pi }_A(\theta ,{\alpha }_A)={\mathcal{P}}_{AB}(\theta ,{\alpha }_A){\mathcal{Z}}_{AB}.$$

(11)

Proceeding with a similar manipulation to Equation (10), marginalising with respect to ${\alpha }_A$, and applying the definition of the nuisance-free likelihood ${\mathcal{L}}_A\left(\theta \right)$ Equation (4) and the marginal prior consistency Equation (7) we recover Equation (6)

$${\mathcal{L}}_A\left(\theta \right){\mathcal{L}}_B\left(\theta \right)\pi \left(\theta \right)={\mathcal{P}}_{AB}\left(\theta \right){\mathcal{Z}}_{AB}.$$

□

2.3. Discussion

Equation (5) represents Bayes theorem for the combined likelihood of both datasets ${\mathcal{L}}_{AB}(\theta ,{\alpha }_A,{\alpha }_B)={\mathcal{L}}_A(\theta ,{\alpha }_A){\mathcal{L}}_B(\theta ,{\alpha }_B)$, using the combined prior ${\pi }_{AB}(\theta ,{\alpha }_A,{\alpha }_B)$. We assume the combined prior is marginally consistent, Equations (7) and (8), which is reasonable, merely demanding that the priors are identical in the parameter spaces where they overlap. In practice, this would usually be achieved by assuming separability between signal and nuisance parameter spaces $\pi (\theta ,\alpha )=\pi \left(\theta \right)\pi \left(\alpha \right)$, but Equations (7) and (8) are a slightly less restrictive requirement and therefore more general.

The upshot of this is that if you have performed inference for two datasets separately, such that you are able to compute the nuisance-free likelihoods with margarine, you may discard the nuisance parameters for the next set of analyses when you combine the datasets.

3. Methods

margarine was first introduced in [9] and uses density estimation to approximate probability distributions such as $\mathcal{P}\left(\theta \right)$ and $\pi \left(\theta \right)$ given sets of representative samples. The code was developed initially to calculate marginal Kullback–Leibler (KL) divergences [10] and Bayesian model dimensionalities (BMD) [11] however as discussed in Section 2 it can be used to calculate the nuisance-free likelihoods. This in turn means that we can use margarine alongside an implementation of the nested sampling algorithm to sample the product ${\mathcal{L}}_A\left(\theta \right){\mathcal{L}}_B\left(\theta \right)$. In this manner, margarine allows us to combine constraints on common parameters across different data sets.

We refer the reader to [9] for a complete discussion of how margarine works; however, we discuss briefly the density estimation here. margarine uses two different types of density estimators to model posterior and prior samples, namely masked autoregressive flows (MAFs, [8]) and kernel density estimators (KDEs, [6,7]).

MAFs transform a multivariate base distribution, the standard normal, into a target distribution via a series of shifts and scaling, which are estimated by autoregressive neural networks. To improve the performance of the MAF the samples representing the target distribution, in our case $\mathcal{P}\left(\theta \right)$ and $\pi \left(\theta \right)$, are transformed into a Gaussianized space. We implement the MAFs using tensorflow and the keras backend [12].

KDEs use a kernel to approximate the multivariate probability density of a series of samples. In our case, the kernel is Gaussian and the probability density is a sum of Gaussians centred on the sample points with a given bandwidth. Again, we transform the target samples into a Gaussianized parameter space allowing the KDE to better capture the distribution. The KDEs are implemented with SciPy in margarine [6,7,13].

Since both types of density estimator build approximations to the target distribution using known distributions, the approximate log probabilities of the target distribution can be easily calculated.

The evaluation of normalised log probabilities for the marginal posterior and marginal prior allows us to calculate the nuisance-free likelihoods, as discussed, along with marginal Kullback–Leibler divergences

$$\mathcal{D}\left(\mathcal{P}\right|\left|\pi \right)=\int \mathcal{P}\left(\theta \right)log\frac{\mathcal{P}\left(\theta \right)}{\pi \left(\theta \right)}d\theta ,$$

(12)

which quantifies the amount of information gained when moving from the marginal prior to posterior.

4. Cosmological Example

It has previously been demonstrated that margarine is capable of replicating complex probability distributions and approximating marginal Bayesian statistics such as the KL divergence and the BMD [9]. Here, we demonstrate the theory discussed in Section 2 by combining samples from the Dark Energy Survey (DES) Year 1 posterior [1] and Planck posterior [3] using margarine to estimate nuisance-free likelihoods. DES surveys supernovae, galaxies, and large-scale cosmic structures in the universe in an effort to measure dark matter and dark energy densities and model the dark energy equation of state. In contrast, Planck mapped the anisotropies in the cosmic microwave background (CMB) and correspondingly provided constraints on key cosmological parameters.

The constraints from DES and Planck have previously been combined using a full nested sampling run over all parameters including a multitude of ‘nuisance’ parameters in a computationally expensive exercise [14]. This corresponds to the flow chart in Figure 1 and the previous analysis gives us access to the combined DES and Planck evidence, which is found to have a value of $log\left(\mathcal{Z}\right)=-5965.7\pm 0.3$. In Figure 3, we show the DES, Planck, and joint posteriors for the six cosmological parameters derived in this work using margarine and the flow chart in Figure 2. The constrained parameters are the baryon and dark matter density parameters, ${\mathsf{\Omega }}_b{h}^2$ and ${\mathsf{\Omega }}_c{h}^2$, the angular size of the sound horizon at recombination, ${\theta }_{MC}$, the CMB optical depth, $\tau$, the amplitude of the power spectrum, $A_s$, and the corresponding spectral index, $n_s$. These make up the set $\theta =({\mathsf{\Omega }}_b{h}^2,{\mathsf{\Omega }}_c{h}^2,{\theta }_{MC},\tau ,A_s,n_s)$. We use the nested sampling algorithm polychord in our analysis [15,16].

We use a uniform prior that is defined to be three sigma around the Planck posterior mean. This is done to improve the efficiency of our nested sampling run. However, we subsequently have to re-weight the samples and correct the evidence for the difference between the priors used here and in the previous full nested sampling run [14] for comparison. If we define

$${\mathcal{Z}}_A=\int \mathcal{L}\left(\theta \right){\pi }_A\left(\theta \right)d\theta ,{\mathcal{Z}}_B=\int \mathcal{L}\left(\theta \right){\pi }_B\left(\theta \right)d\theta ,$$

(13)

where A is our uniform prior space and B is our target prior space from the previous work [14], then

$${\mathcal{Z}}_B=\int \mathcal{L}\left(\theta \right){\pi }_B\left(\theta \right)d\theta =\int \mathcal{L}\left(\theta \right){\pi }_A\left(\theta \right)\frac{{\pi }_B\left(\theta \right)}{{\pi }_A\left(\theta \right)}d\theta ={\mathcal{Z}}_A{〈\frac{{\pi }_B\left(\theta \right)}{{\pi }_A\left(\theta \right)}〉}_{{\mathcal{P}}_A}$$

(14)

giving

$${\mathcal{Z}}_B={\mathcal{Z}}_A{〈\frac{{\pi }_B\left(\theta \right)}{{\pi }_A\left(\theta \right)}〉}_{{\mathcal{P}}_A}.$$

(15)

Following similar arguments, we can transform our posteriors by re-weighting the distributions with the following

$$w_B^{\left(i\right)}=w_A^{\left(i\right)}\frac{{\pi }_B\left({\theta }^{\left(i\right)}\right)}{{\pi }_A\left({\theta }^{\left(i\right)}\right)}.$$

(16)

We see from the figure and corresponding table that with our joint analysis we are able to derive a log-evidence that is approximately consistent with that found in [14] validating the theory discussed and its implementation with margarine. We note that the re-weighting described above relies on the calculations of the two prior log probabilities for which we use margarine, and currently do not have an estimate of the error for. As a result, the error in the combined evidence, $Z_B$, is given by the error in $Z_A$ from the nested samples and is likely underestimated. Using margarine [9] we can also derive the combined KL divergence, also reported in Figure 3, which we find is consistent with the result in the literature of $\mathcal{D}=6.17\pm 0.36$ [11]. Similarly, we derive marginal KL divergences for the DES and Planck cosmological parameters using margarine. A full discussion of the implications of combining the two data sets for our understanding of cosmology can be found in the literature (e.g., [11,14]).

By reducing the number of parameters that need to be sampled, we significantly reduce the nested sampling runtime. For polychord the runtime scales as the cube of the number of dimensions [17]. This can be seen by assessing the time complexity of the algorithm, where $T\propto n_{\mathrm{live}}\times \langle T\left\{\mathcal{L}\left(\theta \right)\right\}\rangle \times \langle T\{Impl.\}\rangle \times \mathcal{D}\left(\mathcal{P}\right|\left|\pi \right)$. Here $n_{\mathrm{live}}$ scales with the number of dimensions, d, as does the Kullback–Leibler divergence. For polychord, the specific implementation time complexity factor, $\langle T\{Impl.\}\rangle$, representing the impact of replacing dead points with higher likelihood live points on the runtime, scales linearly with d. Together this gives $T\propto d^3\times \langle T\left\{\mathcal{L}\left(\theta \right)\right\}\rangle$. Therefore, by using nuisance-free likelihoods and sampling over 6 parameters rather than 41 parameters (cosmological plus 20 nuisance parameters for DES and 15 different nuisance parameters for Planck) we reduce the runtime by a factor of ${(41/6)}^3\approx 319$ with further improvements in $\langle T\left\{\mathcal{L}\right(\theta \left)\right\}\rangle$. Using margarine, $\langle T\left\{\mathcal{L}\right(\theta \left)\right\}\rangle$ is typically reduced since analytic likelihoods are computationally more expensive than emulated likelihoods.

5. Conclusions

In the paper, we demonstrated the consistency between combining constraints from different experiments in a marginal framework using density estimators and the code margarine with a full nested sampling run over all parameters, including those describing ‘nuisance’ components of the data. We have shown this consistency mathematically and with a cosmological example. For the combination of Planck and DES, we found Bayesian evidence and KL divergence that were consistent with previous results [11,14].

The analysis in this paper is only possible because (a) we were able to estimate densities in the (much smaller) cosmological parameter space $\theta$ using margarine, and (b) because we have evidence, $\mathcal{Z}$, from our original nested sampling runs. It is this unique combination that allowed us to compress away or discard nuisance parameters once they have been used. We note also that working in the marginal space results in a compression that is lossless in information on $\theta$ as it recovers an identical marginal posterior and total evidence as is found during a full Bayesian inference. Finally, through the nuisance-free likelihood we can significantly reduce the dimensionality of our problems and since it is faster to emulate a likelihood rather than analytically evaluate, margarine offers a much more computationally efficient path to combined Bayesian analysis.

In principle, our work paves the way for the development of a publicly available library of cosmological density estimators modelled with margarine that can be combined with new data sets using the proposed method in a more efficient manner than currently implemented techniques. However, the work has implications outside of cosmology in any field where multiple experiments probe different aspects of the same physics.