Fast Riemannian-Manifold Hamiltonian Monte Carlo for Hierarchical Gaussian-Process Models

1. Introduction

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]. A GP serves as a prior process for generating nonlinear functions that describe relationships between variables. The goal of Bayesian analysis 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. [9]), 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. Algorithms for inference with hierarchical GP models can be broadly divided into variational algorithms and Monte Carlo (MC) algorithms. Variational algorithms (e.g., Refs.[5,7]) 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.[6,10,11]) allow for precise computation of averages over the posterior, provided the algorithm is run until convergence. The preferred algorithm type depends on the purpose of the computation.

With biometric and econometric applications in mind, we seek an efficient MC algorithm in this study. In these applications, biased results can be misleading and may adversely affect public health or the economy; therefore, the use of variational algorithms is not appropriate. However, to our knowledge, a standard MC algorithm for general hierarchical GP models has not yet been established. Svensson et al. [6] 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. [12] 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. [11] 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 model is that the posterior may be highly stretched or compressed along unknown directions in the parameter space used for sampling. To address this problem, Paquet and Fraccaro [10] 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, we develop an efficient implementation of RMHMC for hierarchical GP models using a soft-absolute Hessian metric introduced by Betancourt [13]. 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 RMHMC with this metric [14] based on PyTorch [15] is available, its performance is poor. We identify the reasons for this, and show that the computational complexity of the implementation can be considerably reduced.

The article is organised as follows. In Section 2.1, we introduce our problem setting in which multiple GPs are used as building blocks for a hierarchical model, employing the representation of GPs introduced by Solin et al. [16]. In Section 2.2, we provide a brief introduction to RMHMC with a soft-absolute Hessian metric. In Section 2.3 and Section 2.4, we identify two sources of redundancy in the calculation of the gradient flow and metric, respectively. 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. 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. 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. Mathematical notations used in this article are summarised in Table A1.

2. Methods and Algorithms

2.1. Hierarchical GP Models and Their Approximate Representations

In this study, we consider hierarchical GP models for which the samplewise likelihood 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}$ 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,\delta +exp{f}_2\left(X_i^{\left(\mathrm{cov}\right)}\right)),\end{array}$$

(2)

with fixed $\delta (>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,\delta +exp{f}_3\left(X_i^{\left(\mathrm{cov}\right)}\right)),s_i\sim \mathrm{Bernoulli}[1/(1+exp{f}_4\left(X_i^{\left(\mathrm{cov}\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 {\mathcal{L}}_{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 {\mathcal{L}}_{jk}}\in {\mathbf{R}}^{|{\mathcal{L}}_{jk}|}$. Thus, each of ${\left\{f_j\right\}}_{j\in \mathcal{J}}$ is a Bayesian multiple-kernel model for which statistical properties were investigated in a previous study [4]. The prior density of the hyperparameters ${\Pi }_{jk}\left({\left\{{\theta }_{jk\ell }\right\}}_{\ell \in {\mathcal{L}}_{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\}}_i$ on ${\left\{{\theta }_{jk\ell }\right\}}_{j,k,\ell }$ hinders efficient sampling. To alleviate this problem, one previous study [16] 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 {\mathcal{L}}_{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 1 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\to \infty$ limit (see Appendix A for the definition of L), the above approximation tends to be exact (Theorem 4 of Ref.[16]). 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 [17]) 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 Eq.(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.2. RMHMC with a Soft-Absolute Hessian Metric

An HMC algorithm [18,19,20] 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 by

$$\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 [21] 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

$$\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

$$\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 Eq.(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 Eq.(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 [13] proposed transforming the Hessian into a positive-definite matrix as

$$\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’s graph. Betancourt [13] further showed how to compute the right-hand sides of Eq.(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 [13] showed that the formula in Eq.(11) allows one to compute the right-hand side of the second equivalence in Eq.(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.3. Structure-Dependent Removal of Redundancy

As we show below, a straightforward application of the above algorithm to the model described by Eqs.(1), (4), (5) and (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}\right\}}_{j,k}\right)$, $N=\lceil d^r\rceil$ and fixed J. Here, the exponent $\omega \left(r\right)$ denotes 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.[22]). 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 Eq.(13) as follows: for $s=1,2$, $a_j=\mathrm{vec}\left({\left\{a_{jkm}\right\}}_{(k,m)}\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 Eq.(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), 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 Eq.(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 Eq.(9) and then suitably assign the forward or reverse mode to the differentiation of each component of the Hamiltonian, according to the model structure.

2.4. Dynamically Programmed Eigendecomposition

The formula given by Betancourt [13] 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 [23] and are suitable for parallelisation [24], 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.5. Numerical Experiments on 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 [14], 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 Eq.(A2) in Appendix A with $d=1$ and $L_1=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 Eq.(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 Eq.(A2) in Appendix A with $d=1$, $L_1=8$ and $M_{1k}=30$. The associated GPs and their hyperparameters are described by Eqs.(4), (5) and (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 0 and the lack of a straightforward implementation in the library we used (NVIDIA Inc., NVHPC23.1 [25]). In addition to the values specified above, we used the following values: (for the comparison of sampling speed) ${\alpha }_{\theta }={\beta }_{\theta }=2.0$ for any hyperparameter $\theta$, $\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 invoke 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 80GB GPU card, for each of five datasets generated with different seeds for the random number generator.

2.6. 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 [26,27,28]. 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 [27], 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 (Eq.(2)). In particular, as the model that best describes the given data is determined by the value of the Bayesian model evidence (BME) [9], we investigate whether the proposed method computes this value within a reasonable computation time (see Appendix B 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 [28]. 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.

Next, we consider the data likelihood described by Eq.(2). The function $f_j$ ($j=1,2$) is 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 use a linear kernel indexed by k for each variable $X_{ik}^{\left(\mathrm{cov}\right)}$, and thus we have ${\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 use 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. Models that describe the conditional mean and variance are determined in this manner. Models that describe only the conditional mean determine $f_1$ in the manner described above, whereas their variance is described by $f_2=b_2$. The priors for GPs and hyperparameters are described by Eqs.(4), (5) and (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), $\delta =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$.

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 Fig.1(A) and confirmed by a Wilcoxon rank-sum test for $\sim 14,000$- and $\sim 60,000$-second simulations ($p=0.31,0.47$), respectively. As shown in Fig.1(A) and (B), our implementation is roughly 10 times faster than the RMHMC based on Hamiltorch. 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 Fig.1(B) 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 Fig.1(B) 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. Examining the number of sweeps in the cyclic Jacobi method in static and dynamic implementations (Fig.1(C)), 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.

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 Fig.2A, NUT-HMC reaches samples of parameter values that have large values of log-posterior density very quickly and continues generating similar samples (Fig.2(A)). This contrasts with the slow convergence of RMHMC to generate samples with lower log-posterior density. We investigate which samplers generate representative samples from the posterior as follows. First, we examine the eigenvalues of the Hessian of the log-posterior density with representative samples from the two samplers (Fig.2(B)). We see that the eigenvalues for the sample from NUT-HMC are 10 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. Indeed, as we run RMHMC using for the initial parameter value a sample obtained using NUT-HMC, we observe that RMHMC swiftly moves to a region with the same lower range of log-density as that of the samples obtained by using the default initial condition $f_{jk}=0$ and ${\theta }_{jk\ell }=1$ (Fig.2(A)).

3.3. Calculation of BME beyond the Laplace approximation with simulated data and NMES data

The existence of negative eigenvalues in Fig.2(B) suggests that the posterior cannot be regarded as being approximately normal. This is not surprising, since hierarchical models are known to be often singular [29]. 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 numerically integrating the Laplace approximation of the likelihood conditioned on each combination of hyperparameter values (see Appendix B for technical details). 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.

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, a propensity function. We successfully calculated 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.

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 the larger models. These findings have non-trivial implications for future modelling based on GP and its implementations, because our results indicate that 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 [10], 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 (Eq.(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 [16] 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.3), 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. Because this effect of parallelism is also observed in the multiplication of much larger matrices [30], we do not attempt to extend the range of dimension of our experiment. The computation time saved by dynamically programmed eigendecomposition was not striking in our simulation study; however, this effect could be much greater for larger model sizes, because eigendecomposition requires $O\left(d^3\right)$ computation, which eventually outweighs the computation of the gradient flow at sufficiently large d. We could not directly observe this phenomenon because tuning the step size and the threshold parameter for the soft-absolute Hessian metric becomes difficult for much larger d. We also note that the scaling of wall time for eigendecomposition depends heavily on the available GPU card. Several years ago, it was shown that the Jacobi method is outperformed by the divide-and-conquer algorithm for matrices larger than $512\times 512$ on a NVIDIA Tesla K40 GPU card [31]; however, a recent study using dynamically programmed eigendecomposition for matrix optimisation showed that one sweep of the cyclic Jacobi method for matrices of size $4,096\times 4,096$ takes much less time than the divide-and-conquer algorithm for matrices of the same size [32]. Because the dynamic programming successfully reduced the number of sweeps to fewer than two, regardless of model size, in our study, we expect that users can also benefit from the dynamically programmed eigendecomposition in larger models if a recent version of a GPU card is used.

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 from which RMHMC swiftly moved away. Such entrapment is a well-known phenomenon and was the main motivation for the development of RMHMC [21]. 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 [33] and RMHMC [34], but it has not yet been properly implemented for the latter. The friction mechanism introduced to Euclidean Langevin MC [35,36] 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.

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 favourable 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 [4]. For this purpose, reversible-jump mechanisms [37,38] 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 those for time-series data must sometimes be used. In this case, RMHMC for GP must be combined with other MC methods, such as sequential MC [39,40]. 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) [41] (see our previous work [42] 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 (Fig.2(D)), whereas DML relies on the asymptotic normality of the estimators. The model is expected to be regular for fixed hyperparameter values and 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 work.