how could I possibly speedup my sampling with a simple reinforcement learning model?

[Lei Zhang]

Hey Stan,

we are recently considering migrating our simple reinforcement learning model from JAGS to Stan. 

However, when we tried the below Stan model, it wasn't really faster than running in JAGS with the same configuration (same samples/chain, burnin, thin, etc.) .

I was wondering if there is a way to speedup / optimize the code or sampling based on the current version?

PS: after reading previous posts, I've already used categorical_logit(), and I already cut the transformed data section to even simply the model, nonetheless it was still slow.

And it might be worth mentioning that, we have 130 subjects and 100 trials. 

Thanks in advance, any suggestion is appreciated!

cheers,

Lei

data {

  int<lower=1> nSubjects;                            // number of subjects

  int<lower=1> nTrials;                              // number of trials 

  int<lower=1,upper=2> choice[nSubjects,nTrials];    // 2nd choices, 1 or 2

  real<lower=-1,upper=1> reward[nSubjects,nTrials];  // outcome, 1 or -1

}

parameters {

  // group-level parameters

  real<lower=0,upper=1> lr_mu;

  real<lower=0,upper=1> lr_kappa0;

  real<lower=0,upper=1> temp0_mu;

  real<lower=0,upper=1> temp0_kappa0;

  // subject-level parameters

  vector<lower=0,upper=1>[nSubjects] lr;

  vector<lower=0,upper=1>[nSubjects] temp0;

}

transformed parameters {

  real<lower=0> lr_kappa;

  real<lower=0> lr_a;

  real<lower=0> lr_b;

  real<lower=0> temp0_kappa;

  real<lower=0> temp0_a;

  real<lower=0> temp0_b;

  vector<lower=0,upper=20>[nSubjects] temp;

  lr_kappa <- 3.0 / (lr_kappa0 * lr_kappa0) - 1.0;

  lr_a <- lr_mu * lr_kappa;

  lr_b <- (1-lr_mu) * lr_kappa;

  temp0_kappa <- 3.0 / (temp0_kappa0 * temp0_kappa0) - 1.0;

  temp0_a <- temp0_mu * temp0_kappa;

  temp0_b <- (1-temp0_mu) * temp0_kappa;

  

  temp <- 20.0 * temp0;

}

model {

  row_vector[2] v[nTrials+1]; // values

  vector[nTrials] pe;         // prediction errors

  // hyperparameters, priors

  lr_mu ~ beta(1,1);

  lr_kappa0 ~ beta(1,1);

  temp0_mu ~ beta(1,1);

  temp0_kappa0 ~ beta(1,1);

  

  // subject loop and trial loop

  for (s in 1:nSubjects) {

    lr[s] ~ beta(lr_a,lr_b);

    temp0[s] ~ beta(temp0_a,temp0_b);

    v[1] <- rep_row_vector(0.0,2);

    for (t in 1:nTrials) {

      //* compute action probs using built-in softmax function and related to choice data */

      choice[s,t] ~ categorical_logit( temp[s] * v[t]' );

      //* prediction error */

      pe[t] <- reward[s,t] - v[t,choice[s,t]];

      //* value updating (learning) */

      v[t+1] <- v[t]; // make a copy of current value into t+1

      v[t,choice[s,t]] <- v[t,choice[s,t]] + lr[s] * pe[t]; // overwrite chosen value with pe update

    }

  }

}

[Daniel Lee]

Hi Lei,

What are the estimated number of effective sample sizes for JAGS and for Stan? Wall time is not a good way of how fast a model is running.

Having transformed data will not slow down the model. That gets executed once only.

Daniel

--
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.

[Bob Carpenter]

Also, you want to use the same method of computing n_eff and Rhat.
The built-ins for R2jags are different and more liberal and from what
I see report effective sample sizes about twice as large as we report
in our more conservative approach in Stan (which uses cross-chain information
on the mean rather than only within chain information which imparts a positive
bias to the n_eff estimates).

I couldn't follow the parameterization --- why are all those variables
constrained to (0,1) and then manipulated with

parameters {
real<lower=0,upper=1> lr_kappa0;

transformed parameters {

lr_kappa <- 3.0 / (lr_kappa0 * lr_kappa0) - 1.0;

model {
lr_kappa0 ~ beta(1,1);

You don't need the beta(1,1) statements --- they're redundant in
Stan, but shouldn't be causing too much slowdown. Variables are by
default uniform on their declared constraints.

You probably want to define all the variables "temp..." as local
variables in the model block so they don't get saved and printed.
It takes both time and space to save them.

You should just declare v as a vector, then you don't have to
transpose.

> row_vector[2] v[nTrials+1]; // values

> vector<lower=0,upper=20>[nSubjects] temp;

> choice[s,t] ~ categorical_logit( temp[s] * v[t]' );

But if it's just size 2, don't you really want to be using a
bernoulli_logit instead of categorical_logit? The choices will
have to be 0:1 rather than 1:2, but otherwise should be the same.

- Bob

[Lei Zhang]

Hi Daniel, 

thanks for your comments. 

I tried to have equal effective sample size per chain on both jags and stan, for 3 chains. On jags, I set sample = 4000, burnin = 2000, thin = 5, and on Stan, I set iter = 22000, warmup = 2000, thin = 5. Both lead to an effective posterior sample size of 4000 per chain if I got them correctly. Running on the same machine with a single core, jags cost roughly 2 hours, however Stan cost nearly 8 hours. 

Alongside the speed difference, I also observed different results of the parameter estimation (in our case, lr and temp) from jags and stan.

I was wondering if I did something wrong..

Thanks,

Lei

[Daniel Lee]

If the parameter estimates are different, then there is no use comparing run time. 

There are a few reasons that your parameter estimates are different. 

- the model was translated incorrectly; this is very common. It could be a mispecified distribution or a typo or a bug or...

- JAGS hasn't really converged; this is less common, but we see this when the posterior distribution is improper. 

I would start by double checking the translation. If you want exact replicability, then make sure the models are identical. Start with simulated data, make sure that fits equally in both. We did that when we started the project to evaluate whether Stan was working properly as an implementation of an MCMC algorithm. 

Then, move away from mathematical conveniences that were helpful in JAGS and start fitting the model you really want to fit. You're not limited by conjugate priors. There are HMC specific things that will make sampling faster.

But first things first. Fix the translation so the estimates are the same. 

Daniel

[Lei Zhang]

Hi Bob,

thanks for your comments and suggestions.

As you must have noticed, we were trying to fit a hierarchical model. 

First, taking the learning rate (lr) as an example. lr must be between 0 and 1, that's why we chose a beta distribution. lr_mu is the population mean of all the lr(s), lr_kappa is the concentration of the beta distribution. 

Our thought of using something like: 

     transformed parameters { 

        lr_kappa <- 3.0 / (lr_kappa0 * lr_kappa0) - 1.0; 

is to prevent bimodal posterior beta distribution (such as beta(0.5,0.5) ). Or perhaps Stan is clever enough to prevent this case automatically?

Then, lr_mu and lr_kappa were used to derive lr_a and lr_b for the beta distribution at the individual level. 

Second, the idea for the temp parameter is basically the same, with temp could be a between 0 and 20. That's why we used the same procedure for temp0 and multiplied with 20 afterwards.

Based on your experience, would you suggest some other distribution for a positive parameter constrained at (0,n], instead of multiplying n with a beta?

Third, I've tried to declare v as a vector to avoid transpose, however there is an indexing error. 

Suppose v is now declared as 'vector[2] v[nTrials+1]', how do I put that into the categorical_logit?

I tried both categorical_logit( temp[s] * v[:,t] ) 

and categorical_logit( temp[s] * v[,t] ), but neither worked. 

Perhaps in more general, how could I extract a particular column?

Thanks in advance.

Lei

[Bob Carpenter]

> On Jun 3, 2015, at 9:49 AM, Lei Zhang <bnuzhan...@gmail.com> wrote:
>
> Hi Bob,
>
> thanks for your comments and suggestions.
>
> As you must have noticed, we were trying to fit a hierarchical model.
>
> First, taking the learning rate (lr) as an example. lr must be between 0 and 1, that's why we chose a beta distribution. lr_mu is the population mean of all the lr(s), lr_kappa is the concentration of the beta distribution.
> Our thought of using something like:
> transformed parameters {
> lr_kappa <- 3.0 / (lr_kappa0 * lr_kappa0) - 1.0;
> is to prevent bimodal posterior beta distribution (such as beta(0.5,0.5) ). Or perhaps Stan is clever enough to prevent this case automatically?
> Then, lr_mu and lr_kappa were used to derive lr_a and lr_b for the beta distribution at the individual level.

You can just declare:

parameters {
real<lower=1> alpha;
real<lower=1> beta;

and then if you have

y ~ beta(alpha, beta);

the alpha and beta will be bounded > 1. The data should pull them
that way anyway if values above 1 make sense.

You used one that looks like this:

real<lower=0, upper=1> theta;
real<lower=0> kappa;

y ~ beta(theta * kappa, (1 - theta) * kappa);

where it's easy to think about theta and kappa, but it's hard to
constrain kappa so that theta * kappa > 1 and (1 - theta) * kappa > 1.
You'd have to do something awkward (from both a readability and
prior specification perspective) like this:

real<lower=0, upper=1> theta;
real<lower=(1 / fmin(theta, 1 - theta))> kappa;

> Second, the idea for the temp parameter is basically the same, with temp could be a between 0 and 20. That's why we used the same procedure for temp0 and multiplied with 20 afterwards.

You could also just constrain temp0 to be between 0 and 20 and then
give it a truncated prior from another family, but beta's good at
bounding it away from the edges (0 and 20). I'm curious what Andrew
thinks about this.

> Based on your experience, would you suggest some other distribution for a positive parameter constrained at (0,n], instead of multiplying n with a beta?

If you just want something vague or weakly informative, it
probably won't matter much. If you have real prior knowledge
about what to expect, then you can use that.

> Third, I've tried to declare v as a vector to avoid transpose, however there is an indexing error.
> Suppose v is now declared as 'vector[2] v[nTrials+1]', how do I put that into the categorical_logit?
> I tried both categorical_logit( temp[s] * v[:,t] )
> and categorical_logit( temp[s] * v[,t] ), but neither worked.
> Perhaps in more general, how could I extract a particular column?

When you declare

vector[2] v[nTrials + 1];

you have an array of 2-vectors. So v[n] is a column vector
already.

- Bob

P.S. Here's the original program for reference:

[Guido Biele]

Steingrove, Wetzels & Wagenmakers (2013) [http://www.ncbi.nlm.nih.gov/pmc/articles/PMC3847555/] used stan to fit a similar model, the PVL model which is kind of an extended RL model.

If I remember correctly, they did not use the beta distribution for the learning rate because they had convergence problems

(Though this was likely been with an earlier version of stan. I a writing this from memory, for details it would be best to contact the Helen Steingrove).

From my experience with fitting RL models a range of 20 for the temperature of the Luce/soft-max choice rule seems large.

An aside about avoiding bimodal learning rate distributions (on the group level): It is not unreasonable to assume a bimodal distribution, because if people learn by slowly accumulating payoff information, you get small learning rates (lets say < .2), but when people simply do win-stay lose-switch this behavior is best captured by high learning rates (close to 1). I have seen data where different participants used both these "strategies".

Guido

[Lei Zhang]

Hi Bob, 

thanks again for your comments and suggestions. They are really helpful. 

I just have one more concern at this moment. For the beta(a,b) distribution, though a>1, b>1 does a good constrain to avoid bimodal shape, however is it too rigid? For instance beta(2.5, 0.7) could also yield a good estimate in some cases (and this indeed happens for my current dataset). That said, either a or b ought be larger than 1, but not both of them. How could I possibly constrain that? Or does it make much sense to you?

btw, an update for other users (if of interest):

I made two changes that really speed up my sampling speed.

1. in the data section, the 'reward' variable was changed as real rather than a two-dimensional vector. Say

           real reward[s,t];

rather than

           row_vector[t] reward[s];

2. vectorizing helps a lot. Say

           lr ~ beta(lr_a, lr_b);

rather than

           for (s in 1:S)  

                lr[s] ~ beta(lr_a, lr_b);

Thanks again.

Lei

--

You received this message because you are subscribed to a topic in the Google Groups "Stan users mailing list" group.
To unsubscribe from this topic, visit https://groups.google.com/d/topic/stan-users/-RB6_X8IMT0/unsubscribe.
To unsubscribe from this group and all its topics, send an email to stan-users+...@googlegroups.com.

[Lei Zhang]

Hi Guido,

thanks for your comments and for sharing your experience. I will contact Steingrove for more details. 

> softmax temperature 20

From a not-so-optimal estimate, I do see this parameter doesn't extend that wide on the parameter space. I'll consider to compress it, and this may save some memory.

> bimodal beta

I agree with what you mentioned. Later, we would plan to employ a transdimensional modeling method to let the model categorize which sub-group a person may belong to. 

Cheers,

Lei

--
You received this message because you are subscribed to a topic in the Google Groups "Stan users mailing list" group.
To unsubscribe from this topic, visit https://groups.google.com/d/topic/stan-users/-RB6_X8IMT0/unsubscribe.
To unsubscribe from this group and all its topics, send an email to stan-users+...@googlegroups.com.

[Young Ahn]

Hi,

You may also want to check this paper by my colleagues and myself. I used the 'Matt trick', which quite dramatically improved the convergence (Rhat) and speed of MCMC sampling.

http://journal.frontiersin.org/article/10.3389/fpsyg.2014.00849/abstract

Ahn, W.-Y., Vasilev, G., Lee, S. H., Busemeyer, J. R., Kruschke, J. K., Bechara, A., & Vassileva, J. (2014). Decision-making in stimulant and opiate addicts in protracted abstinence: evidence from computational modeling with pure users.Frontiers in psychology, 5.

Best,

Young

[Bob Carpenter]

> On Jun 4, 2015, at 12:50 PM, Lei Zhang <bnuzhan...@gmail.com> wrote:
>
> Hi Bob,
>

> thanks again for your comments and suggestions. They are really helpful.
>
> I just have one more concern at this moment. For the beta(a,b) distribution, though a>1, b>1 does a good constrain to avoid bimodal shape, however is it too rigid? For instance beta(2.5, 0.7) could also yield a good estimate in some cases (and this indeed happens for my current dataset).


> That said, either a or b ought be larger than 1, but not both of them. How could I possibly constrain that? Or does it make much sense to you?

You can just replace the fmin with an fmax, then you're
guaranteed at least one is > 1. But...

Have you just tried constraining a > 0 and b > 0 and see
what the fit is? The problem is that constraints can often
lead to boundary behavior; if the fit with only the natural
constraints of a > 0 and b > 0 lead to a = 1/2 and b = 1/2, then
you may just find yourself estimating on the boundary with a constraint
of min(a,b) > 1.

But if you do what Andrew suggests and inverse-logit transform
an unconstrained parameter (essentially an intercept only logistic
regression), then you won't get bimodal behavior.

But if the modes are interesting rather than right at the boundary
of 0 and 1, then you probably want to fit something like a mixture
model.

- 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.

[Lei Zhang]

Hi Young,

thanks for sharing your helpful paper! After checking the hierarchical modeling section, I was just wondering why you used PROBIT rather than LOGIT to reparameterize your priors?

thanks and have a nice weekend.

Lei

Sent using CloudMagic

[Young Ahn]

Hi Lei,

Sorry for the delayed response. As I emailed you, I was out of town. 

I used normal priors and the probit function because it was straightforward to use the Matt trick. The specification of the model followed the model of Wetzels et al. (2010) and Steingroever et al. (2013), which are cited in our paper. 

Best,

Young

[Bob Carpenter]

> On Jun 9, 2015, at 1:14 PM, Young Ahn <wooyou...@gmail.com> wrote:
>
> Hi Lei,
>
> Sorry for the delayed response. As I emailed you, I was out of town.
>
> I used normal priors and the probit function because it was straightforward to use the Matt trick.

It shouldn't matter to the centered (standard) vs. non-centered (Matt tricked)
parameterizations whether the link is probit or logit. It's often helpful in

Gibbs sampling uses probit if you go with the Chib and Albert approach
to sampling with latent variables and want everything to be normal for
conjugacy (you can also take the Chib and Albert approach with logit by
using the logistic density in place of the normal --- it just won't be
conjugate, which really slows down Gibbs samplers like BUGS or JAGS).

- 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.