Faster parallel MCMC: Metropolis adjustment is best served warm

TL;DR

Faster parallel MCMC tackles the wall-clock bottleneck by using a two-phase strategy: an unadjusted warm-up that quickly drives an ensemble of $M$ chains toward the target, followed by Metropolis adjustment for precise convergence. The method, LAPS, automatically tunes hyperparameters via Ensemble Chain Adaptation using a bias proxy from equipartition and an energy-based estimate of asymptotic bias, and automatically switches to adjustment once convergence stalls. Empirically, LAPS outperforms state-of-the-art turnkey ensemble samplers (e.g., MEADS, ChESS) and even rivals sequential methods like NUTS in wall-clock time, often by orders of magnitude, while remaining practical to deploy. The approach is implemented in Blackjax/JAX and is compatible with probabilistic programming frameworks, offering a scalable, out-of-the-box solution for fast Bayesian inference on modern hardware.

Abstract

Despite the enormous success of Hamiltonian Monte Carlo and related Markov Chain Monte Carlo (MCMC) methods, sampling often still represents the computational bottleneck in scientific applications. Availability of parallel resources can significantly speed up MCMC inference by running a large number of chains in parallel, each collecting a single sample. However, the parallel approach converges slowly if the chains are not initialized close to the target distribution (cold start). Theoretically this can be resolved by initially running MCMC without Metropolis-Hastings adjustment to quickly converge to the vicinity of the target distribution and then turn on adjustment to achieve fine convergence. However, no practical scheme uses this strategy, due to the difficulty of automatically selecting the step size during the unadjusted phase. We here develop Late Adjusted Parallel Sampler (LAPS), which is precisely such a scheme and is applicable out of the box, all the hyperparameters are selected automatically. LAPS takes advantage of ensemble-based hyperparameter adaptation to estimate the bias at each iteration and converts it to the appropriate step size. We show that LAPS consistently and significantly outperforms ensemble adjusted methods such as MEADS or ChESS and the optimization-based initializer Pathfinder on a variety of standard benchmark problems. LAPS typically achieves two orders of magnitude lower wall-clock time than the corresponding sequential algorithms such as NUTS.

Figures (9)

• Figure 1: LAPS convergence for the German Credit problem. Left panel: bias of the second moments relative to the ground truth, $b^2_{\mathrm{max}}$ (red) and $b^2_{\mathrm{avg}}$ (blue), equipartition loss (blue dots) and the second-moment fluctuations (dotted lines) are shown as a function of the number of gradient evaluations. The latter two can be estimated when the ground truth is not known. In this example, they reflect the actual bias quite well. The switch from the unadjusted to the adjusted phase is shown by a vertical line. In this case, the unadjusted phase already comes very close to the desired accuracy, but adjusted method further improves on that. Note that the final ensemble is not biased, the residual error is caused by the finite number of chains. Pathfinder's second-moment bias is shown with stars. Right panels: convergence of hyperparameters step size and L is shown in the lower two panels. Step size is tuned based on the quantities derived from the energy error, shown in the upper panel. EEVPD is used in the unadjusted phase, acceptance rate in the adjusted phase. Targeted EEVPD and acceptance rate are shown in grey. Acceptance rate is fixed, while targeted EEVPD is adaptively changed, based on the current bias estimate from the equipartition loss. Note that the step size in the adjusted phase converges within a few steps, thanks to the bisection method.
• Figure 2: Second-moment bias as a function of the number of gradient evaluations on the Stochastic Volatility problem. LAPS converges significantly faster than the state-of-the-art ensemble samplers, which are shown for comparison. Note that in this case, the unadjusted phase does not quite achieve the desired accuracy requirements and the adjusted phase is essential for efficient convergence.
• Figure 3: LAPS posterior on the Banana target, as a function of the number of gradient calls (per chain). The location of chains is shown in red, the target density in black. As can be seen, LAPS converges to the posterior in just around 20 gradient calls.
• Figure 4: Bias as a function of the number of gradient evaluations with the unadjusted MCLMC dynamics for the German Credit problem. Adaptive step size adaptation used in LAPS (yellow) is compared with a fixed step size strategy for various step size choices. Large step size strategy (red) is initially very fast, but is not able to achieve fine convergence. Small step size (blue) is slow to converge but has a smaller asymptotic bias. LAPS adaptively reduces the step size, combining the best of all worlds.
• Figure 5: Schematic of the Ensemble Chain Adaptation (ECA) algorithm that we adopt. $M$ chains are evolved, time goes in the horizontal direction. In each step, the locations of all the chains are used to compute expectation values (orange arrow) which are used to update the hyperparameters (red arrow). Updated hyperparameters specify the kernel (pale purple arrow) which updates the chains in parallel (solid purple arrows).

Theorems & Definitions (2)

• Theorem 1
• proof