Fully Bayesian inference for latent variable Gaussian process models

We consider GP regression models of the form:

Suppose we have a set of $N$ observations $\mathbf{Y}=\left(y_{1},\ldots,y_{N}\right)$ (e.g., from N runs of some complex code that simulates a physical system) at inputs $\textbf{W}=\left[\textbf{w}_{1},\ldots,\textbf{w}_{N}\right]^{\mathsf{T}}$. By the marginalization properties of jointly Gaussian variables, the posterior distribution of $f\left(\textbf{w}_{*}\right)$ at input $\textbf{w}_{*}$ conditioned on these observations, is again Gaussian with mean and variance given by:

The GP hyperparameters $\boldsymbol{\theta}$ are usually unknown in practice. A common strategy is to use a plug-in Bayes approach, also known as Empirical Bayes, where an estimate $\widehat{\boldsymbol{\theta}}_{N}$ is first computed from data, and then plugged into the posterior predictive distribution. Under this plug-in approach, the posterior distribution of $f\left(\mathbf{w}_{*}\right)$ is again Gaussian whose mean and variance are calculated by plugging in $\widehat{\boldsymbol{\theta}}_{N}$ into (3) and (4) respectively. A common strategy for estimation is to optimize $\boldsymbol{\theta}$ with respect to either the log-likelihood

A more rigorous approach to obtain the posterior predictive distributions is to use fully Bayesian inference, where one marginalizes over the posterior distribution of $\boldsymbol{\theta}$ given data. The density of the posterior distribution of $f\left(\mathbf{w}_{*}\right)$ given the data is

There is usually no closed form expression for the above distribution, as the integrals involved are often intractable. In practice, one often uses MCMC methods to draw a finite number of samples from the distribution $p\left(\boldsymbol{\theta}|\mathbf{W},\mathbf{Y}\right)$, which can then be used to compute different quantities of interest. In this work, we use the No-U-Turn-Sampler [20] algorithm for drawing the MCMC samples.

Standard GP correlation functions, such as the squared exponential, Matérn, power exponential, etc. cannot be directly applied to qualitative inputs because there is no inherent distance metric in the qualitative space. The recently proposed LVGP approach [39] provides a natural and convenient way to handle qualitative inputs by mapping the levels of each qualitative factor onto low-dimensional continuous latent spaces. To be more specific, consider a qualitative input $t$ with $L$ levels. The LVGP approach constructs a mapping $\mathbf{z}:\{1,\ldots,L\}\xrightarrow{}\mathbb{R}^{d}$, where $d<L$ is the dimension of the latent variable (LV) space. With this mapping, the LVGP approach defines kernels of the form

The LVGP is based on the premise that for any real physical system, the effects of a qualitative input $t$ must be due to some underlying (and perhaps very high-dimensional) quantitative variables $\{v_{1},v_{2},\ldots\}=\{v_{1}\left(t\right),v_{2}\left(t\right),\ldots\}$. Consider the situation in Figure 1, where the three levels of $t$ are associated with points on the high-dimensional space of $\{v_{1},v_{2},\ldots\}$. The LVGP uses sufficient dimensionality reduction arguments to implicitly construct a low dimensional representation $\mathbf{z}\left(t\right)=g\left(v_{1}\left(t\right),v_{2}\left(t\right),\ldots\right)$ that accounts for most of their effects. Based on empirical studies, [39] argued that $d=2$ was often a good choice.

When there are multiple qualitative inputs, a separate latent space is used for each qualitative input. Let $\mathbf{z}^{(j)}:\{1,\ldots,L_{j}\}\xrightarrow{}\mathbb{R}^{d_{j}}$ denote the $d_{j}$-dimensional LV mapping for the $j^{\mathrm{th}}$ qualitative input $t_{j}$. If the squared exponential kernel is used for both the quantitative variables and the LVs (i.e., $\widetilde{k}\left(\mathbf{z}\left(t\right),\mathbf{z}\left(t^{\prime}\right)\right)=\exp\left[-\frac{1}{2}\left\lVert\mathbf{z}\left(t\right)-\mathbf{z}\left(t^{\prime}\right)\right\rVert^{2}\right]$) for the qualitative inputs, the covariance kernel for the LVGP is given by

The LVs are additional hyperparameters that must be estimated along with the other GP hyperparameters. Existing works in the literature employing the LVGP have all used maximum likelihood estimation followed by the plug-in Bayes approach for generating predictions. The impact of estimation of LVs on the performance of the LVGP model is yet to be studied. In the remainder of the paper, we demonstrate that the impact of the estimation uncertainty in the LVs can be substantial, and that the fully Bayesian approach that we develop for LVGPs appropriately accounts for the uncertainty, improving both the prediction accuracy and the UQ.

For fully Bayesian inference, we need to specify prior distributions for the LVs and for the other GP hyperparameters. In this section, we will discuss our choice of the prior distribution for the LVs. Prior distributions used for the other GP hyperparameters are standard choices. In the following, we drop the superscript $(j)$ from the LVs for notational convenience. For a qualitative variable $t$ with $L$ levels and its LV mapping $\mathbf{z}$, let $\mathbf{z}\left(l\right)=\left(z_{l1},\ldots,z_{ld}\right)$ denote the mapped values for level $l\in\{1,\ldots,L\}$.

The prior on the LVs for a qualitative variable $t$ can be used to encode any domain knowledge (if available) about the similarity between the different levels of $t$ in terms of their effects on the response. This domain knowledge may be available in multiple forms. For example, we could be given an incomplete set of numerical descriptors that are either known to or suspected to have significant effects on the response. The Euclidean distances between these descriptor values could potentially be used to be define some appropriate prior similarity measure for the LVs. However, in this work we focus on the more general setting where no such domain knowledge exists.

In the absence of any domain information, we define the following weakly informative prior over the LVs for $t$:

Alternatively, one can place a completely non-informative prior on the LVs. In practice, one can approximate this via i.i.d. uniform priors with relatively large lower and upper limits (e.g., -10 and 10, respectively). With a completely non-informative prior, there is effectively no regularization on the effects of the qualitative variable. We find that the completely non-informative priors negatively impacted MCMC convergence, irrespective of the choice of the MCMC algorithm. We will, therefore, restrict our attention to the weakly informative prior (10).

With a stationary kernel over the LVs, the correlations between the different levels is dependent only on the pairwise distances between the corresponding LV values of the levels. Therefore, unconstrained LVs are subject to translation and rotation symmetries under the LVGP likelihood (see Figure 2 for illustration in a 2D LV space). While the prior (10) reduces translation symmetries in the LVs under the posterior, rotation symmetries still exist - the two sets of LVs $\{\mathbf{z}\left(l\right):l\in\{1,\ldots,L\}\}$ and $\{\mathbf{Q}^{\mathsf{T}}\mathbf{z}\left(l\right):l\in\{1,\ldots,L\}\}$ have the same posterior probabilities, where $\mathbf{Q}$ is any orthogonal matrix in $\mathbb{R}^{d}$. This rotational symmetry in the LVs can potentially reduce sampling efficiencies of the MCMC scheme, which would necessitate large number of MCMC draws, thereby increasing the computational expense of fully Bayesian inference. To eliminate these symmetries, [39] fix the coordinate frame of reference. For example, in a 2D LV space, they constrain the mapped LVs for the first level, $\mathbf{z}\left(1\right)$, at the origin, and those of the second level, $\mathbf{z}\left(2\right)$, to lie on the horizontal axis. For a general $d$, they set $\mathbf{z}\left(1\right)=\mathbf{0}_{d}$, $\mathbf{z}\left(2\right)=\begin{bmatrix}z_{21}&\mathbf{0}_{d-1}\end{bmatrix}$, and $\mathbf{z}\left(l\right)=\begin{bmatrix}z_{l1}&\ldots&z_{l\left(l-1\right)}&\mathbf{0}_{d-l+1}\end{bmatrix}$ for $2<l\leq d$. This results in $d\left(d-1\right)/2$ fewer free parameters. However, a challenge with this parameterization for fully Bayesian inference is that the prior distribution cannot represent independent and identically distributed LVs, because this parameterization treats the first $d$ levels differently from the other levels. For example, under independent normal or uniform prior distributions with mean at 0, the LVs for any level would be closer on average to the first level than to any other level. This treatment of the LVs is problematic, because the ordering of the levels is arbitrary, and it can affect the quality of the LVGP model, especially with a small number of training observations, or when some of the levels are unobserved in the training data.

To rectify this, we instead modify the prior (10) by first defining “raw” LV parameters $\{\widetilde{\mathbf{z}}\left(l\right)=(\widetilde{z}_{l1},\ldots,\widetilde{z}_{ld}):l\in\{1,\ldots,L\}\}$, which are i.i.d. for all levels, and then transforming them to the coordinate frame of reference used by [39] to obtain the actual LVs. We first consider the case of a 2D LV space. The resulting two-stage prior is as follows:

For $d>2$, the $\mathbf{R}\left(\cdot\right)$ matrix in (12) is replaced by a product of $d(d-1)/2$ Givens rotation matrices in $d$ dimensions. Givens rotation matrices in $d$-dimensions, $\mathbf{R}_{ij}\left(\phi_{ij}\right)$, are $d\times d$ matrices that take the form of the identity matrix except for the $(i,i)$ and $(j,j)$ positions which are replaced by $\cos\left(\phi_{ij}\right)$, and the $(i,j)$ and $(j,i)$ positions which are replaced by $-\sin\left(\phi_{ij}\right)$ and $\sin\left(\phi_{ij}\right)$ respectively. The LVs for a general $d$ are given by

With this formulation, for a qualitative variable with $L$ levels and using a $d$-dimensional LV spave, there are $Ld+1$ parameters to estimate. And so, with $J$ qualitative variables, we have $J(Ld+1)$ parameters to estimate. Algorithm 1 has an $O(Ld^{2})$ computational cost for a qualitative variable with L levels and d number of latent variable dimensions. However, the main computational cost in each posterior evaluation comes from computing the Cholesky factorization, which is $O(N^{3})$.

Finally, for a prior distribution on $\gamma$, we consider the Gamma distribution. The support of the Gamma distribution is the set of positive real numbers. Larger values of $\gamma$ result in larger amounts of regularization on the LVs. The Gamma distribution has two additional hyperparameters, the concentration $\alpha$ and the rate $\beta$. We restrict $\alpha>1$ and set $\beta=\alpha-1$, so that the mode of the distribution is at $\gamma=1$. This results in only a single hyperprior parameter. We find that the convergence and performance of the model is not sensitive to the choice of $\alpha$. Refer to Appendix A for more details. In our numerical experiments in Section 5, we use $\alpha=2$.

In this section, we briefly discuss our choice of the MCMC algorithm for drawing samples from the posterior distribution of the LVGP hyperparameters, and then we state the formulae and procedures for obtaining predictions and confidence intervals on these predictions.

The joint posterior distribution for the LVs and the other GP hyperparameters is typically high dimensional. The MCMC algorithm should have a few desired properties in order to efficiently converge to this posterior distribution. Firstly, the algorithm needs to suppress random walk behavior for the Markov chain to mix well in high dimensions. Otherwise, the chain will take too long to converge to the posterior, and thereby require many evaluations for the posterior. Secondly, if the algorithm has additional tuning parameters, there must be a mechanism to tune them during the burn-in or warm-up phase. Among the various algorithms that we have tested, the No-U-Turn-Sampler [20] algorithm¹¹1as implemented by the NumPyro [28, 7] library in Python satisfies both properties, and is found to work well for our case.

Under MCMC sampling, the posterior predictive distribution of $f\left(\mathbf{w}_{*}\right)$ at input $\mathbf{w}_{*}$ is approximated as a Gaussian mixture distribution

In this work, we define UQ in terms of prediction intervals on the predicted response (as a function of the inputs). Unfortunately, there are no analytical expressions for the quantiles of a Gaussian mixture distribution. So, the variance term (19) cannot be used to directly compute the quantiles. Instead, we estimate the quantiles empirically by drawing samples from this distribution as shown in Algorithm 2. When computing the confidence intervals in our experiments in Section 5, we use about $M=10000$ samples. This is likely overkill, and we could obtain sufficiently accurate confidence intervals using as few as 500-1000 samples. However, the confidence interval computations are very fast, in part because the operations in Algorithm 2 can be easily vectorized. So, we erred on the side of caution and used a large number of samples.

The LV mapping provides an inherent ordering and structure for the different levels of a qualitative input, in terms of their effects on the response. Therefore, a secondary benefit of the LVGP approach is to provide an interpretation of the effects of these different levels. Interpreting the LVs with fully Bayesian inference is, however, less straightforward. As shown in Figure 3, there are multiple latent spaces, each corresponding to a different MCMC sample drawn from the posterior. Since the estimation uncertainty in the LVs can be large, especially with small training data sets, the variation in the MCMC latent space samples can be large. Due to this sample-to-sample variability, the different latent space samples can result in different interpretations of the relations among the levels, thereby making interpretation of the effects of the qualitative variable challenging.

To resolve this ambiguity in interpretation and to help average out sample-to-sample variability, we propose to find a common latent space that reconciles the differences between the different latent spaces by minimizing a discrepancy measure over these different latent spaces. Since the latent variables are ultimately used to compute the covariance matrix over the levels, we define the discrepancy measure in terms of differences in their resulting covariance matrices. In the following, we drop the subscript $j$ from the LVs for notational convenience. For a LV mapping $\mathbf{z}$ for a qualitative input $t$ with $L$ levels, let $\mathbf{Z}=\left[\mathbf{z}\left(1\right)\quad\mathbf{z}\left(2\right)\quad\ldots\quad\mathbf{z}\left(L\right)\right]^{\mathsf{T}}$ denote the corresponding $L\times d$ LV matrix. Let $\widetilde{\mathbf{k}}\left(\mathbf{Z}\right)$ denote the $L\times L$ covariance matrix across levels corresponding to the LV matrix $\mathbf{Z}$, where the element at the row $l_{1}$ and column $l_{2}$ is $\widetilde{k}\left(\mathbf{z}\left(l_{1}\right),\mathbf{z}\left(l_{2}\right)\right)$. We define the discrepancy measure between two LV matrices $\mathbf{Z}$ and $\mathbf{Z}^{\prime}$ (e.g., corresponding to two different MCMC draws of the LVs) as $\left\|\widetilde{\mathbf{k}}\left(\mathbf{Z}\right)-\widetilde{\mathbf{k}}\left(\mathbf{Z}^{\prime}\right)\right\|_{F}$, where $\left\|\cdot\right\|_{F}$ is the Frobenius norm.

Let $\mathbf{Z}_{(1)},\ldots,\mathbf{Z}_{(B)}$ be the $B$ latent variable matrices corresponding to the $B$ different hyperparameter samples from the posterior. We find a representative latent variable matrix $\mathbf{Z}_{*}\in\mathbb{R}^{L\times d}$ as

Once $\mathbf{Z}_{*}$ is obtained, the representative LV space can be interpreted in a similar way as the LV space obtained from point estimates. For systems with multiple qualitative inputs, the optimization problem (20) is run separately for each qualitative input.

For illustration, consider the borehole function, that is commonly used to illustrate GP modeling [26]. The flow rate of water through a borehole, that is drilled from the ground surface through two aquifers is

The estimated representative LV spaces for the two training sets are shown in Figures 4b and 4c, respectively. In both cases, the $z_{1}$ LV corresponds roughly to $H_{l}$, and the $z_{2}$ corresponds fairly closely to $r_{w}$, more so for the set of 64 training observations. Moreover, levels with the same $r_{w}$ value are closer to each other than those with different $r_{w}$ values. Note that the distances between different levels in the (representative) LV space are larger in the case of the training set with 32 observations, and the agreement between the estimated LV space and the true LV space is not as tight as for the case of 64 training observations. This is likely due to larger estimation uncertainty with just 2 observations per level, and hence, greater discrepancies between the different posterior LV samples. Also, the effect of $H_{l}$ is not very clear for the smaller training set, since the orderings of the $H_{l}$ levels are different at different $r_{w}$ values. However, with the larger training set in Figure 4c, the effects of both $r_{w}$ and $H_{l}$ are more clear. In fact, we can conclude that $z_{1}$ approximately represents the effects of $H_{l}$, while the $z_{2}$ approximately represent the effects of $r_{w}$.

The above example shows that LVGP with Bayesian inference can offer useful interpretations even with small training datasets. As we collect more data, we can draw more definitive conclusions from the representative LV space.

We first compare the performance of the two methodologies on a set of engineering models that are commonly used for assessing surrogate models with numerical inputs. In each case, we discretize two numerical inputs and then construct a new qualitative factor $t$ whose levels represent discrete combinations of those two numerical inputs. In addition to the borehole model discussed in Section 3.3, we also consider the following models:

We now compare the estimation methodologies on the two different material design applications that motivated this work. In each case, we train independent LVGP models for multiple response properties of interest as a function of the atomic compositions of the materials, which constitute the purely qualitative inputs. In material design applications, the primary bottleneck is in obtaining physical measurements or first-principles calculations for the properties of interest, as these are typically expensive and time consuming [19]. Therefore, it is important to develop cheap surrogate models that can learn efficiently from (relatively small) available datasets.

We now consider a dataset of DFT computed formation energies of AB$\mathrm{O}_{3}$ perovskites [12]. These compounds are widely used for thermochemical water spitting applications. In these compounds, O is the oxygen atom, while A and B are variables that represent the cations that may occupy these two sites of the compound. There are 73 elements from the periodic table that can occupy either of these sites. The dataset contains a total of 5276 compounds for which formation energies are available. To test the approach, we randomly select 1000 observations as training data, and the remainder of the 5276 compounds are used as test data. We repeated this procedure for 25 replicates, where on each replicate we choose different random subsets to serve as training and test data. In each of the 25 different training sets, there were no levels with missing observations. For this application, fully Bayesian inference with all $N$ = 1000 observations was too prohibitively expensive to run across multiple replicates. Instead, we compare the MAP version of the LVGP model with all $N$ = 1000 observations (which we refer to as the exact MAP) versus both MAP and fully Bayesian versions of the two sparse LVGP approximations discussed in Section 4. For each sparse approximation, we set the number of inducing points to 50.

The RRMSE results are shown in Figure 15a. The first panel corresponds to the exact MAP LVGP, the second to the FITC approximation (MAP and fully Bayesian), and the third to the VFE approximation (MAP and fully Bayesian). Comparing the MAP versions, we can see that the sparse models have much better prediction accuracies than the exact MAP LVGP model, which may seem surprising. However, the optimization of the exact MAP LVGP model suffers from stability issues arising from computing the Cholesky decomposition of a much larger (covariance) matrix than that for the sparse LVGP models, which explains its poorer performance. The VFE model has slightly better prediction accuracy than the FITC model. For both models, there are significant improvements in prediction accuracy with fully Bayesian inference versus MAP.

The MIS results are shown in Figure 15b. Comparing the MAP versions, the VFE model has worse UQ than the exact LVGP model. However, the FITC model has much better UQ than the exact LVGP model. With fully Bayesian inference, there are large improvements in the quality of UQ for both the sparse models. The relative improvement in UQ for the VFE model with the fully Bayesian inference is much larger than that for the FITC models. This can be explained by their respective coverage probabilities, which are shown in Figure 15c. The coverage probability of the MAP FITC model is much larger than the of the MAP VFE model.

These results show that fully Bayesian inference for the LVs and the other hyperparameters results in improved performance even for larger training sets. However, these improvements come with a large increase in computation time. As shown in Figure 15d, the training time for fully Bayesian inference (which include the training time for obtaining the MAP estimates for the inducing points) is about an order of magnitude larger than that for obtaining the corresponding MAP estimates.