Assessment of Gradient-Based Samplers in Standard Cosmological Likelihoods

TL;DR

This paper compares gradient-based samplers like NUTS with traditional methods in cosmological likelihood analyses, finding that NUTS offers efficiency gains in high-dimensional spaces but with increased computational cost.

Contribution

It provides a detailed assessment of NUTS versus Metropolis-Hastings in cosmology, highlighting the efficiency benefits and computational trade-offs involved.

Findings

NUTS achieves about 10 times more effective samples per likelihood evaluation.

The computational cost of gradients makes NUTS roughly twice as slow in real time.

Gradient-based samplers are beneficial for high-dimensional cosmological parameter spaces.

Abstract

We assess the usefulness of gradient-based samplers, such as the No-U-Turn Sampler (NUTS), by comparison with traditional Metropolis-Hastings algorithms, in tomographic $3 \times 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 (Campagne et al. 2023), and we use it to derive parameter constraints from both datasets using the NUTS algorithm as implemented in {\S}4, and the Metropolis-Hastings algorithm as implemented in Cobaya (Lewis 2013). 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 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 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.