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 PR was Claude assisted

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

[stan-buildbot]

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan	0.04	0.04	1.06	5.94% faster
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan	0.16	0.16	1.0	0.22% faster
stat_comp_benchmarks/benchmarks/garch/garch.stan	0.85	0.86	0.99	-0.98% slower
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan	0.01	0.01	1.01	0.5% faster
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan	0.1	0.11	0.99	-0.65% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan	7.96	8.36	0.95	-5.02% slower
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan	3.7	3.71	1.0	-0.07% slower
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan	32.97	33.27	0.99	-0.9% slower
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan	0.38	0.38	0.98	-1.99% slower
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan	6.82	6.79	1.0	0.45% faster
stat_comp_benchmarks/benchmarks/sir/sir.stan	135.45	165.53	0.82	-22.2% slower
stat_comp_benchmarks/benchmarks/arK/arK.stan	3.22	3.23	1.0	-0.26% slower
stat_comp_benchmarks/benchmarks/arma/arma.stan	0.6	0.6	1.0	-0.01% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	25.18	25.83	0.97	-2.57% slower
performance.compilation	406.66	378.62	1.07	6.89% faster
Mean result: 0.9897691373393953

Jenkins Console Log
Jenkins Build Stages
Commit hash: 29e32c4215464fd5550540ce9a448016d22be56f

Machine information

Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           52 bits physical, 57 bits virtual
CPU(s):                                  192
On-line CPU(s) list:                     0-191
Thread(s) per core:                      2
Core(s) per socket:                      48
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              25
Model:                                   17
Model name:                              AMD EPYC 9474F 48-Core Processor
Stepping:                                1
Frequency boost:                         enabled
CPU MHz:                                 1497.097
CPU max MHz:                             4114.4229
CPU min MHz:                             1500.0000
BogoMIPS:                                7189.39
Virtualization:                          AMD-V
L1d cache:                               3 MiB
L1i cache:                               3 MiB
L2 cache:                                96 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-47,96-143
NUMA node1 CPU(s):                       48-95,144-191
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Mitigation; Clear CPU buffers
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d debug_swap

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin