Adaptive multi-stage integration schemes for Hamiltonian Monte Carlo

TL;DR

The paper addresses the sensitivity of Hamiltonian Monte Carlo (HMC) performance to the choice of numerical integrator by introducing s-AIA, an adaptive method that selects the most suitable 2- or 3-stage palindromic splitting integrator based on burn-in data. Extending the original AIA framework from molecular dynamics to Bayesian inference, s-AIA estimates a system-specific stability interval and a fitting factor (S_f or S_omega) to guide integrator selection for any step size within that interval, without adding simulation overhead. Implemented in the HaiCS package, s-AIA demonstrates superior sampling efficiency and energy conservation compared with fixed-parameter multi-stage schemes across Gaussian, German BLR, Musk BLR, and SIR benchmarks, with 3-stage schemes typically outperforming 2-stage ones. The approach offers practical gains in HMC accuracy and convergence by adapting the integrator to the problem, potentially reducing manual tuning and improving scalability in Bayesian analysis. Overall, s-AIA provides a principled, low-overhead mechanism to tailor Hamiltonian integration to the target distribution, enhancing the reliability and efficiency of HMC-based inference.

Abstract

Hamiltonian Monte Carlo (HMC) is a powerful tool for Bayesian statistical inference due to its potential to rapidly explore high dimensional state space, avoiding the random walk behavior typical of many Markov Chain Monte Carlo samplers. The proper choice of the integrator of the Hamiltonian dynamics is key to the efficiency of HMC. It is becoming increasingly clear that multi-stage splitting integrators are a good alternative to the Verlet method, traditionally used in HMC. Here we propose a principled way of finding optimal, problem-specific integration schemes (in terms of the best conservation of energy for harmonic forces/Gaussian targets) within the families of 2- and 3-stage splitting integrators. The method, which we call Adaptive Integration Approach for statistics, or s-AIA, uses a multivariate Gaussian model and simulation data obtained at the HMC burn-in stage to identify a system-specific dimensional stability interval and assigns the most appropriate 2-/3-stage integrator for any user-chosen simulation step size within that interval. s-AIA has been implemented in the in-house software package HaiCS without introducing computational overheads in the simulations. The efficiency of the s-AIA integrators and their impact on the HMC accuracy, sampling performance and convergence are discussed in comparison with known fixed-parameter multi-stage splitting integrators (including Verlet). Numerical experiments on well-known statistical models show that the adaptive schemes reach the best possible performance within the family of 2-, 3-stage splitting schemes.

Figures (15)

• Figure 1: Comparison of the upper bounds $\rho_k (h, b)$, $k = 2$\ref{['eq:rho2stage']}$,3$\ref{['eq:rho3stage']} of the energy error, for fixed-parameter multi-stage splitting integrators --- VV2, VV3, BCSS2, BCSS3, ME2, ME3 (Table \ref{['tab:IntegratorsTable']}) --- and the adaptive integrators AIA and s-AIA$k$. The interval for the step size $h$ is normalized with respect to the number of stages $k$ of the integrator in order to lead to fair comparisons. The zoomed plot in the upper left corner shows the situation for $h/k \in (0, 1.2)$.
• Figure 2: Comparison of the integration coefficient $b$ for fixed-parameter multi-stage splitting integrators ---VV2, VV3, BCSS2, BCSS3, ME2, ME3 (Table \ref{['tab:IntegratorsTable']})--- and the adaptive integrators AIA and s-AIA$k$, $b^2_{\text{opt}}$ and $b^3_{\text{opt}}$\ref{['eq:boptsAIA']}. The interval for the step size $h$ is normalized with respect to the number of stages $k$ of the integrator to lead to fair comparisons.
• Figure 3: Summary of the s-AIA$k$ algorithm. The proposed approach consists of three stages: (i) tuning stage for adjusting the step size $\Delta t_{\text{VV}}$ to get $\text{AR} \approx \alpha_{\text{target}}$ (\ref{['app:ARtarget']}); (ii) burn-in stage; the optimal multi-stage integrator and the HMC simulation parameters are found by combining the simulation data and the analysis provided; (iii) production stage to generate the HMC samples.
• Figure 4: Detailed schematic representation of the s-AIA$k$ algorithm.
• Figure 5: Frequency distributions of the benchmark models.