Leveraging Sparsity to Improve No-U-Turn Sampling Efficiency for Hierarchical Bayesian Models

TL;DR

This paper introduces Sparse NUTS (SNUTS), a preconditioning method that uses a sparse precision matrix to significantly accelerate NUTS sampling in high-dimensional, correlated hierarchical Bayesian models, especially when the posterior is sparse.

Contribution

The paper presents SNUTS, a novel preconditioning approach that leverages sparse precision matrices to improve NUTS efficiency in high-dimensional Bayesian models with correlated parameters.

Findings

SNUTS converges 10-100 times faster than standard NUTS with diagonal or dense preconditioners.

SNUTS outperforms Pathfinder variational inference in several case studies.

SNUTS is most effective in high-dimensional, sparse, and highly correlated models, scalable beyond 10,000 parameters.

Abstract

Analysts routinely use Bayesian hierarchical models to understand natural processes. The no-U-turn sampler (NUTS) is the most widely used algorithm to sample high-dimensional, continuously differentiable models. But NUTS is slowed by high correlations, especially in high dimensions, limiting the complexity of applied analyses. Here we introduce Sparse NUTS (SNUTS), which preconditions (decorrelates and descales) posteriors using a sparse precision matrix ($Q$). We use Template Model Builder (TMB) to efficiently compute $Q$ from the mode of the Laplace approximation to the marginal posterior, then pass the preconditioned posterior to NUTS through the Bayesian software Stan for sampling. We apply SNUTS to seventeen diverse case studies to demonstrate that preconditioning with $Q$ converges one to two orders of magnitude faster than Stan's industry standard diagonal or dense preconditioners. SNUTS also outperforms preconditioning with the inverse of the covariance estimated with Pathfinder variational inference. SNUTS does not improve sampling efficiency for models with the highly varying curvature found in funnels, wide tails, or multiple modes. SNUTS is most advantageous, and can be scaled beyond $10^4$ parameters, in the presence of high dimensionality, sparseness, and high correlations, all of which are widespread in applied statistics. An open-source implementation of SNUTS is provided in the R package SparseNUTS.