---
layout: distill
title: "Dynamic guessing for Hamiltonian Monte Carlo with embedded numerical root-finding"
description: A new algorithm and Python package for a class of difficult Bayesian statistical models
htmlwidgets: true

# Anonymize when submitting
authors:
  - name: Anonymous
    affiliations:
      name: Anonymous

# Only add author names for camera-ready
# authors:
#   - name: Author Name
#     url: "https://[url_of_author]"
#     affiliations:
#       name: Research Center, University Name
#   - name: Another Author Name
#     url: "https://[url_of_author]"
#     affiliations:
#       name: Research Center, University Name

# Must be the same name as your submission. Do not change this name, just use "submission.bib".
bibliography: submission.bib

# Add a table of contents to your submission.
#   - make sure that TOC names match the actual section names
#     for hyperlinks within the submission to work correctly.
toc:
  - name: Abstract
  - name: Introduction
  - name: Related work
  - name: Methods
  - name: Experiments
  - name: Results
  - name: Discussion and conclusions
  - name: "Appendix 1: pseudocode description of dynamic guessing"
  - name: "Appendix 2: implementation of guess-implicit"
  - name: "Appendix 3: software details"
  - name: "Appendix 4: benchmark models"
---

## Abstract

It is possible to fit Bayesian statistical models whose parameters satisfy analytically intractable algebraic conditions by embedding a differentiable numerical root-finder inside a gradient-based sampling algorithm like Hamiltonian Monte Carlo. This technique has enabled important scientific breakthroughs, but is limited by the high computational cost of computing and differentiating large numbers of numerical solutions. We show that dynamically varying the starting guess within a Hamiltonian trajectory 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.

## Introduction

If a modeller knows that some quantities jointly satisfy algebraic constraints, they may want to embed a root-finding problem inside a Bayesian statistical model. For example, biochemical reaction networks are governed by partially-known kinetic parameters and often satisfy steady-state constraints. Typically the root-finding problem has no analytical solution, so that its solution must be approximated using numerical methods.

Statistical inference for this kind of model is possible using gradient-based Markov Chain Monte Carlo algorithms like Hamiltonian Monte Carlo and its variants, as the parameter gradients of root-finding problems can usually be found. Unfortunately, solving and differentiating root-finding problems in the course of gradient-based MCMC imposes a substantial computational overhead, limiting the range of models that can practically be fit.

In this paper we propose to address this problem by dynamically updating the root-finding algorithm's starting guess as the sampler moves along a simulated Hamiltonian trajectory. We propose heuristics for updating the guess and test these on a range of statistical models, showing that dynamic guessing improves performance compared with the state of the art. We also present a Python package `grapevine` containing our implementation of HMC with dynamic guessing, as well as benchmarks and convenience functions that allow users to easily fit their own statistical models using our algorithm.

Our Python package and the code used to perform the experiments reported in this paper are available at [ANONYMISED] and from the Python Package Index. 

## Hamiltonian Monte Carlo

Hamiltonian Monte Carlo (HMC) and its variants sample $N$ parameter vectors $\theta_1, ..., \theta_N \in \mathbb{R}^{k}$ according to a target probability density function $\pi: \mathbb{R}^k \rightarrow [0, 1]$. If the algorithm works, then the statistical properties of the sample will approximately agree with the target distribution, i.e. $\sum_i \frac{1}{N}f(\theta_i) \approx E_{\pi}f(\theta)$ for any function $f$.

See [@betancourtConceptualIntroductionHamiltonian2018] for a full introduction to HMC. The core strategy, shared with the older Metropolis-Hastings-Rosenberg algorithm (MHRMC) is to start from a parameter vector $\theta^{\dagger}$, then generate a proposal vector $\theta^\star$, then accept or reject the proposal randomly, with probability depending on the ratio $\frac{\pi(\theta^{\star})}{\pi(\theta^{\dagger})}$. Whereas MHRMC generates proposal vectors according to a Gaussian random walk, so that $\theta^{\star} \sim N(\theta^{\dagger}, \Sigma)$, HMC constructs an adjoint dynamical system obeying Hamiltonian dynamics, then generates a proposal by numerically simulating the trajectory resulting from randomly perturbing this system. 

Specifically, the adjoint dynamical system maps the parameter vector $\theta\in\mathbb{R}^{k}$ to a particle in $k$-dimensional space with potential energy $V(\theta)-\ln{\pi{\theta}}$ and kinetic energy $K(\theta, \kappa) = -\ln{\pi(\theta\mid \kappa)}$ for auxiliary momentum vector $\kappa\in\mathbb{R}^k$. Perturbations are modelled by choosing an initial $\kappa$ at random, then finding the trajectory where the Hamiltonian $H(\theta, \kappa) = K(\theta, \kappa) + V(\theta)$ is constant.

A Hamiltonian trajectory can be simulated numerically using a symplectic integrator such as Velocity Verlet[@swopeComputerSimulationMethod1982][@blanesNumericalIntegratorsHybrid2014]. These integrators linearise the trajectory into successively computed segments of given "step size" $\epsilon$. To compute each segment the local potential energy $V(\theta)$ and its gradient $\nabla_{\theta}=\frac{\partial V}{\partial \theta}$ must be calculated at every step. Below when we refer to a "simulated Hamiltonian trajectory", we mean the series $\theta^{\dagger}, \theta^{1}, \ldots, \theta^{n}$ resulting from this process.

With carefully chosen hyperparameters, HMC can reliably generate a proposal vector $\theta^{\star}$ that is far away in parameter space from the initial parameter vector $\theta^{\dagger}$, yet has a non-negligible acceptance probability, leading to excellent performance [@mangoubiDoesHamiltonianMonte2018]. This, together with the asymptotic exactness of Markov Chain Monte Carlo and the existence of well-maintained implementations, e.g. [@StanReferenceManual][@abril-plaPyMCModernComprehensive2023][@cabezas2024blackjax], make HMC a popular choice for applications requiring sampling of probability functions, such as Bayesian statistical modelling.

## Embedded root-finding

We consider sampling by Hamiltonian Monte Carlo of a probability density function that embeds a root-finding problem, so that $\pi(\theta) = f(\theta, x)$ where $x$ is a vector such that $g(x, \theta) = \bar{0}$. We assume that solving $g$ for $x$ requires a numerical solver. 

To generate a proposal vector $\theta^{\star}$, HMC simulates a trajectory for the adjoint Hamiltonian system using a symplectic integrator, starting at the current parameter vector $\theta^{\dagger}$. At each step along the simulated trajectory, the numerical solution $x$ is required to calculate the potential energy $V(\theta)$, and the parameter gradient $\frac{\partial x}{\partial \theta}$ is required to calculate the potential energy's gradient $\nabla_{\theta}$. Thus a differentiable numerical root-finding algorithm is needed.

Automatically differentiable numerical root-finders are available that interface with popular HMC implementations[@optimistix2024][@Roots], or are fully integrated[@StanUsersGuide]. These have enabled application of Bayesian statistical modelling with embedded numerical root-finding appears in a wide range of scientific contexts, including ignition chemistry [@najmUncertaintyQuantificationChemical2009], cell biology [@grovesBayesianRegressionFacilitates2024][@linden-santangeliIncreasingCertaintySystems2024] and optimal control [@leeftinkProbabilisticPontryaginsMaximum2025a].

See [@timonenImportanceSamplingApproach2022][@margossianReviewAutomaticDifferentiation2019] for discussion of previous implementations of HMC with embedded root-finding. 

While many numerical root-finding algorithms exist, below we focus on the Newton-Raphson algorithm [@boggsSolutionNonlinearSystems1971]. To find $x$ satifying $g(x, \theta) = \bar{0}$ given $\theta$, this algorithm iteratively updates a starting guess $x_0$ according to the rule

\begin{align*}
x_{i+1} = x_i - J_{x_i}^{-1} g(x_i, \theta)
\end{align*}

where $J_{x_i}=\frac{\partial g(x_i, \theta)}{\partial x_i}$ is the jacobian of the target function with respect to $x_i$.

Once a solution $x$ is found that satisfies $g(x, \theta) = 0$ to within a desired tolerance, its gradients with respect to $\theta$ can be found using the implicit function theorem:

$$
\frac{\partial x}{\partial \theta} = -J_x^{-1} J_\theta
$$

Assuming that calculating and inverting the jacobian term  $J_{x_i}^{-1}$ is approximately equally costly throughout, the computational cost of using the Newton-Raphson algorithm to solve an embedded root-finding problem for a single simulated segment of a Hamiltonian trajectory is approximately proportional to the total number of Newton steps required. Like other numerical root-finding algorithms, this cost is highly sensitive to the starting guess $x_0$: if $x_0$ is in the neighbourhood of the solution then the algorithm converges quadratically, whereas a poor starting guess can prevent convergence altogether.
See [@casellaChoiceInitialGuesses2021] for a general discussion of the effect of the choice of starting guess on numerical root-finding performance.

A natural way to speed up HMC with embedded root-finding is to find the best possible guess for each problem. The main recommendation of the Stan user's guide [@StanUsersGuide] is to hard-code a guess that is reasonable, given the likely values of the parameters. However, previous implementations of HMC with embedded root-finding require the starting guess to be the same for all problems on the same simulated Hamiltonian trajectory. Since the sampling performance of HMC is tied to its ability to propose a parameter vector $\theta^{\star}$ that is distant in parameter space from the current vector $\theta^{\dagger}$, this limitation is problematic. The simulated trajectory from $\theta^{\dagger}$ to a distant $\theta^{\star}$ will include intermediate parameter vectors $\theta^{a}$ and $\theta^{b}$ that are also distant. The solution to problem $a$, i.e. $g(x, \theta^a) = \bar{0}$, will typically be different from that of problem $b$, $g(x, \theta^b) = \bar{0}$. Therefore a starting guess that is optimal for problem $a$ will be sub-optimal for problem $b$.

## Related work

Our approach draws on previous work on numerical continuation and warm starting. Continuation refers to methods that use one numerical solution to generate an initial guess for a perturbed problem, which can then be solved more precisely using another method. In particular, our *guess-implicit* heuristic is a case of the Euler-Newton predictor-corrector method investigated by Allgower and Georg[@allgowerIntroductionNumericalContinuation2003]. Chapter 5 of this work proves that a Euler-Newton predictor-corrector method converges to the correct solution path of a continuously varying root-finding problem under sufficiently small linear perturbations.

In model predictive control it is often useful to "warm start" a numerical solver using a solution from a previous time step [@diehlEfficientNumericalMethods2009]. These ideas have recently been combined in the context of bilevel optimisation of neural network hyperparameters[@pmlr-v108-lorraine20a]. Sambharya et al[@sambharyaLearningWarmstartFixedpoint2024] propose using neural networks to learn the relationship between parameter values and optimal warm starts of numerical fixed-point optimisers and compare this learned warm start approach with a "nearest-neighbour" warm start that is similar to our method. The learned warm start performs better for problems where the problem to be solved is far away from its nearest neighbour, so that the neighbour's parameter and root combination and their gradients are uninformative. In contrast, embedded root-finding problems on adjacent steps of a simulated Hamiltonian trajectory are usually close, and very many warm starts are required for a full MCMC run, motivating our consideration of cheaper linear warm starting procedures. 

## Methods

Consider the simulated Hamiltonian trajectory $\theta^1,\ldots,\theta^{n}$, where each $\theta^{i}$ embeds a root-finding problem $g(x^{i}, \theta^{i}) = \bar{0}$. The state of the art is to use the same starting guess $x_0^{\text{default}}$ to numerically solve each embedded root-finding problem, we propose choosing the starting guess dynamically. Instead, we propose using a default value for the first starting guess $x_0^1$, and choosing subsequent starting guesses $x_0^i$ for $i>1$ using a heuristic function $h$, i.e. $x_0^{i}=h(\text{info}^{i-1})$ where $\text{info}^{i-1}$ is some information available to the sampler at trajectory step $i-1$. This strategy includes the state of the art approach as a trivial heuristic:

$$
\text{guess-static}(\emptyset)=x_0^{\text{default}}
$$

We call our approach the "grapevine method" after the expression "I heard it through the grapevine" and the visual resemblance of a simulated Hamiltonian trajectory to a vine, with the numerical solution at each integrator step representing grapes. 

We implemented this idea by augmenting the velocity Verlet integrator with a dynamic state variable containing information that will be used to calculate an initial guess for the root-finding algorithm when the integrator updates its position. This variable is initialised at a default value, which is used to solve the numerical problem at the first step of any trajectory, and then modified at each position update. See appendix 1 below for a pseudocode description of our algorithm.

Our method does not affect the theoretical validity of MCMC sampling, as it does not change the target probability density function. If all embedded root-finding problems have single roots, the log probability density at any point in parameter space will be the same for a sampler using dynamic guessing as for a static one: the guessing heuristic does not affect what solution the root-finder will find, but only how quickly it converges. If an embedded problem has multiple roots, then the root that the solver finds may depend on the guess, so that the target density is not constant. Users must therefore verify that all embedded root-finding problems have single roots.

We tested two non-trivial guessing heuristics: *guess-previous* and *guess-implicit*. The information for the heuristic *guess-previous* is the solution $x^{i-1}$ of the previous root-finding problem. The heuristic is simply to use the previous solution as the next guess:

$$
\text{guess-previous}(x^{i-1}) = x^{i-1}
$$

The information for the heuristic *guess-implicit* is the solution $x^{i-1}$ of the previous root-finding problem, the current parameter vector $\theta^{i}$ and the previous parameter vector $\theta^{i-1}$. The heuristic is to use implicit differentiation to find the local derivative of the previous solution with respect to the previous parameters, then perturb the previous solution by the product of this derivative and the change in parameter values:

$$
\text{guess-implicit}(x^{i-1}, \theta^{i}, \theta^{i-1}) = x^{i-1} - \frac{\partial x^{i-1}}{\partial \theta^{i-1}}(\theta^i-\theta^{i-1})
$$

See Appendix 2 for details of how we implemented *guess-implicit*.

A limitation of *guess-previous* and *guess-implicit* is that they rely on a smooth relationship between the parameter vector $\theta$ and the embedded root-finding solution vector $x$, so that the solution of the root-finding problem $g(x^i, \theta^i) = \bar{0}$, is informative as to the solution of the problem at the next leapfrog step, i.e. $g(x^{i+1}, \theta^{i+1}) = 0$.

In the context of Hamiltonian Monte Carlo we hypothesise that this kind of smoothness is likely to obtain. Stable HMC sampling requires the target density to be sufficiently smooth for a leapfrog integrator's discretisation error to remain bounded [@bou-rabeeGeometricIntegratorsHamiltonian2018]. If the integrator's step size is too large, or the target density is locally non-smooth, the leapfrog trajectory diverges from the true Hamiltonian flow. Such divergences are a key diagnostic tool: see [@betancourtConceptualIntroductionHamiltonian2018]. Most implementations of HMC detect divergences and tune the integrator specifically to avoid them, for example by adjusting the step size, applying geometric transformations or raising diagnostic errors. This process naturally restricts the sampler to regimes where $\pi(\theta)$ behaves smoothly.

In problems with embedded root-finding, $\pi(\theta)$ is invariably coupled with the root $x$ of the embedded root-finding problem $g(x, \theta) = \bar{0}$. If $\pi(\theta)$ were insensitive to $x$, there would be little reason to calculate $x$ at each leapfrog step. Consequently, we hypothesise, the required smoothness in $\pi(\theta)$ will lead to a correspondingly smooth relationship between the parameter vector $\theta$ and the roots $x$. This argument is difficult to make rigorous due to the complexity of non-linear root-finding so we relied on empirical tests to see whether it tends to hold in practice.

## Illustrative example

As a first test we prepared an example illustrating the benefit of dynamic guessing by comparing behaviour of dynamic and static heuristics along a single Hamiltonian trajectory. Black crosses show the solutions of 11 root-finding problems (finding the minimum of a parametrised 2-dimensional Rosenbrock function) corresponding to points along a simulated Hamiltonian trajectory through parameter space. Coloured lines show the paths through solution space taken by a Newton solver to approximately solve each problem. Note that some of these paths are very small. Coloured dots indicate the guess used for the current problem. For the first problem in the trajectory all heuristics use the default guess at coordinate (1, 1), whereas the dynamic heuristics subsequently use better guesses that lie much closer to the target. As a result they take fewer Newton steps in total. Additionally, the *guess-implicit* heuristic is able to save more steps by exploiting gradient information.

![](assets/img/submission/trajectory.png)

This illustrative example is encouraging, as it is roughly representative of embedded root-finding in practice, but it is not sufficient to demonstrate our method's efficacy. The example only shows a single trajectory, whereas a typical MCMC run may require thousands of trajectories, with root-finding performance varying from one trajectory to another. The example also only considers one root-finding problem, whereas a good heuristic should perform well across a range of problems.

## Performance benchmarks

For a more comprehensive performance comparison, we compared the performance of dynamic vs static guessing for a selection of eleven statistical models with different embedded root-finding problems. These models, described in detail in Appendix 4, fall into three categories.

First, there were seven simple statistical models embedding standard problems used to test optimisation algorithms, reformulated as parametrised root-finding problems: "Easom", "Beale", "Rastrigin3d", "Rosenbrock3d", "Rosenbrock8d", "StyblinskiTang3d" and "Levy3d". These models aimed to test the hypothesis that dynamic guessing would solver performance and robustness on especially difficult embedded problems. We chose the specific embedded problems because they have a range of different difficult features for numerical solvers, vary in dimensions, have global minima so that the associated root-finding problems are well-posed and are straightforward to implement.

Second, we tested two steady state metabolic network models, one relatively small ("Linear network") and one large ("Methionine cycle"). This kind of model is common in many fields, especially biochemistry: see for example [@matosGRASPComputationalPlatform2022][@fiedlerTailoredParameterOptimization2016][@grovesBayesianRegressionFacilitates2024]. These models therefore served to test whether dynamic guessing would yield practically significant benefits in real applications

Third, to test our hypothesis that HMC adaptation would tend to induce a smooth relationship between the solution vector $x$ and the parameter vector $\theta$, we tested two models embedding the root-finding problem $g(x, \theta) = x^3 - x\odot\sin(k\theta)\odot\cos(k\theta)$, with $\odot$ representing element-wise multiplication and scalar hyper-parameter $k$ set to the very high value $1\text{e}8$. The solution to this problem depends non-smoothly on $\theta$ due to the oscillations induced by the trigonometric functions. In one model, "Adversarial Dependent", the target density $\pi(x, \theta)$ depends on the solution $x$ via a Gaussian likelihood function, as would be seen in a typical application. The other model "Adversarial Independent", was exactly the same, but with $\pi$ made independent of $x$ by setting the likelihood to be constant. We expected that, if our hypothesis were correct, then dynamic guessing would outperform static guessing on the dependent model but not on the independent model.

For each model, we simulated 20 sets of measurements based on different, randomly chosen model parameterisations. We then generated samples from each measurement set's posterior distribution using the No-U-Turn sampler, using the Stan adaptation algorithm provided by blackjax to tune the sampler. We compared the performance of the trivial heuristic *guess-static* with our proposed heuristics *guess-previous* and *guess-implicit*. Sampler configurations were and solver tolerances were set per model, and so were the same for all heuristics. In addition, for each measurement set, we initialised each heuristic's sampler at the same point and with the same random seed, so that they explored the same path through parameter space. Appendix 3 below describes our software implementation. 

We quantified performance using the total wall time required to complete the whole MCMC run, including warmup and sampling phases, and by counting the total number of post-warmup Newton steps taken in each run. The wall time metric gives a practical indication of the benefit of our approach, though at the cost of generalisability beyond our software implementation and hardware setup. As argued above, the total number of Newton steps approximately quantifies the computational cost of solving a root finding problem, though at the cost of some specificity. It thus serves as an implementation-independent performance metric.

We diagnosed sampling by verifying that the effective sample size [@StanReferenceManual] was not small compared with the total number of MCMC samples, that there were no post-warmup divergent transitions and that the embedded solver never failed to converge.

The figure below shows, for each of our experiments, the relative performance improvement in number of Newton steps from using dynamic guessing. More specifically, each dot shows the number of Newton steps taken by *guess-static* divided by the number of steps taken to generate the same post-warmup samples by a sampler using dynamic guessing.

![Performance comparison](assets/img/submission/performance.png)

This figure shows the relative improvement in wall time from dynamic guessing. Each dot shows the number of seconds taken by *guess-static*  to complete an MCMC run divided by the number of seconds taken by a dynamic sampler to complete the same run

![Wall time comparison](assets/img/submission/wall-time.png)

The number of failed MCMC runs out of 20 for each heuristic and model was as follows:

| model                   | guess_static | guess_previous | guess_implicit |
|-------------------------|--------------|----------------|----------------|
| Rosenbrock8d            | 19           | 2              | 2              |
| Rosenbrock3d            | 2            | 1              | 1              |
| Methionine cycle        | 1            | 0              | 0              |
| Linear network          | 0            | 0              | 1              |
| Easom                   | 1            | 0              | 0              |
| Levy3d                  | 0            | 0              | 0              |
| Beale                   | 0            | 0              | 0              |
| Adversarial-Dependent   | 0            | 0              | 0              |
| Rastrigin3d             | 0            | 0              | 0              |
| Adversarial-Independent | 0            | 0              | 0              |
| StyblinskiTang3d        | 0            | 0              | 0              |

This table shows the average, minimum (⤓) and maximum (⤒) Newton step counts for each heuristic and model:

| model                   | guess_static                                   | guess_previous                              | guess_implicit                              |
|-------------------------|------------------------------------------------|---------------------------------------------|---------------------------------------------|
| Linear network          |  9718.45 (⤓9111.00, ⤒ 9992.00)          |  6802.25 (⤓6441.00, ⤒ 6976.00)       |  **5589.32** (⤓5417.00, ⤒ 5749.00)       |
| Adversarial-Dependent   |  10935.20 (⤓9525.00, ⤒ 12365.00)        |  9386.20 (⤓7895.00, ⤒ 10862.00)      |  **8740.65** (⤓7381.00, ⤒ 9846.00)       |
| Rosenbrock3d            |  24605.17 (⤓15352.00, ⤒ 45383.00)       |  16335.63 (⤓14063.00, ⤒ 21025.00)    |  **6993.11** (⤓5056.00, ⤒ 10865.00)      |
| Rastrigin3d             |  8392.70 (⤓7854.00, ⤒ 8920.00)          |  8161.55 (⤓7660.00, ⤒ 8661.00)       |  **3755.90** (⤓3590.00, ⤒ 3921.00)       |
| Easom                   |  6527.05 (⤓6216.00, ⤒ 6875.00)          |  6438.80 (⤓6113.00, ⤒ 6810.00)       |  **3169.95** (⤓3057.00, ⤒ 3291.00)       |
| Adversarial-Independent |  **8764.80** (⤓7895.00, ⤒ 9515.00)          |  9594.10 (⤓8608.00, ⤒ 10442.00)      |  44509.45 (⤓38980.00, ⤒ 49076.00)    |
| Methionine cycle        |  2826010.58 (⤓2502717.00, ⤒ 3416417.00) |  427898.65 (⤓409979.00, ⤒ 464347.00) |  **402294.85** (⤓377295.00, ⤒ 429837.00) |
| Beale                   |  8875.40 (⤓8225.00, ⤒ 9670.00)          |  8328.70 (⤓7755.00, ⤒ 9053.00)       |  **3717.00** (⤓3541.00, ⤒ 3919.00)       |
| StyblinskiTang3d        |  12742.40 (⤓10557.00, ⤒ 15834.00)       |  8958.95 (⤓8038.00, ⤒ 9517.00)       |  **4555.60** (⤓4097.00, ⤒ 5147.00)       |
| Levy3d                  |  12450.20 (⤓10387.00, ⤒ 17051.00)       |  10460.25 (⤓9557.00, ⤒ 12248.00)     |  **4478.55** (⤓4072.00, ⤒ 5066.00)       |
| Rosenbrock8d            |  18202.00 (⤓18202.00, ⤒ 18202.00)       |  15868.06 (⤓15210.00, ⤒ 16975.00)    |  **5690.94** (⤓5504.00, ⤒ 5949.00)       |

The next table shows the average, minimum (⤓) and maximum (⤒) wall times for each heuristic and model:

| model                   | guess_static                         | guess_previous                   | guess_implicit                       |
|-------------------------|--------------------------------------|----------------------------------|--------------------------------------|
| Rosenbrock3d            |  0.55 (⤓0.41, ⤒0.87)         |  0.40 (⤓ 0.37, ⤒ 0.45)     |  **0.30** (⤓ 0.27, ⤒ 0.36)     |
| Easom                   |  0.25 (⤓0.23, ⤒0.26)         |  0.24 (⤓ 0.23, ⤒ 0.28)     |  **0.23** (⤓ 0.22, ⤒ 0.26)     |
| Beale                   |  0.29 (⤓0.27, ⤒0.31)         |  0.28 (⤓ 0.26, ⤒ 0.31)     |  **0.24** (⤓ 0.23, ⤒ 0.26)     |
| StyblinskiTang3d        |  0.36 (⤓0.31, ⤒0.42)         |  0.29 (⤓ 0.28, ⤒ 0.32)     |  **0.27** (⤓ 0.25, ⤒ 0.30)     |
| Adversarial-Independent |  **0.04** (⤓0.04, ⤒0.05)     |  0.04 (⤓ 0.04, ⤒ 0.05)     |  0.14 (⤓ 0.13, ⤒ 0.15)         |
| Methionine cycle        |  1218.53 (⤓931.40, ⤒3424.77) |  82.88 (⤓ 74.37, ⤒ 124.10) |  **74.16** (⤓ 63.87, ⤒ 159.66) |
| Linear network          |  0.37 (⤓0.35, ⤒0.40)         |  **0.31** (⤓ 0.30, ⤒ 0.34) |  0.34 (⤓ 0.33, ⤒ 0.36)         |
| Rosenbrock8d            |  0.57 (⤓0.57, ⤒0.57)         |  0.54 (⤓ 0.52, ⤒ 0.60)     |  **0.44** (⤓ 0.41, ⤒ 0.50)     |
| Adversarial-Dependent   |  0.07 (⤓0.05, ⤒0.07)         |  **0.06** (⤓ 0.05, ⤒ 0.07) |  0.07 (⤓ 0.05, ⤒ 0.08)         |
| Levy3d                  |  0.39 (⤓0.36, ⤒0.46)         |  0.35 (⤓ 0.34, ⤒ 0.43)     |  **0.31** (⤓ 0.29, ⤒ 0.37)     |
| Rastrigin3d             |  0.31 (⤓0.30, ⤒0.33)         |  0.30 (⤓ 0.29, ⤒ 0.33)     |  **0.27** (⤓ 0.26, ⤒ 0.29)     |


Dynamic guessing reduced the number of failed MCMC runs and required fewer Newton steps compared with static guessing for all the statistical models, except, as expected, "Adversarial Independent". The heuristic *guess-implicit* generally performed better than *guess-static*: this is also expected considering that *guess-implicit* had access to more information, specifically the previous solution's parameter gradients as well as its value.

The improvement in abstract performance, as indicated by the number of required Newton steps, led to small wall time improvements on the toy models, again with the exception of "Adversarial Independent" and very large improvements for the difficult Methionine cycle model. Interestingly, *guess-implicit* performed worse in wall time than *guess-previous* on the "Linear network" and "Adversarial Dependent" models, despite using fewer Newton steps.

All of our experiments were performed on a MacBook Pro 2024 with Apple M4 Pro processor and 48GB RAM, running macOS 15.7.2. Instructions for reproducing our benchmarks can be found in our code repository in the file `readme.md`. 

## Discussion and conclusions

The dynamic algorithms' lower failure rate is likely because, for these algorithms, the sampler diverged less often when traversing low-probability trajectories at the start of the tuning phase of sampling. Such trajectories are especially unfavourable for static guessing because they have root-finding problem solutions that are far away from any reasonable global guess.

Based on our results, we expect that replacing a static guessing algorithm with dynamic guessing will typically improve sampling performance for MCMC with embedded root-finding, making it possible to fit previously infeasible statistical models. The performance improvements on models embedding optimisation test functions show that dynamic guessing can improve MCMC  performance even when the embedded problem is challenging for the solver. The improvements on the steady state metabolic network models show that this performance benefit extends to realistic cases, and can lead to a practically significant improvement in wall time.

As we expected, dynamic guessing was unhelpful in the "Adversarial Independent" model, which embedded a root-finding problem whose solution varied non-smoothly with the parameter vector, and where there was no coupling of the root-finding problem with the target probability density. In this case *guess-previous* showed similar performance as *guess-static*, whereas *guess-implicit* performed worse due to the local parameter gradient of the solution tending to guide the solver away from the next solution. However, when we coupled the root-finding problem with the target probability density in the "Adversarial Dependent" model, we again saw improved performance from dynamic compared with static guessing. We conclude that HMC adaptation algorithms are able to adjust appropriately for even highly non-smooth parameter-root relationships, provided that the roots are coupled with the target density function via a likelihood function.

In our experiments *guess-implicit* consistently saved Newton steps compared with *guess-previous*, except in the non-representative "Adversarial Independent" case. This improvement in theoretical performance mostly translated to improved wall-time on our hardware and with our software implementation, except in the "Linear network" and "Adversarial Dependent cases", where the *guess-previous* sampler used slightly less time despite performing additional Newton steps. This is likely due to the additional computational overhead imposed by calculating and caching $\frac{\partial x^i}{\partial \theta^i}$ at each leapfrog step. While our implementation calculates this quantity efficiently (see appendices 2 and 3 below for details), re-calculating it at all wastes work as $\frac{\partial x^i}{\partial \theta^i}$ is already needed to obtain the gradient of the potential energy (see section on embedded root-finding above). It is therefore likely possible to reduce the cost of *guess-implicit* by more closely integrating it with the leapfrog integrator. In the meantime, we recommend using *guess-implicit* in preference to *guess-previous* because it will tend to perform better in the worst-case scenario where the cost per Newton step is very high.

Our approach is not strictly limited to Hamiltonian Monte Carlo, and would also work for other MCMC algorithms such as Metropolis-Hastings-Rosenberg or Langevin Monte Carlo [@robnikFluctuationDissipationMicrocanonical2024]. However, those algorithms do not require so many root-finding solutions per MCMC iteration, limiting the benefit of our approach, and do not ensure smoothness of the target density function: consequently, the parameter-root relationship may more often be non-smooth for these algorithms.

An opportunity for further performance improvement would be to use a different method to solve the first root-finding problem in a trajectory than for later problems. Plausibly, a slow but robust solver could be preferable for the first problem, which uses a default guess, whereas a faster but more fragile solver might be preferable for later problems where a potentially better guess is available.

There are several ways in which our approach can fail, beyond those explored in our experiments. For embedded problems with multiple roots, the previous solution may be a bad guess for the next solution. If the jacobian $J_x$ is singular or near-singular then our conjugate gradient based method for finding $\frac{\partial x}{\partial \theta}$ will fail. If the target probability distribution has varying characteristic length scale, so that a too-large leapfrog integrator step size is chosen, then *guess-implicit* may suggest a suboptimal starting guess. These failure modes are shared with non-linear root-finding and Hamiltonian Monte Carlo in general, but dynamic guessing may exacerbate them. 

Additionally, our software implementation must be used with care, as it requires a posterior log density function where the guess variable is only used by the numerical solver, and does not otherwise affect the output. We do not provide an automatic safeguard preventing the user from breaking this requirement, even though doing so risks producing invalid MCMC inference. It may therefore be beneficial to implement dynamic guessing using a stricter interface that makes inappropriate use of the guess variable impossible.

# Appendix 1: pseudocode description of dynamic guessing

We assume we have:

- a function $\text{Solve}$ that takes in a guess and returns an approximate solution to the target root-finding problem
- a scalar-valued function $\text{LogProb}$ that takes in a set of parameter values and a root-finding solution
- a function $\text{GetInfo}$ that takes in a set of parameter values and a root-finding solution, returning information for a heuristic.
- a function $\text{Heuristic}$ that takes in the output of $\text{GetInfo}$ and returns a guess
- initial parameters $\theta_0$
- initial momentum $\kappa_0$
- default information $\text{info}_{default}$
- step size $\epsilon$

The aim is to initialise and then simulate a Hamiltonian trajectory using Leapfrog integration, while passing root-finding information appropriately. For this we use the functions $\text{Heuristic}$, $\text{LogDensityAndInfo}$, $\text{PotentialGradientAndInfo}$, $\text{InitialiseTrajectory}$ and $\text{LeapfrogStep}$:

**Function 1: Generate a new guess ($\text{Heuristic}$)**

- **Input:**
  - Parameters $\theta$
  - Information from previous step, $\text{info}$
- **Output:** A guess for the root-finding problem, $x_0$

For example, for the heuristic *guess-previous*, $\text{info}$ is the solution from previous step, if available, or a dummy value indicating that this is the first problem in the trajectory. The output is the previous solution, if available, or else a default value.

**Algorithm 2: Update log density and information ($\text{LogDensityAndInfo}$)**

- **Input:**
  - Parameters $\theta$
  - Information from previous step, $\text{info}$

1. $x_0 \leftarrow \text{Heuristic}(\theta, \text{info})$
2. $x \leftarrow \text{Solve}(x_0)$
3. $\text{lp} \leftarrow \text{LogProb}(\theta, x)$
4. $\text{info}_{+}\leftarrow \text{GetInfo}(\theta, x)$
5. Return $\text{lp}, \text{info}_{+}$ 

**Algorithm 3: Evaluate potential energy and gradient ($\text{PotentialGradientAndInfo}$)**

- **Input:** Parameters $\theta$, information from previous step, $\text{info}$

1. $lp, \text{info}_{+} \leftarrow \text{LogDensityAndInfo}(\theta, \text{info})$
2. $V \leftarrow -lp$ (Potential energy is the negative log density)
3. $\nabla_{\theta} V \leftarrow -\nabla_{\theta} lp$ (Compute the gradient of the potential energy)
4. Return $V, \nabla_{\theta}V$ and $\text{info}_{+}$

**Algorithm 4: Initialise Trajectory ($\text{InitialiseTrajectory}$)**

1. $V_{0}, \nabla_{\theta} V_{0}, \text{info} \leftarrow \text{PotentialGradientAndInfo}(\theta_{0}, \text{info}_{default})$
2. Return $\theta_{0}, \kappa_{0}, V_{0}, \nabla_{\theta} V_{0}, \text{info}$ 

**Algorithm 5: Leapfrog Integration Step ($\text{LeapfrogStep}$)**

- **Input:**
  - Current state $\theta, \kappa, V, \nabla_{\theta} V, \text{info}$
  - step size $\epsilon$

1. $\kappa_{\text{mo}} \leftarrow \kappa - \frac{\epsilon}{2} \nabla_{\theta} V$ (Update momentum, first half-step)
2. $\theta_{+} \leftarrow \theta + \epsilon \kappa_{\text{mo}}$ (Update parameters)
3. $V_{+}, \nabla_{\theta} V_{+}, \text{info}_{+}  \leftarrow \text{PotentialGradientAndInfo}(\theta_{+}, \text{info})$
4. $\kappa_{+} \leftarrow \kappa_{\text{mo}} - \frac{\epsilon}{2} \nabla_{\theta} V_{+}$ (Update momentum, second half-step)
5. Return $\theta_{+}, \kappa_{+}, V_{+}, \nabla_{\theta} V_{+}, \text{info}_{+}$

The functions $\text{InitialiseTrajectory}$ and $\text{LeapfrogStep}$ are drop-in substitutes for their counterparts in standard Hamiltonian Monte Carlo.

# Appendix 2: implementation of *guess-implicit*

The *guess-implicit* heuristic is defined as follows, given previous solution $x^{i-1}$, previous parameters $\theta^{i-1}$ and current parameters $\theta^{i}$:

$$
\text{guess-implicit}(x^{i-1}, \theta^{i-1}, \theta^{i}) = x^{i-1} + \frac{dx}{d\theta}\Delta\theta
$$

where $\Delta\theta=(\theta^{i}-\theta^{i-1})$. 


To obtain $\frac{dx}{d\theta}$, we use the following consequence of the implicit function theorem[@oliveiraImplicitInverseFunction2014]:

$$
\frac{\partial x}{\partial\theta} = -(\operatorname{jac}_{x}g(x^{i-1}, \theta^{i-1}))^{-1}\operatorname{jac}_{\theta}g(x^{i-1}, \theta^{i-1})
$$

In this expression the term

$$
\operatorname{jac}_{x}g(x^{i-1},\theta^{i-1})
$$

abbreviated below to $J_{x}$, indicates the jacobian with respect to $x$ of $g(x^{i-1}, \theta^{i-1})$. Similarly

$$\operatorname{jac}_{\theta}g(x^{i-1}, \theta^{i-1}) = J_{\theta}$$

is the jacobian with respect to $\theta$ of $g(x^{i-1}, \theta^{i-1})$.

Substituting terms we then have 

$$
\text{guess-implicit}(x^{i-1}, \theta^{i-1}, \theta^{i}) = x^{i-1} - J_{x}^{-1}J_{\theta} \Delta\theta
$$

The *guess-implicit* heuristic can be implemented using the following Python function:

```python
import jax 
from jax import numpy as jnp

def guess_implicit(guess_info, params, f):
    "Guess the next solution using the implicit function theorem."
    old_x, old_p, *_ = guess_info
    delta_p = jax.tree.map(lambda o, n: n - o, old_p, params)
    _, jvpp = jax.jvp(lambda p: f(old_x, p), (old_p,), (delta_p,))
    jacx = jax.jacfwd(f, argnums=0)(old_x, old_p)
    u = -(jnp.linalg.inv(jacx))
    return old_x + u @ jvpp
```

Note that this function avoids materialising the parameter jacobian $J_{\theta}$, instead finding the jacobian vector product $J_{\theta}\nabla_{\theta}$ using the function `jax.jvp`. It is possible to avoid materialising the matrix $J_{x}$ using a similar strategy, as demonstrated by the function `guess_implicit_cg` below.

```python
import jax 

def guess_implicit_cg(guess_info, params, f):
    "Guess the next solution using the implicit function theorem."
    old_x, old_p, *_ = guess_info
    delta_p = jax.tree.map(lambda o, n: n - o, old_p, params)
    _, jvpp = jax.jvp(lambda p: f(old_x, p), (old_p,), (delta_p,))

    def matvec(v):
        "Compute Jx @ v"
        return jax.jvp(lambda x: f(x, old_p), (old_x,), (v,))[1]

    dx = -jax.scipy.sparse.linalg.cg(matvec, jvpp)[0]
    return old_x + dx
```

Which implementation of *guess-implicit* is preferable depends on the relative cost and reliability of directly calculating the matrix inverse $J_{x}^{-1}$ as in the function `guess_implicit`, compared with numerically solving $J_x J_p \nabla_p = 0$ as in `guess_implicit_cg`. In general, this depends on the performance of the characteristics of the numerical solver relative to direct matrix inversion as implemented by the function `jax.numpy.linalg.inv`. For example, if $J_{x}$ is sparse but positive semi-definite, `guess_implicit_cg` will likely perform better as the conjugate gradient method can exploit sparsity [@benTalNemirovskiLectureNotes]. 

# Appendix 3: software details

Using Blackjax [@cabezas2024blackjax], we implemented a version of a No-U-Turn sampler with dynamic guessing, which we call "grapeNUTS". For convenience we provide a Python package `grapevine` containing our implementation, including a utility function `run_grapenuts` with which users can easily test the GrapeNUTS sampler.

Our implementation builds on the popular JAX [@jax2018github] scientific computing ecosystem, allowing users to straightforwardly define statistical models and adapt existing models to work with grapeNUTS. Similarly to Bayeux [@bayeuxdevelopersBayeuxStateArt2025], grapevine requires a model in the form of a function that returns a scalar log probability density given a JAX PyTree of parameters; additionally, in grapevine such a function must also accept and return a PyTree containing information for guessing the answers to embedded root-finding problems. Users can specify root-finding problems using arbitrary JAX-compatible libraries, for example optimistix [@optimistix2024] or diffrax [@kidger2021on].

# Appendix 4: benchmark models

## Optimisation test functions

We compared our four heuristics on a series of variations of the following model:

\begin{align*}
\theta &\sim Normal(0, \sigma_{\theta}) \\
x &\sim Normal(\hat{x}, \sigma_{x})
\end{align*}

where $\hat{x}$ is the root such that $f(\hat{x} + \theta) = 0$.

In this equation $f$ is the gradient of a textbook optimisation test function, $sol$ is the textbook solution and $\theta$ is a vector with the same size as the input to $f$. We tested the following functions from the virtual library of simulation experiments [@simulationlib]:

- Easom function (2 dimensions) 
- 3-dimensional Levy function
- Beale function (2 dimensions)
- 3-dimensional Rastrigin function
- 3-dimensional and 8-dimensional Rosenbrock functions 
- Styblinski-Tang function


## Adversarial example

To test the effect of a non-smooth relationship between the solution vector $x$ and the parameter vector $\theta$, we constructed two adversarial models:
"Adversarial-Dependent" and "Adversarial-Independent". Both of these models embed the non-smooth root-finding problem

$$
g(x, \theta) = x^3 - x\odot\sin(k\theta)\odot\cos(k\theta)
$$

where $\odot$ represents element-wise multiplication and scalar $k=1E8$.

The model "Adversarial-Dependent" was the same as the test function models, i.e.

\begin{align*}
\theta &\sim Normal(0, \sigma_{\theta}) \\
x &\sim Normal(\hat{x}, \sigma_{x})
\end{align*}

In this model the solution $x$ is coupled with the total log probability density $\pi(\theta)$ by the likelihood $x &\sim Normal(\hat{x}, \sigma_{x})$.

The model "Adversarial Independent" was the same, but without the likelihood, i.e.

\begin{align*}
\theta &\sim Normal(0, \sigma_{\theta}) \\
\end{align*}

The code implementing "Adversarial Independent" nonetheless evaluated $x$ at every leapfrog step.

## Steady-state reaction networks

To illustrate our algorithm's practical relevance we constructed two statistical models where evaluating the likelihood $p(y\mid\theta)$ requires solving a steady-state problem, i.e. finding a vector $x$ such that $\frac{dx}{dt} = S\cdot v(x, \theta) = \bar{0}$ for known real-valued matrix $S$ and function $v$. In the context of chemical reaction networks, $S_{ij}\in\mathbb{R}$ can be interpreted as representing the amount of compound $i$ consumed or produced by reaction $j$, $x$ as the abundance of each compound and $v(x, \theta)$ as the rate of each reaction. The condition $\frac{dx}{dt} = \bar{0}$then represents the assumption that the compounds' abundances are constant. 

We tested two similar models with this broad structure, one embedding a small biologically-inspired steady state problem and one a relatively large and well-studied realistic steady-state problem.

The smaller modelled network is a toy model of a linear pathway with three reversible reactions with rates $v_1$, $v_2$ and $v_3$. These reactions affect the internal concentrations $x_A$ and $x_B$ according to the following graph:


![Performance comparison](assets/img/submission/mermaid-diagram.png)

The rates $v_1$, $v_2$ and $v_3$, given internal concentrations

$$
x = x_{A}, x_{B}
$$

and parameters

$$
\theta = k^{m}_{A}, k^{m}_{B}, v^{max}, k^{eq}_1, k^{eq}_2,k^{eq}_3, k^{f}_1, k^{f}_3, x^{ext}_A, x^{ext}_{B}
$$

are calculated as follows:

$$
\begin{aligned}
v_1(x, \theta) &=  k^{f}_1 (x^{ext}_{A} - x_{A} / k^{eq}_1) \\
v_2(x, \theta) &= \frac{\frac{v^{max}}{k^{m}_A} (x_{A} - x_{B} / k^{eq}_2)}{1 + x_{A}/k^{m}_{A} + x_{B}/k^{m}_{B} }  \\
v_3(x, \theta) &=  k^{f}_3 (x^{ext}_{B} - x_{B} / k^{eq}_3)
\end{aligned}
$$

According to these equations, rates $v_1$ and $v_3$ are described by mass-action rate laws: transport reactions are often modelled in this way. Rate $v_2$ is described by the Michaelis-Menten equation that is a popular choice for modelling the rates of enzyme-catalysed reactions. 

The larger network models the mammalian methionine cycle, using equations taken from [@grovesBayesianRegressionFacilitates2024], including highly non-linear regulatory interactions. We selected this model because it describes a real biological system and has a convenient scale, being large and complex enough to test the grapevine method's scalability, but small enough for benchmarking purposes.

For the small linear network, we solved the embedded steady-state problem using the optimistix Newton solver. For the larger model of the methionine cycle we  simulated the evolution of internal concentrations as an initial value problem until a steady-state event occurred, using the steady-state event handler and Kvaerno5 ODE solver provided by diffrax. In this case a guess is still needed in order to provide an initial value. Solving a steady-state problem in this way is often more robust than directly solving the system of algebraic equations; see [@fiedlerTailoredParameterOptimization2016] and [@lakrisenkoEfficientComputationAdjoint2023] for further discussion.

Code used for these two experiments is in the code repository files `benchmarks/methionine.py` and `benchmarks/linear.py`.
