Fast Riemannian-manifold Hamiltonian Monte Carlo for hierarchical Gaussian-process models (2511.06407v1)

Abstract: Hierarchical Bayesian models based on Gaussian processes are considered useful for describing complex nonlinear statistical dependencies among variables in real-world data. However, effective Monte Carlo algorithms for inference with these models have not yet been established, except for several simple cases. In this study, we show that, compared with the slow inference achieved with existing program libraries, the performance of Riemannian-manifold Hamiltonian Monte Carlo (RMHMC) can be drastically improved by optimising the computation order according to the model structure and dynamically programming the eigendecomposition. This improvement cannot be achieved when using an existing library based on a naive automatic differentiator. We numerically demonstrate that RMHMC effectively samples from the posterior, allowing the calculation of model evidence, in a Bayesian logistic regression on simulated data and in the estimation of propensity functions for the American national medical expenditure data using several Bayesian multiple-kernel models. These results lay a foundation for implementing effective Monte Carlo algorithms for analysing real-world data with Gaussian processes, and highlight the need to develop a customisable library set that allows users to incorporate dynamically programmed objects and finely optimises the mode of automatic differentiation depending on the model structure.

• The paper presents a structure-aware RMHMC approach that significantly improves sampling efficiency for hierarchical Gaussian process models.
• It leverages a soft-absolute Hessian metric and dynamic eigendecomposition to reduce computational complexity and enhance posterior exploration.
• Empirical results demonstrate rapid mixing and accurate Bayesian marginal likelihood estimation compared to conventional methods like NUT-HMC.

Efficient RMHMC for Hierarchical Gaussian Process Models

Introduction

The presented work develops a highly efficient implementation of Riemannian-manifold Hamiltonian Monte Carlo (RMHMC) for inference in hierarchical Gaussian process (GP) models, emphasizing practical strategies to overcome computational bottlenecks. Hierarchical GP-based Bayesian models are frequently employed for modeling complex nonlinear statistical dependencies but pose significant challenges for Markov Chain Monte Carlo (MCMC) inference due to non-log-concave posteriors and high-dimensional parameter spaces. Existing RMHMC algorithms, while theoretically promising, suffer from poor scaling and performance in these hierarchical contexts, especially when implemented naively using general-purpose automatic differentiation libraries.

Hierarchical GP Model Representation and Challenges

Hierarchical GP models are formulated such that the likelihood for each data sample is parameterized by multiple nonlinear functions, each modeled as sums of GPs with kernel hyperparameters. Real-world applications—such as modeling conditional means and variances in medical expenditure data—necessitate multiple GPs and non-Gaussian likelihoods, thus presenting posteriors with significant curvature heterogeneity, directionality, and non-log-concavity. Efficient sampling from such posteriors is impractical using vanilla HMC or conventional implementations of RMHMC.

A key technical challenge is efficient representation and sampling when the parameterization includes both function values (often represented via random Fourier features) and kernel hyperparameters. Naive differentiation approaches for the Hamiltonian, especially for metrics involving higher-order derivatives, incur prohibitive $O(Nd^3)$ costs for sample size $N$ and parameter dimension $d$.

RMHMC with Soft-Absolute Hessian Metric

The RMHMC implementation utilizes the soft-absolute transformation of the negative log-posterior Hessian as the metric, enabling applicability to non-log-concave posteriors. Given the Hessian $H(q)$, the metric $G(q)$ is constructed via spectral decomposition, replacing eigenvalues $\lambda_i$ with $g(\lambda_i) \approx |\lambda_i|$, and preserving eigenvectors. This choice guarantees positive-definiteness and aligns leapfrog proposal directions with the intrinsic geometry of the posterior.

Critical to this work is the adoption of Betancourt’s formula for Hamiltonian derivatives, which leverages cached eigensystem computations and the structure of the Hessian to keep per-step complexity to $O(d^3)$. However, the order of computation and caching strategy are essential; automatic differentiation libraries do not exploit this structure, resulting in severe redundancy.

Structure-Aware Redundancy Removal and Dynamic Eigendecomposition

The implementation reorganizes gradient and metric computations according to model structure—identifying where forward or reverse mode differentiation is optimal, exploiting sparsity, batching, and matrix operation efficiencies. For model dimensions $d$ scaling as $N^r$ (with $r$ dependent on the number of GPs and features), the computational complexity is reduced to $O(d^{\omega(r)} + d^{1+r})$ for mathematical exponent $\omega(r)$, compared to the naive $O(Nd^3)$.

Dynamic programming for eigendecomposition further accelerates RMHMC. Rather than decomposing the Hessian from scratch at each leapfrog step, the eigensystem is updated via Jacobi method or related approaches, exploiting the temporal continuity between successive parameter states. For practical high-dimensional models, dynamic eigendecomposition reduces wall-clock time, especially as model size $d$ increases, by minimizing the number of sweeps and leveraging GPU parallelism.

Numerical Results

Empirical comparisons reveal that the proposed implementation is an order of magnitude faster than RMHMC based on Hamiltorch (a PyTorch-based library), confirming the dominance of structure-dependent gradient optimization over eigendecomposition improvements for small and moderate $d$. For large models, dynamic eigendecomposition provides additional speed-ups. The RMHMC implementation achieves rapid mixing, effective exploration, and reliable computation of model evidence (Bayesian marginal likelihood, BME) in both synthetic logistic regression tasks and substantive causal inference applications on the NMES medical expenditure survey data.

When compared to NUT-HMC (No-U-Turn HMC with Euclidean metric), RMHMC avoids getting trapped in narrow, high-density regions—demonstrating that data-adaptive metrics are critical for exploration of hierarchical GP posteriors. Moreover, observed posteriors are singular, with non-Gaussian curvature, invalidating Laplace approximations for BME computation; nevertheless, the MC-integrated RMHMC reliably estimates BME within statistically acceptable margins.

Implications and Future Directions

This work establishes that RMHMC sampling in hierarchical GP models is only viable with structure-aware implementations that leverage both analytical simplifications and hardware-specific optimizations (such as GPU parallelism and dynamic matrix decomposition). The finding that automatic differentiation is insufficient without custom computational ordering has broad implications: future software libraries for probabilistic inference must expose lower-level primitives or optimization strategies to support efficient RMHMC, particularly for models with complex dependencies and non-Gaussian likelihoods.

The demonstrated efficiency also opens the door to the application of RMHMC-based inference in large-scale real-world domains, including biostatistics and econometrics, where unbiased posterior averaging is required. The structure-aware enhancements can be extended further, for instance, by combining reversible jump mechanisms for variable selection or integrating RMHMC within sequential Monte Carlo frameworks for time-series models.

However, the work is primarily restricted to sum-of-low-dimensional GPs and one-dimensional kernels; generalization to deep kernel learning or high-dimensional nonstationary models will require additional advancements, particularly in scalable eigendecomposition and gradient handling. The potential integration with doubly-robust debiased machine learning remains speculative due to singularity concerns in model posteriors.

Conclusion

The paper provides a comprehensive solution to the longstanding computational inefficiency in RMHMC sampling for hierarchical GP models, combining theoretical insight with practical optimization—most notably, structure-dependent computation and dynamic eigensystem maintenance. These innovations directly address scalability and accuracy challenges for MCMC inference in hierarchical Bayesian modeling, forming a foundational basis for future algorithmic and software development in probabilistic machine learning.