Introduction

When computing ELPD-based LOO-CV for a Bayesian model we need to compute the log leave-one-out predictive densities log p(y_i|y_−i) for every response value y_i, i = 1, …, N, where y_−i denotes all response values except observation i. To obtain p(y_i|y_−i), we need to have access to the pointwise likelihood p(y_i | y_−i, θ) and integrate over the model parameters θ:

p(y_i | y_−i) = ∫p(y_i | y_−i, θ) p(θ | y_−i) dθ

Here, p(θ | y_−i) is the leave-one-out posterior distribution for θ, that is, the posterior distribution for θ obtained by fitting the model while holding out the ith observation (we will later show how refitting the model to data y_−i can be avoided).

If the observation model is formulated directly as the product of the pointwise observation models, we call it a factorized model. In this case, the likelihood is also the product of the pointwise likelihood contributions p(y_i | y_−i, θ). To better illustrate possible structures of the observation models, we formally divide θ into two parts, observation-specific latent variables f = (f₁, …, f_N) and hyperparameters ψ, so that p(y_i | y_−i, θ) = p(y_i | y_−i, f_i, ψ). Depending on the model, one of the two parts of θ may also be empty. In very simple models, such as linear regression models, latent variables are not explicitly presented and response values are conditionally independent given ψ, so that p(y_i | y_−i, f_i, ψ) = p(y_i | ψ). The full likelihood can then be written in the familiar form

$$ p(y \,|\, \psi) = \prod_{i=1}^N p(y_i \,|\, \psi), $$

where y = (y₁, …, y_N) denotes the vector of all responses. When the likelihood factorizes this way, the conditional pointwise log-likelihood can be obtained easily by computing p(y_i | ψ) for each i with computational cost O(n).

Yet, there are several reasons why a non-factorized observation model may be necessary or preferred. In non-factorized models, the joint likelihood of the response values p(y | θ) is not factorized into observation-specific components, but rather given directly as one joint expression. For some models, an analytic factorized formulation is simply not available in which case we speak of a non-factorizable model. Even in models whose observation model can be factorized in principle, it may still be preferable to use a non-factorized form for reasons of efficiency and numerical stability (Bürkner et al. 2020).

Whether a non-factorized model is used by necessity or for efficiency and stability, it comes at the cost of having no direct access to the leave-one-out predictive densities and thus to the overall leave-one-out predictive accuracy. In theory, we can express the observation-specific likelihoods in terms of the joint likelihood via

$$ p(y_i \,|\, y_{i-1}, \theta) = \frac{p(y \,|\, \theta)}{p(y_{-i} \,|\, \theta)} = \frac{p(y \,|\, \theta)}{\int p(y \,|\, \theta) \, d y_i}, $$

but the expression on the right-hand side may not always have an analytical solution. Computing log p(y_i | y_−i, θ) for non-factorized models is therefore often impossible, or at least inefficient and numerically unstable. However, there is a large class of multivariate normal and Student-t models for which there are efficient analytical solutions available.

More details can be found in our paper about LOO-CV for non-factorized models (Bürkner, Gabry, & Vehtari, 2020), which is available as a preprint on arXiv (https://arxiv.org/abs/1810.10559).

LOO-CV for multivariate normal models

In this vignette, we will focus on non-factorized multivariate normal models. Based on results of Sundararajan and Keerthi (2001), Bürkner et al. (2020) show that, for multivariate normal models with coriance matrix C, the LOO predictive mean and standard deviation can be computed as follows:

where g_i and c̄_ii are

Using these results, the log predictive density of the ith observation is then computed as

$$ \log p(y_i \,|\, y_{-i},\theta) = - \frac{1}{2}\log(2\pi) - \frac{1}{2}\log \sigma^2_{-i} - \frac{1}{2}\frac{(y_i-\mu_{-i})^2}{\sigma^2_{-i}}. $$

Expressing this same equation in terms of g_i and c̄_ii, the log predictive density becomes:

$$ \log p(y_i \,|\, y_{-i},\theta) = - \frac{1}{2}\log(2\pi) + \frac{1}{2}\log \bar{c}_{ii} - \frac{1}{2}\frac{g_i^2}{\bar{c}_{ii}}. $$ (Note that Vehtari et al. (2016) has a typo in the corresponding Equation 34.)

From these equations we can now derive a recipe for obtaining the conditional pointwise log-likelihood for all models that can be expressed conditionally in terms of a multivariate normal with invertible covariance matrix C.

Lagged SAR models

A common non-factorized multivariate normal model is the simultaneously autoregressive (SAR) model, which is frequently used for spatially correlated data. The lagged SAR model is defined as

y = ρWy + η + ϵ or equivalently (I − ρW)y = η + ϵ, where ρ is the spatial correlation parameter and W is a user-defined weight matrix. The matrix W has entries w_ii = 0 along the diagonal and the off-diagonal entries w_ij are larger when areas i and j are closer to each other. In a linear model, the predictor term η is given by η = Xβ with design matrix X and regression coefficients β. However, since the above equation holds for arbitrary η, these results are not restricted to linear models.

If we have ϵ ∼ N(0, σ²I), it follows that (I − ρW)y ∼ N(η, σ²I), which corresponds to the following log PDF coded in Stan:

/** 
 * Normal log-pdf for spatially lagged responses
 * 
 * @param y Vector of response values.
 * @param mu Mean parameter vector.
 * @param sigma Positive scalar residual standard deviation.
 * @param rho Positive scalar autoregressive parameter.
 * @param W Spatial weight matrix.
 *
 * @return A scalar to be added to the log posterior.
 */
real normal_lagsar_lpdf(vector y, vector mu, real sigma, 
                        real rho, matrix W) {
  int N = rows(y);
  real inv_sigma2 = 1 / square(sigma);
  matrix[N, N] W_tilde = -rho * W;
  vector[N] half_pred;
  
  for (n in 1:N) W_tilde[n,n] += 1;
  
  half_pred = W_tilde * (y - mdivide_left(W_tilde, mu));
  
  return 0.5 * log_determinant(crossprod(W_tilde) * inv_sigma2) -
         0.5 * dot_self(half_pred) * inv_sigma2;
}

For the purpose of computing LOO-CV, it makes sense to rewrite the SAR model in slightly different form. Conditional on ρ, η, and σ, if we write

or more compactly, with W̃ = (I − ρW),

then this has the same form as the zero mean Gaussian process from above. Accordingly, we can compute the leave-one-out predictive densities with the equations from Sundararajan and Keerthi (2001), replacing y with (y − W̃⁻¹η) and taking the covariance matrix C to be σ²(W̃^TW̃)⁻¹.

library("loo")
library("brms")
library("bayesplot")
library("ggplot2")
color_scheme_set("brightblue")
theme_set(theme_default())


SEED <- 10001 
set.seed(SEED) # only sets seed for R (seed for Stan set later)

# loads COL.OLD data frame and COL.nb neighbor list
data(oldcol, package = "spdep") 

str(COL.OLD[, c("CRIME", "HOVAL", "INC")])

fit <- brm(
  CRIME ~ INC + HOVAL + sar(COL.nb, type = "lag"), 
  data = COL.OLD,
  data2 = list(COL.nb = COL.nb),
  chains = 4,
  seed = SEED
)

lagsar <- as.matrix(fit, pars = "lagsar")
estimates <- quantile(lagsar, probs = c(0.25, 0.5, 0.75))
mcmc_hist(lagsar) + 
  vline_at(estimates, linetype = 2, size = 1) +
  ggtitle("lagsar: posterior median and 50% central interval")

posterior <- as.data.frame(fit)
y <- fit$data$CRIME
N <- length(y)
S <- nrow(posterior)
loglik <- yloo <- sdloo <- matrix(nrow = S, ncol = N)

for (s in 1:S) {
  p <- posterior[s, ]
  eta <- p$b_Intercept + p$b_INC * fit$data$INC + p$b_HOVAL * fit$data$HOVAL
  W_tilde <- diag(N) - p$lagsar * spdep::nb2mat(COL.nb)
  Cinv <- t(W_tilde) %*% W_tilde / p$sigma^2
  g <- Cinv %*% (y - solve(W_tilde, eta))
  cbar <- diag(Cinv)
  yloo[s, ] <- y - g / cbar
  sdloo[s, ] <- sqrt(1 / cbar)
  loglik[s, ] <- dnorm(y, yloo[s, ], sdloo[s, ], log = TRUE)
}

# use loo for psis smoothing
log_ratios <- -loglik
psis_result <- psis(log_ratios)

plot(psis_result, label_points = TRUE)

yloo_sub <- yloo[S, ]
sdloo_sub <- sdloo[S, ]
df <- data.frame(
  y = y, 
  yloo = yloo_sub,
  ymin = yloo_sub - sdloo_sub * 2,
  ymax = yloo_sub + sdloo_sub * 2
)
ggplot(data=df, aes(x = y, y = yloo, ymin = ymin, ymax = ymax)) +
  geom_errorbar(
    width = 1, 
    color = "skyblue3", 
    position = position_jitter(width = 0.25)
  ) +
  geom_abline(color = "gray30", size = 1.2) +
  geom_point()

(psis_loo <- loo(loglik))

# see help("mi", "brms") for details on the mi() usage
fit_dummy <- brm(
  CRIME | mi() ~ INC + HOVAL + sar(COL.nb, type = "lag"), 
  data = COL.OLD,
  data2 = list(COL.nb = COL.nb),
  chains = 0
)

S <- 500
res <- vector("list", N)
loglik <- matrix(nrow = S, ncol = N)
for (i in seq_len(N)) {
  dat_mi <- COL.OLD
  dat_mi$CRIME[i] <- NA
  fit_i <- update(fit_dummy, newdata = dat_mi, 
                  # just for vignette
                  chains = 1, iter = S * 2)
  posterior <- as.data.frame(fit_i)
  yloo <- sdloo <- rep(NA, S)
  for (s in seq_len(S)) {
    p <- posterior[s, ]
    y_miss_i <- y
    y_miss_i[i] <- p$Ymi
    eta <- p$b_Intercept + p$b_INC * fit_i$data$INC + p$b_HOVAL * fit_i$data$HOVAL
    W_tilde <- diag(N) - p$lagsar * spdep::nb2mat(COL.nb)
    Cinv <- t(W_tilde) %*% W_tilde / p$sigma^2
    g <- Cinv %*% (y_miss_i - solve(W_tilde, eta))
    cbar <- diag(Cinv);
    yloo[s] <- y_miss_i[i] - g[i] / cbar[i]
    sdloo[s] <- sqrt(1 / cbar[i])
    loglik[s, i] <- dnorm(y[i], yloo[s], sdloo[s], log = TRUE)
  }
  ypred <- rnorm(S, yloo, sdloo)
  res[[i]] <- data.frame(y = c(posterior$Ymi, ypred))
  res[[i]]$type <- rep(c("pp", "loo"), each = S)
  res[[i]]$obs <- i
}
res <- do.call(rbind, res)

res_sub <- res[res$obs %in% 1:4, ]
ggplot(res_sub, aes(y, fill = type)) +
  geom_density(alpha = 0.6) +
  facet_wrap("obs", scales = "fixed", ncol = 4)

log_mean_exp <- function(x) {
  # more stable than log(mean(exp(x)))
  max_x <- max(x)
  max_x + log(sum(exp(x - max_x))) - log(length(x))
}
exact_elpds <- apply(loglik, 2, log_mean_exp)
exact_elpd <- sum(exact_elpds)
round(exact_elpd, 1)

df <- data.frame(
  approx_elpd = psis_loo$pointwise[, "elpd_loo"],
  exact_elpd = exact_elpds
)
ggplot(df, aes(x = approx_elpd, y = exact_elpd)) +
  geom_abline(color = "gray30") +
  geom_point(size = 2) +
  geom_point(data = df[4, ], size = 3, color = "red3") +
  xlab("Approximate elpds") +
  ylab("Exact elpds") +
  coord_fixed(xlim = c(-16, -3), ylim = c(-16, -3))

without_pt_4 <- c(
  approx = sum(psis_loo$pointwise[-4, "elpd_loo"]),
  exact = sum(exact_elpds[-4])  
)
round(without_pt_4, 1)

Working with Stan directly

So far, we have specified the models in brms and only used Stan implicitely behind the scenes. This allowed us to focus on the primary purpose of validating approximate LOO-CV for non-factorized models. However, we would also like to show how everything can be set up in Stan directly. The Stan code brms generates is human readable and so we can use it to learn some of the essential aspects of Stan and the particular model we are implementing. The Stan program below is a slightly modified version of the code extracted via stancode(fit_dummy):

// generated with brms 2.2.0
functions {
/** 
 * Normal log-pdf for spatially lagged responses
 * 
 * @param y Vector of response values.
 * @param mu Mean parameter vector.
 * @param sigma Positive scalar residual standard deviation.
 * @param rho Positive scalar autoregressive parameter.
 * @param W Spatial weight matrix.
 *
 * @return A scalar to be added to the log posterior.
 */
  real normal_lagsar_lpdf(vector y, vector mu, real sigma,
                          real rho, matrix W) {
    int N = rows(y);
    real inv_sigma2 = 1 / square(sigma);
    matrix[N, N] W_tilde = -rho * W;
    vector[N] half_pred;
    for (n in 1:N) W_tilde[n, n] += 1;
    half_pred = W_tilde * (y - mdivide_left(W_tilde, mu));
    return 0.5 * log_determinant(crossprod(W_tilde) * inv_sigma2) -
           0.5 * dot_self(half_pred) * inv_sigma2;
  }
}
data {
  int<lower=1> N;  // total number of observations
  vector[N] Y;  // response variable
  int<lower=0> Nmi;  // number of missings
  int<lower=1> Jmi[Nmi];  // positions of missings
  int<lower=1> K;  // number of population-level effects
  matrix[N, K] X;  // population-level design matrix
  matrix[N, N] W;  // spatial weight matrix
  int prior_only;  // should the likelihood be ignored?
}
transformed data {
  int Kc = K - 1;
  matrix[N, K - 1] Xc;  // centered version of X
  vector[K - 1] means_X;  // column means of X before centering
  for (i in 2:K) {
    means_X[i - 1] = mean(X[, i]);
    Xc[, i - 1] = X[, i] - means_X[i - 1];
  }
}
parameters {
  vector[Nmi] Ymi;  // estimated missings
  vector[Kc] b;  // population-level effects
  real temp_Intercept;  // temporary intercept
  real<lower=0> sigma;  // residual SD
  real<lower=0,upper=1> lagsar;  // SAR parameter
}
transformed parameters {
}
model {
  vector[N] Yl = Y;
  vector[N] mu = Xc * b + temp_Intercept;
  Yl[Jmi] = Ymi;
  // priors including all constants
  target += student_t_lpdf(temp_Intercept | 3, 34, 17);
  target += student_t_lpdf(sigma | 3, 0, 17)
    - 1 * student_t_lccdf(0 | 3, 0, 17);
  // likelihood including all constants
  if (!prior_only) {
    target += normal_lagsar_lpdf(Yl | mu, sigma, lagsar, W);
  }
}
generated quantities {
  // actual population-level intercept
  real b_Intercept = temp_Intercept - dot_product(means_X, b);
}

Here we want to focus on two aspects of the Stan code. First, because there is no built-in function in Stan that calculates the log-likelihood for the lag-SAR model, we define a new normal_lagsar_lpdf function in the functions block of the Stan program. This is the same function we showed earlier in the vignette and it can be used to compute the log-likelihood in an efficient and numerically stable way. The _lpdf suffix used in the function name informs Stan that this is a log probability density function.

Second, this Stan program nicely illustrates how to set up missing value imputation. Instead of just computing the log-likelihood for the observed responses Y, we define a new variable Yl which is equal to Y if the reponse is observed and equal to Ymi if the response is missing. The latter is in turn defined as a parameter and thus estimated along with all other paramters of the model. More details about missing value imputation in Stan can be found in the Missing Data & Partially Known Parameters section of the Stan manual.

The Stan code extracted from brms is not only helpful when learning Stan, but can also drastically speed up the specification of models that are not support by brms. If brms can fit a model similar but not identical to the desired model, we can let brms generate the Stan program for the similar model and then mold it into the program that implements the model we actually want to fit. Rather than calling stancode(), which requires an existing fitted model object, we recommend using make_stancode() and specifying the save_model argument to write the Stan program to a file. The corresponding data can be prepared with make_standata() and then manually amended if needed. Once the code and data have been edited, they can be passed to RStan’s stan() function via the file and data arguments.

Conclusion

In summary, we have shown how to set up and validate approximate and exact LOO-CV for non-factorized multivariate normal models using Stan with the brms and loo packages. Although we focused on the particular example of a spatial SAR model, the presented recipe applies more generally to models that can be expressed in terms of a multivariate normal likelihood.

References

Anselin L. (1988). Spatial econometrics: methods and models. Dordrecht: Kluwer Academic.

Bürkner P. C., Gabry J., & Vehtari A. (2020). Efficient leave-one-out cross-validation for Bayesian non-factorized normal and Student-t models. Computational Statistics, :10.1007/s00180-020-01045-4. ArXiv preprint.

Sundararajan S. & Keerthi S. S. (2001). Predictive approaches for choosing hyperparameters in Gaussian processes. Neural Computation, 13(5), 1103–1118.

Vehtari A., Mononen T., Tolvanen V., Sivula T., & Winther O. (2016). Bayesian leave-one-out cross-validation approximations for Gaussian latent variable models. Journal of Machine Learning Research, 17(103), 1–38. Online.

Vehtari A., Gelman A., & Gabry J. (2017). Practical Bayesian model evaluation using leave-one-out cross-validation and WAIC. Statistics and Computing, 27(5), 1413–1432. :10.1007/s11222-016-9696-4. Online. arXiv preprint arXiv:1507.04544.

Vehtari, A., Simpson, D., Gelman, A., Yao, Y., and Gabry, J. (2024). Pareto smoothed importance sampling. Journal of Machine Learning Research, 25(72):1-58. PDF