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 (${\tt NUTS}$), by comparison with traditional Metropolis-Hastings algorithms, in tomographic 3 × 2 point analyses. Specifically, we use the DES Year 1 data and a simulated future LSST-like 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 datasets using the NUTS algorithm implemented in ${\tt 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 ∼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 ∼4.5 times larger for both experiments. We validate these results making use of analytical multi-variate distributions (a multivariate Gaussian and a Rosenbrock distribution) with increasing dimensionality. Based on these results, we conclude that gradient-based samplers such as ${\tt 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.