Approximate cross-validated mean estimates for Bayesian hierarchical regression models

In linear mixed models (LMM), the continuous response vector, $Y$, follows

$$Y|\beta,\phi\sim N(X\beta,\phi^{2}),$$

(1)

where $X$ are the non-random covariates, $\beta\coloneqq(\beta_{1}^{\prime},\beta_{2}^{\prime})^{\prime}\in I\!\!R^{P}$ includes fixed effects $\beta_{1}$ and random effects $\beta_{2}$, $\phi\in I\!\!R^{+}$ is a positive scalar.

$$\begin{split}\beta\sim N\left(\alpha,\begin{bmatrix}C&0\\ 0&\Sigma\end{bmatrix}\right),\\ \Sigma\sim f_{\Sigma}(\Sigma),\quad\phi\sim f_{\phi}(\phi),\end{split}$$

(2)

where $C\in I\!\!R^{P_{1}\times P_{1}}$ is a fixed positive-definite matrix and often a diagonal matrix with large positive values. The hyperparameters, $\Sigma\in I\!\!R^{P_{2}\times P_{2}}$, can take many forms, as long as it is positive-definite, e.g., models with Gaussian Markov random fields sample from the space of all possible $\Sigma$’s such that a variable in $\beta_{2}$ is independent of all others, given its neighborhood. Note that $f_{\phi}(\phi)$ in our setting is a fixed proper prior. Hence $\phi=\mathcal{O}_{P}\left(1\right)$, which means that for any $\epsilon>0$, there exists $\delta>0$ such that $P(|\phi|\leq\delta)>1-\epsilon$. In addition, from $\phi>0$, we know that $\phi^{-1}=\mathcal{O}_{P}\left(1\right)$. The probabilities for $\phi$ and $\phi^{-1}$ come from their prior distributions following the Bayesian point of view. Finally, we take $\alpha=0$ throughout this paper without loss of generality.

The posterior mean for $\beta$ conditioned on variance parameters has the form

$$E(\beta|\Sigma,\phi,Y)=\phi^{-2}VX^{\prime}Y,\quad V=\left(\phi^{-2}X^{\prime}X+\begin{bmatrix}C^{-1}&0\\ 0&\Sigma^{-1}\end{bmatrix}\right)^{-1},$$

(3)

which is analogous to the form of a Frequentist linear regression. We propose to use the full-data posterior mean estimates $\hat{\Sigma}=E(\Sigma|Y)$ and $\hat{\phi}=E(\phi|Y)$ as plug-in estimators in (3) to derive cross-validated mean estimates. In the linear regression setting, this results in a simple and straightforward closed-form solution for the cross-validated mean estimates:

$$\hat{Y}_{\mathpzc{s}_{j}}^{\text{AXE}}=E(X_{\mathpzc{s}_{j}}\beta|Y_{-\mathpzc{s}_{j}},\hat{\Sigma},\hat{\phi})=\hat{\phi}^{-2}X_{\mathpzc{s}_{j}}\left(\hat{\phi}^{-2}X_{-\mathpzc{s}_{j}}^{\prime}X_{-\mathpzc{s}_{j}}+\begin{bmatrix}C^{-1}&0\\ 0&\hat{\Sigma}^{-1}\end{bmatrix}\right)^{-1}X_{-\mathpzc{s}_{j}}^{\prime}Y_{-\mathpzc{s}_{j}}.$$

(4)

The posterior variance of fixed effects $\beta_{1}$, $C$, is not estimated and is taken as the diagonal matrix of prior variances for $\beta_{1}$. One common practice is to take $C^{-1}$ as the matrix of $0$s, which corresponds to the assumption that the fixed effects have infinite variance.

The AXE procedure in (4) essentially proposes using a Frequentist cross-validation to estimate $E(Y_{\mathpzc{s}_{j}}|Y_{-\mathpzc{s}_{j}})$ in place of a full Bayesian estimation of the posterior predictive density. AXE is likewise $\mathcal{O}\left(N^{2}P+P^{3}\right)$ in time for each CV fold. In comparison, Gibbs sampling of the same problem (when available) is $\mathcal{O}\left(MN^{3}P+MNP^{2}+MP^{3}\right)$ in time for each CV fold, where $M$ is the number of MCMC iterations.

Our approach relies on the approximation that the conditional expectation of the posterior mean evaluated at the full-data posterior means of the variance parameters, $\hat{\Sigma}=E(\Sigma|Y)$ and $\hat{\phi}=E(\phi|Y)$, approximates the desired posterior predictive mean, i.e., $E(X_{\mathpzc{s}_{j}}\beta|Y_{-\mathpzc{s}_{j}},\hat{\Sigma},\hat{\phi})\approx E(X_{\mathpzc{s}_{j}}\beta|Y_{-\mathpzc{s}_{j}})$. This looks similar to the approximation in Kass and Steffey (1989) but here we are in a cross-validation setting and we are plugging in the full-data estimates (as opposed to the training data estimates). In Section 2.4, we show that the approximation holds asymptotically under leave-one-cluster-out cross-validation.

Generalized linear mixed models (GLMMs) assume that $Y|\beta$ is not normally distributed. We follow the approximate inference in GLMM developed by Breslow and Clayton (1993) and derive the corresponding AXE for GLMM. Suppose $Y|\beta$ are independent with some probability density function $\pi$, so that

$$Y_{i}|X_{i}\beta\sim\pi(u_{i},a_{i},\phi^{2},v(u_{i})),$$

(5)

where $E(Y_{i}|X_{i}\beta)=u_{i}$, $\mathrm{var}(Y_{i}|X_{i}\beta)=a_{i}\phi^{2}v(u_{i})$, $a_{i}$ is a known constant, $\phi^{2}$ is a dispersion parameter. The regression model assumes that $g(u_{i})=X_{i}\beta$, where $g(.)$ is the link function and the prior distribution of $\beta$ is specified in (2).

Note that, the normal priors on the coefficients $\beta$ are no longer conjugate, and the analytic solution of (4) is unavailable. Instead, an iterative weighted least square algorithm (IWLS) is commonly used to fit GLM or GLMM (Holland and Welsch, 1977; Green, 1984). Defining a working response,

$$\tilde{Y_{i}}=X_{i}\beta+(Y_{i}-u_{i})g^{\prime}(u_{i}),$$

and a $n\times n$ diagonal weight matrix, $W$, with diagonal terms

$$w_{i}=\{\phi^{2}a_{i}v(u_{i})[g^{\prime}(u_{i})]^{2}\}^{-1},$$

i.e., for the logistic model, $g(u_{i})=\log(u_{i})-\log(1-u_{i})$, $g^{\prime}(u_{i})=\frac{1}{u_{i}(1-u_{i})}$, and $w_{i}=u_{i}(1-u_{i})$; for the Poisson model, $g(u_{i})=\log(u_{i})$, $g^{\prime}(u_{i})=\frac{1}{u_{i}}$, and $w_{i}=u_{i}$. IWLS repeatedly updates the working response vector and the weight matrix until its convergence. Given $W$, Harville (1977) shows that the best linear unbiased estimation (BLUE) of $\beta$ can be obtained from the associated weighted linear model,

$$\tilde{Y}|\beta,\sigma^{2}\sim N(X\beta,W^{-1}).$$

In the cross-validation setting, we approximate the training data weights, $W_{-\mathpzc{s}_{j}}$, by the corresponding weights at the convergence of the full data analysis. It is reasonable under the conditions we present in Section 2.1: 1) $E(X_{\mathpzc{s}_{j}}\beta|\tilde{Y}_{-\mathpzc{s}_{j}},\dot{\Sigma},\dot{\phi})\approx E(X_{\mathpzc{s}_{j}}\beta|\tilde{Y}_{-\mathpzc{s}_{j}})$, where $\dot{\Sigma}=E(\Sigma|\tilde{Y}_{-\mathpzc{s}_{j}})$, $\dot{\phi}=E(\phi|\tilde{Y}_{-\mathpzc{s}_{j}})$ and 2) $E(\Sigma|\tilde{Y})\approx\dot{\Sigma}$, $E(\phi|\tilde{Y})\approx\dot{\phi}$.

We plug in the working response vector, $\tilde{Y}_{-\mathpzc{s}_{j}}$, the weight matrix, $W_{-\mathpzc{s}_{j}}$, and other variance estimates, $\hat{\Sigma}$, in the full data analysis, and obtain the following AXE estimates for the test data working response, $\tilde{Y}_{\mathpzc{s}_{j}}$,

$$\hat{\tilde{Y}}_{\mathpzc{s}_{j}}^{\text{AXE}}=E(X_{\mathpzc{s}_{j}}\beta|\tilde{Y}_{-\mathpzc{s}_{j}},W_{-\mathpzc{s}_{j}},\hat{\Sigma})=X_{\mathpzc{s}_{j}}\left(X_{-\mathpzc{s}_{j}}^{\prime}W_{-\mathpzc{s}_{j}}X_{-\mathpzc{s}_{j}}+\begin{bmatrix}C^{-1}&0\\ 0&\hat{\Sigma}^{-1}\end{bmatrix}\right)^{-1}X_{-\mathpzc{s}_{j}}^{\prime}W_{-\mathpzc{s}_{j}}\tilde{Y}_{-\mathpzc{s}_{j}}.$$

(6)

Finally, we transform the test data working response to the original scale by $\hat{{Y}}_{\mathpzc{s}_{j}}^{\text{AXE}}=g^{-1}(\hat{\tilde{Y}}_{\mathpzc{s}_{j}}^{\text{AXE}})$. If there is any concern that the weight matrix might differ between the full-data analysis and the training data analysis, then we can run a few more iterations of the IWLS algorithm. We start from the AXE estimates of the working responses and iteratively update the training data weights and the working responses. In practice, $\hat{\tilde{Y}}_{\mathpzc{s}_{j}}^{\text{AXE}}$ without additional IWLS iterations works well.

One can apply AXE to any CV design, of which leave-one-out (LOO-CV) and leave-one-cluster-out (LCO-CV) are popular choices. For both LOO-CV and LCO-CV, the number of CV folds can grow with the amount of data, making the computation expensive. Thus they typically benefit the most from CV approximation methods. We focus on proving AXE convergence under LCO-CV, of which LOO-CV is a special case. We first define LCO-CV and related notation. Then, we introduce and discuss existing LCO-CV approximation methods.

In LOO-CV, a single observation is held out as test data for each CV fold, and the number of CV folds equals the number of data points, $N$. LOO-CV is commonly used because it maintains as much similarity to the full data model posterior as possible while evaluating the fit to new data. LCO-CV is often used in models with clustered data and cluster-specific parameters, e.g., the patient-specific random effects in a longitudinal model with $J$ patients. When the data are sampled uniformly at random to form CV folds, data from one patient can be split between the training and test data. It introduces correlation between the training and test data and can result in overfitting to the patient-specific parameters (Arlot et al., 2010; Opsomer et al., 2001). To fairly evaluate the predictive capability for a new cluster, LCO-CV defines $J$ CV folds corresponding to $J$ clusters and withholds all observations informing the estimation of a cluster-specific random effect.

LCO-CV has been a challenging CV design for approximate inference. Existing CV approximation methods use the full-data posterior density $f(\beta|Y)$ to approximate the training data posterior density $f(\beta|Y_{-\mathpzc{s}_{j}})$. Under LCO-CV, these approximations can be inaccurate for two reasons. First, since the test data corresponds to the data in a cluster, the test data sample size can be large. It is shown that the larger the sample size of the test data, the more likely it is that the difference in densities $f(\beta|Y)$ and $f(\beta|Y_{-\mathpzc{s}_{j}})$ is large (Gelfand et al., 1992). Second, under cross-validation, the cluster-specific random effects, $\theta_{j}$, cannot be estimated from the training data, relying instead on samples drawn from $\theta_{j}|\theta_{-j},\Sigma$.

To focus on the random effects which pertain to the LCO-CV design, we split $\beta$ into $\theta$ and the remaining coefficients in $\beta$:

$$\beta=(\beta_{/\theta}^{\prime},\hskip 5.0pt\theta^{\prime})^{\prime},$$

(7)

where $\theta$ refers to the part of random effects that pertain to the CV design, and $\beta_{/\theta}$ refers to the components in $\beta$ other than $\theta$, which contain all fixed effects and the part of random effects that do not pertain to the CV design. $\theta\coloneqq(\theta_{1},\theta_{2},\ldots,\theta_{J})^{\prime}\in I\!\!R^{K}$ ($J\leq K\leq P_{2}$). Similarly. we split $X$ column-wise into

$$X=[X_{\beta_{/\theta}}\hskip 5.0ptX_{\theta}],$$

(8)

where $X_{\theta}$ contains the subset of covariates that correspond to $\theta$ and $X_{\beta_{/\theta}}$ contains the remaining columns of $X$.

For the linear mixed models (LMM) case, we first show that the full-data posterior means for variance parameters $\Sigma$ and $\phi$ approximate the training data posterior means; see Theorem 2.1. We then show the convergence of the AXE approximation in Corollary 2.2.

For proof, see Appendix A, included in Supplemental Materials.

As the AXE estimate is a straightforward function of $\Sigma$ and $\phi$, we can use the result of Theorem 2.1 above to derive an overall error bound for $E(X_{\mathpzc{s}_{j}}\beta|Y_{-\mathpzc{s}_{j}},\hat{\Sigma},\hat{\phi})$.

For proof, see Appendix A.

As the conditional expectation approximates the posterior predictive mean with error rate $\mathcal{O}_{P}\left(n_{j}/N\right)$, AXE approximates the posterior predictive mean with overall error rate $\mathcal{O}_{P}\left(n_{j}/N\right)$. In the LMM case, we have generally found that the accuracy of AXE largely depends on the accuracy of using the posterior means of the variance parameters plug-in estimators. In the GLMM case, AXE introduces the same approximations to the IWLS algorithm.

In general, when $\Sigma$ and $\phi$ are estimated with narrow posterior densities, $E(\Sigma|Y_{-\mathpzc{s}_{j}})$ and $E(\phi|Y_{-\mathpzc{s}_{j}})$ will not deviate much from their full-data posterior estimates and AXE will be accurate. This is also when LCO-CV is the most expensive with a large number of clusters and AXE provides the most benefit. In our experience with AXE, even when posterior densities are wide, we have found that AXE can perform well (e.g., Radon subsets and Eight schools examples, see Figure 1). The cases where AXE is inaccurate are typically those with a low number of clusters or severe data imbalance such that removal of the test data severely impacts the estimation of $\Sigma$ or $\phi$.

Intuitively, when the number of clusters is low, running manual cross-validation (MCV) may not be of concern. If additional assurance is needed, we suggest running MCV on a randomly selected subset of CV folds, stratified by clusters. By comparing the MCV point estimates to the AXE approximation, we can determine how similar the conclusions derived from the AXE approximation are to those using the MCV estimates. One common model evaluation criteria is the root mean square error (RMSE), where the error is the difference between $\hat{Y}_{\mathpzc{s}_{j}}^{(\text{MCV})}=E(Y_{\mathpzc{s}_{j}}|Y_{-\mathpzc{s}_{j}})$ and $Y_{\mathpzc{s}_{j}}$. We take the log ratio of AXE-approximated RMSE to ground-truth MCV RMSE,

$$\text{LRR}_{j}=\log\left(\frac{\sum_{i=1}^{n_{j}}(\hat{Y}_{i}^{(\text{AXE})}-Y_{i})^{2}}{\sum_{i=1}^{n_{j}}(\hat{Y}_{i}^{(\text{MCV})}-Y_{i})^{2}}\right),$$

(9)

where LRR_j is calculated separately for each chosen CV fold $j$ to obtain more fine-grained comparisons. A low $|LRR|$ represents high similarity between the approximate RMSE and ground-truth MCV RMSE. We use LRR as our criterion rather than $\sum(\hat{Y}^{\text{AXE}}-\hat{Y}^{MCV})^{2}$, because the magnitude of the latter depends on the standard error for $\hat{Y}^{\text{AXE}}$; it is possible that the approximate RMSE is close to the true RMSE while $\sum(\hat{Y}^{\text{AXE}}-\hat{Y}^{MCV})^{2}$ is large. We also use LRR to compare AXE approximations to other LCO methods in Section 4 by replacing $\hat{Y}_{i}^{(\text{AXE})}$ in (9) with the alternative LCO method’s point estimate for $Y_{i}$.

If exchangeability is assumed between clusters, then the sample mean and variance of $\text{LRR}_{j}$’s can provide inference for the expected accuracy of the approximation. When the variance parameters $\Sigma$ and $\phi$ are not well-estimated, they are less likely to be consistent across CV folds; in these cases, we expect a large standard deviation among the AXE $\text{LRR}_{j}$’s across CV folds due to the instability of the variance parameter posterior means. In practice, we recommend selecting those clusters with particularly large $n_{j}$, or for whose estimated random effects are more extreme, and if the mean or standard deviation of $\text{LRR}_{j}$’s is over some threshold $\delta$, we recommend running MCV. The choice of $\delta$ represents the tolerable degree of error in approximating MCV RMSE. For our examples, we use $\delta=0.25$ as our threshold value, which translates to a ratio of AXE RMSE to MCV RMSE between 0.78 and 1.28. We refer to this additional validation step in our calculations of time cost as AXE+ in Section 4.

We compare AXE to the CV approximation methods described in Section 3 by calculating LRR_j for each method and CV fold $j$, as in (9). For IJ and NS, we calculate MCV RMSE using the posterior mode.

Figure 1 shows the proportion of CV folds with $|\text{LRR}|\leq x$ against different cutoff values, $x$. We call these line plots “LRR percentage curves”. A perfectly accurate method would have 100% of $|\text{LRR}|\text{s}=0$. An $|\text{LRR}|$ larger than $\log(2)$ corresponds to an approximate RMSE that is over 100% different from the ground truth MCV RMSE, making it a poor approximation; thus we focus on $x\in[0,\log(2)]$ on the x-axis. Figure 1 presents three top-performing methods: AXE, Ghost, and iIS. Comparisons across all methods are provided in Appendix E. AXE is consistently one of the better-performing methods across our examples, which is not the case for any other method. We examine each panel in detail below.

Figure 1A shows an example where $n_{j}=1$ and LCO-CV is the same as LOO-CV. AXE and both iIS methods outperform all others, with roughly 70% of $|\text{LRR}|\leq 0.1$ compared to GHOST’s 25%. As GHOST assumes $f(\beta_{/\theta}|Y)=f(\beta_{/\theta}|Y_{-\mathpzc{s}_{j}})$, the Ghost mean estimate for this example is $\bar{Y}$, resulting in larger LRRs, while AXE and the iIS methods have mean estimate $\bar{Y}_{-\mathpzc{s}_{j}}$ (see, e.g., Eight schools in Figure 2B).

Figure 1B contains results for the full Radon data, which is a relatively rich data set compared to the number of parameters to estimate. As such, AXE, iIS, and GHOST all have over $97$% of $|\text{LRR}|$s under $0.1$. A few CV folds do have relatively large $|\text{LRR}|$; these correspond to the same two CV folds across all LCO methods and are cases where the data are particularly informative for the random effect variance.

Figure 1C contains results for the Radon subsets data, where the training data are a randomly selected subset of counties. AXE is the best-performing method with over $99$% of $|\text{LRR}|\text{s}\leq 0.1$. GHOST and iIS perform worse, with $71$% and $39$% of $|\text{LRR}|\text{s}\leq 0.1$, respectively. Each of those three methods assumes $f(\beta_{/\theta}|Y)$ and $f(\beta_{/\theta}|Y_{-\mathpzc{s}_{j}})$ are similar, but the low number of counties in the training data means that estimation of the fixed effect coefficient that describes the county-level heterogeneity can be quite different under cross-validation. We also note that in a closer examination of these results, we found that GHOST tended to underestimate RMSE for the Radon subsets data, while iIS overestimated RMSE. As GHOST uses the full-data posteriors directly, it follows that it produces over-optimistic results.

In Figure 1D, iIS is able to produce stable estimates because it incorporates information from $\beta_{/\theta}|Y$, but it is outperformed by AXE which has a nearly perfect LRR percentage curve. The good performance of AXE suggests that posterior expectations for $\beta_{/\theta}$ and $\Sigma$ were relatively stable across CV folds. In panel E, LRRs are in general larger for all methods than preceding panels (see Table LABEL:tab:lrrtable, in Appendix F). This is partly due to small sample sizes ($n_{j}=1$). In panel F, AXE yields satisfactory results with $92$% of $|\text{LRR}|\text{s}\leq 0.1$ while the remaining LCO approximation methods perform poorly.

We summarize each method’s performance using the area under the LRR percentage curves (AUC-LRRP) of Figure 1. An ideal method would have AUC-LRRP of $\log(2)$, and a method with all $|\text{LRR}|$s $>\log(2)$ would have AUC-LRRP of 0. We normalize AUC-LRRP and divide its value by $\log(2)$ so that the maximum is 1 and the minimum is 0. Table 4 gives AUC-LRRP$/\log(2)$ for each method and example. Further summary metrics for each method and example, consisting of mean LRR and the standard deviation of LRR, are provided in Appendix F.

Table 4 shows that AXE most frequently has the highest AUC-LRRP, except for Eight schools where it had the second highest. The main weakness of Vehtari and iIS is the computational expense. They are the most expensive of all the methods and, in some of our examples, require more time than manual cross-validation, as discussed in Section 4.3.

Both NS and IJ performed worse than the majority of other CV approximation methods. We note that our examples have fewer observations in comparison to the number of model parameters than is typical for NS or IJ, which is usually applied to data large enough such that cross-validation using standard optimization methods is prohibitively expensive. In such situations, the large amount of data typically means that posterior densities are very narrow and $f(\beta_{/\theta},\Sigma,\phi|Y)$ is more similar to $f(\beta_{/\theta},\Sigma,\phi|Y_{-\mathpzc{s}_{j}})$. When there are fewer observations and, subsequently, more uncertainty in the model, a small change in $w$ can result in a larger change in the log-likelihood. Both methods are often limited to LOO-CV, as in Wang et al. (2018); Rad and Maleki (2020); Giordano et al. (2019) and Stephenson and Broderick (2020), which typically results in smaller changes to the log-likelihood over CV folds. We note that all of the CV approximation methods included in this paper improve as the amount of data increases; the main advantage for IJ and NS is that their computing times may scale more easily to LOO-CV with large data sets.

We provide graphical point-by-point comparisons of $\hat{Y}^{(\text{method})}$ to ground-truth cross-validated $\hat{Y}^{(\text{MCV})}$. Figure 2A contains scatter plots of the AXE approximation for each data point $\hat{Y}$ against actual MCV values for all examples. The vast majority of points lie on or near the 45-degree line. In many cases, the AXE approximation is point-by-point equivalent to MCV estimates. It is in contrast to the scatter plots in Figure 2B, which contain AXE approximations as well as GHOST and iIS, both of which performed relatively well, based on Figure 1. iIS has a higher variance than AXE for the Radon subsets, ESP, and SRD data sets, with approximations farther from the 45-degree line. GHOST has a higher variance in the Eight schools and Radon subsets data sets. All three methods perform similarly in the Radon data sets.

We found that the accuracy of AXE was not particularly impacted as the proportion of test data increased, but it improved greatly as the number of clusters increased. For instance, some of the largest deviations from the 45-degree line are in row B, the Radon subsets data, where the number of clusters ranges from 3 to 12. It is intuitive, as the estimation of $\Sigma$ depends largely on cluster-specific intercepts, so when there are fewer clusters, $E(\Sigma|Y_{-\mathpzc{s}_{j}})$ is less stable across CV folds.

Table 5 provides the total computation time on an Intel i7-8700k CPU with 40 Gb of memory, to aid with running examples in parallel; note a standard 8 Gb of memory would be sufficient for the examples. Except for the SRD data, models were fit using STAN or the rstanarm package, with four chains of 2000 samples and a 1000-sample burn-in, resulting in 4000 total samples. For the SRD data, the CAR.ar() function from the CARBayesST package was used, for four chains of 220,000 samples each, with a 20,000 sample burn-in and thinned to produce 4000 total posterior samples. Calculations for IJ were conducted in Python using the autograd package. AXE took a reasonable amount of time for all examples: a few seconds for Eight schools, Radon, Radon subsets, and SLC, about one minute for ESP, and 10 minutes for SRD. To date, Vehtari has been used only in the LOO-CV case (Vehtari et al., 2016; Vanhatalo et al., 2013). When CV folds contained multiple data points, Vehtari took much longer than other approximation methods.