Vectorization Problems

Modeling
,
1

Hey folks :)

I know there are plenty of threads on vectorization, but I couldn’t figure out the solution for my issue.
I want to vectorize this model:

data {
  int <lower=0> N;  // number of subjects
  int <lower=0> K;  // categories 
  int R[K];         // number of responses per category
  int count[N,K];   // observed data
  real scale_b;     // set scaling for background noise
}

parameters {
  // subject parameters
  real <lower=0> a[N];
  real <lower=0> c[N];
  real log_f[N];



  // Mu & Sigma for hyper distributions
  real <lower=0> mu_a;
  real <lower=0> sig_a;
  real <lower=0> mu_c;
  real <lower=0> sig_c;
  real  logMu_f;
  real <lower=0> logSig_f;
 
}

transformed parameters{
  // Transform f Parameter
  
  real f[N] = inv_logit(log_f);
  real mu_f = inv_logit(logMu_f);
  real sig_f = sd(f);
  
  // activations
  vector[K] acts[N];
  real SummedActs[N];
  
  // probabilities
  
  vector[K] probs[N];
  
  // loop over subjects to compute activations and probabilites
  
  for (i in 1:N){
    acts[i,1] = scale_b + a[i] + c[i]; // Item in Position
    acts[i,2] = scale_b + a[i];        // Item in Other Position
    acts[i,3] = scale_b + f[i]*(a[i]+c[i]);// Distractor in Position
    acts[i,4] = scale_b + f[i] *a[i]; // Distractor in other Position
    acts[i,5] = scale_b; // non presented Lure
    
    SummedActs[i] = R[1] * acts[i,1] + R[2] * acts[i,2] + R[3] * acts[i,3]+ R[4] * acts[i,4]+ R[5] * acts[i,5];
    
    probs[i,1] = (R[1] * acts[i,1]) ./ (SummedActs[i]);  
    probs[i,2] = (R[2] * acts[i,2]) ./ (SummedActs[i]);
    probs[i,3] = (R[3] * acts[i,3]) ./ (SummedActs[i]);
    probs[i,4] = (R[4] * acts[i,4]) ./ (SummedActs[i]);
    probs[i,5] = (R[5] * acts[i,5]) ./ (SummedActs[i]);
  
  }
}

model{
  
  // priors for hyper parameters
  mu_c ~ normal(10,10);
  sig_c ~ gamma(1,0.01);
  
  mu_a ~ normal(10,10);
  sig_a ~ gamma(1,0.01);
  
  logMu_f ~ normal(0,1);
  logSig_f ~ gamma(1,0.01);
  
  // Loop over subjects
  for(i in 1:N){
    // Draw subject parameters from truncated normal
    
    c[i] ~ normal(mu_c, sig_c);
    a[i] ~ normal(mu_a, sig_a);
    log_f[i] ~ normal(logMu_f, logSig_f);


    // draw data from probabilities determined by MMM parms
    count[i,]  ~ multinomial(probs[i,]);
  }
}

As a fiirst try, I changed just the sampling statement for the subject parameters, and pulled it out of the loop:


    c ~ normal(mu_c, sig_c);
    a  ~ normal(mu_a, sig_a);
    log_f ~ normal(logMu_f, logSig_f);

// Loop over subjects
  for(i in 1:N){
    // draw data from probabilities determined by MMM parms
    count[i,]  ~ multinomial(probs[i,]);
  }
}

This worked, but there was not really a gain in speed.

However, I didn’t figure out how to vectorize the count statement, as well as the transformed parameters block. I tried to vectorize the transformation of the acts in a matrix way, but this leads to a wrong calculation of the probailities (“not a valid simplex”). Here my try:

 real f[N] = inv_logit(log_f);
  real mu_f = inv_logit(logMu_f);
  real sig_f = sd(f);
  
  // activations
  vector[K] acts[N];
  real SummedActs[N];
  
  // probabilities
  
  vector[K] probs[N];
  
    acts[,1] = scale_b + a+ c; // Item in Position
    acts[,2] = scale_b + a;        // Item in Other Position
    acts[,3] = scale_b + f*(a+c);// Distractor in Position
    acts[,4] = scale_b + f *a; // Distractor in other Position
    acts[,5] = scale_b; // non presented Lure
    
    SummedActs = R[1] * acts[,1] + R[2] * acts[,2] + R[3] * acts[,3]+ R[4] * acts[,4]+ R[5] * acts[,5];
    
    probs[,1] = (R[1] * acts[,1]) ./ (SummedActs);  
    probs[,2] = (R[2] * acts[,2]) ./ (SummedActs);
    probs[,3] = (R[3] * acts[,3]) ./ (SummedActs);
    probs[,4] = (R[4] * acts[,4]) ./ (SummedActs);
    probs[,5] = (R[5] * acts[,5]) ./ (SummedActs);
   

  }

I tried to define the acts and probabilities as matrix [N,K], but that also didn’t work out. I think it’s either a problem with the acts calculation or with the probs. Putting it step by step out of the loop produced the same error (no valid simplex) for the probabilites, even if I just unlooped the acts block. I’m really frustrated, maybe somebody could give an advice for vector / matrix operations, as the reference manual was not that helpful to me.

Thanks and stay healthy !

*Jan

2

Unfortunately, multinomial cannot currently be vectorized in this way.

That’s weird, because probs should contain valid simplices regardless of the values of acts (as long as those are positivie and finite).

The only problem I see immediately is that scale_b + f*(a+c) should probably be scale_b + f .* (a+c) to make sure you are doing component-wise multiplication and not a dot product (I am actually surprised this was not a compile-time error). Similarly for the f * a part.

Note also that you can use print statements in your code to help you debug.

Other than that your approach to vectorization seems good.

Unfortunately, that can easily happen. Vectorization gives huge speedups only if there are computations to be shared and the loops are large, e.g.:

int N = 1000;
real sigma = ...;
vector[N] mu = ... ;
vector[N] y = ... ;
for(i in 1:N) {
   y[i] ~ normal(mu[i], sigma); // 1/sigma and 1/sigma^2 evaluated for each iteration to compute the density
}

y ~ normal(mu, sigma); // 1/sigma and 1/sigma^2 evaluated just once and shared

the benefits of vectorizing arithmetic operations are in my experience only marginal.

Best of luck with your model!

3

Hey Martin, thanks for your review ! As there are no further speed gains, I left it the way it worked. Hopefully, multinomial vectorization will be supported soon!