Hiearchical reinforcement-learning model (for IGT) with mixing - check & optimization

[Martin Šmíra]

Hi Stan,

I implemented a hierarchical reinforcement-learning (PVL-delta for Iowa Gambling Task) model in Stan. Now I'm interested in using the model to compare two groups. In particular, I’m interested in whether the individual parameters of both group come from the same group-level distribution, or from two different group-level distributions. I would like to implement this comparison by including a binary mixture component mix that controls which model (either the model with identical group-level distributions or the model with different group-level distributions) has generated the data of the two groups.  

In the original model file, I specified one group-level distribution for each parameter of interest and there are 4 such parameters. Implementing mixing here will introduce a second group-level distribution for each of the four parameters. I assume that the individual-level parameters of the subjects of group 1 are always drawn from the first group-level distribution, but the individual-level parameters of the subjects of group 2 can either be drawn from the first or the second group-level distribution. 

Below is are the crucial parts of my solution; the complete model file can be downloaded here: https://github.com/neekro/Stan_translate/blob/master/pvl_d_a_michael_mix_all_optimized.stan (in the repository is also an R file, data and graphical model of PVL-Delta without my mixture extension). 

The loop (1:4) goes through our 4 parameters of interest. The variable “mix” is used as mixing parameter (i.e., a binary model index); “ind_pr” is a vector of individual parameters. Here, I used the function tail to extract the individual parameters of group 2; “mu_pr” and “std” are parameters of the group-level distributions. 

First, in the model block, there is a normal sampling statement for first group’s individual parameters. Sampling for the second group is done in transformed parameters block, so I can later use lp_parts variable to compute mixing proportion probability and use it for sampling binary variable z[i], which is the main interest of the model. It indicates whether group 2 is having same group distributions for each parameter of interest as group 1. 

transformed parameters {

  for (i in 1:4) {

    lp_parts[i,1] <- log1m(mix[i]) + normal_log(tail(ind_pr[i], n_s_2),  mu_pr[i,1], std[i,1]);

    lp_parts[i,2] <- log(mix[i]) + normal_log(tail(ind_pr[i], n_s_2), mu_pr[i,2], std[i,2]);

}

model {

  for (i in 1:4) {  

    for (s in 1:n_s_1) {  

      # Group 1

      ind_pr[i,s] ~ normal(mu_pr[i,1], std[i,1]);     

    }

    # Group 2 

    increment_log_prob(log_sum_exp(lp_parts[i]));

  }

}

generated quantities {

  int<lower=0,upper=1> z[4];  // high values of z indicate group distribution to be different

  for (i in 1:4) {

    real prob;

    prob <- exp(lp_parts[i,1]) / (exp(lp_parts[i,1]) + exp(lp_parts[i,2]));

    z[i] <- bernoulli_rng(prob); 

  }

}

My questions are:

1)  Is the mixing part implemented correctly?  - The obtained posterior distributions seem appropriate so far, but I want to be sure about it. 

2)  Any suggestions for optimizing the model file?  - With the current version of the model file, there are a few parameters with n_eff = 20 even after an hour of sampling.

3)  If I want to assign prior information to the mixing proportion, is it valid to just exchange the “mix” variable with a given numerical probability – e.g.,  0.9 for one part and 0.1 for the second? And if I put there 0.5 and 0.5, will it lead to same behavior as using a parameter mix and (1 - mix) while mix ~ uniform(0, 1)? What is preferred here?

4)  Can I used (exp(lp_parts[i,1]) + exp(lp_parts[i,2])) interchangeable with sum(exp(lp_parts[i]))?  -  I tried to use the second version in a few models and in some of those it resulted in "Error in function stan::prob::bernoulli_rng(d): Probability parameter is 1, but must be between (0, 1)". 

Thank you for any suggestions.

Martin

[Bob Carpenter]

These models are getting complicated!

Comment 1: Don't use tabs; you can see the kind of trouble they
cause by viewing lines 51 and 53 of your code on GitHub. If you
use emacs, you can set it up to get rid of tabs automatically (untabify
or you can use the elisp config we give from our developer tools wiki
page on stan-dev/stan). Emacs has a Stan mode (downloadable from stan-dev)!

Comment 2: Probit is suboptimal compared to logit computationally in
Stan because Phi() is so much harder to compute than inv_logit.
The reverse was true in BUGS/JAGS, I think.

Comment 3: This is suboptimal because the max and min will cut the relation
to the gradients:

dummy[d] <- exp(fmax(fmin(Ev[d] * theta, 450), -450));

Instead, you should use the built-in softmax function, or even better,
just use the softmax-parameterized categorical distribution, which cuts

choice[s, t + 1] ~ categorical_softmax(Ev);

Comment 4: I wouldn't put the lp_parts in transformed parameters unless
you need to save them for some reason. There's a two-argument log_sum_exp
that you can use directly. What we really need are some mixture functions
to make this all easier.

Comment 5: I still think it's easier to just work in expectation
than sample bernoulli data from a known probability only to estimate
that probabilty. The Rao-Blackwell theorem tells you this in theory, but
you can also see it in practice because there's sampling error in estimating
by sampling, whereas there's none in just using the expectation. But if you
really need the samples for some reason, then what you're doing is the right
way to get them.


Comment 6: Use vector-based distributions wherever possible, even at the cost
of some data manipulation. So this (mangled by tabs!):

for (i in 1:4) {

# Prior on the group-level mean parameters
# probit scale [-Inf, Inf]
mu_pr[i] ~ normal(0, 1);

# Individual-level paramters

for (s in 1:n_s_1) {
# Group 1
ind_pr[i,s] ~ normal(mu_pr[i,1], std[i,1]);
}
# Group 2
increment_log_prob(log_sum_exp(lp_parts[i]));
}

should be

for (i in 1:4)

mu_pr[i] ~ normal(0,1);

for (i in 1:4)

ind_pr[i] ~ normal(mu_pr[i,1], std[i,1]);

for (i in 1:4)

increment_log_prob(log_sum_exp(lp_parts[i]));

You can put them all in the same loop, but it's usually
clearer with a loop per variable (and just as fast, if
not faster in the case of long arrays).

To answer your questions.

1) Yes, this looks right for the mixture.

2) See above --- the move to categorical_softmax should be a big boost.

3) The usual approach in a mixture model is to estimate the mixing proportion
lambda constrained to [0,1]. You can also set it if you know it ahead of time,
or you can take the in-between route and give it an informative prior.

4) Yes, they're interchangeable, but see above on exp(). Any time you find yourself
using exp() on scalars or inverse() on matrices, you want to step back and
reformulate wherever possible for better arithmetic stability. But the + version
may be faster because the vector-based sum() isn't saving you anything with
only two operands.

- 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.
> For more options, visit https://groups.google.com/d/optout.

[Martin Šmíra]

Hi Bob,

thank you for the comments and answers. This model is actually not from the book.

 

Comment 3:  This is suboptimal because the max and min will cut the relation
to the gradients:

        dummy[d] <- exp(fmax(fmin(Ev[d] * theta, 450), -450));      

Instead, you should use the built-in softmax function, or even better,
just use the softmax-parameterized categorical distribution, which cuts

   choice[s, t + 1] ~ categorical_softmax(Ev);

1) Is there the end of the sentence missing? 

2) I didn't find categorical_softmax function. Did you mean categorical_logit(Ev * theta) ? 

Comment 4:  I wouldn't put the lp_parts in transformed parameters unless
you need to save them for some reason.  There's a two-argument log_sum_exp
that you can use directly.  What we really need are some mixture functions
to make this all easier.  

I need it for inferring mixing proportion. So I think I need lp_parts in transformed parameters so I can access it in generated quantities, correct?

 

Comment 6:  Use vector-based distributions wherever possible, even at the cost
of some data manipulation.  So this (mangled by tabs!):

 for (i in 1:4) {  
        # Prior on the group-level mean parameters
    # probit scale [-Inf, Inf]
        mu_pr[i] ~ normal(0, 1);
        
    # Individual-level paramters
        for (s in 1:n_s_1) {  
          # Group 1
      ind_pr[i,s] ~ normal(mu_pr[i,1], std[i,1]);    
    }
        # Group 2
        increment_log_prob(log_sum_exp(lp_parts[i]));
  }


should be

  for (i in 1:4)
    mu_pr[i] ~ normal(0,1);

  for (i in 1:4)
    ind_pr[i] ~ normal(mu_pr[i,1], std[i,1]);

  for (i in 1:4)
    increment_log_prob(log_sum_exp(lp_parts[i]));

You can put them all in the same loop, but it's usually
clearer with a loop per variable (and just as fast, if
not faster in the case of long arrays).

Here is a catch, the loop is actually going from 1:n_s_1, which is just a first half of the vector. I was thinking of using head() function instead of the loop, would it be faster to compute? 

 

To answer your questions.

4) Yes, they're interchangeable, but see above on exp().  Any time you find yourself
   using exp() on scalars or inverse() on matrices, you want to step back and
   reformulate wherever possible for better arithmetic stability.  But the + version
   may be faster because the vector-based sum() isn't saving you anything with
   only two operands.

But still, if it's interchangeable, why do I get an error if I use the sum there? ... as I mentioned before.

Thank you,

Martin

[Martin Šmíra]

Hi Bob (and Stan),

I've tried several ways to improve this model, but I wasn't successful. 

I implemented your suggestions - softmax and inv_logit, but it didn't really help the performance.  Additionally using inv_logit yield quite different parameter estimations. 

I also tried to implement Matt trick here, but I wasn't able to do it - here is a thread about it https://groups.google.com/forum/?hl=cs#!topic/stan-users/TqEdlLPbDoE .

Recently I run the model for 35 hours - 50 000 iterations. And yet there is a parameter with n_eff as small as 79. From traceplot (link below) you can also see some weird behavior - especially with parameter mu_c[2] and its deviation std[4,2].  

I'm asking for any suggestion for optimization of the model. 

Here is model code https://github.com/neekro/Stan_translate/tree/master/long%20run%20optim_2_phi   and also some traceplots. 

Inference for Stan model: pvl_d_mix_all_optimized_2_phi.

5 chains, each with iter=12000; warmup=2000; thin=1; 

post-warmup draws per chain=10000, total post-warmup draws=50000.

              mean se_mean     sd      2.5%       25%       50%       75%     97.5% n_eff  Rhat

z[1]         0.696   0.016  0.460     0.000     0.000     1.000     1.000     1.000   794 1.005

z[2]         0.048   0.004  0.215     0.000     0.000     0.000     0.000     1.000  3026 1.003

z[3]         0.490   0.028  0.500     0.000     0.000     0.000     1.000     1.000   311 1.016

z[4]         0.690   0.031  0.462     0.000     0.000     1.000     1.000     1.000   227 1.016

mu_A[1]      0.660   0.006  0.168     0.304     0.540     0.680     0.794     0.926   789 1.008

mu_A[2]      0.767   0.006  0.246     0.089     0.729     0.864     0.929     0.992  1940 1.002

mu_w[1]      0.138   0.002  0.104     0.024     0.067     0.113     0.179     0.398  1938 1.003

mu_w[2]      2.357   0.046  1.493     0.065     1.037     2.381     3.618     4.864  1041 1.005

mu_a[1]      0.152   0.010  0.088     0.031     0.085     0.136     0.205     0.362    79 1.052

mu_a[2]      0.432   0.003  0.235     0.050     0.257     0.397     0.573     0.948  5099 1.002

mu_c[1]      0.841   0.015  0.189     0.498     0.722     0.820     0.946     1.273   151 1.030

mu_c[2]      1.259   0.061  1.155     0.438     0.661     0.718     1.014     4.565   359 1.009

std[1,1]     1.171   0.008  0.236     0.639     1.014     1.215     1.363     1.487   919 1.005

std[1,2]     0.725   0.016  0.400     0.059     0.399     0.708     1.045     1.446   642 1.014

std[2,1]     1.317   0.004  0.147     0.963     1.236     1.351     1.435     1.494  1686 1.004

std[2,2]     0.784   0.019  0.453     0.015     0.389     0.811     1.195     1.467   554 1.008

std[3,1]     1.160   0.015  0.227     0.657     1.007     1.197     1.344     1.482   228 1.021

std[3,2]     1.002   0.010  0.412     0.075     0.750     1.134     1.333     1.484  1751 1.004

std[4,1]     0.415   0.006  0.153     0.214     0.307     0.379     0.489     0.809   615 1.003

std[4,2]     0.299   0.022  0.388     0.023     0.061     0.112     0.326     1.377   313 1.011

lp__     -5941.816   1.526 25.740 -5988.523 -5960.085 -5942.995 -5924.574 -5890.403   285 1.034

Cheers,

Martin