Dynamic guessing for Hamiltonian Monte Carlo with embedded numerical root-finding

Teddy Groves, Nicholas Luke Cowie, Lars Keld Nielsen

Keywords: MCMC, Computational Biology, implicit differentiation, Hamiltonian Monte Carlo, JAX, root-finding

TL;DR: We used dynamic guessing to improve performance of Hamiltonian style Monte Carlo for target functions that embed numerical root-finding and optimisation problems.

Abstract: Thanks to scientific machine learning, it is possible to fit Bayesian statistical models whose parameters satisfy analytically intractable algebraic conditions like steady-state state constraints. This is often done by embedding a differentiable numerical root-finder inside a gradient-based sampling algorithm like Hamiltonian Monte Carlo. However, computing and differentiating large numbers of numerical solutions comes at a high computational cost. We demonstrate that dynamically varying the starting guess within Hamiltonian trajectorie can improve performance. To choose a good guess we propose two heuristics: *guess-previous* reuses the previous solution as the guess and *guess-implicit* extrapolates the previous solution using implicit differentiation. We benchmark these heuristics on a range of representative models. We also present a JAX-based Python package providing easy access to a performant sampler augmented with dynamic guessing.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 12