Adaptive tuning of Hamiltonian Monte Carlo methods

Hamiltonian Monte Carlo [1] is a Markov Chain Monte Carlo (MCMC) method for obtaining $N$ correlated samples $\{\bm{\theta}_{n}\}_{n=1,...,N}$ from a target probability distribution $\pi(\bm{\theta})$ in $\mathbb{R}^{D}$, $D$ being the dimension of $\bm{\theta}$. HMC generates a Markov chain in the joint phase space $\Omega\subseteq\mathbb{R}^{2D}$ with invariant distribution

and recovers $\pi(\bm{\theta})$ by marginalizing out the auxiliary momentum variable $\bm{p}$ (distributed according to $n(\bm{p})$).

In (1), $H(\bm{\theta},\bm{p})$ is the Hamiltonian

where $M$ is the symmetric positive definite mass matrix, and $U(\bm{\theta})$ is the potential energy, which contains information about the target distribution $\pi(\bm{\theta})$:

HMC generates new proposals $(\bm{\theta}^{\prime},\bm{p}^{\prime})$ by alternating momentum update steps, where $\bm{p}$ is drawn from a Gaussian $\mathcal{N}(0,M)$, with the steps that update position $\bm{\theta}$ and momenta $\bm{p}$ through the numerical integration of the Hamiltonian equations

performed $L$ times using an explicit symplectic and reversible integrator $\Psi_{h}$ [43]. Thus,

Symplecticity and reversibility of $\Psi_{h}$ ensure that the generated Markov chain is ergodic, thus that $\pi(\bm{\theta},\bm{p})$ is an invariant measure for the generated Markov chain.

Due to numerical integration errors, the Hamiltonian energy and hence the target density (1) are not exactly preserved. To address this issue, the invariance of the target density is secured through a Metropolis test with the acceptance probability

where

is the energy error resulting from the numerical integration (4). In case of acceptance, $\bm{\theta}^{\prime}$ becomes the starting point for the subsequent iteration. Conversely, if the proposal is rejected, the previous state $\bm{\theta}$ is retained for the next iteration. Regardless of acceptance or rejection, the momentum variable is discarded, and a new momentum $\bm{p}$ is drawn from $\mathcal{N}(0,M)$.

Generalized Hamiltonian Monte Carlo [30] incorporates a Partial Momentum Update (PMU) [29] in standard Monte Carlo to replace the full refreshment of auxiliary momenta. The PMU reduces random-walk behavior by partially updating the auxiliary momenta $\bm{p}$ through mixing with an independent and identically distributed (i.i.d) Gaussian noise $\bm{u}\sim\mathcal{N}(0,M)$:

Here, $\varphi\in(0,1]$ is a random noise parameter which determines the extent to which the momenta are refreshed. When $\varphi=1$, the update (7) reduces to the standard HMC momentum resampling scheme.

Notice that the PMU produces configurations that preserve the joint target distribution, as the orthogonal transformation in (7) maintains the distributions of $\bm{p}\sim\mathcal{N}(0,M)$. Since the momenta are not fully discarded, GHMC performs a momentum flip after each rejected proposal. As shown by Neal [2] and formally proved by Fang et al. [37], the three core steps of GHMC (Hamiltonian dynamics combined with the Metropolis test, the PMU, and momentum flip) each individually satisfy detailed balance and maintain reversibility, while their combination does not. This means that the resulting Markov chains are not reversible. Nevertheless, Fang et al. [37] demonstrated that GHMC fulfils the weaker modified detailed balance condition, ensuring the stationarity of the generated Markov chains.

The most commonly used numerical scheme for solving the Hamiltonian equations (3) is the Velocity Verlet integrator [6, 7]. This scheme is widely favored due to its simplicity, good stability, and energy conservation properties. The algorithm reads

Here, $h$ represents the integration step size. The updates of momenta and position in (8) are often called momentum kick and position drift, respectively.

More sophisticated integration schemes have been proposed in literature, including multi-stage palindromic splitting integrators [44, 10, 45]. These methods perform multiple splits of the separable Hamiltonian system (3), applying a sequence of kicks and drifts (8) with arbitrary chosen step sizes, in contrast to the half-kick, drift, half-kick structure of the standard Verlet method. By denoting a position drift and a momentum kick, respectively, with

the family of $k$-stage splitting integrators is defined as [16]

if $k=2k^{\prime}$, and

if $k=2k^{\prime}-1$. The coefficients $b_{i}$, $a_{j}$ in (10)–(11) have to satisfy the conditions $2\sum_{i=1}^{k^{\prime}}b_{i}+b_{k^{\prime}+1}=2\sum_{j=1}^{k^{\prime}}a_{j}=1$, and $2\sum_{i=1}^{k^{\prime}}b_{i}=2\sum_{j=1}^{k^{\prime}-1}a_{j}+a_{k^{\prime}}=1$, respectively, to ensure an integration step of length $h$. A term stage ($k$) refers to the number of gradient evaluations per integration. It is straightforward to verify that the only 1-stage splitting integrator of the form (11) is the Velocity Verlet scheme (8), which can be expressed as

Various choices of integration coefficients $b_{i}$, $a_{j}$ in (10)–(11) yield integrators that reach their pick performance at different specific values of a simulation step size $h$. To address the step size dependency, adaptive integration approaches AIA [11] and s-AIA [14] were developed for molecular simulations and Bayesian inference applications, respectively. For a given system and simulation step size $h$, these methods select the most suitable 2- (AIA, s-AIA2) and 3- (s-AIA3) stage splitting integrator ensuring optimal energy conservation for harmonic forces. Examples of 2- and 3-stage integrators, including those used in this study, are reviewed in [10, 45, 14].

To identify endpoints of a randomization interval for an optimal step size we refer to the upper bound of the expected energy error for 3-stage integrators $\rho_{3}(h,b)$ derived in [14] and plotted for the range of integrators considered in this study in Figure 2.

When searching for an optimal randomization interval for $h_{\text{CoLSI}}$ the following considerations have to be taken into account. First, the interval should include the step sizes that guarantee the high accuracy of the numerical integration, i.e. $\rho_{3}(h,b)\ll 1$. Then, one should keep interval narrow enough to stay close to an estimated optimal value of a step size and to avoid including too small step sizes, which usually satisfy the aforementioned inequality but lead to correlated samples in HMC. Finally, the approximate analysis proposed in [14] may likely result in the overestimated values of the stability limit and thus $\Delta{t_{\text{CoLSI}}}$ (see A). This implies a reasonable choice for the right endpoint to be $h_{\text{CoLSI}}$.

On the other hand, we recall that BCSS3 was derived to minimize the maximum of $\rho_{3}(h,b)$ in $0<h<h_{\text{CoLSI}}$ and thus to achieve the best performance near $h_{\text{CoLSI}}$. The inspection of Figure 2 reveals that $\rho_{3}(h,b)$ with $b=b_{\text{BCSS3}}$ exhibits a local maximum $h_{\text{BCSS3max}}$ on the left side of $h_{\text{CoLSI}}$. Within the interval ($h_{\text{BCSS3max}}$, $h_{\text{CoLSI}}$) both s-AIA3 and BCSS3 demonstrate very high accuracy, which deteriorates after $h$ passing the center of the stability interval. Thus, such an interval obeys all requirements on an optimal randomization interval summarized above.

The local maxima of $\rho_{\text{BCSS3}}(h)$, i.e. $h_{\text{BCSS3max}}$, can be obtained by setting $\frac{\partial\rho_{3}(h,b)}{\partial h}=0$ and $b=b_{\text{BCSS3}}$ [10] (see B) which yields $h_{\text{lower}}\equiv h_{\text{BCSS3max}}\approx 2.0772$.

Therefore, the proposed randomization interval for the nondimensional optimal integration step size is

And the corresponding interval for the dimensional step size is

where $\Delta{t_{\text{lower}}}$ and $\Delta{t_{\text{CoLSI}}}$ are calculated as proposed in [14] and briefly described in C.

To validate the proposed optimal randomized step size and to choose an optimal integrator among the 3-stage integrators, we run HMC simulations with s-AIA3, BCSS3, ME3 and VV at $\Delta{t}$ in (14) (in 1-stage units for VV) for all benchmarks introduced in Section 3 and listed in Table 2. The remaining HMC parameters, namely, a number of integration steps per Monte Carlo iteration, $L$, and numbers of burn-in and production iterations, are kept the same as in the numerical experiments in [14]. In addition, we conduct GHMC simulations to evaluate the efficacy of the proposed randomization interval with this sampler in comparison with HMC. For the PMU in Eq. (7), we set $\varphi$ to be uniformly selected from the interval $(0,0.5)$ (a standard recommendation) and use twice shorter trajectories than for HMC. We remark that this choice of GHMC parameters is based on heuristic intuition. The optimal selection of these parameters are investigated in the following Sections.

The simulation parameters for HMC and GHMC used in this section are summarized in Table 3.

Figure 3 compares performance (in terms of gESS) of HMC and GHMC using $\Delta{t_{\text{CoLSI}}}$ and a randomization interval (14) with $\text{gESS}^{\bigstar}$ in Table 2 through $\text{gESS}^{\bigstar}$/gESS relations.

First, we recognize that, as expected, for all benchmarks and for both HMC and GHMC, the adaptive s-AIA3 integrator yields either better or similar sampling performance (in terms of gESS) compared with other tested integrators. Moreover, for GHMC, this performance is visibly higher (up to 2x) than gESS observed in HMC with the newly proposed settings (see Table 3) and with the previously reported setup (Table 2). We notice that the new settings for HMC with s-AIA3 do not necessarily improve the results in Table 2 for each benchmark, but they always stay close to them. We recall that the $\text{gESS}^{\bigstar}$ values in Table 2 for each benchmark were achieved with various integrators and at different step sizes by trial and error, whereas our new settings were generated by the procedure universal for all simulated systems.

In summary, the numerical experiments not only justify the proposed optimal choices for a numerical integrator (s-AIA3), a step size ($h_{\text{CoLSI}}$ / $\Delta{t_{\text{CoLSI}}}$), and a randomization interval ((13) / (14)), which we will consider for the rest of this study, but also reveal the performance superiority of GHMC over HMC.

We remark that, in contrast to the HMC experiments which relied on the recommended choices of $L$, the results for GHMC in Figure 3 were obtained without a proper tuning of $L$ and its additional parameter $\varphi$ for the PMU. Obviously, refining these two parameters can further improve performance of GHMC.

Our next objective is to search for optimal, or close to optimal choices of the remaining GHMC parameters.

The irreversibility of GHMC may allow it to reach convergence faster, compared to reversible samplers, like HMC. To check this, we choose the performance metric that includes information about the speed of convergence of the Markov chains. The metric is briefly introduced in Section 3. Here we provide more details on the metric and extend it to the alternative approaches for computing effective sample size (ESS).

Following the recommendation in [50], we consider a chain to be converged if $\max\text{PSRF}<1.01$. The Potential Scale Reduction Factor (PSRF) [48] can be calculated as described in [49].

Starting from that, let us define the number of samples in the production chains, required for reaching $\max\text{PSRF}<1.01$, as $N_{1.01}$, and evaluate the ratio grad/ESS between the number of gradient computations and ESS for production chains of the length of $N_{1.01}+1000$ iterations. Such a metric combines both the information about convergence and sampling efficiency.

The number of gradient evaluations reads as

where $\overline{L}$ is the mean of the number of integration steps per Monte Carlo iteration (it varies depending on a chosen randomization scheme for $L$) and $k$ is the number of stages of an integrator in use (in our case, $k=3$). We recall that the number of stages $k$ indicates the number of times the algorithm performs an evaluation of gradients per step size. Therefore, $\overline{L}k$ represents the theoretical average number of gradient evaluations per Monte Carlo iteration.

In Section 3, we calculate ESS using the effectiveSize function of the coda package of R [51] and consider the minimum ESS across variates, i.e. minESS. We want to include two additional ESS metrics in our list, which now appears like:

• 1.

  minESS (coda) – the minimum ESS across variates

• 2.

  meanESS (coda) – the average ESS over the sampled variables calculated through the effectiveSize function of the coda package of R [51]

• 3.

  multiESS – the metric calculated in the multiESS function of the mcmcse package [56] using the approach proposed in [57].

To summarize, we propose three metrics for assessing convergence and sampling performance, such as grad/minESS, grad/meanESS, and grad/multiESS. For each metric, a lower gradient value (grad) indicates better convergence, while a higher ESS suggests improved sampling efficiency. Consequently, the smallest grad/ESS ratio represents the best overall performance.

To enable comparison of efficiency between two samplers, a Relative Efficiency Factor (REF) is introduced as:

The REF value quantifies the extent to which $\text{Sampler}_{1}$ outperforms $\text{Sampler}_{2}$ for a particular $\mathcal{M}=\text{grad}/\text{ESS}$ metric, where ESS is picked from the set {minESS, meanESS, multiESS}.

We use the the metrics described in 8.1 for the full set of simulation benchmarks (Table 4). For the biggest benchmark (G2000), we consider the ESS metrics evaluated over the first $N_{1.01}+2000$ iterations, in order to take into account enough iterations ($>D=2000$) for the multiESS calculation. We also remark that for Gaussian benchmarks the grad/ESS metrics for the simulations with $L=1$ were calculated over the first $\tilde{N}_{1.01}+1000$ ($2000$ for G2000), where $\tilde{N}_{1.01}$ is the number of iterations required to achieve $\text{avgPSRF}<1.01$. The relaxation of the PSRF threshold helped to avoid extremely long simulations ($N_{\text{prod}}\gg 10^{6}$) needed for reaching convergence in terms of $\max\text{PSRF}<1.01$.

In order to evaluate the optimal settings proposed in this study, we compare our recommended choice $\varphi\sim$ $\mathcal{U}(\varphi_{\text{lower}},\varphi_{\text{upper}})$ (21) with the fixed choice of $\varphi=1$ (which corresponds to standard HMC) and with three randomization intervals $\varphi\sim$ $\mathcal{U}(0,0.1)$, $\mathcal{U}(0,0.5)$, $\mathcal{U}(0,0.9)$. For each choice of $\varphi$, in addition to the proposed optimal randomization interval for $L$ (27), we consider five different $\mathcal{U}\{1,L_{\text{upper}}\}$ intervals (four for Banana), where the values of $L_{\text{upper}}$ are proportional to the dimension of the simulated system. The optimal randomization intervals for $\varphi$ and $L$ are summarized in Table 5.

As put forward at the end of Section 4.2, we use s-AIA3 as a numerical integrator for Hamiltonian dynamics and the step size randomization interval proposed in Section 4.1 (Eq. (14)).

Table 6 presents the performance comparison between GHMC with $\varphi\sim$ $\mathcal{U}(\varphi_{\text{lower}},\varphi_{\text{upper}})$ and standard HMC ($\varphi=1$) across the tested benchmarks, with varying values of $L_{\text{upper}}$. For any benchmark and for each metric, the best performance is consistently achieved by GHMC (shown in red). Combining GHMC with our proposed choice of $L_{\text{opt}}$ leads to the best GHMC results for the average metrics, i.e. grad/meanESS and grad/multiESS. Moreover, for non-Gaussian benchmarks (also, with smaller dimensions), grad/minESS shows the best results with such a choice of $L$. For Gaussian benchmarks, the best results for grad/minESS are reached for larger values of $L_{\text{upper}}$ – they are lower for GHMC (around $D$) and higher for HMC ($>\text{2}D$).

The results obtained with other tested randomization intervals for $\varphi$, $\mathcal{U}(0,0.1)$, $\mathcal{U}(0,0.5)$, $\mathcal{U}(0,0.9)$, support our choice for the optimal randomization interval ($\varphi_{\text{lower}}$, $\varphi_{\text{upper}}$) and can be examined in E.

Finally, in Figure 6, we provide performance comparison in terms of the REF (29) between GHMC with our optimal settings, i.e. AT-GHMC, and HMC with its best settings found heuristically (see Table 6) using all proposed ESS metrics and all tested benchmarks. For the Gaussian benchmarks, we also present the comparison at the $L_{\text{upper}}$ values that optimize the minESS metric in HMC and GHMC (see Table 6). Both plots in Figure 6 confirm the improvements (up to 34x and 10x, respectively) gained with our proposed GHMC settings across all metrics and benchmarks. They also demonstrate that greater improvements are achieved with the average metrics, i.e. meanESS and multiESS.

Estrogen-positive (ER+) breast cancer is the most common subtype of breast cancer, accounting for over 70% of all breast tumors diagnosed. Hormone therapies such as tamoxifen are the gold-standard treatment. However, a significant amount of patients develops resistance to the treatment, which constitutes a serious clinical problem. The development of this resistance is not yet well understood, and plenty of clinical and laboratory research aims to find which genes may be relevant to the problem.

For one such clinical study presented in [60], ER+ breast cancer patients treated with tamoxifen for 5 years were tested for the presence of a specific gene marker, SOX2. The objective of the study was to analyze the connection between SOX2 and the emergence of resistance to tamoxifen treatment. An overexpression of this gene was previously identified in bioinformatic analyses of resistant breast cancer cells as a key contributor for resistance, but clinical evidence was presented in [60] for the first time.

The dataset used in our work stems from that clinical trial. It contains numerical data for SOX2 values and 7 other gene-related covariates taken from 75 tumors samples. The recorded response variable was the event of relapse after 5 years. Notably, the SOX2 variable perfectly separates the dataset for this response variable, as all cases with SOX2 values below 3 respond well to the treatment and all those above this value present a resistant response. The appearance of separation implies that maximum likelihood estimates do not exist for some binary classification models such as logistic regression [61] in the classical frequentist interpretation. The perfect separation makes the Bayesian framework natural to model such a dataset, as it can provide estimates for the parameters involved without the need to use regularization techniques.

The model of choice is a Bayesian Logistic Regression, where the dataset has 75 observations and dimension $D=8$ (including intercept parameter). We test the effect of three priors, informative $\mathcal{N}(0,1)$, weakly informative $\mathcal{N}(0,2.5)$ (according to [62]) and low informative $\mathcal{N}(0,5)$.

We compare the performance of AT-GHMC proposed in this study against the No-U-Turn-Sampler (NUTS) [5], the optimized HMC sampler implemented in Stan [63] using its R interface rstan [64]. NUTS in Stan uses the leapfrog/Verlet integrator [6] and the original optimization routine to determine both a step size, $\Delta t$, and a number of Hamiltonian dynamics integration steps per iterations, $L$. Additionally, we perform numerical experiments with HMC using $L_{\text{opt}}$ from Eq. (27) (AT-HMC) and setting $L=L_{\text{NUTS}}$ (HMC${}_{L_{\text{NUTS}}}$). The latter is determined such that the integration leg $\Delta tL$ matches that found by NUTS.

All samplers are run with the same number of burn-in iterations, $N_{\text{burn-in}}$, to ensure a fair comparison. We assess performance using grad/ESS metrics introduced in Section 8.1. Small values of $N_{1.01}$ in NUTS are expected, given its optimization routine aims to bring PSRF close to 1. The simulation settings are summarized in Table 7.

Figure 7 compares the sampling performance of AT-GHMC, AT-HMC, and HMC${}_{L_{\text{NUTS}}}$ with NUTS using REF metrics (29) for $\text{ESS}\in\{\text{minESS},\text{meanESS},\text{multiESS}\}$. AT-GHMC consistently outperforms NUTS across all metrics and priors, achieving more than 8 times improvement (e.g., multiESS for the $\mathcal{N}(0,2.5)$ prior).

We emphasize that, the ATune algorithm does not offer an optimal choice of $L$ for HMC. The values of $L_{\text{opt}}$ (27) used in AT-HMC are the recommended values for GHMC. However, the strong meanESS and multiESS performance of AT-HMC in all experiments confirms their applicability in HMC. On the other hand, the relative efficiency of HMC${}_{L_{\text{NUTS}}}$ with respect to NUTS shows its strong dependence on a chosen prior. Specifically, HMC${}_{L_{\text{NUTS}}}$ demonstrates visible improvements over NUTS with the low informative prior $\mathcal{N}(0,5)$; the comparable performance with NUTS if the chosen prior is $\mathcal{N}(0,2.5)$, and the poor sampling efficiency with the highly informative prior $\mathcal{N}(0,1)$. For $\mathcal{N}(0,5)$, the improved performance can be attributed to the higher accuracy and longer optimal step size of the 3-stage integrator (s-AIA3) employed in HMC${}_{L_{\text{NUTS}}}$ compared to those achievable with the standard 1-stage Verlet integrator used in NUTS. On the contrary, the poor performance of HMC${}_{L_{\text{NUTS}}}$ with $\mathcal{N}(0,1)$ is likely caused by the underestimation of $\Delta t_{\text{CoLSI}}$ hinted by the comparison of its value with $\overline{\Delta t}$ (see Table 7). Indeed, the values are close, which is not what one expects when comparing optimal step sizes for 1- and 3- stage integrators. It is also not the case for $\mathcal{N}(0,5)$ and $\mathcal{N}(0,2.5)$, where optimal step sizes are longer than the NUTS optimal step sizes (less pronounced for $\mathcal{N}(0,2.5)$ though), and where AT-GHMC and AT-HMC have never been outperformed by NUTS. These observations are encouraging. They mean that even in the worst scenario, when the estimated optimal parameters are less accurate, they still guarantee the good sampling performance (see AT-GHMC and AT-HMC for $\mathcal{N}(0,1))$ and are not inferior to the optimal parameters offered by the state-of-the-art NUTS approach. Moreover, the sampling advantages of GHMC over HMC, enhanced by the stochastic procedure for finding its optimal settings, disregard the inaccuracy in the estimation of an optimal step size and, as a result, AT-GHMC demonstrates a significant improvement over NUTS, AT-HMC and HMC${}_{L_{\text{NUTS}}}$ as happened in the case of $\mathcal{N}(0,1)$. The fact that the performance of AT-HMC in this case is comparable with, or better than the NUTS performance also confirms the robustness of the proposed methodology.

Many phenomena on the cellular level such as tissue formation or cell sorting can be explained by introducing differential cell-cell adhesion [65, 66, 67, 68, 69, 70]. This phenomena has been used to explain pattern formation for zebrafish lateral lines [71, 72], cancer invasion modeling [66], and cell sorting in early fly brain [70] and early mice brain development [73]. When modeling such biological systems using PDEs, the adhesion can be represented by means of a non-local interaction potential [66, 65, 68, 69]. To enhance predictive power of such models, a careful calibration of the model parameters according to experimental lab data is required. This task can become increasingly difficult due to sparsity of available data or increasing complexity of the model (such as, for example, by the addition of an interaction term). The application of state-of-the-art methodologies for performing parameter estimation on PDE models is receiving lots of attention in the last years [74, 75, 76], and, in this context, our goal is to investigate the feasibility of using the methodology introduced in this work for the task.

To achieve this goal, we propose a benchmark PDE model:

which describes the time evolution of the density of cells $u(\mathbf{x},t)>0$, $\mathbf{x}\in\mathbb{R}^{d},t>0$ ($d=1,2,3$ is the spatial dimension), as governed by a series of potentials. In (30), $\Pi(u)$ is a density of the internal energy and $W(\mathbf{x})$ is an interaction potential (also called an interaction kernel) describing cell adhesion. The convolution ($W*u$) gives rise to non-local interactions. Depending on the explicit form of the potentials, they can describe a wide range of biological phenomena such as the formation of ligands through long philopodia, modeled by nonlocal attractive potentials, and volume constraints, modeled by strong repulsive nonlocal potentials or nonlinear diffusive terms. We refer to [77] for more information about the applications of these aggregation-diffusion models in the sciences. We consider the following explicit forms for $\Pi^{\prime}(u)$ and $W(\mathbf{x})$: