% 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{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}

Thanks to scientific machine learning, it is possible to fit Bayesian
statistical models whose parameters satisfy analytically intractable
algebraic conditions like steady-state constraints. This is often done
by embedding a differentiable numerical root-finder inside a
gradient-based sampling algorithm like Hamiltonian Monte Carlo. However,
computing and differentiating large numbers of numerical solutions comes
at a high computational cost. We 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 partially-known 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
known to satisfy steady-state constraints. Often the root-finding
problem's solution must be found 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.

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 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, 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{Related work}\label{related-work}

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}.

We are not aware of any previous work that developed specialised
inference methods for Bayesian models with embedded root finding.
Instead, previous work has embedded numerical root-finding within
generic inference algorithms such as Hamiltonian Monte Carlo
{[}\citet{nealMCMCUsingHamiltonian2011}{]}\citep{betancourtConceptualIntroductionHamiltonian2018}
or a related algorithm such as the No-U-Turn sampler
\citep{hoffmanNoUTurnSamplerAdaptively2014} (below we include such
variants under ``HMC''). See
{[}\citet{timonenImportanceSamplingApproach2022}{]}\citep{margossianReviewAutomaticDifferentiation2019}
for discussion of this approach. HMC is often preferable to alternative
inference algorithms because of its good performance
\citep{mangoubiDoesHamiltonianMonte2018}, asymptotic exactness and the
existence of well-maintained implementations,
e.g.~{[}\citet{StanReferenceManual}{]}\citep{abril-plaPyMCModernComprehensive2023}.

Hamiltonian Monte Carlo couples the target probability distribution with
a dynamical system representing a particle lying on a surface with one
dimension per parameter of the target distribution and a potential
energy derived from the target density. To generate a sample, a
symplectic integrator simulates the particle's trajectory along the
surface after a perturbation. Ideally, the trajectory takes the particle
far from its starting point, leading to efficient sampling. The
symplectic integrator proceeds by linearising the trajectory in small
segments, evaluating the target density and its parameter gradients on
log scale at the end of each segment. For targets with embedded
numerical root-finding, the roots and their local parameter gradients
must also be found at every step.

Many numerical root-finding algorithms are iterative, generating a
series of numbers that start with an initial guess and converge towards
the true solution. These algorithms tend to perform better, the closer
the initial guess is to the true solution: see
\citep{casellaChoiceInitialGuesses2021} for discussion of this topic.
Thus a natural way to speed up HMC with embedded root-finding is to find
the best possible guess for each problem; indeed, the main
recommendation of the Stan user's guide \citep{StanUsersGuide} is to
choose a guess that is reasonable, given the likely values of the
parameters. However, previous implementations of HMC with embedded
root-finding have required the starting guess to be the same for all
problems that lie on the same simulated Hamiltonian trajectory. Since
HMC trajectories aim to traverse a large distance in parameter space,
and embedded root-finding problems will typically have different
solutions depending on the parameters, the solution is likely to vary
for different points on the trajectory. As a result, a guess that is
optimal at one point on the trajectory will be sub-optimal elsewhere.

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. For example, an
Euler predictor, which perturbs the previous solution using local
gradient information, is known to work well with Newton
solvers\citep{allgowerIntroductionNumericalContinuation2003}. Similarly,
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}.

\subsection{Methods}\label{methods}

Instead of using the same starting guess for every root-finding problem
on one simulated Hamiltonian trajectory, we propose choosing the guess
dynamically, based on the previous integrator state. We call this
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
{[}\citet{swopeComputerSimulationMethod1982}{]}\citep{blanesNumericalIntegratorsHybrid2014}
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. In
this way all steps except the first have access to non-default guessing
information. See appendix 1 below for a pseudocode description of the
grapevine algorithm.

For HMC sampling to be valid, it is important that the target
probability density function is constant. If all embedded root-finding
problems have single roots, this will be the case up to numerical error
for the grapevine algorithm, as the guessing heuristic does not affect
what solution the root-finder will find, but only how quickly it
converges. Therefore in this case there is no difference in MCMC
validity between a grapevine sampler and an equivalent sampler with
static guessing. 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 heuristics for generating an initial guess:
\emph{guess-previous} and \emph{guess-implicit}. The information for the
heuristic \emph{guess-previous} is the solution of the previous
root-finding problem. The heuristic is simply to use the previous
solution as the next guess. The information for the heuristic
\emph{guess-implicit} is the solution of the previous root-finding
problem, and also the parameter values at the previous step. The
heuristic is to use implicit differentiation to find the local
derivative of the previous solution with respect to the previous
parameters, then obtain a guess by perturbing the previous solution by
the product of this derivative and the change in parameter values. See
Appendix 2 for details of how \emph{guess-implicit} is defined and
implemented.

\subsection{Experiments}\label{experiments}

First 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.

For a more comprehensive performance comparison, we therefore compared
the performance of dynamic vs static guessing for a selection of models,
each embedding a different root-finding problem. To test the dynamic
heuristics we augmented the No-U-Turn sampler with dynamic guessing,
calling the resulting sampler ``grapeNUTS''. For each model, we randomly
generated 6 parameter sets, each of which we used to randomly simulate
an observation set. We then sampled from the resulting posterior
distribution using unaugmented NUTS (\emph{guess-static}), and grapeNUTS
with our heuristics \emph{guess-previous} and \emph{guess-implicit}. The
sampler configurations were the same for all heuristics and solver
tolerances were set per problem. Appendix 3 below describes our software
implementation and appendix 4 details the models.

We quantified performance by dividing the total number of Newton steps
taken in each run by the effective sample size. See
\citep{StanReferenceManual} for definition of effective sample size and
\citep{kumarArviZUnifiedLibrary2019} for the software we used to
calculate it. We diagnosed sampling by verifying that the effective
sample size was not small compared with the total number of MCMC samples
and that there were no post-warmup divergent transitions. We also
recorded MCMC runs containing at least one solver convergence failure:
in these cases we counted the whole run as a failure.

\subsection{Results}\label{results}

The figure below shows performance comparisons for three dynamic
guessing heuristics over nine statistical models with embedded
root-finding problems. For each point, we randomly selected a true
parameter set, then generated a simulated dataset consistently with the
parametrised model. We then used the simulated data as measurements and
performed posterior sampling. We quantified sampler performance as the
number of solver steps divided by the number of effective samples
generated, as plotted on logarithmic scale on the y axis. Lower values
indicate better performance. Vertical lines denote MCMC runs that were
unsuccessful because of a numerical solver failure.

\pandocbounded{\includegraphics[keepaspectratio]{assets/img/submission/performance.png}}

Dynamic guessing tended to improve MCMC performance compared with static
guessing for all the statistical models that we tested. The heuristic
\emph{guess-previous} showed similar or better performance compared with
\emph{guess-static}, whereas \emph{guess-implicit} performed
substantially better than \emph{guess-static} on every benchmark. It is
also notable that the dynamic algorithms failed less frequently than
\emph{guess-static} on the difficult \texttt{Rosenbrock8d} and
\texttt{Easom} benchmarks.

All of our results 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 adaptation phase. 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 a grapevine-augmented algorithm like grapeNUTS will
typically improve sampling performance for MCMC with embedded
root-finding, making it possible to fit previously infeasible
statistical models.

While dynamic guessing generally outperformed static guessing, the
relative performance of \emph{guess-previous} compared with
\emph{guess-implicit} varied between benchmarks. We expect that this
variation was caused by differences in how smoothly the solution of the
embedded problem changes with changes in parameters. The smoother this
relationship, the more likely that the embedded problems at adjacent
points in an HMC trajectory will have similar solutions, leading to
better relative performance of the \emph{guess-previous} heuristic.

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.

While our implementation of the grapevine method is performant and
flexible, it 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. In our implementation
there is no 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 a stricter
grapevine 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 \(p_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_{\text{guess}}\)
\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_{\text{guess}} \leftarrow \text{Heuristic}(\theta, \text{info})\)
\item
  \(x \leftarrow \text{Solve}(x_{\text{guess}})\)
\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
  \(U \leftarrow -lp\) (Potential energy is the negative log density)
\item
  \(\nabla_{\theta} U \leftarrow -\nabla_{\theta} lp\) (Compute the
  gradient of the potential energy)
\item
  Return \(U, \nabla_{\theta}U\) and \(\text{info}_{+}\)
\end{enumerate}

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

\begin{enumerate}
\def\labelenumi{\arabic{enumi}.}
\tightlist
\item
  \(U_{0}, \nabla_{\theta} U_{0}, \text{info} \leftarrow \text{PotentialGradientAndInfo}(\theta_{0}, \text{info}_{default})\)
\item
  Return
  \(\theta_{0}, p_{0}, U_{0}, \nabla_{\theta} U_{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, p, U, \nabla_{\theta} U, \text{info}\)
  \item
    step size \(\epsilon\)
  \end{itemize}
\end{itemize}

\begin{enumerate}
\def\labelenumi{\arabic{enumi}.}
\tightlist
\item
  \(p_{\text{mo}} \leftarrow p - \frac{\epsilon}{2} \nabla_{\theta} U\)
  (Update momentum, first half-step)
\item
  \(\theta_{+} \leftarrow \theta + \epsilon p_{\text{mo}}\) (Update
  parameters)
\item
  \(U_{+}, \nabla_{\theta} U_{+}, \text{info}_{+}  \leftarrow \text{PotentialGradientAndInfo}(\theta_{+}, \text{info})\)
\item
  \(p_{+} \leftarrow p_{\text{mo}} - \frac{\epsilon}{2} \nabla_{\theta} U_{+}\)
  (Update momentum, second half-step)
\item
  Return
  \(\theta_{+}, p_{+}, U_{+}, \nabla_{\theta} U_{+}, \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_{prev}\), previous parameters \(\theta_{prev}\)
and current parameters \(\theta_{next}\):

\[
\text{guess-implicit}(x_{prev}, \theta_{prev}, \theta_{next}) = x_{prev} + \frac{dx}{d\theta}\Delta\theta
\]

where \(\Delta\theta=(\theta_{next}-\theta_{prev})\).

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}f(x_{prev}, \theta_{prev}))^{-1}\operatorname{jac}_{\theta}f(x_{prev}, \theta_{prev})
\]

In this expression the term

\[
\operatorname{jac}_{x}f(x_{prev},\theta_{prev})
\]

abbreviated below to \(J_{x}\), indicates the jacobian with respect to
\(x\) of \(f(x_{prev}, \theta_{prev})\). Similarly

\[\operatorname{jac}_{\theta}f(x_{prev}, \theta_{prev}) = J_{\theta}\]

is the jacobian with respect to \(\theta\) of
\(f(x_{prev}, \theta_{prev})\).

Substituting terms we then have

\[
\text{guess-implicit}(x_{prev}, \theta_{prev}, \theta_{next}) = x_{prev} - 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}

We chose these functions 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.

\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. This kind of model is common in many
fields, especially biochemistry: see for example
{[}\citet{matosGRASPComputationalPlatform2022}{]}\citep{fiedlerTailoredParameterOptimization2016}\citep{grovesBayesianRegressionFacilitates2024}.

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:

\pandocbounded{\includegraphics[keepaspectratio]{assets/img/submission/mermaid-diagram-2026-01-15-123253.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 \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}
