Rapid Parameter Estimation for Pulsar-Timing-Array Datasets with Variational Inference and Normalizing Flows

In the gravitational-wave analysis of pulsar-timing-array datasets, parameter estimation is usually performed using Markov chain Monte Carlo methods to explore posterior probability densities. We introduce an alternative procedure that instead relies on stochastic gradient-descent Bayesian variational inference, whereby we obtain the weights of a neural-network-based approximation of the posterior by minimizing the Kullback-Leibler divergence of the approximation from the exact posterior. This technique is distinct from simulation-based inference with normalizing flows since we train the network for a single dataset, rather than the population of all possible datasets, and we require the computation of the data likelihood and its gradient. Unlike Markov chain methods, our technique can trivially exploit highly parallel computing platforms. This makes it extremely fast on modern graphical processing units, on which it can analyze the NANOGrav 15-yr dataset in a few tens of minutes, depending on the probabilistic model, compared to hours or days with the analysis codes used so far. We expect that this speed will unlock new astrophysical and cosmological explorations of pulsar-timing-array datasets with statistical models that are currently too computationally expensive. Furthermore, this kind of variational inference is viable in other contexts of gravitational-wave data analysis, as long as differentiable and parallelizable likelihoods are available.