Laplace Bug when passing Eigen::Map in tuple of functor arguments

@WardBrian

I'm trying to simplify the result here but have not been successful yet. i.e. this gives correct values. @WardBrian can you see anything here that is meaningfully different than what the compiler generates?

https://github.com/stan-dev/math/blob/e53e95b7fcdfad20e1d0e7df24333e4d5a8b5510/test/unit/math/laplace/aki_ex_test.cpp

@SteveBronder

@SteveBronder I think you misread Aki's post, the data you've hardcoded in here is what works.

The other data, like datat <- list(N=1, y=183, mu=0.5, sigmaz=2.5, run_laplace=0), causes problems

Oh I did misread it my apologies

Changing the data to

  const std::vector<int>    y{183};
  const std::vector<double> mu{0.5};
  const double sigmaz = 2.5;

does indeed lead to the test failing with an exception thrown

C++ exception with description "Error in laplace_marginal_density: theta contains NaN values for all values." thrown in the test body.

@avehtari

By adding some prints to the likelihood it seems the issue here is that theta is shooting off to massive values in the optimization:

[0]
[0]
[100.265]
[100.265]
[100.265]
[4.78573e+30]
[4.78573e+30]
[4.78573e+30]

Exception: Error in laplace_marginal_density: theta contains NaN values for all values.

Inside integrand, the value of theta never exceeds ~7000 it seems

If you play around with theta_0, you can find values that work (for example, { "N": 1, "y": [183], "mu": [0.5], "sigmaz": 2.5, "run_laplace": 1, "theta_0": [3] }) seems fine! Theta never gets larger than 7.5 ), but most of them shoot off into either massively positive or massively negative numbers and then you get the exception

cc @avehtari

I added print also for the log likelihood. Running with

datat <- list(N=1, y=183, mu=0.5, sigmaz=2.5, run_laplace=1)

interestingly theta does come back from e+31 range

Chain 1 [0] 
Chain 1 -684.009 
Chain 1 [0] 
Chain 1 -684.009 
Chain 1 [100.265] 
Chain 1 -5.77624e+43 
Chain 1 [100.265] 
Chain 1 -5.77624e+43 
Chain 1 [100.265] 
Chain 1 -5.77624e+43 
Chain 1 [-1.58456e+31] 
Chain 1 -2.89975e+33 
Chain 1 [-1.58456e+31] 
Chain 1 -2.89975e+33 
Chain 1 [-1.58456e+31] 
Chain 1 -2.89975e+33 
Chain 1 [1143.75] 
Chain 1 -inf 
Chain 1 [1143.75] 
Chain 1 -inf 
Chain 1 [1143.75] 
Chain 1 -inf

One thing the optimization should do is that when ever the log likelihood is evaluated to be -inf or nan, the step should be halved. I'll try to have a look at the current code.

I had a quick look at the code in laplace_marginal_density.hpp. It could be made safer by checking that every time when laplace_likelihood::log_likelihood() or laplace_likelihood::diff() are called, the return value is finite (not -inf or nan) and if a non-finite return value is encountered then the step size is halved. Now it is possible to end up with theta which has non-finite log likelihood and/or derivative.

@avehtari

the step should be halved.

Reading the current code, it looks like the stepsize is always 0.5 -- is that what you were referring to @avehtari?

There used to be a comment from Charles about changing this but it looks like it got deleted at some point

math/stan/math/mix/functor/laplace_marginal_density.hpp

Line 219 in a23525a

= (a + a_prev) * 0.5; // TODO(Charles) -- generalize for any factor

@avehtari

@avehtari @WardBrian I have a branch here showing what is happening. From the debugging log below we can see that with a mean of 0.5 the gradient for theta shoots off to 180'ish which then causes a and b for updating theta to become huge. That's an issue because in the poisson_log function we take an exponential of that value.

__________
ITER: 0 / 100
diff theta gradient: 181.351
curr theta: 
	0
W:  
	Nonzero entries:
	(1.64872,0) 
	Outer pointers:
	0  $
	1.64872

W_r:  
	1.28403
theta_grad:  
	181.351
a: 
	16.0424
b: 
	181.351
--------------
theta: 
	100.265
--------------
objective_old: 
	-1.79769e+308
objective_new: 
	-5.77624e+43
__________
ITER: 1 / 100
diff theta gradient: 
	-5.77624e+43
curr theta: 
	100.265
W:  
Nonzero entries:
(5.77624e+43,0) 
Outer pointers:
0  $
5.77624e+43

W_r:  
	7.60016e+21
theta_grad:  
	-5.77624e+43
a: 
	-2.5353e+30
b: 
	5.73378e+45

This ran went of to infinity in the objective, but other runs would just go off to some huge but wrong number.

I changed laplace here so that we take a partial step in the direction of the new a instead of a full step. We try a new partial step and if that fails to give us a better theta we halve the step size and keep trying until we to a step size of 1e-8 or find a better theta. This seems to be better and the new version can solve this problem with a step size of 0.9. @avehtari do you know of a better way to fix this?

If we do this step approach do you both think we should expose step_size to the user in tol so they can control it?

When you say "halve the step size and keep trying", from where do you keep going? The original point, or the point immediately before failure?

The point immediately before failure. So same a and b, but a smaller step size for the update.

I would expect it to perform better if it was trying again from the original point, because then you're not having to compensate for the initial too-large-but-not-yet-failing steps. Is that hard to try?

@charlesm93

For this problem that would not work because it can hit a very large number that is still wrong but not infinite.***

***EDIT: Let me try this as it's not that hard to change. But I think it will not catch the edge case^. Just to be clear, you mean completely restarting the algorithm with a smaller stepsize?

@charlesm93 do you have any thoughts on this? Is there a way we can add something like the Armijo rule to dynamically change the step size?

@mitzimorris

@mitzimorris when I was working on GPs a lot, I mostly use log-normal for length-scale as it's easy to define with desired mass on desired interval, and it makes the inference stable as the tails are Gaussian (which is great for HMC/NUTS, too). Only paper I know investigating GP length scale priors is this workshop paper https://drive.google.com/file/d/0B3WHb3BabixAYlptTVBWUGdyVEE/view?resourcekey=0-mj7f4AZQ-UN1Rvd9NrRlHg. I'm curious, can you ask Claude for the reference for its recommendation?

Reading the current code, it looks like the stepsize is always 0.5 -- is that what you were referring to @avehtari?

The "0.5" indicates that the stepsize is halved. The step size itself is given by the inverted Hessian. See page 9 at https://arxiv.org/pdf/2306.14976. If you picked 0.1 instead of 0.5, you would reduce the stepsize by 10. Maybe something worth playing with.

I don't know optimization INLA is using, but all GP packages doing Laplace are using something else than Newton, so maybe we should also use something else. Ad L-BFGS is already there, how hard it would be to try that?

We could try solving the underlying optimization problem without worrying about the gradient computation and see if that improves the result.

The gradient calculation (adjoint-differentiation) would need some refactoring if we switch from a Newton solver to an L-BFGS. While there is no technical challenge that I foresee, it will require a bit of work.

@avehtari

@avehtari try out the latest version

cd stan/lib/stan_math
git pull origin fix/laplace-line-search
git checkout 21bd94183b2006082c750ec88c6bce9fa7e2274f
cd ../../../
make clean-all
make -j4 build

From your script this seems to be running okay now. Though in your example there are some cases I'm seeing where the error relative to integrate_1d increases quite a bit as mu becomes more negative and sigma grows. Since laplace_marginal no longer has lpdf at the end I wonder if it would be useful to return back our final estimated theta values? I think knowing how mu and theta interacted here would be nice.

For everyone else, the graph below is comparing the error in log likelihood of embedded laplace with the likelihood from using integrate 1d on the roach data. @avehtari would it be okay if I posted the R and stan code here?

Code for above

worst_iter = which.max(max(abs(dr$log_lik - dr$log_lik2)))
graph_err = data.frame(sigma=draws_of(dr$sigma),
  mu=draws_of(dr$mu[worst_iter]),
  error=abs(draws_of(dr$log_lik[worst_iter])-draws_of(dr$log_lik2[worst_iter]))) |>
  ggplot(aes(x=mu,y=sigma,color=error)) +
  geom_hline(aes(yintercept = mean(sigma))) +
  geom_vline(aes(xintercept = mean(mu))) +
  geom_point() + ggtitle(paste0("Iter: ", worst_iter, "\tyline: mean of sigma\txline: mean of mu")) +
  scale_color_gradientn(colors = c("blue", "yellow", "red")) 
print(graph_err)
for (i in 1:length(dr$mu)) {
  graph_err = data.frame(sigma=draws_of(dr$sigma),
    mu=draws_of(dr$mu[i]),
    error=abs(draws_of(dr$log_lik[i])-draws_of(dr$log_lik2[i]))) |>
    ggplot(aes(x=mu,y=sigma,color=error)) +
    geom_point() + ggtitle(paste0("Iter: ", i))
  print(graph_err)
  readline(prompt="Press [enter] to continue")
}

This is a graph of all of the errors between the likelihoods of laplace vs integrate_1d by mu and sigma estimates for all y. Very weird U shape when mu is around -2.5?

@avehtari

@avehtari try out the latest version

Sorry, I missed this due to weekend. I tested, but it prints out so much diagnostic messages that it's very slow. CmdStanR allows suppressing the messages, but it seems they still cause significant overhead and sampling for the roaches did not finish in 3 hours. Can you make a version without all that output?

the graph below is comparing the error in log likelihood of embedded laplace with the likelihood from using integrate 1d on the roach data.

Errors seem to be small enough for practical purposes

@avehtari would it be okay if I posted the R and stan code here?

Here's the code

#' **Load packages**
library(cmdstanr)
options(mc.cores = 4)
library(loo)
library(ggplot2)
theme_set(bayesplot::theme_default(base_family = "sans"))
library(posterior)

#' Load data
data(roaches, package="rstanarm")
# Roach1 is very skewed and we take a square root
roaches$sqrt_roach1 <- sqrt(roaches$roach1)

#' Compile model
poisson_re_int <- "poisson_re_laplace.stan"
writeLines(readLines(poisson_re_int))
modpri <- cmdstan_model(stan_file = poisson_re_int, force_recompile=TRUE)

#' Stan data
datap <- list(N = dim(roaches)[1],
              P = 3,
              offsett = log(roaches$exposure2),
              X = roaches[,c('sqrt_roach1','treatment','senior')],
              y = roaches$y,
              integrate_1d_reltol = 1e-6)

#' Sample
fitpri <- modpri$sample(data = datap, refresh = 0, show_messages = FALSE,
                        chains = 4, parallel_chains = 4)

#' LOO with Laplace integrated log_lik
(loopri <- fitpri$loo(save_psis=TRUE, variables="log_lik"))

#' LOO with Quadrature integrated log_lik
(loopri <- fitpri$loo(save_psis=TRUE, variables="log_lik2"))

#' Plot the abs(log_lik - log_lik2) with different mu and sigmaz values
dr <- as_draws_rvars(fitpri$draws(variables=c('log_lik','log_lik2','mu','sigmaz'), format="df"))

data.frame(sigma=draws_of(dr$sigma),
           mu=draws_of(dr$mu[93]),
           error=abs(draws_of(dr$log_lik[93])-draws_of(dr$log_lik2[93]))) |>
  ggplot(aes(x=mu,y=sigma,color=error)) +
  geom_point()

and the Stan code

// Poisson regression with hierarchical intercept ("random effect")
functions {
  real integrand(real z, real notused, array[] real theta,
               array[] real X_i, array[] int y_i) {
    real sigmaz = theta[1];
    real mu_i = theta[2];
    real p = exp(normal_lpdf(z | 0, sigmaz) + poisson_log_lpmf(y_i | z + mu_i));
    return (is_inf(p) || is_nan(p)) ? 0 : p;
  }
  real poisson_re_log_ll(vector theta, data int y, real mu) {
    return poisson_log_lpmf(y | mu + theta);
  }
  matrix cov_fun(real sigma, data int N) {
    return diag_matrix(rep_vector(sigma^2, N));
  }
}
data {
  int<lower=0> N;           // number of data points
  int<lower=0> P;           // number of covariates
  matrix[N,P] X;            // covariates
  array[N] int<lower=0> y;  // target
  vector[N] offsett;        // offset (offset variable name is reserved)
  real integrate_1d_reltol;
}
parameters {
  real alpha;               // intercept
  vector[P] beta;           // slope
  vector[N] z;              // individual intercept ("random effect")
  real<lower=0> sigmaz;     // prior scale for z
}
model {
  // priors
  alpha ~ normal(0, 3);  
  beta ~ normal(0, 3);    
  z ~ normal(0, sigmaz);  
  sigmaz ~ normal(0, 1);
  // observation model
  y ~ poisson_log_glm(X, z+offsett+alpha, beta); 
}
generated quantities {
  // log_lik for PSIS-LOO
  vector[N] log_lik;
  vector[N] log_lik2;
  vector[N] y_rep;
  vector[N] y_loorep;
  vector[N] mu = offsett + alpha + X*beta;
  vector[N] s = sqrt(diagonal(cov_fun(sigmaz, N)));
  for (i in 1:N) {
    // z as posterior draws, this would be challenging for PSIS-LOO (and WAIC)
    // log_lik[i] = poisson_log_glm_lpmf({y[i]} | X[i,], z[i]+offsett[i]+alpha, beta);
    // posterior predictive replicates conditional on p(z[i] | y[i])
    // in theory these could be used with above (non-integrated) log_lik and PSIS-LOO
    // to get draws from LOO predictive distribution, but the distribution of the
    // ratios from above lok_lik is bad
    real mu_i = offsett[i] + alpha + X[i,]*beta;
    y_rep[i] = poisson_log_rng(z[i] + mu_i);
    // embedded Laplace
    log_lik[i] =  laplace_marginal(poisson_re_log_ll, (y[i], mu_i), cov_fun, (sigmaz, 1));
    // we can integrate each z[i] out with 1D adaptive quadrature to get more
    // stable log_lik and corresponding importance ratios
    log_lik2[i] = log(integrate_1d(integrand,
                              negative_infinity(),
                              positive_infinity(),
                              append_array({sigmaz}, {mu_i}),
			      {0}, // not used, but an empty array not allowed
			      {y[i]},
			      integrate_1d_reltol));
    // conditional LOO predictive replicates conditional on p(z[i] | sigmaz, mu_i)
    // these combined with integrated log_lik and PSIS-LOO provide
    // more stable LOO predictive distributions
    y_loorep[i] = poisson_log_rng(normal_rng(0, sigmaz) + mu_i);
  }
}

Sorry try this from a cmdstan repo. I switched to a different branch

cd stan/lib/stan_math
git pull origin fix/laplace-wolfe
git checkout 21bd94183b2006082c750ec88c6bce9fa7e2274f
cd ../../../
make clean-all
make -j4 build

That should have all the error messages removed

Works now! The error in Laplace integrated LOO compared to quadrature integrated LOO is about the same order as Monte Carlo error with 4000 posterior draws (with quite low ESS for some parameters)

With Laplace

Computed from 4000 by 262 log-likelihood matrix.

         Estimate   SE
elpd_loo   -878.9 38.4
p_loo         5.1  0.5
looic      1757.7 76.8
------
MCSE of elpd_loo is 0.2.
MCSE and ESS estimates assume MCMC draws (r_eff in [0.0, 0.3]).

All Pareto k estimates are good (k < 0.7).
See help('pareto-k-diagnostic') for details.

With integrated LOO

Computed from 4000 by 262 log-likelihood matrix.

         Estimate   SE
elpd_loo   -878.6 38.3
p_loo         5.1  0.5
looic      1757.3 76.6
------
MCSE of elpd_loo is 0.2.
MCSE and ESS estimates assume MCMC draws (r_eff in [0.0, 0.3]).

Excellent!!

I'm going to cleanup the code today and tomorrow and put up a pr hopefully by the end of the week

@avehtari

@avehtari and @charlesm93 I have a branch up and running that gives results that look nicer. Can you try out cmdstan with this math branch?

cd stan/lib/stan_math
git pull origin debug/laplace-wolfe
git checkout 99307e22b3b46b93b3ce0450b49775f5ddca33d9
cd ../../../
make clean-all
make -j4 build

The results from the tests look nice with this branch. I plan to open up a branch on Monday with a PR for these changes

git pull origin debug/laplace-wolfe
git checkout 99307e2

I tested Roaches Poisson varying intercept integrated LOO and it worked fine. I'll try to get it break with some other example

Navigation Menu

Repository navigation

Laplace Bug when passing Eigen::Map in tuple of functor arguments #3205

25 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Laplace Bug when passing Eigen::Map in tuple of functor arguments #3205

Description

Activity

WardBrian commented on Jun 13, 2025

SteveBronder commented on Jun 13, 2025

WardBrian commented on Jun 16, 2025

WardBrian commented on Jun 16, 2025

avehtari commented on Jun 22, 2025

avehtari commented on Jun 23, 2025

WardBrian commented on Jul 9, 2025

SteveBronder commented on Jul 17, 2025

WardBrian commented on Jul 17, 2025

SteveBronder commented on Jul 17, 2025

WardBrian commented on Jul 17, 2025

SteveBronder commented on Jul 17, 2025

25 remaining items

avehtari commented on Jul 30, 2025

charlesm93 commented on Aug 6, 2025

charlesm93 commented on Aug 6, 2025

SteveBronder commented on Aug 7, 2025

SteveBronder commented on Aug 7, 2025

SteveBronder commented on Aug 11, 2025

avehtari commented on Aug 12, 2025

SteveBronder commented on Aug 12, 2025

avehtari commented on Aug 12, 2025

SteveBronder commented on Aug 12, 2025

SteveBronder commented on Aug 12, 2025

SteveBronder commented on Aug 22, 2025

avehtari commented on Aug 29, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions