Improve floating point accuracy

NEWS.md

@@ -1,6 +1,9 @@
 # loo (development version)
 
-* Fix `loo_compare()` when used with subsampling: compute model comparison by comparison-model minus reference-model by @florence-bockting in #391
+* Improve numerical stability in `loo()`, `psis()`, model weighting, subsampling, 
+and moment matching in #395 
+* Fix `loo_compare()` when used with subsampling: compute model comparison by 
+comparison-model minus reference-model by @florence-bockting in #391
 * Update user messages in `print()` by @ishaan-arora-1, @florence-bockting in 
 #328.
 

R/E_loo.R

@@ -219,7 +219,10 @@ E_loo.matrix <-
   # sample size ESS is estimated with the generic target quantity invariant
   # estimate 1/sum(w^2), see e.g. "Monte Carlo theory, methods and examples"
   # by Owen (2013).
-  (sum(.wmean(x^2, w)) - sum(.wmean(x, w)^2)) / (1 - sum(w^2))
+  # The two-pass form avoids the cancellation in E[x^2] - E[x]^2 and is
+  # equivalent to it only because `w` sums to one.
+  weighted_mean <- .wmean(x, w)
+  sum(w * (x - weighted_mean)^2) / (1 - sum(w^2))
 }
 .wsd <- function(x, w, ...) {
   sqrt(.wvar(x, w))

R/helpers.R

@@ -26,6 +26,52 @@ colLogMeanExps <- function(x) {
   matrixStats::colLogSumExps(x) - logS
 }
 
+#' More stable version of `exp(x) - exp(y)`
+#'
+#' @noRd
+#' @param x A numeric vector.
+#' @param y A numeric scalar or vector recycled to the length of `x`.
+#'   Must satisfy `x >= y` elementwise.
+#' @return A numeric vector equal to `exp(x) - exp(y)`.
+#'
+exp_x_minus_exp_y <- function(x, y) {
+  out <- -exp(x) * expm1(y - x)
+  # which() drops the NA comparisons that NA or NaN inputs would produce
+  out[which(x == y)] <- 0
+  out
+}
+
+#' More stable version of `x^2 - y^2`
+#'
+#' @noRd
+#' @param x,y Numeric vectors of the same length.
+#' @return A numeric vector equal to `x^2 - y^2`.
+#'
+difference_of_squares <- function(x, y) {
+  (x - y) * (x + y)
+}
+
+#' More stable version of `(exp(a) - exp(b)) / exp(c)`
+#'
+#' @noRd
+#' @param a,b,c Numeric vectors of the same length.
+#' @return A numeric vector equal to `(exp(a) - exp(b)) / exp(c)`. Elements
+#'   with `a == b` are returned as an exact zero regardless of `c`; elsewhere
+#'   `NA` and `NaN` inputs propagate.
+#'
+exp_diff_over_exp <- function(a, b, c) {
+  # `a >= b` is NA if `a` or `b` is NA or NaN, and R silently ignores NA
+  # indices in `[<-`. Seed the result from the inputs and index with which()
+  # so that missing values propagate instead of leaving a zero behind.
+  out <- a + b + c
+  larger <- which(a >= b)
+  smaller <- which(a < b)
+  out[larger] <- exp(a[larger] - c[larger]) * -expm1(b[larger] - a[larger])
+  out[smaller] <- exp(b[smaller] - c[smaller]) * expm1(a[smaller] - b[smaller])
+  out[which(a == b)] <- 0
+  out
+}
+
 #' Compute point estimates and standard errors from pointwise vectors
 #'
 #' @noRd
@@ -63,6 +109,43 @@ validate_ll <- function(x) {
   invisible(x)
 }
 
+#' Check that a log-likelihood array/matrix/vector is finite
+#'
+#' `loo()` requires finite log-likelihood values, for two unrelated reasons.
+#'
+#' A `-Inf` log likelihood is meaningful on its own — the observation has zero
+#' likelihood under that draw — but the leave-one-out importance ratio is
+#' `1 / p(y_i | theta)`, which is then infinite, so the PSIS estimate does not
+#' exist. Because `loo()` negates the log likelihood before importance
+#' sampling, this used to surface as [validate_ll()]'s `+Inf` message, with the
+#' polarity reversed.
+#'
+#' A `+Inf` log likelihood is not meaningful, and it passes the log-ratio check
+#' as a `-Inf` ratio. `ll + lw` is then `Inf + -Inf`, so a single such value
+#' made every estimate for the model `NA`.
+#'
+#' This is deliberately stricter than [validate_ll()], which is also used for
+#' log ratios, where `-Inf` is a valid zero importance weight.
+#'
+#' @noRd
+#' @param x Array/matrix/vector of log-likelihood values.
+#' @return `x`, invisibly, if no error is thrown.
+#'
+validate_log_lik <- function(x) {
+  if (is.list(x)) {
+    stop("List not allowed as input.")
+  }
+  # single pass covering NA, NaN and both infinities; the more specific checks
+  # below only run when something is already known to be wrong
+  if (!all(is.finite(x))) {
+    if (anyNA(x)) {
+      stop("NAs not allowed in input.")
+    }
+    stop("All log-likelihood values must be finite.")
+  }
+  invisible(x)
+}
+
 #' Convert iter by chain by obs array to (iter * chain) by obs matrix
 #'
 #' @noRd

R/importance_sampling.R

@@ -12,6 +12,16 @@ importance_sampling <- function(log_ratios, method, ...) {
   UseMethod("importance_sampling")
 }
 
+validate_log_ratios <- function(x) {
+  validate_ll(x)
+  # validate_ll() has already ruled out NA and +Inf, so a column without a
+  # finite value is exactly a column whose maximum is -Inf
+  if (any(matrixStats::colMaxs(x) == -Inf)) {
+    stop("Each column of log ratios must contain at least one finite value.")
+  }
+  invisible(x)
+}
+
 
 #' @rdname importance_sampling
 #' @inheritParams psis
@@ -24,8 +34,8 @@ importance_sampling.array <-
     cores <- loo_cores(cores)
     stopifnot(length(dim(log_ratios)) == 3)
     assert_importance_sampling_method_is_implemented(method)
-    log_ratios <- validate_ll(log_ratios)
     log_ratios <- llarray_to_matrix(log_ratios)
+    log_ratios <- validate_log_ratios(log_ratios)
     r_eff <- prepare_psis_r_eff(r_eff, len = ncol(log_ratios))
     do_importance_sampling(log_ratios, r_eff = r_eff, cores = cores, method = method)
   }
@@ -40,7 +50,7 @@ importance_sampling.matrix <-
            cores = getOption("mc.cores", 1)) {
     cores <- loo_cores(cores)
     assert_importance_sampling_method_is_implemented(method)
-    log_ratios <- validate_ll(log_ratios)
+    log_ratios <- validate_log_ratios(log_ratios)
     r_eff <- prepare_psis_r_eff(r_eff, len = ncol(log_ratios))
     do_importance_sampling(log_ratios, r_eff = r_eff, cores = cores, method = method)
   }

R/loo.R

@@ -198,6 +198,7 @@ loo.array <-
            cores = getOption("mc.cores", 1),
            is_method = c("psis", "tis", "sis")) {
     is_method <- match.arg(is_method)
+    validate_log_lik(x)
     psis_out <- importance_sampling.array(log_ratios = -x, r_eff = r_eff, cores = cores, method = is_method)
     ll <- llarray_to_matrix(x)
     pointwise <- pointwise_loo_calcs(ll, psis_out)
@@ -222,6 +223,7 @@ loo.matrix <-
            cores = getOption("mc.cores", 1),
            is_method = c("psis", "tis", "sis")) {
     is_method <- match.arg(is_method)
+    validate_log_lik(x)
     psis_out <-
       importance_sampling.matrix(
         log_ratios = -x,
@@ -371,6 +373,7 @@ loo_i <-
     if (!is.matrix(ll_i)) {
       ll_i <- as.matrix(ll_i)
     }
+    validate_log_lik(ll_i)
     psis_out <-
       importance_sampling.matrix(
         log_ratios = -ll_i,
@@ -488,29 +491,64 @@ importance_sampling_loo_object <- function(pointwise, diagnostics, dims,
 #' @return Vector of standard error estimates.
 #'
 mcse_elpd <- function(ll, lw, E_elpd, r_eff, n_samples = NULL) {
-  lik <- exp(ll)
-  w2 <- exp(lw)^2
-  E_epd <- exp(E_elpd)
-  if (length(r_eff) == 1 && !is.null(ncol(ll))) {
+  if (!is.matrix(ll)) {
+    ll <- as.matrix(ll)
+  }
+  if (!is.matrix(lw)) {
+    lw <- as.matrix(lw)
+  }
+  S <- nrow(ll)
+  if (length(r_eff) == 1) {
     r_eff <- rep(r_eff, ncol(ll))
   }
-  var_elpd <-
-    vapply(
-      seq_len(ncol(w2)),
-      FUN.VALUE = numeric(1),
-      FUN = function(i) {
-        # Variance in linear scale
-        # Equation (6) in Vehtari et al. (2024)
-        var_epd_i <- sum(w2[, i] * (lik[, i] - E_epd[i]) ^ 2) / r_eff[i]
-        # Compute variance in log scale by match the variance of a
-        # log-normal approximation
-        # https://en.wikipedia.org/wiki/Log-normal_distribution#Arithmetic_moments
-        log(1 + var_epd_i / E_epd[i]^2)
-      }
+  # Everything is computed relative to the loo predictive density, so that
+  # 1) exp() of the log likelihood never over- or underflows, and
+  # 2) expm1() avoids the cancellation in `exp(ll) - exp(E_elpd)`.
+  # `ll - E_elpd` is bounded above by `-lw`, so the product below cannot
+  # overflow for consistent (ll, lw, E_elpd); the fallback covers the rest.
+  #
+  # Variance in linear scale, relative to E_epd^2.
+  # Equation (6) in Vehtari et al. (2024)
+  var_epd_ratio <-
+    matrixStats::colSums2((exp(lw) * expm1(ll - rep(E_elpd, each = S)))^2) /
+      r_eff
+  # Variance in log scale by matching the variance of a log-normal
+  # https://en.wikipedia.org/wiki/Log-normal_distribution#Arithmetic_moments
+  var_elpd <- log1p(var_epd_ratio)
+  undefined <- is.infinite(E_elpd) & E_elpd < 0
+  overflow <- !is.finite(var_epd_ratio) & !undefined
+  if (any(overflow)) {
+    lvr <- log_var_epd_ratio(
+      ll[, overflow, drop = FALSE] - rep(E_elpd[overflow], each = S),
+      lw[, overflow, drop = FALSE],
+      r_eff[overflow]
     )
+    var_elpd[overflow] <-
+      ifelse(lvr > 0, lvr + log1p(exp(-lvr)), log1p(exp(lvr)))
+  }
+  var_elpd[undefined] <- NA_real_
   sqrt(var_elpd)
 }
 
+#' Log of the relative linear-scale ELPD variance, for the rare case where
+#' `exp(lw) * expm1(log_lik_ratio)` over- or underflows
+#'
+#' @noRd
+#' @param log_lik_ratio Matrix of `ll - E_elpd` values.
+#' @param lw Matrix of normalized log weights.
+#' @param r_eff Vector of relative effective sample sizes.
+#' @return Vector of `log(var_epd / E_epd^2)` values.
+#'
+log_var_epd_ratio <- function(log_lik_ratio, lw, r_eff) {
+  log_abs_diff <- log(abs(expm1(log_lik_ratio)))
+  big <- which(log_lik_ratio > 700)
+  if (length(big)) {
+    log_abs_diff[big] <-
+      log_lik_ratio[big] + log(-expm1(-log_lik_ratio[big]))
+  }
+  matrixStats::colLogSumExps(2 * (lw + log_abs_diff)) - log(r_eff)
+}
+
 
 #' Warning message if r_eff not specified
 #' @noRd

R/loo_model_weights.R

@@ -257,6 +257,10 @@ stacking_weights <-
     if (K < 2) {
       stop("At least two models are required for stacking weights.")
     }
+    validate_lpd_point(lpd_point)
+    if (any(rowSums(is.finite(lpd_point)) == 0)) {
+      stop("Each observation must have a finite predictive density for at least one model.")
+    }
 
     negative_log_score_loo <- function(w) {
       # objective function: log score
@@ -272,11 +276,15 @@ stacking_weights <-
       stopifnot(length(w) == K - 1)
       w_full <- c(w, 1 - sum(w))
       grad <- rep(0, K - 1)
-      # avoid over- and underflows using log weights, rowLogSumExps,
-      # and by subtracting the row maximum of lpd_point
-      mlpd <- matrixStats::rowMaxs(lpd_point)
+      mixture_lpd <- matrixStats::rowLogSumExps(
+        sweep(lpd_point, 2, log(w_full), "+")
+      )
       for (k in 1:(K - 1)) {
-        grad[k] <- sum((exp(lpd_point[, k] - mlpd) - exp(lpd_point[, K] - mlpd)) / exp(matrixStats::rowLogSumExps(sweep(lpd_point, 2, log(w_full), '+')) - mlpd))
+        grad[k] <- sum(exp_diff_over_exp(
+          lpd_point[, k],
+          lpd_point[, K],
+          mixture_lpd
+        ))
       }
       return(-grad)
     }
@@ -317,9 +325,13 @@ pseudobma_weights <-
     if (K < 2) {
       stop("At least two models are required for pseudo-BMA weights.")
     }
+    validate_lpd_point(lpd_point)
+    elpd <- colSums2(lpd_point)
+    if (!any(is.finite(elpd))) {
+      stop("At least one model must have a finite total predictive density.")
+    }
 
     if (!BB) {
-      elpd <- colSums2(lpd_point)
       uwts <- exp(elpd - max(elpd))
       wts <- structure(
         uwts / sum(uwts),
@@ -345,15 +357,6 @@ pseudobma_weights <-
   }
 
 
-#' Generate dirichlet simulations, rewritten version
-#' @importFrom stats rgamma
-#' @noRd
-dirichlet_rng <- function(n, alpha) {
-  K <- length(alpha)
-  gamma_sim <- matrix(rgamma(K * n, alpha), ncol = K, byrow = TRUE)
-  gamma_sim / rowSums(gamma_sim)
-}
-
 #' @export
 print.stacking_weights <- function(x, digits = 3, ...) {
   cat("Method: stacking\n------\n")
@@ -372,6 +375,27 @@ print.pseudobma_bb_weights <- function(x, digits = 3, ...) {
   print_weight_vector(x, digits = digits)
 }
 
+
+
+# internal ----------------------------------------------------------------
+# `-Inf` is a valid zero predictive density; `NA`, `NaN` and `+Inf` make
+# stacking fail in the optimizer and pseudo-BMA return invalid weights.
+validate_lpd_point <- function(lpd_point) {
+  if (anyNA(lpd_point) || any(lpd_point == Inf)) {
+    stop("All values in 'lpd_point' must be finite or -Inf.")
+  }
+  invisible(lpd_point)
+}
+
+#' Generate dirichlet simulations, rewritten version
+#' @importFrom stats rgamma
+#' @noRd
+dirichlet_rng <- function(n, alpha) {
+  K <- length(alpha)
+  gamma_sim <- matrix(rgamma(K * n, alpha), ncol = K, byrow = TRUE)
+  gamma_sim / rowSums(gamma_sim)
+}
+
 print_weight_vector <- function(x, digits) {
   z <- cbind(x)
   colnames(z) <- "weight"

R/loo_moment_matching.R

@@ -408,7 +408,7 @@ loo_moment_match_i <- function(i,
   elpd_loo_i <- matrixStats::logSumExp(log_liki + lwi)
   mcse_elpd_loo <- mcse_elpd(
     ll = as.matrix(log_liki), lw = as.matrix(lwi),
-    E_elpd = exp(elpd_loo_i), r_eff = r_eff_i
+    E_elpd = elpd_loo_i, r_eff = r_eff_i
   )
 
   list(elpd_loo_i = elpd_loo_i,
@@ -531,11 +531,14 @@ shift_and_scale <- function(x, upars, lwi) {
   # compute moments using log weights
   S <- dim(upars)[1]
   mean_original <- colMeans(upars)
-  mean_weighted <- colSums(exp(lwi) * upars)
+  weights <- exp(lwi)
+  mean_weighted <- colSums(weights * upars)
   shift <- mean_weighted - mean_original
-  mii <- exp(lwi)* upars^2
-  mii <- colSums(mii) - mean_weighted^2
-  mii <- mii*S/(S-1)
+  # The two-pass form avoids the cancellation in E[x^2] - E[x]^2 and is
+  # equivalent to it only because `weights` sums to one.
+  centered <- sweep(upars, 2, mean_weighted)
+  mii <- colSums(weights * centered^2)
+  mii <- mii * S / (S - 1)
   scaling <- sqrt(mii / matrixStats::colVars(upars))
   # transform posterior draws
   upars_new <- sweep(upars, 2, mean_original, "-")

R/loo_subsample.R

@@ -1185,7 +1185,7 @@ srs_diff_est <- function(y_approx, y, y_idx) {
   t_pi_tilde <- sum(y_approx)
   t_pi2_tilde <- sum(y_approx^2)
   t_e <- N * mean(e_i)
-  t_hat_epsilon <- N * mean(y^2 - y_approx_m^2)
+  t_hat_epsilon <- N * mean(difference_of_squares(y, y_approx_m))
 
   est_list <- list(m = length(y), N = N)
   # eq (7)

R/psis.R

@@ -254,7 +254,7 @@ psis_smooth_tail <- function(x, cutoff) {
   exp_cutoff <- exp(cutoff)
 
   # save time not sorting since x already sorted
-  fit <- posterior::gpdfit(exp(x) - exp_cutoff, sort_x = FALSE)
+  fit <- posterior::gpdfit(exp_x_minus_exp_y(x, cutoff), sort_x = FALSE)
   k <- fit$k
   sigma <- fit$sigma
   if (is.na(k)) {