Fast Riemannian Manifold Hamiltonian Monte Carlo for Hierarchical Gaussian Process Models

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 applying the chain rule for the differentiation of the Hamiltonian in the optimal order determined by the model structure, and by dynamically programming the eigendecomposition of the Riemannian metric with the recursive update of the eigenvectors at the previous move. This improvement cannot be achieved when using a naive automatic differentiator included in commonly used libraries. 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 to finely optimise the mode of automatic differentiation depending on the model structure.

1. Introduction

In machine learning, Bayesian analysis based on Gaussian processes (GPs) associated with positive semidefinite kernels has been studied extensively from modelling, algorithmic, and theoretical points of view [1,2,3,4,5,6,7,8,9,10,11,12]. A GP serves as a prior process for generating nonlinear functions that describe relationships between variables. The goal of Bayesian analysis with GPs is to infer the posterior process for such nonlinear functions conditioned on given data or to estimate the average of quantities of interest over such a posterior process. Although the posterior process can be represented as a closed-form solution for the simplest problems (e.g., Chapter 6 of Ref. [13]), such a solution cannot be expected when GPs are used as building blocks for hierarchical models that capture complex statistical dependencies among variables in real-world data (e.g., GPLVM, GPDM, iWMM, and DGP [1,2,7,8]). Algorithms for inference with hierarchical GP models can be broadly divided into variational algorithms and Monte Carlo (MC) algorithms. Variational algorithms (e.g., Refs. [9,11]) are usually much faster than MC algorithms but introduce unavoidable bias due to discrepancies between the original and variational models. In contrast, MC algorithms (e.g., Refs. [10,14,15]) allow for precise computation of averages over the posterior, provided the algorithm is run until convergence, namely, asymptotically exact computation. The preferred algorithm type depends on the purpose of the computation. At the moment, however, all algorithms included in the commonly used program libraries designed for the inference with GP-based hierarchical Bayesian models (GPyTorch and GPFlow [16,17]) are variational or a hybrid of variational and MC approaches.

Apart from the developments in machine learning, closely related frameworks for modelling with GPs have been proposed in the other branches of statistics as well. With the early recognition of the correspondences between smoothing by splines and Bayesian estimation with Gaussian processes [18], researchers have proposed estimation methods for the latter framework. Since a stricter guarantee on the quality of estimation results tends to be required outside the machine learning community, MC algorithms are preferred to variational algorithms. For example, in biometric and econometric applications, biased results can be misleading and may adversely affect public health or the economy; therefore, the use of variational algorithms is not appropriate. Therefore, MC algorithms or a deterministic method termed “iterative nested Laplace approximation (INLA)” [19] are used for the estimation. Different types of Markov-chain MC (MCMC) algorithms, INLAs, and softwares for their easy implementation are available (for MCMC, different derivatives of BUGS [20] and Stan [21]; for INLA, R-INLA [22]). Depending on the model class, INLA has been shown to perform inference very fast (see Ref. [23] for the comparison between INLA and MCMC).

However, it is important to note that these GPs used in statistics are different from those employed in machine learning. A strong condition is imposed on these GPs in order to keep the correlation of the prior Gaussian process short-range [24]. The combination of such short-range GPs and independent observations results in a short-range posterior which is much easier to sample from with the currently available MC algorithms than the posteriors of the GPs in machine learning that are strongly correlated over data points. INLA also exploits this short-range correlation. INLA assumes that the posterior conditioned on a small number of hyperparameters is approximated well with a Gaussian distribution (i.e., the Laplace approximation), which is justified by the central limit theorem. INLA takes advantage of the sparseness of the precision matrix of this approximating Gaussian distribution due to the short-range correlation. These assumptions made in the statistical analysis with MCMC and INLA reveal their limitations. First, as we discuss further below, the posterior of the hierarchical GP models with long-range correlation can be neither efficiently sampled with commonly used MCMC nor efficiently approximated with the integral of Gaussian conditional distributions. Second, with regard to INLA, the Laplace approximation is not valid if the conditional posterior is singular or nearly singular. Singular models are known to be ubiquitous, ranging from simple Gaussian models where both the mean and variance are parameterised, to large complex hierarchical models [25]. Although INLA was extended to near-Gaussian cases [26], the use of INLA is not appropriate for the analysis of such potentially singular hierarchical Bayesian models and is severely limited to simple models such as the latent Gaussian model discussed in Ref. [24].

As we have discussed above, the literature on the application of MC algorithms to hierarchical GP-based models is limited in both machine learning and statistics. In this study, we seek an efficient MC algorithm for hierarchical GP models used in machine learning, in order to enhance the applicability of the models in the other statistical fields, especially in biometrics and econometrics. To our knowledge, a standard MC algorithm for this purpose has not yet been established. Svensson et al. [10] sampled a function from the posterior process using a closed-form solution for the GP component of their model conditioned on other parameters. Their formula can be applied only when the likelihood term is Gaussian and is thus not applicable to, for example, binary classification. Hensman et al. [27] introduced a non-Gaussian variational approximation of the posterior with the aid of inducing points sampled using Hamiltonian MC (HMC). Their method, however, is variational and not guaranteed to describe the posterior precisely, and its scalability with increasing numbers or dimensions of GPs remains unclear. Pandita et al. [15] introduced an adaptive sequential MC for GPs to solve problems in mechanical engineering. However, the authors used hundreds of CPU cores for inference, and the efficiency of their sampling scheme itself was not clearly demonstrated.

One fundamental issue in MC sampling from the posterior in a GP-based model is that the posterior is highly stretched or compressed along unknown directions in the parameter space used for sampling, namely, an unknown complex long-range correlation. To address this problem, Paquet and Fraccaro [14] implemented Riemannian manifold HMC (RMHMC) for GP models. With the Hessian metric they employed, RMHMC adjusts the direction of MC moves according to the curvature of the log posterior density surface. They demonstrated that RMHMC is much more efficient than the ordinary HMC method based on the Euclidean metric. However, their method is applicable only to cases with log-concave posterior densities, whereas hierarchical models that describe complex dependencies among variables typically have non-log-concave posterior densities.

To overcome this limitation, in this study, we aim to obtain an efficient implementation of RMHMC for hierarchical GP models using a soft-absolute Hessian metric introduced by Betancourt [28]. This metric transforms an indefinite Hessian into a positive semidefinite matrix that shares eigenvectors with the Hessian but has associated eigenvalues set close to the absolute values of the original Hessian eigenvalues. Although a general-purpose implementation of this algorithm [29] based on PyTorch  [30] is available, its performance is too poor to be used for sampling from the posterior of hierarchical Bayesian models with multiple GPs. We study the cause of this poor performance and identify two sources of redundancy. First, we find that the computation order of the differentiation of the Hamiltonian according to the chain rule crucially affects the performance. This implies that the performance of an implementation of the algorithm cannot be optimised as long as it relies on the automatic differentiator provided by high-level programming languages designed for performing machine learning with graphical processing units (GPUs), such as PyTorch and TensorFlow [30,31]. This limitation is serious when we take into account the current standard practice for developers of libraries and their users to apply these automatic differentiators to manually coded functions, in order to avoid manually coding the (higher-order) derivatives, a major source of program bugs. Second, we find that redundancy can be removed by dynamically programming eigendecomposition required at each step of RMHMC. Dynamic eigendecomposition has recently been shown to improve performance in matrix optimisation [32], and we find RMHMC with the modified Hessian metric to be as suitable for its application as this case. With these findings, we implemented the RMHMC based on the soft-absolute Hessian metric by coding programs with low-level languages, C++ and OpenACC [33,34], and we confirmed that this implementation performs the algorithm more than ten times faster than the existing library based on the high-level languages [29].

The article is organised as follows. In Section 2.1, we provide preliminaries. In the Section 2.1.1, we introduce our problem setting in which multiple GPs are used as building blocks for a hierarchical Bayesian model, employing the representation of GPs introduced by Solin et al. [35]. In Section 2.1.2, we provide a brief introduction to RMHMC with a soft-absolute Hessian metric. In Section 2.2, we identify two sources of redundancy described above. We show how performance deteriorates when reverse-mode automatic differentiation and a divide-and-conquer algorithm for the eigendecomposition of symmetric matrices in PyTorch are naively applied. In Section 3, we numerically demonstrate that our implementation considerably outperforms an implementation based on a general-purpose library. The details of the experimental methods are provided in Section 2.3. Specifically, in Section 3.1, we compare the performance of different implementations using toy examples of varying sizes and show that both sources of redundancy contribute to the improvement. In Section 3.2, we compare the performance of our implementation with that of no-U-turn (NUT)-HMC based on the Euclidean metric. We show that NUT-HMC becomes trapped within a very narrow region and fails to converge to the posterior, suggesting that the use of a data-adaptive metric is essential for efficient sampling. In Section 3.3, we show, using a toy example and a real-world dataset, that our implementation successfully calculates the marginal likelihood, (that is, the Bayesian model evidence) within a reasonable computation time. Although INLA performs this calculation faster on the toy examples, it is too slow to analyse real-world data using nonlinear GP models with non-log-concave posteriors. Finally, in Section 4, we discuss, based on these results, the advantages and limitations of the proposed method compared with existing approaches and suggest directions for future work. The mathematical notations used in this article are summarised in Table A1.

2. Methods and Algorithms

2.1. Preliminaries

2.1.1. Hierarchical GP Models and Their Approximate Representations

As we have stated in the Introduction, we aim to obtain an implementation of RMHMC that is applicable to hierarchical GP models with non-log-concave (conditional) posteriors. Although numerous such models have been proposed, the following simplest class of models that slightly extend basic GP models with log-concave (conditional) posteriors [6,36] would be most suitable for testing the performance of samplers.

Let us consider hierarchical GP models for which the samplewise likelihood $P\left(X_i\right|{\left\{f_j\right\}}_{1\le j\le J})$ of i.i.d. data $X_i\in {\mathbf{R}}^d$ (with sample index $i\in \{1,2,\dots ,N\}$) is parameterised by J nonlinear functions ${\left\{f_j\left(X_i\right)\right\}}_{1\le j\le J}$ generated from GPs through function U as

$$\begin{array}{c}\hfill P\left(X_i\right|{\left\{f_j\right\}}_{1\le j\le J})=exp(-U\left({\left\{f_j\left(X_i\right)\right\}}_{1\le j\le J}\right)).\end{array}$$

(1)

For instance, the simplest examples describing the relationship between a target variable $X_i^{\left(\mathrm{tgt}\right)}$ and covariates $X_i^{\left(\mathrm{cov}\right)}$ are as follows.

Example 1.

Covariate-dependent target mean and variance ($J=2$):

$$\begin{array}{c}\hfill X_i^{\left(\mathrm{tgt}\right)}=f_1\left(X_i^{\left(\mathrm{cov}\right)}\right)+{ϵ}_i,{ϵ}_i\sim \mathcal{N}(0,{\sigma }_0^2+exp\left(f_2\left(X_i^{\left(\mathrm{cov}\right)}\right)\right)),\end{array}$$

(2)

with fixed ${\sigma }_0^2(>0)$.

Example 2.

Nonlinear Gaussian mixture model ($J=4$):

$$\begin{array}{c}\hfill X_i^{\left(\mathrm{tgt}\right)}=(1-s_i)f_1\left(X_i^{\left(\mathrm{cov}\right)}\right)+s_i{f}_2\left(X_i^{\left(\mathrm{cov}\right)}\right)+{ϵ}_i,\\ \hfill {ϵ}_i\sim \mathcal{N}(0,{\sigma }_0^2+exp\left(f_3\left(X_i^{\left(\mathrm{cov}\right)}\right)\right)),s_i\sim \mathrm{Bernoulli}[1/(1+exp\left(f_4\left(X_i^{\left(\mathrm{cov}\right)}\right)\right))].\end{array}$$

(3)

In the above modelling, we assume that each of ${\left\{f_j\right\}}_{1\le j\le J}$ is the sum of nonlinear functions generated by GP priors ${\left\{{\mathcal{G}}_{jk}\right\}}_{1\le j\le J,k\in {\mathcal{K}}_j}$:

$$\begin{array}{c}\hfill f_j=\sum _{k\in {\mathcal{K}}_j}{f}_{jk}+b_j,f_{jk}\sim {\mathcal{G}}_{jk}\left({\left\{{\theta }_{jk\ell }\right\}}_{\ell \in {\Theta }_{jk}}\right),b_j\sim \mathcal{N}(0,\Sigma ).\end{array}$$

(4)

Each GP ${\mathcal{G}}_{jk}$ is associated with a positive semidefinite kernel $k(·,·)$ that describes the prior covariance of the generated functions and is parameterised by scalar hyperparameters ${\left\{{\theta }_{jk\ell }\right\}}_{\ell \in {\Theta }_{jk}}\in {\mathbf{R}}^{|{\Theta }_{jk}|}$. Thus, each of ${\left\{f_j\right\}}_{j\in \mathcal{J}}$ is a Bayesian multiple-kernel model, the statistical properties of which were investigated in a previous study [6]. The prior density of the hyperparameters ${\Pi }_{jk}\left({\left\{{\theta }_{jk\ell }\right\}}_{\ell \in {\Theta }_{jk}}\right)$ is assumed to be given.

As we jointly sample functions ${\left\{f_{jk}\right\}}_{j,k}$ and hyperparameters ${\left\{{\theta }_{jk\ell }\right\}}_{j,k,\ell }$, the complicated dependence of the prior density of ${\left\{f_{jk}\left(X_i\right)\right\}}_{j,k,i}$ on ${\left\{{\theta }_{jk\ell }\right\}}_{j,k,\ell }$ hinders efficient sampling. To alleviate this problem, one previous study [35] introduced a useful approximation scheme that took the following form:

$$\begin{array}{c}\hfill f_{jk}\left(X\right)\approx \sum _{m=1}^{M_{jk}}{a}_{jkm}{\varphi }_{jkm}\left(X\right),a_{jkm}\sim \mathcal{N}(0,c_{jk}{V}_{jkm}\left({\sigma }_{jk}\right)),\end{array}$$

(5)

where the generated function was represented as the sum of sinusoidal feature functions ${\varphi }_{jkm}$ weighted by the normally distributed coefficients $a_{jkm}$. For the Gaussian kernels that we mainly use in this study, hyperparameters ${\left\{{\theta }_{jk\ell }\right\}}_{\ell \in {\Theta }_{jk}}=\{c_{jk},{\sigma }_{jk}\}$ tune the amplitude and bandwidth, respectively, of the GP through the above equation. Concrete representations of ${\varphi }_{jkm}$ and $V_{jkm}$ are given in Appendix A. In this study, for simplicity, we restrict ourselves to the use of one-dimensional Gaussian kernels and linear kernels, as well as the use of a single set of hyperparameters $\{c_g,{\sigma }_g\}$ for all the GPs associated with Gaussian kernels, and a single hyperparameter $c_{\ell }$ for all the GPs associated with linear kernels (equivalently Gaussian random variables). Note that $V_{jkm}$ can be set to one for linear kernels. We assign inverse Gamma priors to the hyperparameters as follows:

$$\begin{array}{c}\hfill \theta \sim \mathrm{InvGamma}({\alpha }_{\theta },{\beta }_{\theta }),(\theta =c_g,{\sigma }_g\mathrm{or}{c}_{\ell }).\end{array}$$

(6)

In the $L,M_{jk}\to \infty$ limit (see Appendix A for the definition of L), the above approximation tends to be exact (Theorem 4 of Ref. [35]). The advantage of this approximation is the use of features (${\varphi }_{jkm}$) that are independent of the bandwidth parameter ${\sigma }_{jk}$, and the use of a tractably computed prior density for ${\left\{a_{jkm}\right\}}_{j,k,m}$. However, the prefixed sinusoidal features are less flexible than are data-adaptive features (obtained, e.g., by the incomplete Cholesky decomposition [37]) and are not suitable for approximating high-dimensional functions. In biometric and econometric applications (e.g., epidaemiology), modelling a target variable as the sum of the effects of several factors is preferred for the sake of interpretability and robustness. Thus, modelling with the sum of several low-dimensional functions in Equation (4) is justifiable.

The above modelling motivates our investigation of efficient MC algorithms. The variance $c_{jk}{V}_{jkm}$ varies with ${\sigma }_{jk}$ over a few orders of magnitude. Combined with the structure of the data likelihood, the posterior distribution is often extremely stretched or compressed along unknown directions. Thus, a data-adaptive sampling scheme, such as RMHMC, is required.

2.1.2. RMHMC with a Soft-Absolute Hessian Metric

An HMC algorithm [38,39,40] obtains samples of parameters $q\in {\mathbf{R}}^d$ from a target density $P\left(q\right)$ by regarding q as the position of particles and simulating its time evolution together with that of the associated momentum vector $p\in {\mathbf{R}}^d$ according to the Hamilton equation of motion with a suitably designed Hamiltonian $\mathcal{H}(q,p)$. The time evolution over discretised timesteps is determined as follows:

$$\begin{array}{ccc}\hfill p\left(t+{\displaystyle \frac{ϵ}{2}}\right)& =& p\left(t\right)-{\displaystyle \frac{ϵ}{2}}{\displaystyle \frac{\partial }{\partial q}}\mathcal{H}\left(q\left(t\right),p\left(t+{\displaystyle \frac{ϵ}{2}}\right)\right),\hfill \\ \hfill q(t+ϵ)& =& q\left(t\right)+{\displaystyle \frac{ϵ}{2}}\left\{{\displaystyle \frac{\partial }{\partial p}}\mathcal{H}\left(q\left(t\right),p\left(t+{\displaystyle \frac{ϵ}{2}}\right)\right)+{\displaystyle \frac{\partial }{\partial p}}\mathcal{H}\left(q(t+ϵ),p\left(t+{\displaystyle \frac{ϵ}{2}}\right)\right)\right\},\hfill \\ \hfill p(t+ϵ)& =& p\left(t+{\displaystyle \frac{ϵ}{2}}\right)-{\displaystyle \frac{ϵ}{2}}{\displaystyle \frac{\partial }{\partial q}}\mathcal{H}\left(q(t+ϵ),p\left(t+{\displaystyle \frac{ϵ}{2}}\right)\right).\hfill \end{array}$$

(7)

This time evolution for a time period of duration $ϵ$ is called a leapfrog. The mapping from $\left(q\right(t),p(t\left)\right)$ to $\left(q\right(t+ϵ),p(t+ϵ\left)\right)$ preserves the volume and approximately preserves the value of $\mathcal{H}$. As we show in Algorithm 1, HMC repeats the combination of C leapfrogs and the sampling of a new value for p from a suitable q-dependent Gaussian distribution. At the end of every C leapfrogs, the change in $\mathcal{H}$ due to discretisation is adjusted via the Metropolis–Hastings procedure. Samples obtained from a long-time simulation after a burn-in period approximate the target distribution, and the statistics calculated with the samples converge to the average over the target at the infinite time limit. Girolami and Calderhead [41] introduced the use of a nontrivial Riemannian metric $q\mapsto G\left(q\right)\in {\mathbf{R}}^d\times {\mathbf{R}}^d$ for the parameter space and showed that, in this case, the Hamiltonian should be defined as follows:

$$\begin{array}{c}\hfill \mathcal{H}(q,p)\stackrel{\mathrm{def}}{=}-lnP\left(q\right)+{\displaystyle \frac{1}{2}}ln\left\{{\left(2\pi \right)}^d\left|G\left(q\right)\right|\right\}+{\displaystyle \frac{1}{2}}{p}^{T}G{\left(q\right)}^{-1}p.\end{array}$$

(8)

For this version of Hamiltonian, their derivatives with respect to q and p read as

$$\begin{array}{ccc}\hfill {\displaystyle \frac{\partial }{\partial p_i}}\mathcal{H}& =& -{\left(G{\left(q\right)}^{-1}p\right)}_i,\hfill \\ \hfill {\displaystyle \frac{\partial }{\partial q_i}}\mathcal{H}& =& {\displaystyle \frac{\partial }{\partial q_i}}lnP\left(q\right)-{\displaystyle \frac{1}{2}}{p}^{T}G{\left(q\right)}^{-1}{\displaystyle \frac{\partial G\left(q\right)}{\partial q_i}}G{\left(q\right)}^{-1}p+{\displaystyle \frac{1}{2}}tr\left(G{\left(q\right)}^{-1}{\displaystyle \frac{\partial G\left(q\right)}{\partial q_i}}\right).\hfill \end{array}$$

(9)

The choice of $G\left(q\right)=1$ recovers the ordinary HMC based on the Euclidean metric.

Algorithm 1 (RM)HMC

Input: Target density $P\left(q\right)$, initial value $q\left(0\right)\in {\mathbf{R}}^d$, number of total and burn-in MC moves A and $A_0$, stepsize $ϵ$, number of leapfrogs C in a single move, metric $G:q\mapsto G\left(q\right)\in {\mathbf{R}}^d\times {\mathbf{R}}^d$ and Hamiltonian $\mathcal{H}:(q,p)\mapsto \mathcal{H}(q,p)\in \mathbf{R}$ defined in Equation (8). Output: A collection of parameter values ${\left\{q(ϵ\ell C)\right\}}_{A_0\le \ell \le A}$ useful for approximating $P\left(q\right)$. Generate an initial value for the momentum vector: $p\left(0\right)\sim \mathcal{N}(0,G(q\left(0\right)\left)\right)$. for ℓ in $\{0\dots A-1\}$ do    for s in $\{0\dots C-1\}$ do      Carry out a leapfrog step according to Equation (7)    end for    Generate $r_{\ell }\sim \mathrm{Uniform}(0,1)$ for the following Metropolis–Hastings procedure:    if $r_{\ell }>exp(-\mathcal{H}(q\left(ϵ(\ell +1)C\right),p\left(ϵ(\ell +1)C\right))+\mathcal{H}(q(ϵ\ell C),p(ϵ\ell C)))$ then      $q\left(ϵ\right(\ell +1\left)C\right)\leftarrow q(ϵ\ell C)$    end if    Generate an updated value for the momentum vector: $p\left(ϵ\right(\ell +1\left)C\right)\sim \mathcal{N}(0,G(q\left(ϵ\right(\ell +1\left)C\right)\left)\right)$. end for

The Hessian of $-lnP\left(q\right)$ is a natural choice for $G\left(q\right)$, as we explained in the previous section. However, the Hessian of the negative log posterior density for a hierarchical probabilistic model is often not positive definite and therefore cannot be used as a metric. For this problem, Betancourt [28] proposed transforming the Hessian into a positive definite matrix as follows:

$$\begin{array}{c}\hfill H\left(q\right)=\sum _{1\le i\le d}{\lambda }_i{\psi }_i{\psi }_i^T\mapsto G\left(q\right)=\sum _{1\le i\le d}g\left({\lambda }_i\right){\psi }_i{\psi }_i^T,\end{array}$$

(10)

where the differentiable function g approximates the absolute value, i.e., $g\left(\lambda \right)\approx \left|\lambda \right|$. In the above, ${\left\{{\lambda }_i\right\}}_{1\le i\le d}$ and ${\left\{{\psi }_i\right\}}_{1\le i\le d}$ are the eigenvalues and eigenvectors of the Hessian H. Employing this transformed metric is equivalent to rescaling the motion of particles according to the absolute value of the curvature of the log density graph. Betancourt [28] further showed how to compute the right-hand sides of Equation (9) for the soft-absolute Hessian, as follows, for $\Psi =\mathrm{mat}\left({\left\{{\left({\psi }_i\right)}_j\right\}}_{j,i}\right)$:

$$\begin{array}{ccc}\hfill {\Psi }^T{\partial }_{q_i}G\left(q\right)\Psi & =& T\odot ({\Psi }^T{\partial }_{q_i}H\left(q\right)\Psi ),\hfill \\ \hfill T_{j\ell }& =& \left\{\begin{array}{cc}{\displaystyle \frac{g\left({\lambda }_j\right)-g\left({\lambda }_{\ell }\right)}{{\lambda }_j-{\lambda }_{\ell }}}& ({\lambda }_j\ne {\lambda }_{\ell })\\ g^{\prime }\left({\lambda }_j\right)& ({\lambda }_j={\lambda }_{\ell })\end{array}\right.,(1\le i,j,\ell ,\le d),\hfill \end{array}$$

(11)

where ⊙ denotes the Hadamard product. Given that eigenvectors and eigenvalues were already obtained, Betancourt [28] showed that the formula in Equation (11) allows one to compute the right-hand side of the second equivalence in Equation (9) for $1\le i\le d$ with an $O\left(d^3\right)$ computational cost by first caching

$$\begin{array}{c}\hfill W_1=\Psi BTB{\Psi }^T,W_2=\Psi (R\odot T){\Psi }^T\end{array}$$

(12)

for $B=\mathrm{diag}{({\left({\Psi }^{T}p\right)}_i/g\left({\lambda }_i\right))}_i$ and $R=\mathrm{diag}{\left(g{\left({\lambda }_i\right)}^{-1}\right)}_i$, and then calculating

$$\begin{array}{ccc}\hfill p^{T}G{\left(q\right)}^{-1}{\displaystyle \frac{\partial G\left(q\right)}{\partial q_i}}G{\left(q\right)}^{-1}p& =& tr\left(W_1{\displaystyle \frac{\partial }{\partial q_i}}H\left(q\right)\right),\hfill \\ \hfill tr\left(G{\left(q\right)}^{-1}{\displaystyle \frac{\partial G\left(q\right)}{\partial q_i}}\right)& =& tr\left(W_2{\displaystyle \frac{\partial }{\partial q_i}}H\left(q\right)\right).\hfill \end{array}$$

(13)

Here, $W_1$ and $W_2$ are obtained by $O\left(d^{2.38}\right)$ matrix multiplications, the full gradient ${\displaystyle \frac{\partial }{\partial q}}H\left(q\right)$ is obtained by an $O\left(d^3\right)$ computation, and the $2d$ trace operations are carried out by taking the sum of the elementwise products $2d$ times; hence, the overall complexity of the computation is $O\left(d^3\right)$.

2.2. Two Sources of Redundancy in the Previous Implementations

2.2.1. Order of Computations in the Calculation of Gradient Flow

As we show below, a straightforward application of the above algorithm to the model described by Equations (1), (4)–(6) is time consuming. We show that the computational cost can be reduced to $O(d^{\omega \left(r\right)}+d^{1+r})$ for $d=\dim \left(\{{\left\{a_{jkm}\right\}}_{j,k,m},{\left\{b_j\right\}}_j\}\right)+\dim \left({\left\{{\theta }_{jk\ell }\right\}}_{j,k,\ell }\right)$, $N=\lceil d^r\rceil$, and fixed J. Here, the exponent $\omega \left(r\right)$ describes the computational complexity of multiplying matrices of sizes $d^r\times d$ and $d\times d$ (equivalently matrices of sizes $d\times d^r$ and $d^r\times d$) (see Table 2 and Footnote 1 of Ref. [42]). This can be verified by substituting a concrete representation of the third derivatives of the negative log posterior density for ${\displaystyle \frac{\partial }{\partial q}}H\left(q\right)$ in Equation (13) as follows: for $s=1,2$, $a_j=\mathrm{vec}\left({\left\{a_{jkm}\right\}}_{(k,m)}\right)$, $U_i=U\left({\left\{f_j\left(X_i\right)\right\}}_{j\in \mathcal{J}}\right)$, $f_{j,i}=f_j\left(X_i\right)$, ${\varphi }_{j,i}=\mathrm{vec}\left({\left\{{\varphi }_{jkm}\left(X_i\right)\right\}}_{(k,m)}\right)$, ${\Phi }_j=\mathrm{mat}\left({\left\{{\varphi }_{jkm}\left(X_i\right)\right\}}_{i,(k,m)}\right)$, and $W_{s{j}_1{j}_2}=\mathrm{mat}\left({\left\{W_{s,j_1{k}_1{m}_1,j_2{k}_2{m}_2}\right\}}_{(k_1,m_1),(k_2,m_2)}\right)$.

$$\begin{array}{ccc}& & tr\left(W_s{\displaystyle \frac{\partial }{\partial a_j}}H\left(q\right)\right)\hfill \\ & =& \sum _{i,j_1,j_2,k_1,k_2,m_1,m_2}{\displaystyle \frac{{\partial }^3{U}_i}{{\partial }_{f_{j_1,i}}{\partial }_{f_{j_2,i}}{\partial }_{f_{j,i}}}}{\varphi }_{j_1{k}_1{m}_1,i}{\varphi }_{j_2{k}_2{m}_2,i}{\varphi }_{j,i}{W}_{s,\left(j_1{k}_1{m}_1\right),\left(j_2{k}_2{m}_2\right)}+\dots \hfill \\ & =& \sum _{j_1,j_2}{\Phi }_j\left[\mathrm{diag}{\left({\displaystyle \frac{{\partial }^3{U}_i}{{\partial }_{f_{j_1,i}}{\partial }_{f_{j_2,i}}{\partial }_{f_{j,i}}}}\right)}_i\left\{\left(\left({\Phi }_{j_1}{W}_{s{j}_1{j}_2}\right)\odot {\Phi }_{j_2}\right)\mathrm{vec}\left(\left\{1\right\}\right)\right\}\right]+\dots ,\hfill \end{array}$$

(14)

where we omit writing down the terms that involve the derivatives with respect to the bias term ${\left\{b_j\right\}}_j$ and the hyperparameters ${\left\{{\theta }_{jk\ell }\right\}}_{j,k,\ell }$ whose contribution can be computed in the same manner as those of the derivatives with respect to ${\left\{a_{jkm}\right\}}_{j,k,m}$. For each combination of $(j_1,j_2)$, the suitably ordered computation in the second line requires $O\left(d^{\omega \left(r\right)}\right)$, $O\left(d^{1+r}\right)$, $O\left(d^{1+r}\right)$, $O\left(d^r\right)$, and $O\left(d^{1+r}\right)$ computations for calculating ${\Phi }_{j_1}{W}_{s,j_1,j_2}$, calculating the Hadamard product, calculating the matrix vector product with vector $\mathrm{vec}\left(\right\{1\left\}\right)$ (a vector all of whose entries are 1), multiplying the diagonal matrix from the left, and calculating the matrix vector product with ${\Phi }_j$, respectively.

The important point that should be noted here is that the standard library based on the automatic differentiation of the Hamiltonian does not follow the above order of computations. First, the direct derivatives of the Hamiltonian cannot be automatically computed when the soft-absolute metric is used. If the formula in Equation (9) is naively implemented by computing the third-order derivatives of the log likelihood with the aid of an automatic differentiator in the reverse mode (the default mode in PyTorch 2.8.0), at least an $O\left(N{d}^3\right)=O\left(d^{3+r}\right)$ computation is required. This differs greatly from the efficient computation described above. Even if the joint likelihood is log concave and the formula in Equation (9) can be avoided, a naive application of a reverse-mode automatic differentiator to the Hamiltonian results in the same computational cost. In order to obtain an efficient implementation using an automatic differentiator, one must first implement a differentiator that recognises the formula in Equation (9) and then suitably assign the forward or reverse mode to the differentiations of the components of the Hamiltonian according to the chain rule in a model-structure-dependent manner.

2.2.2. Dynamically Programmed Eigendecomposition

The formula given by Betancourt [28] relies on the availability of the complete set of eigenvectors $\Psi$. The fact that $H\left(q\right(t\left)\right)$ and $H\left(q\right(t+ϵ\left)\right)$ are close to each other motivates us to reduce computational costs by dynamically computing $\Psi \left(q\right(t+ϵ\left)\right)$ to take advantage of $\Psi \left(q\right(t\left)\right)$. The boundedness of the gradient of $H\left(q\right)$ and p in the region where the joint probability for $(q,p)$ is concentrated implies that

$$\begin{array}{c}\hfill \Psi {\left(q\left(t\right)\right)}^{T}H\left(q(t+ϵ)\right)\Psi \left(q\left(t\right)\right)=\mathrm{diag}\left({\lambda }_i\left(q\left(t\right)\right)\right)+O\left(ϵ\right)\end{array}$$

(15)

holds with a high probability in terms of the Frobenius norm and can be efficiently eigendecomposed as

$$\begin{array}{c}\hfill \Psi {\left(q\left(t\right)\right)}^{T}H\left(q(t+ϵ)\right)\Psi \left(q\left(t\right)\right)=Q(t+ϵ,t)\mathrm{diag}\left({\lambda }_i\left(q(t+ϵ)\right)\right)Q{(t+ϵ,t)}^T\end{array}$$

(16)

using the Jacobi method. Note that cyclic versions of the Jacobi method quadratically converge [43] and are suitable for parallelisation [44], and thus are expected to carry out the above decomposition very quickly with a small error tolerance $\zeta$. With this decomposition, we update the eigenvectors as

$$\begin{array}{c}\hfill \Psi \left(q\right(t+ϵ\left)\right)=\Psi \left(q\right(t\left)\right)Q(t+ϵ,t).\end{array}$$

(17)

In practice, we perform the Gram–Schmidt orthogonalisation of $\Psi \left(q\right(t\left)\right)$ every ten steps before applying it to $H\left(q\right(t+ϵ\left)\right)$ to remove accumulated numerical errors.

2.3. Methods of Numerical Experiments

2.3.1. Examination of Computational Complexity

We investigate the efficiency of the proposed algorithm by performing Bayesian logistic regression with artificially generated data and measuring its computation time. For comparison, we also perform posterior sampling using Hamiltorch [29], a publicly available library for different types of HMCs based on PyTorch that works on CUDA devices.

First, we generate standardised explanatory variables $X_i^{\left(\exp \right)}\in {\mathbf{R}}^D\stackrel{i.i.d.}{\sim }\mathcal{N}(0,1)$ for each sample i ($1\le i\le 500$). Then, we calculate the true log odds ratio $f_1^{*}\left(X_i^{\left(\exp \right)}\right)$ according to

$$\begin{array}{c}\hfill f_1^{*}\left(X_i^{\left(\exp \right)}\right)=\sum _{1\le k\le D}\sum _{m=1}^{M_{1k}}{a}_{1km}^{*}{\varphi }_{1km}\left(X_{ik}^{\left(\exp \right)}\right)+b_1^{*},\end{array}$$

(18)

with $b_1^{*}=-0.5$, $a_{1km}^{*}\stackrel{i.i.d.}{\sim }\mathcal{N}(0,c^{*}{m}^{-1})$ for $4\le m\le 16$ and $a_{1km}^{*}=0$ otherwise. For any k, the feature function ${\varphi }_{1mk}$ is given by Equation (A2) in Appendix A with $\delta =1$ and $L=8$. Here, the value of constant $c^{*}$ is determined so that the empirical standard deviation $\widehat{\mathrm{SD}}\left[f_1\left(X^{\left(\exp \right)}\right)\right]$ is 1.5.

We then generate the sample label according to

$$\begin{array}{c}\hfill X_i^{\left(\mathrm{lab}\right)}=\left\{\begin{array}{cc}+1& \mathrm{prob}.{\displaystyle \frac{1}{1+e^{-f_1^{*}\left(X_i^{\left(\exp \right)}\right)}}}\\ -1& \mathrm{otherwise}\end{array}\right..\end{array}$$

(19)

For this dataset, we perform posterior sampling, where the model is determined by Equation (1) with $J=1$ and

$$\begin{array}{c}\hfill U\left(f_1\left(\{X_i^{\left(\exp \right)},X_i^{\left(\mathrm{lab}\right)}\}\right)\right)=ln(1+exp(-X_i^{\left(\mathrm{lab}\right)}{f}_1\left(X_i^{\left(\exp \right)}\right))),\end{array}$$

(20)

where function $f_1$ is further decomposed into functions associated with one-dimensional Gaussian kernels:

$$\begin{array}{c}\hfill f_1\left(X_i^{\left(\exp \right)}\right)=\sum _{1\le k\le D}\sum _{1\le m\le M_{1k}}{a}_{1km}{\varphi }_{1km}\left(X_{ik}^{\left(\exp \right)}\right)+b_1.\end{array}$$

(21)

Here, we abuse a notation for brevity by identifying index k with a Gaussian kernel function $k(X_i,X_i^{\prime })=exp(-{(X_{ik}-X_{ik}^{\prime })}^2/{\sigma }_{1k})$, for which, again, the feature function ${\varphi }_{1mk}$ is given by Equation (A2) in Appendix A with $\delta =1$, $L=8$ and $M_{1k}=30$. The associated GPs and their hyperparameters are described by Equations (4)–(6). For the transformation into the soft-absolute Hessian in the proposed method, we used $g\left(\lambda \right)=\sqrt{{\kappa }^2+{\lambda }^2}$, which takes values close to $g\left(\lambda \right)=\lambda \coth \left({\kappa }^{-1}\lambda \right)$ employed in Hamiltorch. We did not implement the latter because of the singularity of $\coth$ at zero and the lack of a straightforward implementation in the library that we used (NVIDIA Inc. (Santa Clara, CA, USA), NVHPC23.1 [34]).

In addition to the values specified above, we used the following values: (for the comparison of sampling speed) ${\alpha }_{c_g}={\alpha }_{c_{\ell }}=5,{\beta }_{c_g}={\beta }_{c_{\ell }}=0.5,{\alpha }_{{\sigma }_g}={\beta }_{{\sigma }_g}=1$, $\Sigma =1$, $A=9600$, $A_0=2400$, $C=100$, $ϵ=0.001$, $\kappa =1$, $\zeta =1.0\times {10}^{-13}$; (for the calculation of BME (see Appendix B)) ${\alpha }_{c_g}={\alpha }_{c_{\ell }}=5,{\beta }_{c_g}={\beta }_{c_{\ell }}=0.5,{\alpha }_{{\sigma }_g}={\beta }_{{\sigma }_g}=1$, $\Sigma =1$, $A=50$, $C=100$, $ϵ=0.001$, $\kappa =1$, $\zeta =1.0\times {10}^{-13}$, $Z=50$, ${\tau }_s=1-0.02(s-1)$ ($1\le s\le 41$), ${\tau }_s=0.2-0.005(s-41)$ ($41<s\le 71$), ${\tau }_s=0.05-0.002(s-71)$ ($71<s\le 91$), and ${\tau }_s=0.01-0.001(s-91)$ ($91<s\le 101$). The initial parameter values for MC we used were $a_{jkm}=0$, $b_j=0$ and $\theta =1$ for any hyperparameter $\theta$.

We confirm the convergence of the proposed algorithm by performing Wilcoxon’s rank sum test for the values of the log posterior density in the first and second halves of the trajectory after an initial burn-in period. The proposed method was implemented by writing codes in C++ and OpenACC that offload most of the calculation to a CUDA device retaining all necessary data in its own memory, and by invoking cuBLAS and cuSolver for linear algebra routines as much as possible. For comparison, we wrote a code in PyTorch for the same model that invokes Hamiltorch for carrying out RMHMC or NUT-HMC with a manually coded likelihood. To conduct a comparison of different implementations, we measure the wall time spent on MC moves run on a single CPU core (a recent version of Intel Xeon processor) connected to a single NVIDIA Tesla A100 PCIe 80 GB GPU card, for each of the five datasets generated with different seeds for the random number generator.

All of the source codes used for the analysis described above are provided as Supplementary Materials.

2.3.2. Analysis of 1987 National Medical Expenditure Survey (NMES)

To demonstrate that our implementation is acceptably efficient in a real-world application, we perform Bayesian estimation with excerpted data from the 1987 national medical expenditure survey (NMES) of the United States. A few groups of authors performed causal inference with this dataset about the effects of smoking on medical expenditures [45,46,47]. In the present analysis, we attempt to increase the precision of the propensity function used for causal inference. As shown by Imai and van Dyk [46], the identification of a set of parameters f that characterise the conditional distribution $p_f\left(X^{\left(\mathrm{tgt}\right)}\right|X^{\left(\mathrm{cov}\right)})$ of the actual treatment variable $X^{\left(\mathrm{tgt}\right)}$ for the given values of covariates $X^{\left(\mathrm{cov}\right)}$ reduces the dimensionality of subsequent causal analysis.

In this dataset, $X^{\left(\mathrm{tgt}\right)}$ denotes the packyear of smoking (i.e., the product of the number of packs of cigarette consumed by the subject and the duration of smoking measured in years). The covariates $X^{\left(\mathrm{cov}\right)}$ include the age at the time of the survey, the age at the initiation of smoking, gender, race, marital status, education level, census region, poverty status, and seat belt usage. The previous study focused mainly on the estimation of the mean of the conditional distribution as f and not its variance. We compare the performance of linear and nonlinear models that describe only the conditional mean, and the performance of linear and nonlinear models that describe both the conditional mean and variance (Equation (2)). In particular, as the model that best describes the given data is determined by the value of the Bayesian model evidence (BME) [13], we investigate whether the proposed method (and INLA) computes this value within a reasonable computation time (see Appendix B and Appendix C for the technical details of the calculation of BME). Carrying out causal inference with the estimated propensity functions requires further theoretical development and is not within the scope of the present study. Thus, we restrict our analysis to the estimation of propensity functions and discuss the further development that is needed in Section 4.

We obtained the dataset included in a library for causal inference [47]. The categorical covariates in this dataset were transformed to a set of binary covariates that retained the original information. The continuous covariates were standardised to have a zero mean and unit variance.

For this dataset, we considered the data likelihood described by Equation (2). The function $f_j$ ($j=1,2$) was described as

$$\begin{array}{c}\hfill f_j\left(X_i^{\left(\mathrm{cov}\right)}\right)=\sum _{1\le k\le D}\sum _{1\le m\le M_{jk}}{a}_{jkm}{\varphi }_{jkm}\left(X_{ik}^{\left(\mathrm{cov}\right)}\right)+b_j.\end{array}$$

(22)

For the linear models, we used a linear kernel indexed by k for each variable $X_{ik}^{\left(\mathrm{cov}\right)}$, and thus we had ${\varphi }_{jk1}\left(X_i^{\left(\mathrm{cov}\right)}\right)=X_{ik}^{\left(\mathrm{cov}\right)}$ with $M_{jk}=1$. For the nonlinear models, we used a one-dimensional Gaussian kernel for each continuous variable with $L_1=8$ and $M_{jk}=30$, and a linear kernel for each binary variable. The models that describe the conditional mean and variance were determined in this manner. The models that describe only the conditional mean determined $f_1$ in the manner described above, whereas their variance was described by $f_2=b_2$. The priors for GPs and hyperparameters were described by Equations (4)–(6) in the same manner as for the simulated data.

In addition to the values specified above, we used the following values: ${\alpha }_{c_g}={\alpha }_{c_{\ell }}=5,$${\beta }_{c_g}={\beta }_{c_{\ell }}=0.5,{\alpha }_{{\sigma }_g}={\beta }_{{\sigma }_g}=1$, $\Sigma =1$, $A=50$, $C=400$, $ϵ=0.0001$ ($ϵ=0.00008$ for the nonlinear model describing both of the conditional mean and variance), ${\sigma }_0^2=0.2$, $\kappa =1$, $\zeta =1.0\times {10}^{-13}$, $Z=10$, ${\tau }_s=1-0.02(s-1)$ ($1\le s\le 41$), ${\tau }_s=0.2-0.005(s-41)$ ($41<s\le 71$), ${\tau }_s=0.05-0.002(s-71)$ ($71<s\le 91$), and ${\tau }_s=0.01-0.001(s-91)$ ($91<s\le 101$). The initial parameter values for MC we used were $a_{jkm}=0$, $b_j=0$, and $\theta =1$ for any hyperparameter $\theta$.

All of the source codes used for the analysis described above are provided as Supplementary Materials.

3. Results

3.1. Comparison Among Different Implementations of RMHMC with a Soft-Absolute Hessian Metric

The proposed implementation is fast enough to converge to equilibrium in both small and moderately large models, as seen from the trajectories in Figure 1A and confirmed by a Wilcoxon rank sum test for ~14,000- and ~60,000-s simulations ($p=0.31,0.47$), respectively. As shown in Figure 1A,B, our implementation is roughly ten times faster than the RMHMC based on Hamiltorch.

Figure 1. RMHMC for Bayesian logistic regression on simulated data. Comparison between our implementations based on the dynamic and static eigendecomposition of the Hessian and the implementation based on Hamiltorch. (A) Representative trajectories of log posterior density for $D=1$ ($d=34$) and $D=16$ ($d=484$). (B) Mean and standard deviation of the wall time spent on computing 100 blocks of 100 leapfrogs are shown on the linear (main panel) and logarithmic (inset) scales for different model sizes. (C) Mean and standard deviation of the number of sweeps in the cyclic Jacobi method for a single dynamic and static eigendecomposition of the metric are shown on the logarithmic scale for different model sizes. In (B,C), model sizes are shown on the logarithmic scales. (Abbreviations) prop.(dyn.): the proposed method based on dynamic eigendecomposition, prop.(syevd): the proposed method based on static eigendecomposition with the divide-and-conquer algorithm, and prop.(syevj): the proposed method based on static eigendecomposition with the Jacobi method.

To investigate the relative impact of the order of computations compared with that of the dynamically programmed eigendecomposition, we also measured the wall time for implementations in which the eigendecomposition was replaced by static ones based on either the Jacobi method or the divide-and-conquer algorithm. The inset in Figure 1B suggests that the difference in computation time between our implementation and the implementation based on Hamiltorch is mainly attributed to the computation of the gradient flow rather than eigendecomposition.

However, a close examination of Figure 1B also shows differences among the algorithms for eigendecomposition. Although little difference is observed in computation time among the examined algorithms for small D, the dynamic eigendecomposition saves substantial computation time for large D. However, we also observe that the faster increase in the computational cost of the Jacobi method compared to that of the divide-and-conquer algorithm reduces this saving for $D=64$ ($d=1924$) and $D=96$ ($d=2884$). Examining the number of sweeps in the cyclic Jacobi method in static and dynamic implementations (Figure 1C), we confirm the advantage of the dynamic implementation, which requires only one to two sweeps at each step regardless of the model dimensionality, whereas the static implementation requires an increasing number of sweeps for greater model dimensionality on average. In the current model setting, the performance with $D\ge 128$($d\ge 3844$) could not be examined because of numerical instability.

3.2. Comparison with NUT-HMC Sampler

We also confirm the advantage of using RMHMC by comparing its performance with that of NUT-HMC. As shown in Figure 2A, NUT-HMC reaches samples of parameter values that have large values of log posterior density very quickly and continues generating similar samples (Figure 2A). This contrasts with the slow convergence of RMHMC to generate samples with lower log posterior density. To obtain a clue as to which samplers generate representative samples from the posterior, we examine the eigenvalues of the Hessian of the log posterior density with representative samples from the two samplers (Figure 2B). We see that the eigenvalues for the sample from NUT-HMC are ten times larger than those for the sample from RMHMC, except for several small positive or negative eigenvalues. This suggests that the samples from NUT-HMC have been taken from a much narrower region having a larger density. We therefore doubt that the total posterior probability for this narrow region is smaller than that for the region explored by RMHMC. We will see that this is the case below.

Figure 2. Comparison of proposed implementation of RMHMC and NUT-HMC in Hamiltorch. (A) Representative trajectories of the log posterior density for $D=1$ ($d=34$) obtained using the proposed implementation of RMHMC based on dynamic eigendecomposition (prop.(dyn)) and NUT-HMC in Hamiltorch. (B) Eigenvalues of the Hessians of the negative log posterior density for the samples of parameter values obtained after running the proposed implementation of RMHMC and NUT-HMC in Hamiltorch for 2000 s. The eigenvalues are plotted in descending order on the logarithmic scale. For negative eigenvalues, their absolute values are plotted. The absolute values of the 32nd to 34th eigenvalues for the sample from NUT-HMC were extremely small and are therefore out of the range in the plot.

3.3. Calculation of BME Beyond the Laplace Approximation with Simulated Data and NMES Data

The existence of negative eigenvalues in Figure 2B suggests that the posterior cannot be regarded as being approximately normal. This is not surprising, since hierarchical models are known to be often singular [25]. Thus, the BME of the proposed model cannot be calculated by simply applying the Laplace approximation to the posterior density. Even in this case, we successfully carried out the calculation of BME through MC integration with reasonably high precision, as confirmed by the standard error of the estimate determined with multiple sequences of RMHMC (Table 1). The obtained value was in good agreement with the BME value obtained by performing INLA, namely, numerically integrating the Laplace approximation of the likelihood conditioned on each combination of hyperparameter values (see Appendix B for technical details). As expected from the nature of the two computation methods, INLA carried out the calculation within a shorter time. The latter calculation based on the Laplace approximation is justified by the fact that the model has log-concave posterior density and is regular when conditioned on hyperparameters.

Table 1. Calculation of BME using the proposed implementation of RMHMC and INLA. The mean (and standard error) of the estimated model evidence and the mean (and standard deviation) of the wall time for its computation in a single RMHMC (among Z RMHMCs) or in INLA are shown. For NMES data, linear (L) and nonlinear (NL) models that describe only the conditional mean (mean) or both of the conditional mean and variance (mean/var) were used.

Data	Model	Algorithm	Estimated BME	Wall Time (s)	# MC (Z)
Simulation	NL-logistic	prop. method	$-279.81(\pm 0.40)$	16,243 $(\pm 1102)$	50
Simulation	NL-logistic	INLA	$-282.31$	11,884	-
NMES	L-mean	prop. method	−11,716.78 $(\pm 2.11)$	43,314 $(\pm 5120)$	10
NMES	L-mean	INLA	−11,746.10	$42.81$	-
NMES	NL-mean	prop.method	−11,602.37 $(\pm 2.65)$	49,648 $(\pm 3613)$	10
NMES	NL-mean	INLA	−11,606.00	1,407,963	-
NMES	L-mean/var	prop. method	$-9540.63(\pm 1.87)$	49,424 $(\pm 4509)$	10
NMES	L-mean/var	INLA	$-9587.66$	$132.17$	-
NMES	NL-mean/var	prop. method	$-9163.53(\pm 3.43)$	52,284 $(\pm 3501)$	10
NMES	NL-mean/var	INLA	-	≥6,000,000	-

In the calculation with INLA, it is also confirmed that the samples obtained with RMHMC and NUT-HMC in Figure 2 are positioned near the center and periphery of the posterior probability mass, respectively (Figure 3A), while they are positioned near the lower/broader and higher/narrower peaks of the posterior density, respectively (Figure 3B).

Figure 3. Colour plot of (A) marginal posterior density and (B) peak value of joint posterior ($-lnP\left(q\right)$ used in RMHMC) for different values of amplitude $c_g$ and bandwidth ${\sigma }_g$. Representative samples obtained after ~2000-s simulations with RMHMC and NUT-HMC are also plotted with the same symbols as those used in Figure 2.

To demonstrate the usefulness of the proposed method for analysing real-world data, we also calculate the BME for the NMES dataset using linear or nonlinear models that estimate the mean (and variance) of the distribution of the treatment variable conditioned on the covariates, namely, the propensity functions. We successfully calculated the BME values for this estimation problem within an acceptable computation time (Table 1). The computed BME values showed improved estimation of the propensity function with the nonlinear model for both the mean and variance of the conditional treatment density.

The computation of BME with these models was also carried out using INLA. The repeated alternate application of the l-BFGS algorithm to the non-log-concave likelihood term of the model required more computation time than was required for INLA on simulated data. For the linear model of the conditional mean or the conditional mean and variance, INLA requires the integration with respect to only one hyperparameter, which was carried out within a very short time. For the nonlinear models, however, INLA requires the integration with respect to three hyperparameters and hence the evaluation of approximate Gaussian conditional distributions on much more points in the hyperparameter space. Furthermore, in these models, the evaluation of each combination of hyperparameter values took several tens of seconds. As a result, INLA took far more time to compute the BME values for nonlinear models than RMHMC.

4. Discussion

In the present study, we have investigated MC methods for the posterior sampling of functions and hyperparameters in hierarchical GP models, and we have shown that a straightforward implementation of RMHMC based on currently available general-purpose libraries is highly redundant and its performance can be greatly improved. The main source of redundancy was the order of computation in the calculation of the gradient flow on the Riemannian manifold, whereas the eigendecomposition of the metrics was also a major source of redundancy for larger models. These findings have non-trivial implications for future modelling based on GP and its implementations, because our results indicate that in the current standard practice of coding only the likelihood of the model and allowing libraries to carry out inference with the aid of automatic differentiation results in poor performance, unless an intelligent library that optimises overall computational complexity is developed.

This problem was not recognised in the previous study of RMHMC for GP [14], because they sampled in the space for dual variables (${\left\{f_{j,i}\right\}}_{1\le j\le J,1\le i\le N}$ with $J=1$ in our notation) for which the entries of the third-order derivative tensor were sparse. In a hierarchical model with multiple GPs, the use of dual variables incurs a large computational cost. In this case, sampling in the space for multiple sets of dual variables ${\left\{f_{jki}\right\}}_{1\le j\le J,k\in {\mathcal{K}}_j,1\le i\le N}$ ($J>1$) is very high-dimensional. In this case, the eigendecomposition of the metric required for the Betancourt’s formula (Equation (9)) is essentially impossible to compute. Furthermore, the derivatives of posterior density with respect to both of the dual variables and the hyperparameters ($c_{jk}$ and ${\sigma }_{jk}$ in our notations) are also difficult to compute. The previous study avoided these difficulties by using only a single GP and fixing the hyperparameter values. In contrast, we use a representation [35] with a reduced number of primal variables, which makes the dependence on hyperparameters easier to compute.

In our numerical study, the dependence of computation time on model size does not precisely agree with theoretical expectations (Section 2.2.1), presumably because the GPU accelerator carries out many arithmetic operations in parallel, and the effect of the size of multiplied matrices in the examined range is masked by this parallelism. As we have discussed in Section 2.2.1, with an optimised computation order, the gradient flow for RMHMC can be calculated with costs of the same complexity as the multiplication of two matrices. Although this is a power of model size with exponent between two and three, the growth of the measured computation time was approximately linear. The effect of parallelism is well known for the multiplication of matrices [48] and observed for the multiplication of much larger matrices than those in our computation. Because the simulations with model sizes greater than 3844 suffered numerical instability, we could not attempt to extend the range of dimension of our experiment to see how the computational cost grows in very large models.

In numerical simulations, by dynamically programming eigendecomposition, in comparison with static one based on the divide-and-conquer algorithm, 25 to 30 percent of the computation time was saved for large models ($d=484$–1924), whereas this proportion reduced for $d=2884$. A similar phenomenon was observed in the previous study of dynamically programmed eigendecomposition for matrix optimisation [32]. Although the efforts to speed up the eigendecomposition resulted in this relatively small improvement, the $O\left(d^3\right)$ eigendecomposition is expected to eventually outweigh the computation of the gradient flow at sufficiently large d. Again, we could not directly observe this phenomenon because of the numerical instability for $d\ge 3844$. Theoretically, less than 1.5 sweeps of the Jacobi method are faster than the divide-and-conquer algorithm in the large d limit (number of sweeps $\times (6n^3+81n^2-87n)$ vs $9n^3+(30n+n+6n^2)\times$ number of iterations [49]). Since the average number of sweeps in the dynamically programmed eigendecomposition was around 1.5 regardless of the model size, we could not conclude which method for eigendecomposition is preferred for a very large model. Since the Jacobi method is more suitable for parallel computing, we speculate its advantage.

For the theoretical estimation of the computational costs, the tradeoff between the increased number of sweeps in the Jacobi method due to a large stepsize and the increased leapfrogs due to a small stepsize also needs to be considered. Since the computational cost of one sweep in the Jacobi method is high, taking the largest stepsize that does not increase the average number of sweeps in a step would be the best choice. However, we could not validate this point with increasing stepsizes in our numerical simulations, because the acceptance rate was decreased by the deviation from the conservation of the Hamiltonian before the optimal point of the tradeoff was observed.

In the numerical study, we showed that RMHMC outperforms NUT-HMC. The latter apparently reaches a parameter region with higher log posterior density more quickly, but this region has turned out to be a narrow spurious region to which a small posterior probability is assigned. Such entrapment is a well-known phenomenon and was the main motivation for the development of RMHMC [41]. Thus, the rescaling of MC moves according to the curvature of log posterior density has been suggested to play a crucial part in reaching the centre of the posterior, and the MC methods commonly used in statistics (e.g., metropolis-adjusted Langevin MC used in Ref. [23]) are not expected to allow for effective sampling in hierarchical GP models used in machine learning. However, various elaborations have been introduced to Euclidean HMC and related Langevin algorithms but not yet to RMHMC. For example, avoiding random walk behavior with a no-U-turn mechanism was proposed for Euclidean HMC [50] and RMHMC [51], but it has not yet been properly implemented for the latter. The friction mechanism introduced to Euclidean Langevin MC [52,53] suppresses inefficient oscillatory behaviour and could also be beneficial if applied to RMHMC. One factor that might have hindered elaboration of RMHMC in this regard is the previously poor performance of the plain RMHMC. Now that we have demonstrated improved RMHMC performance, further development in this direction should be encouraged.

In comparison with INLA, we observed that INLA is faster in the computation of BME for the nonlinear Bayesian logistic regression with simulated data. This calculation might be made even faster, if the joint posterior for hyperparameters is successfully approximated with the CCD strategy [19]. However, this approach must first identify the peak of the joint posterior for hyperparameters, which is not trivial in our example, unlike in the latent Gaussian model studied by INLA. With the simple grid-based integration, the computation time is much longer than expected from the performance of INLA for the hierarchical GP model with short-range correlation [23]. While the number of hyperparameters is two in our simplest example, the computation time grows exponentially with this number.

Indeed, the computation time was intolerably large for INLA of NMES dataset with the nonlinear model of the conditional mean (and variance), as we used the simple grid-based numerical integration. This difficulty was also due to the non-log concavity of the likelihood term. For the non-log concave, highly stretched/compressed posterior density, only repeated alternate applications of l-BFGS effectively worked. Straightforward applications of l-BFGS and ordinary gradient descent algorithms failed to identify the solution within a reasonable computation time, presumably because of the non-log concavity and the highly stretched/compressed landscape of the density, respectively. The scalable performance of RMHMC shown in Table 1 was in stark contrast with INLA in these respects. Here, note that the issue concerning the non-log concavity may be alleviated by the use of a trust region method [54]. Since the trust region method needs fine tuning and a choice of suitable subroutines and is still a theoretical subject for active research, we did not attempt to improve this point.

Although we have successfully shown that RMHMC can be greatly accelerated, it is fair to note that we restricted our investigation to simple model settings. Bayesian multiple-kernel models have been shown to have favorable statistical properties when irrelevant GP is plugged out with sparsifying mechanisms such as a prior that imposes a penalty according to the number of included GPs [6]. For this purpose, reversible jump mechanisms [55,56] must be introduced to RMHMC. Roughly speaking, this amounts to performing a shorter version of the MC integration in the calculation of BME multiple times in the simulation. For real-world applications, more elaborate structured models, such as GPDM for time-series data [2], must sometimes be used. In this case, RMHMC for GP must be combined with other MC methods, such as sequential MC [57,58]. Whether the proposed method works within a reasonable computation time when it is tailored to the models described above needs to be investigated.

Although the present work was intended to solve a biometrical and econometrical problem, unbiased effect estimation with Bayesian models requires further theoretical development, not just an accelerated implementation. The posterior we inferred with NMES data was asymptotically biased, as are essentially all machine learning estimators. Therefore, in order to complete the causal inference with the biased propensity function, one needs to construct a corrected estimator for the treatment effect. This may be carried out by using the framework of doubly-robust debiased machine learning (DML) [59] (see our previous work [60] for the application of DML to a model based on multiple reproducing kernel Hilbert spaces). However, its application is not straightforward, because the hierarchical Bayesian model is apparently singular, whereas DML relies on the asymptotic normality of the estimators. Since the model can be made regular by fixing parts of the model, the framework of DML may be applied to a regular submodel conditioned on suitable hyperparameter values. As the focus of the present study is on the efficiency of RMHMC, we leave this development to a future works.