SuperNest: Accelerated Nested Sampling Applied to Astrophysics and Cosmology ^†

Abstract

We present a method for improving the performance of nested sampling as well as its accuracy. Building on previous work we show that posterior repartitioning may be used to reduce the amount of time nested sampling spends in compressing from prior to posterior if a suitable “proposal” distribution is supplied. We showcase this on a cosmological example with a Gaussian posterior, and release the code as an LGPL licensed, extensible Python package supernest.

1. Introduction

Nested sampling is a multi-purpose tool invented by John Skilling [1] for simultaneous probabilistic integration of high-dimensional distributions and drawing weighted samples. These are essential tasks for the performance of Bayesian model comparison and parameter estimation [2]. It has become widely used across the physical sciences from astrophysics and cosmology through biochemistry and materials science to high energy physics [3].

The performance of nested sampling algorithms, however, suffer in comparison with alternative posterior samplers such as Metropolis–Hastings [4], in particular, when the latter are provided with a pre-trained proposal distribution and start point. This paper allows an apples-to-apples comparison, by providing a mechanism for accelerating nested sampling using proposal distributions. This is achieved using an extension of the posterior repartitioning techniques developed by Chen et al. [5].

In Section 2, we establish notation and definitions. In Section 3, we review and extend posterior repartitioning. In Section 4, we present our chosen approach and in Section 5, we showcase this on a cosmology example before concluding in Section 6.

2. Background

2.1. Notation

In this section, we establish notation and recall the key terms in nested sampling and Bayesian inference. A probabilistic inference model $\mathcal{M}$ with a vector of tunable parameters $\theta$ for data D is specified by likelihood $\mathcal{L}\left(\theta \right)\equiv P\left(D\right|\theta ,\mathcal{M})$ and prior $\pi \left(\theta \right)=P\left(\theta \right|\mathcal{M})$. The two-step challenge of inference is to compute the posterior $\mathcal{P}\left(\theta \right)\equiv P\left(\theta \right|D,\mathcal{M})$ and evidence $\mathcal{Z}=P\left(D\right|\mathcal{M})$. Constructing a representation of the posterior is termed parameter estimation, and evaluating an estimate of the evidence comprises model comparison. The four distributions are linked via Bayes theorem:   

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

(1)

where both sides of the equation equal to the joint distribution $P(D,\theta )$.

It is useful to define the Kullback-Leibler divergence $\mathcal{D}$ [6] between the prior $\pi$ and the posterior $\mathcal{P}$:

$${\mathcal{D}}_{\pi }\left\{\mathcal{P}\right\}=\int \mathcal{P}\left(\theta \right)ln\frac{\mathcal{P}\left(\theta \right)}{\pi \left(\theta \right)}d\theta .$$

(2)

One can understand the properties of the KL divergence as a logarithmic ratio of the volume contained in the prior distribution to the volume contained in the posterior distribution:

$${\mathcal{D}}_{\pi }\left\{\mathcal{P}\right\}\approx ln\frac{V_{\pi }}{V_{\mathcal{P}}}.$$

(3)

When both the prior and posterior are uniform (i.e., constant value where they are non-zero), this relation is exact (one can think of the posterior averaging process in Equation (2) as an information theory inspired “smoothing” for non top-hat distributions).

2.2. Nested Sampling

In this section, we will provide a brief primer on the nested sampling algorithm. For a survey of established nested sampling packages and applications, readers are recommended the Nature primer [3]. We begin by initialising $n_{\mathrm{live}}$ initial “live points” drawn from the prior $\pi$. Then at each iteration, the lowest likelihood live point is marked “dead” and replaced with a new live point drawn from the prior, that is at a higher likelihood. The compression procedure allows one to estimate the density of states, which given the recorded locations of live and dead points leads to estimates of posterior and evidence with respective errors.

Crucially, the algorithm terminates once the running estimate of the error in the evidence is greater than a set fraction of the total estimated evidence. Thus, the time complexity of nested sampling is:

$$T\propto n_{\mathrm{live}}\times \langle \mathcal{T}\left\{\mathcal{L}\left(\theta \right)\right\}\rangle \times \langle \mathcal{T}\left\{\mathrm{Impl}.\right\}\rangle \times {\mathcal{D}}_{\pi }\left\{\mathcal{P}\right\},$$

(4)

where $\mathcal{T}\left\{\mathcal{L}\right(\theta \left)\right\}$ is the time complexity of evaluating $\mathcal{L}\left(\theta \right)$ and $\mathcal{T}\left\{\mathrm{Impl}.\right\}$ is the implementation-specific factor representing the time complexity of replacing a dead point with a live point with a higher likelihood, and is the dominant factor for nested sampling over parameter spaces with many dimensions [3]. For rejection nested sampling (e.g., MultiNest [7]) this factor is exponential in the number of tunable parameters d (with a turnover constant in tens of dimensions), and linear in d for chain-based approaches (e.g., PolyChord [8]). The Kullback–Leibler divergence ${\mathcal{D}}_{\pi }\left\{\mathcal{P}\right\}$ is linear in the d for parameters that are well constrained, as can be seen from Equation (3) by noticing that volumes in parameter space scale with a power of d.

From Equation (4), it follows that the run time of nested sampling, given a resolution fixed by $n_{\mathrm{live}}$ can be reduced by (i) speeding up the likelihood evaluation; (ii) improving the replacement efficiency; (iii) speeding up the prior to posterior compression. The final approach is the one adopted for the repartitioning techniques in this paper.

The accuracy of the evidence is associated with the accumulated Poisson noise, such that the error in $ln\mathcal{Z}$ is [1]:

$$\sigma \propto \sqrt{{\mathcal{D}}_{\pi }\left\{\mathcal{P}\right\}/n_{\mathrm{live}}}.$$

(5)

Thus, faster compression leads to better accuracy and precision. For a fixed accuracy $\sigma$, we can adjust $n_{\mathrm{live}}\propto \mathcal{D}/{\sigma }^2$, which means $T\propto {\mathcal{D}}^2/{\sigma }^2$, given Equation (4). In precision-normalised cases, reducing the effect of the KL divergence term results in quadratic improvements to the speed of nested sampling. Note that the number of live points can be also adjusted. The minimum number is dictated by the quantity of modes in the multi-modal posterior, and implementation dependent considerations: for example, PolyChord needs $n_{\mathrm{live}}\sim \mathcal{O}\left(d\right)$ to generate enough dead point phantoms for an accurate covariance computation. MultiNest needs at least $n_{\mathrm{live}}\sim \mathcal{O}\left(100\right)$ to compute ellipsoidal decompositions.

2.3. Historical Overview of Acceleration Attempts

One approach to partially achieve this form of speed-up is to perform two passes of nested sampling. First, a low resolution pass (few live points) over the entire original prior volume $V_{\pi ,1}$ finds the location of the highest concentration of posterior volume, and an evidence estimate ${\mathcal{Z}}_1$. Next, one constructs a much tighter bounding prior ${\pi }_2$ with volume $V_{\pi ,2}$ around the largest concentration of the posterior. Trivially:

$${\mathcal{D}}_{\pi ,1}\left\{\mathcal{P}\right\}<{\mathcal{D}}_{\pi ,2}\left\{\mathcal{P}\right\}$$

(6)

hence, latter inference over the tighter box will finish faster finding the evidence ${\mathcal{Z}}_2$. With a changed prior, this evidence will not be the same as the original. However, for box-priors we can compensate the volumes geometrically, using the fact that the overall evidence satisfies:

$$\mathcal{Z}={\mathcal{Z}}_1={\mathcal{Z}}_2\frac{V_{\pi ,2}}{V_{\pi ,1}}$$

(7)

See Anstey et al. [9] for examples of this method’s application in real astronomical cases.

3. Posterior Repartitioning

Nested sampling is unusual amongst Bayesian numerical algorithms, in that it distinguishes between the prior $\pi \left(\theta \right)$ and likelihood $\mathcal{L}\left(\theta \right)$. Metropolis–Hastings, and Hamiltonian Monte Carlo, for example, are techniques which are only sensitive; $\mathcal{L}\times \pi$ (i.e., the joint or unnormalised posterior) is considered for the purposes of acceptance/rejection. Nested sampling, on the other hand, by “sampling from the prior $\pi$, subject to a hard likelihood constraint $\mathcal{L}$”, separates the two functions.

Chen et al. [5] pioneered an approach which exploits this split, noting that we have the freedom to define a new prior $\tilde{\pi }$ and likelihood $\tilde{\mathcal{L}}$, which providing they obey the relation:

$$\tilde{\mathcal{L}}\left(\theta \right)\tilde{\pi }\left(\theta \right)=\mathcal{L}\left(\theta \right)\pi \left(\theta \right)$$

(8)

will recover the same posterior and evidence. Conceptually, one can think of this as transferring portions of the likelihood into the prior, or vice versa. A cosmological example of this subtle rearranging of portions of likelihood into prior can be found in many codes which implement nested sampling such as cobaya [10], cosmomc [11,12], MontePython [13] or cosmosis [14] in the manner in which they consider Gaussian priors. Codes can either make the choice to incorporate this as part of the prior (typically as an inverse-error-function-based unit hypercube transformation), or as a density term added to the rest of the log-likelihood calculations.

This rearranging affects the Kullback–Leibler divergence. The difference can be found using a posterior average:

$${\mathcal{D}}_{\tilde{\pi }}\left\{\mathcal{P}\right\}={\mathcal{D}}_{\pi }\left\{\mathcal{P}\right\}+{〈ln\frac{\tilde{\pi }}{\pi }〉}_{\mathcal{P}}$$

(9)

If the final term in Equation (9) is negative, repartitioning provides a mechanism for reducing the time to convergence by transferring a reasonable “guess”: the proposal distribution, from likelihood to prior, reducing the Kullback–Leibler divergence between the prior $\tilde{\pi }$ and the posterior $\tilde{\mathcal{P}}=\mathcal{P}$. As we have seen before, this both improves performance and increases precision.

3.1. Example 1: Power Posterior Repartitioning

The scheme considered in Chen et al. [5] is power posterior repartitioning (PPR):

$$\tilde{\pi }\left(\theta \right|\beta )=\frac{{\pi \left(\theta \right)}^{\beta }}{z_{\beta }},\tilde{\mathcal{L}}\left(\theta \right|\beta )=\mathcal{L}\left(\theta \right){\pi \left(\theta \right)}^{1-\beta }{z}_{\beta },z_{\beta }=\int {\pi \left(\theta \right)}^{\beta }d\theta .$$

(10)

For $\beta =0$, any distribution $\pi$ becomes uniform. If a distribution $\pi$ is not uniform, raising it to a power of $\beta$ serves to sharpen or spread the peaks, depending on whether $\beta$ is greater or smaller than 1, in analogy with the thermodynamic inverse temperature of a Boltzmann distribution [3]. For any given value of $\beta$, a nested sampling run with likelihood $\tilde{\mathcal{L}}$ and prior $\tilde{\pi }$ recovers the same posterior samples and evidence as a run with $\mathcal{L}$ and prior $\pi$.

Chen et al. [5] considered this in the context of a geophysical example. The original intention behind repartitioning was not to improve performance, but rather as a way of making nested sampling more resilient to over-zealously specified priors (e.g., if a user specifies a Gaussian prior which is narrow and/or distant from the peak of the likelihood). Still, the performance implications are valuable. Substituting Equation (10) into Equation (9) shows that varying $\beta$ will change the time to convergence as:

$$T\left(\beta \right)\propto {\mathcal{D}}_{\tilde{\pi }}\left\{\mathcal{P}\right\}={\mathcal{D}}_{\pi }\left\{\mathcal{P}\right\}+{〈ln\frac{{\pi }^{\beta -1}}{z_{\beta }}〉}_{\mathcal{P}}$$

(11)

Whether or not this increases or decreases the speed of nested sampling depends on the choice of $\beta$ and to what extent the sharpened prior overlaps with the posterior. Equation (10) is valid for any choice of $\beta$, and while any specific choice can be useful, Chen et al. [5] demonstrated that still greater power comes from considering it as an additional parameter to sample over, effectively extending the parameter space to $\tilde{\theta }=(\theta ,\beta )$ with a prior on $\beta$. A posterior value of the parameter $\beta$ close to 0 therefore indicates that the prior was poorly chosen, i.e., that there is significant tension between prior and posterior.

One could consider this for the purposes of acceleration by allowing $\beta$ to take a range wider than $[0,1]$, for example, by using an exponential prior $P\left(\beta \right)=e^{-\beta }[0\le \beta <\infty ]$. There is, however, no need to restrict ourselves to the case of power posterior repartitioning, or indeed to a priors $\tilde{\pi }$, which are direct transformations of the original $\pi$.

3.2. Example 0: Replacement Repartitioning

The simplest repartitioning scheme is to simply use a different prior ${\pi }_{*}$, but compensate for this with the adjusted likelihood:

$$\tilde{\pi }\left(\theta \right)={\pi }_{*}\left(\theta \right),\tilde{\mathcal{L}}\left(\theta \right)=\frac{\mathcal{L}\left(\theta \right)\pi \left(\theta \right)}{{\pi }_{*}\left(\theta \right)}.$$

(12)

From Equation (9), we can see that if ${\pi }_{*}$ is close to $\mathcal{P}$ as quantified by the KL divergence, nested sampling terminates sooner (with ${\mathcal{D}}_{\tilde{\pi }}\left\{\mathcal{P}\right\}=0$ in the limiting case $\tilde{\pi }=\mathcal{P}$). However, a poorly specified proposal ${\pi }_{*}$ will pessimise the choices of new live points, thus making the problem harder. Indeed, a poorly chosen proposal can prevent us gaining any usable information.

3.3. Example 2: Additive Superpositional Repartitioning

Another repartitioning scheme is to consider an additive superposition of the original prior $\pi$ with another normalised distribution ${\pi }_{*}$:

$$\tilde{\pi }\left(\theta \right|\beta )=(1-\beta )\pi \left(\theta \right)+\beta {\pi }_{*}\left(\theta \right),\tilde{\mathcal{L}}\left(\theta \right|\beta )=\frac{\mathcal{L}\left(\theta \right)\pi \left(\theta \right)}{(1-\beta )\pi \left(\theta \right)+\beta {\pi }_{*}\left(\theta \right)}$$

(13)

Here, ${\pi }_{*}$ does not need to be related to the original prior, and could, for example, be a Gaussian proposal distribution (often supplied with cosmological sampling packages for common likelihood combinations), or a tighter box prior as suggested in Section 2.3.

This scheme can be generalised to superpositions of N alternative distributions ${\pi }_1,\dots {\pi }_N$. Consider:

$$\tilde{\pi }\left(\theta \right|\beta )=\left(1-{\sum }_{i=1}^N{\beta }_i\right)\pi \left(\theta \right)+{\sum }_{i=1}^N{\beta }_i{\pi }_i\left(\theta \right),\tilde{\mathcal{L}}\left(\theta \right|\beta )=\frac{\mathcal{L}\left(\theta \right)\pi \left(\theta \right)}{(1-{\sum }_{i=1}^N{\beta }_i)\pi \left(\theta \right)+{\sum }_{i=1}^N{\beta }_i{\pi }_i\left(\theta \right)}$$

(14)

However, additive superpositional repartitioning in practice has the issue of poor generalisation to higher dimensions and/or complex priors due to the difficulty of implementing an analytic unit hypercube transformation. Further work is required to fully explore this using nested samplers that do not rely on prior point–percent functions which produce said transformations. We thank Xy Wang for exploring the possibility of these models in the first instance as part of her University of Cambridge master’s dissertation.

4. Supernest: Stochastic Superpositional Repartitioning

The above examples lead us to the proposed repartitioning scheme. Just as with additive repartitioning, we allow, arbitrarily, many unrelated priors to be superimposed. However, we side-step the multivariate point-percent function correction problem by sampling from each of the priors with some probability. We achieve this by adding an integer selection parameter $\beta$, such that:

$$\tilde{\pi }\left(\theta \right|\beta )=\left\{\begin{array}{ccc}\pi \left(\theta \right)& :& \beta =0\\ {\pi }_1\left(\theta \right)& :& \beta =1\\ ⋮& & ⋮\\ {\pi }_N\left(\theta \right)& :& \beta =N\end{array}\right.\tilde{\mathcal{L}}\left(\theta \right|\beta )=\left\{\begin{array}{ccc}\mathcal{L}\left(\theta \right)& :& \beta =0\\ \mathcal{L}\left(\theta \right)\pi \left(\theta \right)/{\pi }_1\left(\theta \right)& :& \beta =1\\ ⋮& & ⋮\\ \mathcal{L}\left(\theta \right)\pi \left(\theta \right)/{\pi }_N\left(\theta \right)& :& \beta =N\end{array}\right.$$

(15)

and then $\beta$ is sampled over with a discrete (typically uniform) prior $P\left(\beta \right)$. The choice of said prior depends on the confidence in each of the proposals. When applied to nested sampling, at the end of the run, one will recover a (joint) posterior on $\beta$, which will indicate which proposal was chosen more frequently, as the usual Bayesian balance between the corresponding modified likelihood $\tilde{\mathcal{L}}$ and KL divergence ${\mathcal{D}}_{\tilde{\pi }}\left\{\mathcal{P}\right\}$ of the favoured bin.

In the additive case from Section 3.3, the primary difficulty is in obtaining the point-percent function of the repartitioned prior. The more priors are introduced, the harder the calculation. For the stochastic case, the prior of the superposition is the superposition of the individual priors. Thus, if one knows the point percent functions of the constituent priors, simply copying them is the “calculation”.

Additionally, only one of the prior quantile/likelihood pairs is evaluated for each point. Thus, the time complexity of the likelihood evaluation of a superposition of multiple likelihoods is $\mathcal{O}(\langle \mathcal{T}(\mathcal{L})\rangle )$. By contrast, the equivalent time complexity for the additive superposition is the sum of the time complexities: $\mathcal{O}\left(\sum \mathcal{T}\left(\mathcal{L}\right)\right)$. For comparison, if we fed 100 identical models to the stochastic mixture, we would not see a performance regression, while the additive superposition would take 100 times longer to compute.

5. Cosmological Example

For the purposes of demonstration, we choose a likelihood $\mathcal{L}$ which is Gaussian in the parameters, with parameter means and parameter covariances as given by the Planck legacy archive chains [15], and prior given by the default prior widths as specified in those chains. This has six cosmological parameters and 21 nuisance parameters associated with galactic foregrounds and satellite calibration, giving 27 parameters to constrain. In its default mode, this therefore gives $n_{\mathrm{live}}=675$ and $n_{\mathrm{repeats}}=135$, and for this example an $n_{\mathrm{dead}}\sim 72,000$. This provides a proxy for a nested sampling problem of interests to cosmologists, whilst being fast enough numerically to explore a variety of configurations.

We run PolyChord wrapped with supernest, both with their default settings. PolyChord, in addition, is set with a precision criterion of $10\%$, i.e., stopping when the live points contain 10% of the remaining evidence. As a representative example, we first choose a Gaussian proposal with the same covariance as the parameters, inflated by a factor of $3^d$ in volume, and offset by a random vector drawn from the posterior. This gives a proposal which is conservative (much wider than the posterior), and approximately correctly centered.

The performance of this example is shown in Figure 1. We can see that the algorithm terminates approximately three times faster, with an error bar that is approximately two times smaller despite the perturbation. Usually to reduce the error bar by a factor of 2, one would need to increase the number of live points by a factor of 4. Thus, when precision-normalised supernest achieved an order-of-magnitude $12\times$ improvement over the traditional nested sampling run. Further improvements to this speed are achievable if a more compact (and possibly larger number of) proposal distributions are used.

As a second example, where the covariance of the proposal does not align with the covariance of the posterior, we draw an inverse Wishart distributed covariance matrix, with $d=27$ degrees of freedom and scale matrix equal to the covariance, once again inflating the result in volume by $3^d$, and offsetting the result by a random vector drawn from the posterior. The performance is unaffected to within sampling fluctuation, with the recovered evidence equal to within error.

6. Conclusions

We presented a method for improving the performance and accuracy of nested sampling using the recent innovations in posterior repartitioning of Chen et al. [5]. The approach functions by reducing the effective prior-to-posterior compression that nested sampling “sees”, and hence reducing both the runtime and accumulated Poisson error in the evidence estimate. We discussed a variety of alternative posterior repartitioning strategies before settling on a stochastic superpositional scheme, and demonstrated that on a cosmological example even conservative proposals have the capacity to speed up nested sampling runs by an order of magnitude.

This performance improvement is achieved by trimming off the long prior-to-posterior compression phase seen in the blue but not orange curve of the lower panel of Figure 1. The remaining majority of the time is spent in traversing the typical set. Accelerating the pace at which nested sampling does this is a matter of ongoing research.

This work will be followed by a re-examination of the repartitioning, in particular, by considering more general cases than the ones considered here. In a follow up study we shall provide a more mathematical look into how further performance improvements can be achieved via modifications to nested samplers, hysteretic priors on the choice parameter, as well as exploring the connection between prior spaces and Hilbert spaces. Finally, we shall consider a reformulation of Bayesian methods using functionals, and explore the implications of consistent partitioning (a generalisation of posterior repartitioning) applied to Bayesian inference in functional space.