% Options for packages loaded elsewhere
\PassOptionsToPackage{unicode}{hyperref}
\PassOptionsToPackage{hyphens}{url}
\documentclass[
]{article}
\usepackage{xcolor}
\usepackage{amsmath,amssymb}
\setcounter{secnumdepth}{-\maxdimen} % remove section numbering
\usepackage{iftex}
\ifPDFTeX
  \usepackage[T1]{fontenc}
  \usepackage[utf8]{inputenc}
  \usepackage{textcomp} % provide euro and other symbols
\else % if luatex or xetex
  \usepackage{unicode-math} % this also loads fontspec
  \defaultfontfeatures{Scale=MatchLowercase}
  \defaultfontfeatures[\rmfamily]{Ligatures=TeX,Scale=1}
\fi
\usepackage{lmodern}
\ifPDFTeX\else
  % xetex/luatex font selection
\fi
% Use upquote if available, for straight quotes in verbatim environments
\IfFileExists{upquote.sty}{\usepackage{upquote}}{}
\IfFileExists{microtype.sty}{% use microtype if available
  \usepackage[]{microtype}
  \UseMicrotypeSet[protrusion]{basicmath} % disable protrusion for tt fonts
}{}
\makeatletter
\@ifundefined{KOMAClassName}{% if non-KOMA class
  \IfFileExists{parskip.sty}{%
    \usepackage{parskip}
  }{% else
    \setlength{\parindent}{0pt}
    \setlength{\parskip}{6pt plus 2pt minus 1pt}}
}{% if KOMA class
  \KOMAoptions{parskip=half}}
\makeatother
\usepackage{color}
\usepackage{fancyvrb}
\newcommand{\VerbBar}{|}
\newcommand{\VERB}{\Verb[commandchars=\\\{\}]}
\DefineVerbatimEnvironment{Highlighting}{Verbatim}{commandchars=\\\{\}}
% Add ',fontsize=\small' for more characters per line
\newenvironment{Shaded}{}{}
\newcommand{\AlertTok}[1]{\textcolor[rgb]{1.00,0.00,0.00}{\textbf{#1}}}
\newcommand{\AnnotationTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
\newcommand{\AttributeTok}[1]{\textcolor[rgb]{0.49,0.56,0.16}{#1}}
\newcommand{\BaseNTok}[1]{\textcolor[rgb]{0.25,0.63,0.44}{#1}}
\newcommand{\BuiltInTok}[1]{\textcolor[rgb]{0.00,0.50,0.00}{#1}}
\newcommand{\CharTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
\newcommand{\CommentTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textit{#1}}}
\newcommand{\CommentVarTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
\newcommand{\ConstantTok}[1]{\textcolor[rgb]{0.53,0.00,0.00}{#1}}
\newcommand{\ControlFlowTok}[1]{\textcolor[rgb]{0.00,0.44,0.13}{\textbf{#1}}}
\newcommand{\DataTypeTok}[1]{\textcolor[rgb]{0.56,0.13,0.00}{#1}}
\newcommand{\DecValTok}[1]{\textcolor[rgb]{0.25,0.63,0.44}{#1}}
\newcommand{\DocumentationTok}[1]{\textcolor[rgb]{0.73,0.13,0.13}{\textit{#1}}}
\newcommand{\ErrorTok}[1]{\textcolor[rgb]{1.00,0.00,0.00}{\textbf{#1}}}
\newcommand{\ExtensionTok}[1]{#1}
\newcommand{\FloatTok}[1]{\textcolor[rgb]{0.25,0.63,0.44}{#1}}
\newcommand{\FunctionTok}[1]{\textcolor[rgb]{0.02,0.16,0.49}{#1}}
\newcommand{\ImportTok}[1]{\textcolor[rgb]{0.00,0.50,0.00}{\textbf{#1}}}
\newcommand{\InformationTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
\newcommand{\KeywordTok}[1]{\textcolor[rgb]{0.00,0.44,0.13}{\textbf{#1}}}
\newcommand{\NormalTok}[1]{#1}
\newcommand{\OperatorTok}[1]{\textcolor[rgb]{0.40,0.40,0.40}{#1}}
\newcommand{\OtherTok}[1]{\textcolor[rgb]{0.00,0.44,0.13}{#1}}
\newcommand{\PreprocessorTok}[1]{\textcolor[rgb]{0.74,0.48,0.00}{#1}}
\newcommand{\RegionMarkerTok}[1]{#1}
\newcommand{\SpecialCharTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
\newcommand{\SpecialStringTok}[1]{\textcolor[rgb]{0.73,0.40,0.53}{#1}}
\newcommand{\StringTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
\newcommand{\VariableTok}[1]{\textcolor[rgb]{0.10,0.09,0.49}{#1}}
\newcommand{\VerbatimStringTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
\newcommand{\WarningTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
\usepackage{longtable,booktabs,array}
\usepackage{calc} % for calculating minipage widths
% Correct order of tables after \paragraph or \subparagraph
\usepackage{etoolbox}
\makeatletter
\patchcmd\longtable{\par}{\if@noskipsec\mbox{}\fi\par}{}{}
\makeatother
% Allow footnotes in longtable head/foot
\IfFileExists{footnotehyper.sty}{\usepackage{footnotehyper}}{\usepackage{footnote}}
\makesavenoteenv{longtable}
\usepackage{graphicx}
\makeatletter
\newsavebox\pandoc@box
\newcommand*\pandocbounded[1]{% scales image to fit in text height/width
  \sbox\pandoc@box{#1}%
  \Gscale@div\@tempa{\textheight}{\dimexpr\ht\pandoc@box+\dp\pandoc@box\relax}%
  \Gscale@div\@tempb{\linewidth}{\wd\pandoc@box}%
  \ifdim\@tempb\p@<\@tempa\p@\let\@tempa\@tempb\fi% select the smaller of both
  \ifdim\@tempa\p@<\p@\scalebox{\@tempa}{\usebox\pandoc@box}%
  \else\usebox{\pandoc@box}%
  \fi%
}
% Set default figure placement to htbp
\def\fps@figure{htbp}
\makeatother
\setlength{\emergencystretch}{3em} % prevent overfull lines
\providecommand{\tightlist}{%
  \setlength{\itemsep}{0pt}\setlength{\parskip}{0pt}}
\usepackage[]{natbib}
\bibliographystyle{plainnat}
\usepackage{bookmark}
\IfFileExists{xurl.sty}{\usepackage{xurl}}{} % add URL line breaks if available
\urlstyle{same}
\hypersetup{
  pdftitle={Dynamic guessing for Hamiltonian Monte Carlo with embedded numerical root-finding},
  hidelinks,
  pdfcreator={LaTeX via pandoc}}

\title{Dynamic guessing for Hamiltonian Monte Carlo with embedded
numerical root-finding}
\author{}
\date{}

\begin{document}
\maketitle

{
\setcounter{tocdepth}{3}
\tableofcontents
}
\subsection{Abstract}\label{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: \emph{guess-previous} reuses the
previous solution as the guess and \emph{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.

\subsection{Introduction}\label{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 \texttt{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.

\subsection{Hamiltonian Monte Carlo}\label{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 \citep{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{[}\citet{swopeComputerSimulationMethod1982}{]}\citep{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
\citep{mangoubiDoesHamiltonianMonte2018}. This, together with the
asymptotic exactness of Markov Chain Monte Carlo and the existence of
well-maintained implementations,
e.g.~{[}\citet{StanReferenceManual}{]}\citep{abril-plaPyMCModernComprehensive2023}\citep{cabezas2024blackjax},
make HMC a popular choice for applications requiring sampling of
probability functions, such as Bayesian statistical modelling.

\subsection{Embedded root-finding}\label{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{[}\citet{optimistix2024}{]}\citep{Roots}, or are fully
integrated\citep{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 \citep{najmUncertaintyQuantificationChemical2009}, cell
biology
{[}\citet{grovesBayesianRegressionFacilitates2024}{]}\citep{linden-santangeliIncreasingCertaintySystems2024}
and optimal control
\citep{leeftinkProbabilisticPontryaginsMaximum2025a}.

See
{[}\citet{timonenImportanceSamplingApproach2022}{]}\citep{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 \citep{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 \citep{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 \citep{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\).

\subsection{Related work}\label{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
\emph{guess-implicit} heuristic is a case of the Euler-Newton
predictor-corrector method investigated by Allgower and
Georg\citep{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
\citep{diehlEfficientNumericalMethods2009}. These ideas have recently
been combined in the context of bilevel optimisation of neural network
hyperparameters\citep{pmlr-v108-lorraine20a}. Sambharya et
al\citep{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.

\subsection{Methods}\label{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: \emph{guess-previous} and
\emph{guess-implicit}. The information for the heuristic
\emph{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 \emph{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 \emph{guess-implicit}.

A limitation of \emph{guess-previous} and \emph{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
\citep{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
\citep{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.

\subsection{Illustrative example}\label{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
\emph{guess-implicit} heuristic is able to save more steps by exploiting
gradient information.

\pandocbounded{\includegraphics[keepaspectratio]{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.

\subsection{Performance benchmarks}\label{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
{[}\citet{matosGRASPComputationalPlatform2022}{]}\citep{fiedlerTailoredParameterOptimization2016}\citep{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
\emph{guess-static} with our proposed heuristics \emph{guess-previous}
and \emph{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
\citep{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 \emph{guess-static} divided by the number of steps taken to
generate the same post-warmup samples by a sampler using dynamic
guessing.

\begin{figure}
\centering
\pandocbounded{\includegraphics[keepaspectratio]{assets/img/submission/performance.png}}
\caption{Performance comparison}
\end{figure}

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

\begin{figure}
\centering
\pandocbounded{\includegraphics[keepaspectratio]{assets/img/submission/wall-time.png}}
\caption{Wall time comparison}
\end{figure}

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

\begin{longtable}[]{@{}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.3521}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.1972}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2254}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2254}}@{}}
\toprule\noalign{}
\begin{minipage}[b]{\linewidth}\raggedright
model
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_static
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_previous
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_implicit
\end{minipage} \\
\midrule\noalign{}
\endhead
\bottomrule\noalign{}
\endlastfoot
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 \\
\end{longtable}

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

\begin{longtable}[]{@{}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.1534}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2945}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2761}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2761}}@{}}
\toprule\noalign{}
\begin{minipage}[b]{\linewidth}\raggedright
model
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_static
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_previous
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_implicit
\end{minipage} \\
\midrule\noalign{}
\endhead
\bottomrule\noalign{}
\endlastfoot
Linear network & 9718.45 (⤓9111.00, ⤒ 9992.00) & 6802.25 (⤓6441.00, ⤒
6976.00) & \textbf{5589.32} (⤓5417.00, ⤒ 5749.00) \\
Adversarial-Dependent & 10935.20 (⤓9525.00, ⤒ 12365.00) & 9386.20
(⤓7895.00, ⤒ 10862.00) & \textbf{8740.65} (⤓7381.00, ⤒ 9846.00) \\
Rosenbrock3d & 24605.17 (⤓15352.00, ⤒ 45383.00) & 16335.63 (⤓14063.00, ⤒
21025.00) & \textbf{6993.11} (⤓5056.00, ⤒ 10865.00) \\
Rastrigin3d & 8392.70 (⤓7854.00, ⤒ 8920.00) & 8161.55 (⤓7660.00, ⤒
8661.00) & \textbf{3755.90} (⤓3590.00, ⤒ 3921.00) \\
Easom & 6527.05 (⤓6216.00, ⤒ 6875.00) & 6438.80 (⤓6113.00, ⤒ 6810.00) &
\textbf{3169.95} (⤓3057.00, ⤒ 3291.00) \\
Adversarial-Independent & \textbf{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) & \textbf{402294.85} (⤓377295.00, ⤒
429837.00) \\
Beale & 8875.40 (⤓8225.00, ⤒ 9670.00) & 8328.70 (⤓7755.00, ⤒ 9053.00) &
\textbf{3717.00} (⤓3541.00, ⤒ 3919.00) \\
StyblinskiTang3d & 12742.40 (⤓10557.00, ⤒ 15834.00) & 8958.95 (⤓8038.00,
⤒ 9517.00) & \textbf{4555.60} (⤓4097.00, ⤒ 5147.00) \\
Levy3d & 12450.20 (⤓10387.00, ⤒ 17051.00) & 10460.25 (⤓9557.00, ⤒
12248.00) & \textbf{4478.55} (⤓4072.00, ⤒ 5066.00) \\
Rosenbrock8d & 18202.00 (⤓18202.00, ⤒ 18202.00) & 15868.06 (⤓15210.00, ⤒
16975.00) & \textbf{5690.94} (⤓5504.00, ⤒ 5949.00) \\
\end{longtable}

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

\begin{longtable}[]{@{}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.1852}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2815}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2519}}
  >{\raggedright\arraybackslash}p{(\linewidth - 6\tabcolsep) * \real{0.2815}}@{}}
\toprule\noalign{}
\begin{minipage}[b]{\linewidth}\raggedright
model
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_static
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_previous
\end{minipage} & \begin{minipage}[b]{\linewidth}\raggedright
guess\_implicit
\end{minipage} \\
\midrule\noalign{}
\endhead
\bottomrule\noalign{}
\endlastfoot
Rosenbrock3d & 0.55 (⤓0.41, ⤒0.87) & 0.40 (⤓ 0.37, ⤒ 0.45) &
\textbf{0.30} (⤓ 0.27, ⤒ 0.36) \\
Easom & 0.25 (⤓0.23, ⤒0.26) & 0.24 (⤓ 0.23, ⤒ 0.28) & \textbf{0.23} (⤓
0.22, ⤒ 0.26) \\
Beale & 0.29 (⤓0.27, ⤒0.31) & 0.28 (⤓ 0.26, ⤒ 0.31) & \textbf{0.24} (⤓
0.23, ⤒ 0.26) \\
StyblinskiTang3d & 0.36 (⤓0.31, ⤒0.42) & 0.29 (⤓ 0.28, ⤒ 0.32) &
\textbf{0.27} (⤓ 0.25, ⤒ 0.30) \\
Adversarial-Independent & \textbf{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) & \textbf{74.16} (⤓ 63.87, ⤒ 159.66) \\
Linear network & 0.37 (⤓0.35, ⤒0.40) & \textbf{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) &
\textbf{0.44} (⤓ 0.41, ⤒ 0.50) \\
Adversarial-Dependent & 0.07 (⤓0.05, ⤒0.07) & \textbf{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) & \textbf{0.31} (⤓
0.29, ⤒ 0.37) \\
Rastrigin3d & 0.31 (⤓0.30, ⤒0.33) & 0.30 (⤓ 0.29, ⤒ 0.33) &
\textbf{0.27} (⤓ 0.26, ⤒ 0.29) \\
\end{longtable}

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
\emph{guess-implicit} generally performed better than
\emph{guess-static}: this is also expected considering that
\emph{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, \emph{guess-implicit} performed worse in wall time than
\emph{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 \texttt{readme.md}.

\subsection{Discussion and
conclusions}\label{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 \emph{guess-previous} showed similar performance
as \emph{guess-static}, whereas \emph{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 \emph{guess-implicit} consistently saved Newton steps
compared with \emph{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 \emph{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 \emph{guess-implicit} by more closely integrating
it with the leapfrog integrator. In the meantime, we recommend using
\emph{guess-implicit} in preference to \emph{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
\citep{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 \emph{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.

\section{Appendix 1: pseudocode description of dynamic
guessing}\label{appendix-1-pseudocode-description-of-dynamic-guessing}

We assume we have:

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

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}\):

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

\begin{itemize}
\tightlist
\item
  \textbf{Input:}

  \begin{itemize}
  \tightlist
  \item
    Parameters \(\theta\)
  \item
    Information from previous step, \(\text{info}\)
  \end{itemize}
\item
  \textbf{Output:} A guess for the root-finding problem, \(x_0\)
\end{itemize}

For example, for the heuristic \emph{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.

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

\begin{itemize}
\tightlist
\item
  \textbf{Input:}

  \begin{itemize}
  \tightlist
  \item
    Parameters \(\theta\)
  \item
    Information from previous step, \(\text{info}\)
  \end{itemize}
\end{itemize}

\begin{enumerate}
\def\labelenumi{\arabic{enumi}.}
\tightlist
\item
  \(x_0 \leftarrow \text{Heuristic}(\theta, \text{info})\)
\item
  \(x \leftarrow \text{Solve}(x_0)\)
\item
  \(\text{lp} \leftarrow \text{LogProb}(\theta, x)\)
\item
  \(\text{info}_{+}\leftarrow \text{GetInfo}(\theta, x)\)
\item
  Return \(\text{lp}, \text{info}_{+}\)
\end{enumerate}

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

\begin{itemize}
\tightlist
\item
  \textbf{Input:} Parameters \(\theta\), information from previous step,
  \(\text{info}\)
\end{itemize}

\begin{enumerate}
\def\labelenumi{\arabic{enumi}.}
\tightlist
\item
  \(lp, \text{info}_{+} \leftarrow \text{LogDensityAndInfo}(\theta, \text{info})\)
\item
  \(V \leftarrow -lp\) (Potential energy is the negative log density)
\item
  \(\nabla_{\theta} V \leftarrow -\nabla_{\theta} lp\) (Compute the
  gradient of the potential energy)
\item
  Return \(V, \nabla_{\theta}V\) and \(\text{info}_{+}\)
\end{enumerate}

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

\begin{enumerate}
\def\labelenumi{\arabic{enumi}.}
\tightlist
\item
  \(V_{0}, \nabla_{\theta} V_{0}, \text{info} \leftarrow \text{PotentialGradientAndInfo}(\theta_{0}, \text{info}_{default})\)
\item
  Return
  \(\theta_{0}, \kappa_{0}, V_{0}, \nabla_{\theta} V_{0}, \text{info}\)
\end{enumerate}

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

\begin{itemize}
\tightlist
\item
  \textbf{Input:}

  \begin{itemize}
  \tightlist
  \item
    Current state \(\theta, \kappa, V, \nabla_{\theta} V, \text{info}\)
  \item
    step size \(\epsilon\)
  \end{itemize}
\end{itemize}

\begin{enumerate}
\def\labelenumi{\arabic{enumi}.}
\tightlist
\item
  \(\kappa_{\text{mo}} \leftarrow \kappa - \frac{\epsilon}{2} \nabla_{\theta} V\)
  (Update momentum, first half-step)
\item
  \(\theta_{+} \leftarrow \theta + \epsilon \kappa_{\text{mo}}\) (Update
  parameters)
\item
  \(V_{+}, \nabla_{\theta} V_{+}, \text{info}_{+}  \leftarrow \text{PotentialGradientAndInfo}(\theta_{+}, \text{info})\)
\item
  \(\kappa_{+} \leftarrow \kappa_{\text{mo}} - \frac{\epsilon}{2} \nabla_{\theta} V_{+}\)
  (Update momentum, second half-step)
\item
  Return
  \(\theta_{+}, \kappa_{+}, V_{+}, \nabla_{\theta} V_{+}, \text{info}_{+}\)
\end{enumerate}

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

\section{\texorpdfstring{Appendix 2: implementation of
\emph{guess-implicit}}{Appendix 2: implementation of guess-implicit}}\label{appendix-2-implementation-of-guess-implicit}

The \emph{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\citep{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 \emph{guess-implicit} heuristic can be implemented using the
following Python function:

\begin{Shaded}
\begin{Highlighting}[]
\ImportTok{import}\NormalTok{ jax }
\ImportTok{from}\NormalTok{ jax }\ImportTok{import}\NormalTok{ numpy }\ImportTok{as}\NormalTok{ jnp}

\KeywordTok{def}\NormalTok{ guess\_implicit(guess\_info, params, f):}
    \CommentTok{"Guess the next solution using the implicit function theorem."}
\NormalTok{    old\_x, old\_p, }\OperatorTok{*}\NormalTok{\_ }\OperatorTok{=}\NormalTok{ guess\_info}
\NormalTok{    delta\_p }\OperatorTok{=}\NormalTok{ jax.tree.}\BuiltInTok{map}\NormalTok{(}\KeywordTok{lambda}\NormalTok{ o, n: n }\OperatorTok{{-}}\NormalTok{ o, old\_p, params)}
\NormalTok{    \_, jvpp }\OperatorTok{=}\NormalTok{ jax.jvp(}\KeywordTok{lambda}\NormalTok{ p: f(old\_x, p), (old\_p,), (delta\_p,))}
\NormalTok{    jacx }\OperatorTok{=}\NormalTok{ jax.jacfwd(f, argnums}\OperatorTok{=}\DecValTok{0}\NormalTok{)(old\_x, old\_p)}
\NormalTok{    u }\OperatorTok{=} \OperatorTok{{-}}\NormalTok{(jnp.linalg.inv(jacx))}
    \ControlFlowTok{return}\NormalTok{ old\_x }\OperatorTok{+}\NormalTok{ u }\OperatorTok{@}\NormalTok{ jvpp}
\end{Highlighting}
\end{Shaded}

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

\begin{Shaded}
\begin{Highlighting}[]
\ImportTok{import}\NormalTok{ jax }

\KeywordTok{def}\NormalTok{ guess\_implicit\_cg(guess\_info, params, f):}
    \CommentTok{"Guess the next solution using the implicit function theorem."}
\NormalTok{    old\_x, old\_p, }\OperatorTok{*}\NormalTok{\_ }\OperatorTok{=}\NormalTok{ guess\_info}
\NormalTok{    delta\_p }\OperatorTok{=}\NormalTok{ jax.tree.}\BuiltInTok{map}\NormalTok{(}\KeywordTok{lambda}\NormalTok{ o, n: n }\OperatorTok{{-}}\NormalTok{ o, old\_p, params)}
\NormalTok{    \_, jvpp }\OperatorTok{=}\NormalTok{ jax.jvp(}\KeywordTok{lambda}\NormalTok{ p: f(old\_x, p), (old\_p,), (delta\_p,))}

    \KeywordTok{def}\NormalTok{ matvec(v):}
        \CommentTok{"Compute Jx @ v"}
        \ControlFlowTok{return}\NormalTok{ jax.jvp(}\KeywordTok{lambda}\NormalTok{ x: f(x, old\_p), (old\_x,), (v,))[}\DecValTok{1}\NormalTok{]}

\NormalTok{    dx }\OperatorTok{=} \OperatorTok{{-}}\NormalTok{jax.scipy.sparse.linalg.cg(matvec, jvpp)[}\DecValTok{0}\NormalTok{]}
    \ControlFlowTok{return}\NormalTok{ old\_x }\OperatorTok{+}\NormalTok{ dx}
\end{Highlighting}
\end{Shaded}

Which implementation of \emph{guess-implicit} is preferable depends on
the relative cost and reliability of directly calculating the matrix
inverse \(J_{x}^{-1}\) as in the function \texttt{guess\_implicit},
compared with numerically solving \(J_x J_p \nabla_p = 0\) as in
\texttt{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
\texttt{jax.numpy.linalg.inv}. For example, if \(J_{x}\) is sparse but
positive semi-definite, \texttt{guess\_implicit\_cg} will likely perform
better as the conjugate gradient method can exploit sparsity
\citep{benTalNemirovskiLectureNotes}.

\section{Appendix 3: software
details}\label{appendix-3-software-details}

Using Blackjax \citep{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 \texttt{grapevine}
containing our implementation, including a utility function
\texttt{run\_grapenuts} with which users can easily test the GrapeNUTS
sampler.

Our implementation builds on the popular JAX \citep{jax2018github}
scientific computing ecosystem, allowing users to straightforwardly
define statistical models and adapt existing models to work with
grapeNUTS. Similarly to Bayeux
\citep{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 \citep{optimistix2024} or diffrax
\citep{kidger2021on}.

\section{Appendix 4: benchmark
models}\label{appendix-4-benchmark-models}

\subsection{Optimisation test
functions}\label{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
\citep{simulationlib}:

\begin{itemize}
\tightlist
\item
  Easom function (2 dimensions)
\item
  3-dimensional Levy function
\item
  Beale function (2 dimensions)
\item
  3-dimensional Rastrigin function
\item
  3-dimensional and 8-dimensional Rosenbrock functions
\item
  Styblinski-Tang function
\end{itemize}

\subsection{Adversarial example}\label{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.

\subsection{Steady-state reaction
networks}\label{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:

\begin{figure}
\centering
\pandocbounded{\includegraphics[keepaspectratio]{assets/img/submission/mermaid-diagram.png}}
\caption{Performance comparison}
\end{figure}

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 \citep{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
\citep{fiedlerTailoredParameterOptimization2016} and
\citep{lakrisenkoEfficientComputationAdjoint2023} for further
discussion.

Code used for these two experiments is in the code repository files
\texttt{benchmarks/methionine.py} and \texttt{benchmarks/linear.py}.

\bibliography{submission.bib}

\end{document}
