Embedded Laplace, solver 3: the Newton iteration converges to saddle points and converges slowly

The Newton loop is run_newton_loop(), with flip_direction() and the Wolfe line search that starts from the Barzilai-Borwein step; the saddle-point check and the stopping rule are in the same path. Run the R script in the issue (issue_laplace_solver3.R), which uses laplace_latent_solve_tol() with solver 3 for example 1 (saddle point) and example 2 (iteration count). Done means example 1 reaches a maximum from every start or fails with a named cause, and example 2 needs about 10 iterations at each tolerance.

Description

Solver 3 (LU) is the only solver of the embedded Laplace that accepts a log-likelihood that is not concave, that is, when the negative Hessian W of the log-likelihood is indefinite at some Newton iterate. I used it for a nonlinear joint model with 4 latent variables per group. Its log-likelihood is not concave at the mode for about half of the groups. The Newton iteration of solver 3 has two problems there. The examples below are small; both problems occur with CmdStan 2.40.0 and with develop (43bc5684a3).

Notation: f(theta) = log p(y | theta) + log N(theta | 0, K) is the objective of the Newton iteration. At a maximum, K^-1 + W is positive definite.

1. The Newton iteration can converge to a saddle point of f. In run_newton_loop(), flip_direction() reverses the Newton direction only when it is not an ascent direction, and nothing checks that K^-1 + W is positive definite at the result. A Newton step can be an ascent step in total while it moves to a saddle point along a direction in which f is convex. At the saddle point the gradient is zero, and the stopping rule accepts the point. The Laplace approximation is not defined there: laplace_latent_solve*() and laplace_latent*_rng() fail, and laplace_marginal*() returns NaN or fails. In example 1, the starts (+-0.01, 3) go to the saddle point (0, 0) in one step. In the joint model, 41 of 1000 posterior draws failed when the Newton iteration started at 0, and 163 of 1000 when it started at the posterior draw of the latent variables. In the failing draw that I traced, one group converged to a saddle point (eigenvalues of -Hessian of f: 104.4, 1.27, 0.85, -1.6). With #, the error message says that det(I + Sigma * W) is not positive at the final iterate, but the solve still fails. That check finds only an odd number of negative eigenvalues.

2. The convergence is linear, not quadratic. The line search starts from the Barzilai-Borwein step size, and the strong Wolfe curvature condition with c2 = 0.9 accepts every step alpha with |1 - alpha| <= 0.9 along a Newton direction of a nearly quadratic objective. In example 2 (4 latent variables, unimodal posterior), a trace of the iteration shows that alpha = 1 is never used: 26 of the 37 iterations at tolerance 1e-6 use alpha = 1.855. The error then shrinks by the factor |1 - alpha| = 0.855 per iteration. The solve needs 37, 50 and 65 iterations at tolerance 1e-6, 1.49e-8 (default) and 1e-10. As an experiment, I started the line search at alpha = 1 (the Wolfe conditions unchanged): the same solve took 10 iterations at each of the three tolerances and ended at the same mode. Also, the stopping rule checks only the change of the objective. With linear convergence it stops while |K * grad f| is still 0.093 (tolerance 1e-6) or 0.012 (default tolerance).

Possible changes:

• Item 1: check that K^-1 + W is positive definite during the iteration, for example with a Cholesky factorization of the symmetric I + L_K^T W L_K (K = L_K L_K^T). Where it is not, use a modified Newton step (Levenberg-Marquardt damping, or a modified Hessian) so that directions of negative curvature are ascent directions. At the end, a failed check could move the iterate along a direction of negative curvature and restart, or give an error that names the cause.
• Item 2: try alpha = 1 first for a Newton direction, and use the Barzilai-Borwein step only when alpha = 1 fails the line-search conditions. Stop on a gradient-based criterion as well, for example the Newton decrement grad f^T (K^-1 + W)^-1 grad f.

Example

The R script below (cmdstanr; the Stan programs are inline) runs both examples.

• Example 1 (item 1): 2 latent variables, K = [4 1; 1 1]. theta[1] has two Cauchy observations, y = (-3, 3), s = 0.3, and theta[2] has one Gaussian observation, z = 0, sz = 0.1. f has maxima at (+-2.9385, +-0.0097) and a saddle point at (0, 0). laplace_latent_solve_tol() with solver 3, tolerance 1e-10, from 5 start points.
• Example 2 (item 2): 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. BFGS from 50 start points finds one mode. laplace_latent_solve_tol() with solver 3 from one start point. The script finds the smallest max_num_steps that gives no "max number of iterations" warning.

# issue_laplace_solver3.R
#
# Reproduction for a stan-dev/math issue on the Newton iteration of the
# embedded Laplace with solver 3 (LU), the solver that accepts a
# log-likelihood that is not concave.  Self-contained (the Stan programs are
# inline).
#
#   1. The Newton iteration converges to a saddle point of
#      log p(theta | y), where the Laplace approximation is not defined.
#   2. The iteration needs many steps: the line search starts from the
#      Barzilai-Borwein step and accepts steps up to about 1.9 times the
#      Newton step, so the convergence is linear, not quadratic.
#      Shown by the smallest max_num_steps that gives no
#      "max number of iterations" warning.
#
#   Rscript tests/issue_laplace_solver3.R
#   CMDSTAN_PATH=~/proj/cmdstan STAN_MATH=~/proj/stan_math_laplace_mean/ \
#     Rscript tests/issue_laplace_solver3.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 output lines.
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) rep(NA_real_, 2))
  list(vals = vals, out = out)
}
exception_of <- function(out) {
  unique(sub(".*Exception: (.*?) \\(in '.*", "\\1",
             grep("Exception", out, value = TRUE), perl = TRUE))
}

# ── 1. 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;
}
generated quantities {
  vector[2] mode;
  {
    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;
  }
}
"
y <- c(-3, 3); s <- 0.3; z <- 0; sz <- 0.1
K <- matrix(c(4, 1, 1, 1), 2)
log_f1 <- function(th) {
  sum(dcauchy(y, th[1], s, log = TRUE)) + dnorm(z, th[2], sz, log = TRUE) -
    0.5 * sum(th * solve(K, th))
}
cat("\n1. Saddle point.  Stationary points of log p(theta | y):\n")
for (p in list(c(0, 0), c(2.9385, 0.0097), c(-2.9385, -0.0097))) {
  q <- optim(p, function(t) sum(numDeriv::grad(log_f1, t)^2), method = "BFGS",
             control = list(reltol = 1e-20))$par
  cat(sprintf("   (%7.4f, %7.4f): eigenvalues of -Hessian %s\n", q[1], q[2],
              paste(sprintf("%.3f", eigen(-numDeriv::hessian(log_f1, q),
                                          symmetric = TRUE)$values),
                    collapse = " ")))
}
mod1 <- compile(saddle_code)
for (start in list(c(0.3, 3), c(0.05, 3), c(0.01, 3), c(-0.01, 3),
                   c(-0.05, 3))) {
  r <- run_once(mod1, list(y = y, s = s, z = z, sz = sz, K = K,
                           theta0 = start), "mode")
  ex <- exception_of(r$out)
  cat(sprintf("   start (%5.2f, %4.1f): %s\n", start[1], start[2],
              if (length(ex)) paste("exception:", ex)
              else sprintf("mode (%.4f, %.4f)", r$vals[1], r$vals[2])))
}

# ── 2. Number of Newton iterations ─────────────────────────────────────────
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;
  int<lower=1> max_steps;
}
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, max_steps, 3, 100, 0));
    mode = lap.1;
  }
}
"
data2 <- list(n = 12L,
              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,
              theta0 = c(0.3636, -0.3385, 1.9794, -1.9322))
mod2 <- compile(tumor_code)
hits_limit <- function(tol, max_steps) {
  r <- run_once(mod2, c(data2, list(tol = tol, max_steps = as.integer(max_steps))),
                "mode")
  any(grepl("max number of iterations", r$out))
}
cat("\n2. Newton iterations, 4 latent variables (smallest max_num_steps without the\n",
    "   'max number of iterations' warning):\n")
for (tol in c(1e-6, 1.49012e-8, 1e-10)) {
  lo <- 1L
  hi <- 200L
  while (lo < hi) {
    mid <- (lo + hi) %/% 2L
    if (hits_limit(tol, mid)) lo <- mid + 1L else hi <- mid
  }
  cat(sprintf("   tol %-11g max_num_steps %d\n", tol, lo))
}

Expected Output

Current output (CmdStan 2.40.0):

1. Saddle point.  Stationary points of log p(theta | y):
   ( 0.0000,  0.0000): eigenvalues of -Hessian 101.334 -0.099
   ( 2.9385,  0.0097): eigenvalues of -Hessian 101.335 19.879
   (-2.9385, -0.0097): eigenvalues of -Hessian 101.335 19.879
   start ( 0.30,  3.0): mode (2.9385, 0.0097)
   start ( 0.05,  3.0): mode (2.9385, 0.0097)
   start ( 0.01,  3.0): exception: cholesky_decompose: Matrix m is not positive definite
   start (-0.01,  3.0): exception: cholesky_decompose: Matrix m is not positive definite
   start (-0.05,  3.0): mode (-2.9385, -0.0097)

2. Newton iterations, 4 latent variables (smallest max_num_steps without the
    'max number of iterations' warning):
   tol 1e-06       max_num_steps 37
   tol 1.49012e-08 max_num_steps 50
   tol 1e-10       max_num_steps 65

A trace of example 1 (debug prints in run_newton_loop()) for the start (0.01, 3): the first step (alpha = 1, directional derivative 912, so no flip) goes to (1.8e-6, 5.9e-9), the saddle point; the next step changes the objective by 1.2e-13, and the iteration stops.

Expected:

1. Example 1: a maximum from every start point, or, if the iteration cannot leave the saddle point, an error that names the cause.
2. Example 2: quadratic convergence near the mode, that is, about 10 iterations at each tolerance, as in the experiment with alpha = 1.

Current Version:

v5.4.0 (CmdStan 2.40.0). The same code is in develop at 43bc5684a3.