Speeding up a hierarchical negative binomial model

[Patrick Miller]

Hi,

I am new to STAN, but have been finding it very powerful for a couple of problems that I have needed to model. The most recent model that I am trying to build has run into speed issues. It is a hierarchical model with a negative binomial response variable that has 2 parameters (plus an intercept) that vary by group. After a while trying to tweak it to fit the data correctly, I decided to do what I should have done in the first place, simplify it and use simulated data.

However, even after doing this, I am still running into what I find to be slow (small ESS/time). Perhaps, it is just the nature of the problem, but I figured I would ask for help here. From my limited knowledge about how to debug, what I have been seeing is very large leapfrog jumps with small stepsizes (I increased maximum tree depth to be safe). Is this indicative of a problem with my model?

I would like to speed it up as much as possible because the actual data set is orders of magnitude larger and has more groups.

Thank you for any advice ahead of time,

Patrick

// STAN MODEL

data {

    int<lower=1> n_train;

    int<lower=1> n_clusters;

    int<lower=0, upper=n_clusters> cluster_train[n_train];

    int<lower=0> y[n_train];

    vector[n_train] log_x1;

    vector[n_train] x2;    

}

parameters {

    real<lower=0> r;

    

    real mu_b0;

    real mu_b1;

    real mu_b2;

        

    real<lower=0> sigma_b0;

    real<lower=0> sigma_b1;

    real<lower=0> sigma_b2;

    

    vector[n_clusters] b0_raw;

    vector[n_clusters] b1_raw;

    vector[n_clusters] b2_raw;

}

transformed parameters {

    vector[n_clusters] b0;

    vector[n_clusters] b1;

    vector[n_clusters] b2;

    

    vector[n_train] eta_train;

    

    b0 <- mu_b0 + sigma_b0 * b0_raw;

    b1 <- mu_b1 + sigma_b1 * b1_raw;

    b2 <- mu_b2 + sigma_b2 * b2_raw;

    

    for(i in 1:n_train){

      eta_train[i] <- 

        b0[cluster_train[i]] + 

        b1[cluster_train[i]] * log_x1[i] + 

        b2[cluster_train[i]] * x2[i];

    }

}

model {

    mu_b0 ~ normal(3, 1);

    mu_b1 ~ normal(-0.2, 0.1);

    mu_b2 ~ normal(-0.15, 0.5);

    

    sigma_b0 ~ cauchy(0, 1);

    sigma_b1 ~ cauchy(0, 0.1);

    sigma_b2 ~ cauchy(0, 0.1);

      

    b0_raw ~ normal(0, 1);

    b1_raw ~ normal(0, 1);

    b2_raw ~ normal(0, 1);

    

    r ~ cauchy(0, 2.5);

    

    y ~ neg_binomial_2_log(eta_train, r);

}

[Bob Carpenter]

I'm not sure why it's slow, though the overdispersion can be problematic,
as can problems where things have the kind of varying scales your priors
seem to be indicating. And it's easy to make a mistake on negative
binomial parameterizations, but if it fits fake data at a smaller scale,
it's probably OK in terms of parameterization.

I'm also not sure whether adding a covariance prior over
random effects would help --- are they highly correlated? If so,
that might help control the sampling, though it gives you a more
complex estimation problem.

- Bob

On Jun 25, 2015, at 1:55 PM, Patrick Miller <mille...@gmail.com> wrote:
>
> Hi,
>
> I am new to STAN, but have been finding it very powerful for a couple of problems that I have needed to model. The most recent model that I am trying to build has run into speed issues. It is a hierarchical model with a negative binomial response variable that has 2 parameters (plus an intercept) that vary by group. After a while trying to tweak it to fit the data correctly, I decided to do what I should have done in the first place, simplify it and use simulated data.
>

> However, even after doing this, I am still running into what I consider slowness. Perhaps, it is just the nature of the problem, but I figured I would ask for help here. From my limited knowledge about how to debug, what I have been seeing is very large leapfrog jumps with small stepsizes (I increased maximum tree depth to be safe). Is this indicative of a problem with my model?

> --
> You received this message because you are subscribed to the Google Groups "Stan users mailing list" group.
> To unsubscribe from this group and stop receiving emails from it, send an email to stan-users+...@googlegroups.com.
> For more options, visit https://groups.google.com/d/optout.
> <stan_dat.rda><stan_negbin_model.txt><stan_r.r>

[Patrick Miller]

Thanks for the reply. That makes sense that overdispersion would make the sampling slow with the varying scales. I'll try adding a covariance prior -- I had just figured that the added complexity would slow it down even more.

[Bob Carpenter]

> On Jun 26, 2015, at 8:29 AM, Patrick Miller <mille...@gmail.com> wrote:
>
> Thanks for the reply. That makes sense that overdispersion would make the sampling slow with the varying scales. I'll try adding a covariance prior -- I had just figured that the added complexity would slow it down even more.

It will definitely be slower per log density and gradient eval.

But it's very possible it'll be faster in terms of effective
sample size per second if it controls the chains better.

I don't have good intuition for this problem. We should do some
testing on our side since it's a common situation for us. We
typically start the way you did --- without the covariance priors,
then gradually add more structure.

- Bob

[Patrick Miller]

I have added the covariance matrix parameterization, although I'm not sure I did it properly. Over a very small subset of the data it seems to be speeding it up a little bit (80% of the time to get an effective sample), but the leapfrog steps are the same size. I'll run it on a larger subset of the data over the weekend.

// Model with covariance priors

data {

    int<lower=0> n_train;

    int<lower=1> n_clusters;

    int<lower=0, upper=n_clusters> cluster_train[n_train];

    int<lower=0> y[n_train];

    

    int K; // number of parameters (3)

    

    matrix[n_train, K] x; // individual predictors

}

parameters {

    real<lower=0> r;

    

    vector[K] mu_beta;

    

    matrix[K, n_clusters] z;

    cholesky_factor_corr[K] L_Omega;

    vector<lower=0>[K] tau; // prior scale on the parameters

}

transformed parameters {

    matrix[n_clusters, K] beta; // parameter matrix

    matrix[n_clusters, K] mu_beta_mat;

    

    for(cl in 1:n_clusters)

      mu_beta_mat[cl] <- mu_beta';

    

    beta <- mu_beta_mat + (diag_pre_multiply(tau, L_Omega) * z)';

}

model {

  vector[n_train] x_beta_cluster;

  for(n in 1:n_train)

    x_beta_cluster[n] <- x[n] * beta[cluster_train[n]]';

  

  mu_beta[1] ~ normal(3, 1);

  mu_beta[2] ~ normal(-0.2, 0.1);

  mu_beta[3] ~ normal(-0.15, 0.1);

  

  L_Omega ~ lkj_corr_cholesky(2);

  

  tau[1] ~ normal(0, 1);

  tau[2] ~ normal(0, 0.1);

  tau[3] ~ normal(0, 0.1);

  

  to_vector(z) ~ normal(0, 1);

  

  r ~ normal(0, 3);

  

  y ~ neg_binomial_2_log(x_beta_cluster, r);

[Bob Carpenter]

The leapfrog steps are going to take a bit longer, but
it might not even be measurable at low dimensions.

The coding looks OK to me, but Ben's the expert here.

- Bob

[Patrick Miller]

Just reporting back...adding the covariance prior definitely helped with speeding up the model, especially for the full data set. I have gotten the warm-up portion to be particularly fast, while the sampling portion takes ~12 minutes per effective sample (down from 30+ minutes). Thanks for the help!

Patrick