Return the final Newton iterate as the Laplace latent mean, and fix the solver-3 log determinant

[avehtari]

Summary

Two independent fixes in the embedded Laplace approximation. This PS is based on #3436 which should merged first, because both change laplace_base_rng().

1. The latent mean is the final Newton iterate (stan/math/mix/functor/laplace_base_rng.hpp). laplace_base_rng() computed the mean as mean_train = covariance_train * md_est.theta_grad, that is K * grad log p(y | theta_hat). At an exact mode this equals theta_hat, but before exact convergence the two differ by K * grad log p(theta | y) at theta_hat. The stopping rule limits the change of the objective, not this difference, and the factor K makes the difference large when the posterior is much narrower than the prior. In example A (tolerance 1e-6), the final iterate is 0.0012 posterior sd from the mode, but the returned mean was 0.53 posterior sd from it, at a log density 7.1 lower. The mean is now md_est.theta, the final iterate. This affects laplace_latent_solve(), laplace_latent_solve_tol(), and the mean of the draws of laplace_latent_rng(), laplace_latent_tol_rng() and laplace_latent[_tol]_{bernoulli_logit,poisson_log,neg_binomial_2_log}_rng().

2. The log determinant of solver 3 uses |U_ii| and checks the sign (LUSolver::compute_log_determinant() in stan/math/mix/functor/laplace_marginal_density_estimator.hpp). Solver 3 factorizes the non-symmetric B = I + K W with partial pivoting. The diagonal of U can then have negative entries also when det(B) > 0, and the old sum(log(diag(U))) was NaN: laplace_marginal*() returned NaN without a message, with solver 3 and after a fallthrough to solver 3. This happens also for a Gaussian likelihood (example B). The log determinant is now sum(log(abs(diag(U)))), and the sign of det(B) is the sign of the permutation times the signs of the U_ii. Because det(B) = det(K) det(K^-1 + W) and det(K) > 0, a non-positive det(B) means that K^-1 + W is not positive definite at the final iterate, for example at a saddle point, so the Laplace approximation is not defined there. The function then throws std::domain_error with a message that says so. (A positive det(B) does not prove positive definiteness: the check finds an odd number of negative eigenvalues.) compute_log_determinant() also runs for laplace_latent_solve*() and laplace_latent*_rng(). So at such a point (example C), laplace_marginal_tol() now throws instead of returning NaN, and laplace_latent_solve_tol() throws this message instead of cholesky_decompose: Matrix m is not positive definite.

Examples

• A: the longitudinal part of the joint model for one group: 12 observations y ~ normal(m, sigma * m) with m(t) = bsld * (phi * exp(-d * t) + (1 - phi) * exp(g * t)), and 4 latent variables on the log or logit scale of (bsld, d, g, phi) with diagonal K. The posterior sd is 5.6 to 9.4 times smaller than the prior sd. laplace_latent_solve_tol() with solver 3. The reference mode comes from optim() with Newton polishing.
• B: Gaussian likelihood y ~ normal(theta, s), y = (0.3, -0.2), s = sqrt(0.1), K = [1 1.5; 1.5 4]. B = I + K W has det(B) = 226, but its LU diagonal is (15, -15.07) because the rows are swapped. The exact log marginal is multi_normal_lpdf(y | 0, K + s^2 I).
• C: 2 latent variables, K = [4 1; 1 1]; theta[1] has two Cauchy observations y = (-3, 3), s = 0.3, and theta[2] one Gaussian observation z = 0, sz = 0.1. The log posterior has two maxima and a saddle point at (0, 0). From the start (0.01, 3), the Newton iteration of solver 3 converges to the saddle point.

example	call	develop	this PR
A, tolerance 1e-6	laplace_latent_solve_tol()	0.53 posterior sd from the mode, log density 7.12 lower	0.0012 sd, 2.3e-6 lower
A, default tolerance	laplace_latent_solve_tol()	0.068 sd, 0.135 lower	0.00015 sd, 3.8e-8 lower
B	laplace_marginal_tol(), solver 3	NaN	-2.376754 (exact)
B	laplace_marginal_tol(), solvers 1 and 2	-2.376754	-2.376754
C	laplace_marginal_tol()	NaN	exception: det(I + Sigma * W) is not positive at the final iterate, ...
C	laplace_latent_solve_tol()	exception: cholesky_decompose: Matrix m is not positive definite	exception: det(I + Sigma * W) is not positive at the final iterate, ...

The develop column is from CmdStan 2.40.0; examples A and B give the same values with develop 43bc568. The column "this PR" is from CmdStan develop with this branch.

R script (cmdstanr; the Stan programs are inline)

# pr_laplace_mean_lu_logdet.R
#
# Examples for a stan-dev/math PR (self-contained; the Stan programs are
# inline):
#
#   A. laplace_latent_solve_tol() returned K * grad log p(y | theta_hat)
#      instead of the final Newton iterate theta_hat.  Before exact
#      convergence the two differ by K * grad log p(theta | y), which is
#      large when the posterior is much narrower than the prior.
#   B. laplace_marginal_tol() with solver 3 returned NaN when the LU
#      factorization of B = I + K W has a negative pivot (also for a
#      Gaussian likelihood).
#   C. At a saddle point of log p(theta | y), the solver-3 log determinant
#      now reports that det(B) is not positive.
#
#   Rscript tests/pr_laplace_mean_lu_logdet.R
#   CMDSTAN_PATH=~/proj/cmdstan STAN_MATH=~/proj/stan_math_laplace_mean/ \
#     Rscript tests/pr_laplace_mean_lu_logdet.R

library(cmdstanr)
library(posterior)
if (nzchar(Sys.getenv("CMDSTAN_PATH"))) {
  set_cmdstan_path(Sys.getenv("CMDSTAN_PATH"))
}
cpp <- list()
if (nzchar(Sys.getenv("STAN_MATH"))) {
  cpp$MATH <- Sys.getenv("STAN_MATH")
}
compile <- function(code) {
  cmdstan_model(write_stan_file(code), cpp_options = cpp,
                force_recompile = length(cpp) > 0)
}
cat(sprintf("CmdStan %s at %s, Stan Math %s\n", cmdstan_version(),
            cmdstan_path(), if (is.null(cpp$MATH)) "of the CmdStan" else cpp$MATH))

# One fixed_param draw; returns the values of `vars` (NA if the draw failed)
# and the "Exception" lines of the output.
run_once <- function(mod, data, vars) {
  out <- capture.output(
    fit <- mod$sample(data = data, fixed_param = TRUE, chains = 1,
                      iter_sampling = 1, refresh = 0, show_messages = TRUE,
                      sig_figs = 12),
    type = "output")
  vals <- tryCatch(as.numeric(as_draws_matrix(fit$draws(vars))),
                   error = function(e) NA_real_)
  list(vals = vals,
       exception = unique(sub(".*Exception: (.*?) \\(in '.*", "\\1",
                              grep("Exception", out, value = TRUE),
                              perl = TRUE)))
}

# ── A. Returned mean ───────────────────────────────────────────────────────
tumor_code <- "
functions {
  real ll(vector theta, vector y, vector t, vector mu, real sigma) {
    real bsld = mu[1] * exp(theta[1]);
    real d = mu[2] * exp(theta[2]);
    real g = mu[3] * exp(theta[3]);
    real phi = inv_logit(logit(mu[4]) + theta[4]);
    vector[num_elements(y)] m
        = bsld * (phi * exp(-d * t) + (1 - phi) * exp(g * t));
    return normal_lpdf(y | m, sigma * m);
  }
  matrix k_fun(vector omega) {
    return diag_matrix(square(omega));
  }
}
data {
  int<lower=1> n;
  vector[n] y;
  vector[n] t;
  vector[4] mu;
  vector<lower=0>[4] omega;
  real<lower=0> sigma;
  vector[4] theta0;
  real<lower=0> tol;
}
generated quantities {
  vector[4] mode;
  {
    tuple(vector[4], matrix[4, 4]) lap = laplace_latent_solve_tol(
      ll, (y, t, mu, sigma), 4, k_fun, (omega,),
      (theta0, tol, 100, 3, 100, 0));
    mode = lap.1;
  }
}
"
y <- c(201.04, 78.30, 48.01, 44.39, 52.91, 50.49, 48.63, 67.29, 52.21,
       72.43, 102.43, 105.06)
t <- seq(0, 693, by = 63)
mu <- c(117.4, 9.121e-3, 1.913e-3, 0.3781)
omega <- c(0.4925, 0.9217, 1.120, 1.299)
sigma <- 0.09826
log_f <- function(th) {
  p <- plogis(qlogis(mu[4]) + th[4])
  m <- mu[1] * exp(th[1]) * (p * exp(-mu[2] * exp(th[2]) * t) +
                               (1 - p) * exp(mu[3] * exp(th[3]) * t))
  sum(dnorm(y, m, sigma * m, log = TRUE)) + sum(dnorm(th, 0, omega, log = TRUE))
}
mode_ref <- optim(rep(0, 4), function(z) -log_f(z), method = "BFGS",
                  control = list(reltol = 1e-15, maxit = 5000))$par
for (it in 1:5) {
  mode_ref <- mode_ref - solve(numDeriv::hessian(log_f, mode_ref),
                               numDeriv::grad(log_f, mode_ref))
}
post_sd <- sqrt(diag(solve(-numDeriv::hessian(log_f, mode_ref))))
cat("\nA. Returned mean.  Mode (optim + Newton):",
    sprintf("%.4f", mode_ref), "\n")
cat("   posterior sd:", sprintf("%.3f", post_sd), " prior sd:", omega, "\n")
mod_a <- compile(tumor_code)
for (tol in c(1e-6, 1.49012e-8, 1e-10)) {
  r <- run_once(mod_a, list(n = length(y), y = y, t = t, mu = mu,
                            omega = omega, sigma = sigma,
                            theta0 = c(0.3636, -0.3385, 1.9794, -1.9322),
                            tol = tol), "mode")
  cat(sprintf("   tol %-9g returned %s: %.2g posterior sd from the mode, log density %.3g lower\n",
              tol, paste(sprintf("%.4f", r$vals), collapse = " "),
              max(abs(r$vals - mode_ref) / post_sd),
              log_f(mode_ref) - log_f(r$vals)))
}

# ── B. Solver 3 with a negative LU pivot ───────────────────────────────────
normal_code <- "
functions {
  real ll(vector theta, vector y, real s) {
    return normal_lpdf(y | theta, s);
  }
  matrix k_fun(matrix K) {
    return K;
  }
}
data {
  vector[2] y;
  real<lower=0> s;
  matrix[2, 2] K;
  int<lower=1, upper=3> solver;
}
generated quantities {
  real lml = laplace_marginal_tol(ll, (y, s), 2, k_fun, (K,),
                                  (rep_vector(0.0, 2), 1e-10, 100, solver,
                                   100, 0));
  real lml_exact = multi_normal_lpdf(y | rep_vector(0.0, 2),
                                     K + diag_matrix(rep_vector(square(s), 2)));
}
"
K_b <- matrix(c(1, 1.5, 1.5, 4), 2)
s_b <- sqrt(0.1)
B <- diag(2) + K_b %*% diag(1 / s_b^2, 2)
cat(sprintf("\nB. Negative LU pivot.  B = I + K W: det %.4g, LU diagonal (partial pivoting) %s\n",
            det(B), paste(sprintf("%.4g", diag(as.matrix(Matrix::expand(Matrix::lu(B))$U))),
                          collapse = " ")))
mod_b <- compile(normal_code)
for (sv in 1:3) {
  r <- run_once(mod_b, list(y = c(0.3, -0.2), s = s_b, K = K_b, solver = sv),
                c("lml", "lml_exact"))
  cat(sprintf("   solver %d: laplace_marginal_tol %.6f, exact %.6f\n", sv,
              r$vals[1], r$vals[2]))
}

# ── C. Saddle point ────────────────────────────────────────────────────────
saddle_code <- "
functions {
  real ll(vector theta, vector y, real s, real z, real sz) {
    return cauchy_lpdf(y | theta[1], s) + normal_lpdf(z | theta[2], sz);
  }
  matrix k_fun(matrix K) {
    return K;
  }
}
data {
  vector[2] y;
  real<lower=0> s;
  real z;
  real<lower=0> sz;
  matrix[2, 2] K;
  vector[2] theta0;
  int<lower=0, upper=1> marginal;
}
generated quantities {
  real lml = not_a_number();
  vector[2] mode = rep_vector(not_a_number(), 2);
  if (marginal) {
    lml = laplace_marginal_tol(ll, (y, s, z, sz), 2, k_fun, (K,),
                               (theta0, 1e-10, 100, 3, 100, 0));
  } else {
    tuple(vector[2], matrix[2, 2]) lap = laplace_latent_solve_tol(
      ll, (y, s, z, sz), 2, k_fun, (K,), (theta0, 1e-10, 100, 3, 100, 0));
    mode = lap.1;
  }
}
"
mod_c <- compile(saddle_code)
data_c <- list(y = c(-3, 3), s = 0.3, z = 0, sz = 0.1,
               K = matrix(c(4, 1, 1, 1), 2), theta0 = c(0.01, 3))
cat("\nC. Saddle point at (0, 0); start (0.01, 3), solver 3\n")
r <- run_once(mod_c, c(data_c, list(marginal = 1L)), "lml")
cat(sprintf("   laplace_marginal_tol: %s\n",
            if (length(r$exception)) paste("exception:", r$exception)
            else sprintf("returned %.6g", r$vals)))
r <- run_once(mod_c, c(data_c, list(marginal = 0L)), "mode")
cat(sprintf("   laplace_latent_solve_tol: %s\n",
            if (length(r$exception)) paste("exception:", r$exception)
            else sprintf("returned %s", paste(sprintf("%.4g", r$vals), collapse = " "))))

Tests

New tests:

• laplace_latent_solve.mean_is_final_iterate_tumor_size (in test/unit/math/laplace/laplace_latent_solve_test.cpp): example A. The reference is a solve with tolerance 1e-14. With tolerance 1e-6, the returned mean must be within 0.01 posterior sd of the reference.
• laplace_marginal_tol.all_solvers_exact_with_negative_lu_pivot (new file test/unit/math/laplace/laplace_lu_log_determinant_test.cpp): example B. The test asserts that the LU factorization of B has a negative diagonal entry and that det(B) > 0. Solvers 1, 2 and 3 must return the exact log marginal.
• laplace_marginal_tol.solver_3_saddle_point_throws (same new file): example C. From (0.01, 3), laplace_marginal_tol() and laplace_latent_solve_tol() must throw std::domain_error with the new message. From (0.3, 3), solver 3 must return the maximum (2.9385, 0.0097).

Without the changes in stan/, all three tests fail: the returned mean is 0.53 posterior sd from the reference, solver 3 returns NaN, and laplace_marginal_tol() does not throw at the saddle point.

With the changes, the new tests pass, and so do the other test files in test/unit/math/laplace/. I stopped the slow laplace_disease_map_test suite in laplace_marginal_lpdf_test.cpp after its first case, which passed. clang-format and cpplint pass on the changed files.

Side Effects

• laplace_latent_solve*() and laplace_latent*_rng(): the mean moves from K * grad log p(y | theta_hat) to theta_hat. The two agree up to the convergence error, so results change only where the iteration stopped before exact convergence, and they move toward the mode.
• laplace_marginal*() with solver 3, or after a fallthrough to solver 3: where it returned NaN because of a negative pivot, it now returns the log marginal. Where det(B) <= 0, it now throws std::domain_error; before, it returned NaN, or a finite value when the permutation was odd and all pivots were positive. In the model block, both a NaN and an exception reject the proposal; in generated quantities, the exception replaces a silent NaN or a wrong value.
• The log-determinant change affects only solver 3. Solvers 1 and 2 compute the log determinant from the Cholesky factor of a symmetric positive definite matrix.
• The Newton iteration can still converge to a saddle point (example C); this PR only makes the error name the cause.

Release notes

The embedded Laplace functions laplace_latent_solve(), laplace_latent_rng() and their variants now return the final Newton iterate as the mean, which is more accurate when the iteration stops before exact convergence. laplace_marginal() and its variants with solver 3 no longer return NaN when the LU factorization has a negative pivot, and they report an error when the Newton iteration ends at a point where the negative Hessian of the log density is not positive definite.

Checklist

• Copyright holder: Aalto University

  The copyright holder is typically you or your assignee, such as a university or company. By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
  - Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
  - Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested