Adaptive tuning of Hamiltonian Monte Carlo methods

TL;DR

This work tackles the persistent challenge of tuning Hamiltonian Monte Carlo (HMC) hyperparameters for efficient Bayesian inference. It introduces Adaptive Tuning (ATune), which automatically discovers system-specific integrators and hyperparameters to produce AT-HMC and Generalized AT-HMC (AT-GHMC) with no production overhead, leveraging burn-in data and s-AIA-based integration. A key innovation is the joint optimization of the integration scheme, step-size randomization, and partial momentum update, including principled randomization intervals for both $\Delta t$ and $\varphi$, plus an automated selection of trajectory lengths. Empirical results across Gaussian benchmarks and three real-world applications (breast cancer BLR, cell-cell adhesion PDEs, and Influenza A outbreaks) show that AT-GHMC substantially outperforms standard HMC and state-of-the-art NUTS in stability and sampling efficiency, demonstrating broad practical impact for complex Bayesian models.

Abstract

With the recently increased interest in probabilistic models, the efficiency of an underlying sampler becomes a crucial consideration. A Hamiltonian Monte Carlo (HMC) is one popular option for models of this kind. Performance of HMC, however, strongly relies on a choice of parameters associated with an integration method for Hamiltonian equations, which up to date remains mainly heuristic or introduces time complexity. We propose a novel computationally inexpensive and flexible approach (we call it Adaptive Tuning or ATune) that, by combining a theoretical analysis of the multivariate Gaussian model with simulation data generated during a burn-in stage of HMC, detects a system specific splitting integrator with a set of reliable HMC hyperparameters, including their credible randomization intervals, to be readily used in a production simulation. The method automatically eliminates those values of simulation parameters which could cause undesired extreme scenarios, such as resonance artefacts, low accuracy or poor sampling. The new approach is implemented in the in-house software package HaiCS, with no computational overheads introduced in a production simulation, and can be easily incorporated in any package for Bayesian inference with HMC. The tests on popular statistical models reveal the superiority of adaptively tuned HMC and generalized Hamiltonian Monte Carlo (GHMC) in terms of stability, performance and accuracy over conventional HMC tuned heuristically and coupled with the well-established integrators. We also claim that the generalized formulation of HMC, i.e. GHMC, is preferable for achieving high sampling performance. The efficiency of the new methodology is assessed in comparison with state-of-the-art samplers, e.g. NUTS, in real-world applications, such as endocrine therapy resistance in cancer, modeling of cell-cell adhesion dynamics and influenza A epidemic outbreak.

Figures (12)

• Figure 1: Best grad/minESS performance ($\text{gESS}^{\bigstar}$) obtained for the tested benchmarks -- G1000 (mauve), German (dark red) and Musk (turquoise) (Table \ref{['tab:MetricsAroundHSL']}). $k$ is the number of stages of an integrator. For any benchmark, top performance is achieved around the CoLSI, i.e. $h = k$.
• Figure 2: The upper bound $\rho_{\text{3-stage}} (h, b)$ of the expected energy error for 3-stage integrators.
• Figure 3: Comparison of $\text{gESS}^{\bigstar}$ performance (Table \ref{['tab:MetricsAroundHSL']}) with $\text{gESS}$ observed in HMC (dashed bars) and GHMC (filled bars) with $\Delta {t} \sim \mathcal{U} (\Delta {t_\text{lower}}, \Delta {t_\text{ColSI}})$ for G1000 (left), German (center) and Musk (right) benchmarks using s-AIA3 (green), BCSS3 (blue), ME3 (orange), VV (red) integrators. GHMC combined with s-AIA3 (green filled bars) always outperforms $\text{gESS}^{\bigstar}$ in Table \ref{['tab:MetricsAroundHSL']} and demonstrates best performance in all three benchmarks. HMC (green dashed bars) also shows the best performance with s-AIA3, which is close to $\text{gESS}^{\bigstar}$ in Table \ref{['tab:MetricsAroundHSL']} but not necessarily better.
• Figure 4: Plot of $\mathcal{K}(h)$ (Eq. \ref{['eq:Kh']}) for $h$$\in$$(h_\text{lower}, h_\text{CoLSI})$\ref{['eq:h_BCSS3maxtoHSL']}.
• Figure 5: Schematic representation of the ATune algorithm for generating a set of optimal parameters for a GHMC/HMC simulation. The features specific to AT-GHMC only are highlighted in blue, whereas the features optional for AT-HMC are displayed in orange. Optimal parameters are shown in magenta.