Nested sampling with normalizing flows for gravitational-wave inference

We present a novel method for sampling iso-likelihood contours in nested sampling using a type of machine learning algorithm known as normalizing flows and incorporate it into our sampler nessai. nessai is designed for problems where computing the likelihood is computationally expensive and therefore the cost of training a normalizing flow is offset by the overall reduction in the number of likelihood evaluations. We validate our sampler on 128 simulated gravitational wave signals from compact binary coalescence and show that it produces unbiased estimates of the system parameters. Subsequently, we compare our results to those obtained with dynesty and find good agreement between the computed log-evidences while requiring 2.07 times fewer likelihood evaluations. We also highlight how the likelihood evaluation can be parallelized in nessai without any modifications to the algorithm. Finally, we outline diagnostics included in nessai and how these can be used to tune the sampler’s settings.

Figure 1

Diagram of a normalizing flow $f(x)$ composed of four coupling transforms which maps an $n$-dimensional input vector x to an $n$-dimensional latent vector z. Each transform splits $x$ in two $[x_{1:m},x_{m+1:n}]$ and updates one part conditioned on the other. In the first and third transforms $x_{1:m}$ is used as the input to a neural network (NN) which then produces the scale $s$ and translation $t$ vectors of length $m$. The element-wise product ($\odot$) is then computed between $x_{1:m}$ and $\exp (s)$ followed by the sum of the output and $t$. This is shown in the left transform. In the second and fourth transforms $x_{1:m}$ is updated conditioned on $x_{m+1:n}$ as shown in the right transform.

Figure 2

Example of how a normalizing flow trained on a set of live points can produce samples within current iso-likelihood contour for simple two-dimensional parameter space. Top: example of training samples in the physical space $\mathcal{X}$ and learned mapping to the latent space $\mathcal{Z}$ with the iso-likelihood contour for the current worst point shown in orange. Middle: samples drawn from a truncated Guassian within the iso-likelihood contour in $\mathcal{Z}$ and mapped to $\mathcal{X}$ using the inverse mapping. Bottom: pool of accepted samples after applying rejection sampling until 1000 points are obtained shown in both $\mathcal{Z}$ and $\mathcal{X}$.

Figure 3

Distribution of the optimal network SNR for the 128 simulated gravitational-wave injections in simulated Gaussian noise. The priors on chirp mass and luminosity distance Appendix pp3 were chosen such that the ringdown frequency [61] does not exceed the Nyquist frequency and that the majority of signals have detectable optimal network SNRs.

Figure 4

Probability-probability (P-P) plot showing the confidence interval versus the fraction of the events within that confidence interval for the posterior distributions obtained using our analysis nessai for 128 simulated compact binary coalescence signals produced with bilby and bilby_pipe. The 1-, 2- and 3-$\sigma$ confidence intervals are indicated by the shaded regions and $p$-values are shown for each of the parameters and the combined $p$-value is also shown.

Figure 5

Difference between the log-evidences $\mathrm{\Delta }\ln Z$ obtained using dynesty and nessai for all 128 injections with distance marginalization (dashed line) and without distance marginalization (solid line).

Figure 6

Distribution of the total number of likelihood evaluations required to reach convergence for and total run time for dynesty (blue) and nessai (orange) when applied 128 simulated signals from compact binary coalescence with the priors and sampler settings described in Sec. 5 and Appendix pp4. The results with distance marginalization disabled are shown with solid lines and those with distance marginalization enabled are shown with dashed lines.

Figure 7

Comparison of the total time (in hours) spent on each stage of the algorithm for increasing number of threads for a single injection with a fixed noise seed. The time spent evaluating the likelihood decreases as the number of threads increases, the theoretical reduction is shown in black. The costs of training and population stages remain approximately constant, as such they act as a lower bound on the minimum run-time. The sum of the time spent on likelihood evaluation, training and population is approximately equal to the total time spent sampling, indicating minimal overhead.

Figure 8

Example of the distribution of insertion indices for two nested sampling runs with 2000 live points. Uniformly distributed indices indicate no under- or overconstraining and deviations from uniformity indicate the opposite. The result with an orange dashed line shows overconstraining and the result with a solid blue line shows the correctly converged run. The shaded region indicates the 2-$\sigma$ errors on the expected distribution.

Figure 9

Example of the statistics that are tracked in our sampler as a function of sampling iteration: (a) minimum (blue solid) and maximum (orange dashed) log-likelihood, (b) cumulative number of likelihood evaluations, (c) log evidence $\log Z$ (blue solid) and fractional change in evidence $\mathrm{d}Z$ (orange dashed), (d) proposal (blue solid) and population (orange dashed) acceptance and (e) $p$-value for cross-checks of every $K$ live points. The iterations at which the normalizing flow is trained are indicated with vertical lines, for this injection these total 87.

Figure 10

Probability-probability (P-P) plot showing the confidence interval versus the fraction of the events within that confidence interval for the posterior distributions obtained using dynesty for 128 simulated compact binary coalescence signals produced with bilby and bilby_pipe. The 1-, 2-, and 3-$\sigma$ confidence intervals are indicated by the shaded regions and $p$-values are shown for each of the parameters and the combined $p$-value is also shown. We use the settings described in [36] with the exception of the number of live points which we increase to 2000.

Figure 11

Corner plot comparing the posterior distributions produced with dynesty (blue) and our sampler nessai (orange) for an injection with an optimal network SNR of 15.54. The phase is marginalized and remaining 14 parameters are shown, see Appendix pp3 for details about the parameters. The injected value is indicated by the cross-hairs in each subplot and the respective 16% and 84% percentiles are also shown in the 1-dimensional marginalized posteriors.