Assessment of gradient-based samplers in standard cosmological likelihoods

ABSTRACT We assess the usefulness of gradient-based samplers, such as the no-U-turn sampler (NUTS), by comparison with traditional Metropolis–Hastings (MH) algorithms, in tomographic $3\times 2$ point analyses. Specifically, we use the Dark Energy Survey (DES) Year 1 data and a simulated dataset for the Large Synoptic Survey Telescope (LSST) survey as representative examples of these studies, containing a significant number of nuisance parameters (20 and 32, respectively) that affect the performance of rejection-based samplers. To do so, we implement a differentiable forward model using jax-cosmo, and we use it to derive parameter constraints from both data sets using the nuts algorithm implemented in numpyro, and the Metropolis–Hastings algorithm as implemented in cobaya. When quantified in terms of the number of effective number of samples taken per likelihood evaluation, we find a relative efficiency gain of $\mathcal {O}(10)$ in favour of NUTS. However, this efficiency is reduced to a factor $\sim 2$ when quantified in terms of computational time, since we find the cost of the gradient computation (needed by nuts) relative to the likelihood to be $\sim 4.5$ times larger for both experiments. We validate these results making use of analytical multivariate distributions (a multivariate Gaussian and a Rosenbrock distribution) with increasing dimensionality. Based on these results, we conclude that gradient-based samplers such as NUTS can be leveraged to sample high-dimensional parameter spaces in Cosmology, although the efficiency improvement is relatively mild for moderate ($\mathcal {O}(50)$) dimension numbers, typical of tomographic large-scale structure analyses.