speeding up zero inflation model through vectorization?

[Guido Biele]

Hello,

in the stan language reference a zero inflation model is described as follows:

for (n in 1:N) {

  if (y[n] == 0)

    increment_log_prob(log_sum_exp(bernoulli_log(1,theta),

                                                     bernoulli_log(0,theta)

                                                     + poisson_log(y[n],lambda)));

  else

    increment_log_prob(bernoulli_log(0,theta)

                                 + poisson_log(y[n],lambda));

}

I was thinking of improving the speed of the model by

- splitting the array y in two parts, 

  one where y = 0 and the other where y != 0

- putting the results of

  log_sum_exp(bernoulli_log(1,theta),

                      bernoulli_log(0,theta)

                      + poisson_log(y[n],lambda))

  and

  bernoulli_log(0,theta) + poisson_log(y[n],lambda)

  into two vectors log_prob_y and log_prob_y0

- and finally pass the vectors log_prob_y 

  and log_prob_y0 to increment_log_prob

The underlying idea is to vectorize as much as possible. 

For example when calculating log_prob_y and log_prob_y0, 

I would "vectorize" for bernoulli_log and poisson_log or 

neg_binomial_2_lo_log. The only place where one would

still have to loop is for log_sum_exp.

I am unsure if such a modification would would improve 

the speed noticeable, because if I understand it correctly 

calculating the gradients is the most time consuming operation,

and I do not understand if the changes I suggested have any 

influence on the calculation of gradients.

Thanks in advance!

Guido

PS: In my particular case the data set is moderately large

(N = 50 000 rows and 20 predictors) and I am using the 

negative binomial model. 

An improvement in speed of 10% would already be useful for me!

[Bob Carpenter]

We discussed this on the list recently and I added code examples
of how to do exactly what you're suggesting in the 2.9.0 manual in
discussion of the hurdle model (p. 115).

The key is to use algebra to reduce the size of the expression
graph.

- Bob

> --
> 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.
> To post to this group, send email to stan-...@googlegroups.com.
> For more options, visit https://groups.google.com/d/optout.

[Guido Biele]

[after a longer holiday-hiatus]

Bob,

thanks for the pointers (and sorry for not remembering, checking the manual better).

I also looked at the hurdle model, and analog to what you 

were showing in the manual for the poisson model there is a 

faster version to calculate the cdf for the negative binomial distribution. 

I'll post this when I am sure I got it right.

I had troubles implementing this because, i believe, the example 

for the hurdle model on page 115 has a small but consequential error:

When implementing the truncation, I think the code increments the log posterior by

 -log(poissonCDF(0,lambda))      instead of the correct

 -log(1-poissonCDF(0,lambda)).

Alternatively, log(1-exp(-lambda)) is already the CCDF, and not the CDF as 

it sais on page 115.

Or am I misunderstanding something/not reading carefully enough?

Guido

in 

[Guido Biele]

I am following up with 2 points:

 

1) My belief that the reference manual describes

and incorrect optimized hurdle model is strengthened

2) Code for a faster zero inflated negative binomial model.

 

### 1) The reference manual describes an ###

###      incorrect optimized hurdle model ###

 

Following up on my last message:

I checked and, as described in the manual,

poissonCDF(0 | lambda) = log(1-exp(-lambda)).

However, to calculate p(y[n] | lambda) for the hurdle model, 

one has to divide (or subtract on the log scale) by the

complementary poisson CDF. poissonCCDF.

Instead, the code on page 115 of the 

reference manual uses the poissonCDF.

 

So it should not be 

increment_log_prob(-log1m_exp(-lambda))

but

increment_log_prob(-log1m_exp(log1m_exp(-lambda))).

[Alternatively, I am making a fool out of myself ...]

 

### 2) Code for a faster zero inflated ###

###      negative binomial model. ###

 

Replacing the model section from page 114 in the 

stan reference the following makes gradient evaluation

ca. 20% faster.

 

This code is for a hierarchical model involving multiple

groups, using following variables:

y: vector of length N with observed counts.

phi: vector with over dispersion parameters for groups

etas: vector of length N containing log(mu) for the negative binomial distribution

thetas: logit of p(y == 0)

sx: design matrix

sx0: design matrix for zero inflation

g_start: index of the first entry for a group in y, sx, or sx0

g_nzstart: index of the first nonzero entry for a group in y, sx, or sx0

g_end: index of the last entry for a group in y, sx, or sx0

betas: matrix with regression weights to obtain log(mu), by group

alphas: regressions weight for zero inflation part (same for all groups)

### begin model ###

...

thetas <- to_vector(sx0 * to_matrix(alpha));

 

for (gr in 1:M) {

    log_phi <- log(phi[gr]);

     etas[g_start[gr]:g_end[gr]] <- sx[g_start[gr]:g_end[gr],1:K] * beta[gr,1:K]';

 

     // y > 0  cases

     increment_log_prob(bernoulli_logit_log(0,thetas[g_nzstart[gr]:g_end[gr]])

                                   + neg_binomial_2_log_log(y[g_nzstart[gr]:g_end[gr]],

                                                                         etas[g_nzstart[gr]:g_end[gr]],

                                                                         phi[gr]));

     // y == 0 cases

     for (k in g_start[gr]:(g_nzstart[gr]-1)) {

       increment_log_prob(log_sum_exp(bernoulli_logit_log(1,thetas[k]),

                                                         bernoulli_logit_log(0,thetas[k]) 

                                                         + neg_binomial_2_log_log(0,etas[k],phi[gr])));            

     }

   }

 

### end model ####

Following additional modifications did make the model

 slightly faster (around 5%), but appeared to be less stable 

(i.e. produced NaNs or infinite values where the above model did not):

bernoulli_logit_log(1,thetas[k]) can be obtained from log_inv_logit(thetas[k])

bernoulli_logit_log(0,thetas[k]) can be obtained from log1m__inv_logit(thetas[k])

neg_binomial_2_log_log(0,etas[k],phi[gr]) can be obtained from

(log1m_exp(etas[k] - log_sum_exp(etas[k],log_phi)))*phi[gr])).

Best, Guido

PS: getting the count of zero values to then use the binomial 

distribution does not work in my case, because every participants

has a different theta value.

[Bob Carpenter]

Thanks --- I think someone else (or you before) may have mentioned this.
I'd thought I'd put it in the manual issues, but apparently not, so here
goes:

https://github.com/stan-dev/stan/issues/1709#issuecomment-171032587

- Bob

[Bob Carpenter]

Even better, it can be vectorized as bernoulli_logit_log(1, theta) if you do
it for all of the k values. It's useful to spend the time to put themtogether
that way if you want.

I don't want to get too carried away optimizing the code, but would you mind
sending a complete working model? I can check it into our example-models
repo and credit you with it after I validate that it works.

- Bob

[Guido Biele]

Hi Bob,

thanks for your feedback and help!

I am attaching 3 files:

1) HZAxNB2Regression.stan 

is a hierarchical zero inflated negative binomial model 

(hierachical only for the count part, not the the zero inflation part).

2) sim_nbinom_data.R 

simulates zero inflated data and returns those 

and design matrix and parameters used to 

created these data

3) test_HZAxNB2Reg.R

used files 1) and 2) to do a parameter recovery 

test and plot the results.

I think the estimated are close enough to the 

generative parameters, expect that 

I am not able to fully recover the zero alteration 

(which is more a hurdle model ...).

I am also sometimes getting messages about a few 

divergent iterations. But the other diagnostic parameters

(Rhat, neff, msce) look fine. I am not sure what the problem here is .

Best - Guido

PS:

The reason I am so interested in optimizing the model

is that my models already run 20 or more hours (40000+ cases 

and 20+ variables), and because I use them in a 

multiple imputation context,I need at least 10 chains per model.

So a 20% faster model is worth something.