Speeding up a Binomial model with count likelihood

I have a binomial model with ~ 15k observations that takes a few hours to sample.

Today, I saw this post by @mike-lawrence

“Note that if you have lots of data and encounter slow sampling, there are tricks to speed up binomial outcomes by having the likelihood evaluated as counts…”

I’m wondering if anyone can point me toward some resources regarding these tricks. I’d like to see if I can speed things up.

See here (code here).

Thanks!

My data are aggregated (number of successes and number of trials) on each row. Any chance this approach will assist in that context?

If you’re already using a likelihood akin to

num_succusses ~ binomial( num_trials, theta)

Then you’re already doing the sufficient stats trick. Other avenues for speed up may be present though, but you’d have to post your model for me to discern.

Thanks!

That is my current likelihood, but I’ll get the code cleaned up a bit and post here for any input.

Hi @mike-lawrence

Below is my stan model.

These data are for a pooled analysis of data from several labs across several years, where some subjects have participated in more than one study.


functions {
 
  matrix scale_r_cor(matrix z, vector SD, matrix L) {
    
    return transpose(diag_pre_multiply(SD, L) * z);
  }
  
}

data {
  
  int<lower=1> N; // total number of observations
  array[N] int nCorr; // response variable
  array[N] int nTrials;  // number of trials
  int<lower=1> K; // number of population-level effects
  matrix[N, K] X; // population-level design matrix
  // Lab-level data
  int<lower=1> nLab; // Number of labs
  array[N] int<lower=1> Lab; // Lab indicator
  // Year-level data
  int<lower=1> nYear; // Number of years
  array[N] int<lower=1> Year; // Year indicator
  // Subject-level data
  int<lower=1> nSubj; // Number of subjects
  int<lower=1> mSubj; // Number of per-subject coefficients 
  array[N] int<lower=1> Subj; // Subject indicator
  // Subject-level predictor values
  vector[N] Z_SNR0;
  vector[N] Z_SNR2;
  vector[N] Z_SNR4;
  vector[N] Z_SNR5;
  vector[N] Z_SNR6;
  vector[N] Z_SNR8;
  vector[N] Z_SNR10;
  vector[N] Z_CepCorr;
  vector[N] Z_Fast;
  vector[N] Z_Slow;
  vector[N] Z_FlOn;
  vector[N] Z_FlLow;
  vector[N] Z_FlHigh;
  int<lower=1> nSubjCorr; // Sumber of subject-level correlations
  // Year:Subject data
  int<lower=1> nYearSubj; // Number of Year:Subject terms
  array[N] int<lower=1> YearSubj; // Year:Subj indicator
  
}

parameters {
  
  vector[K] b; // Coefficients
  vector<lower=0>[1] sd_lab; // Lab-level standard deviation
  array[1] vector[nLab] z_lab; // Lab z-scores
  vector<lower=0>[1] sd_year; // Year-level standard deviation
  array[1] vector[nYear] z_year; // Year z-scores
  vector<lower=0>[1] sd_yearsubj; // Year-level standard deviation
  array[1] vector[nYearSubj] z_yearsubj; // Year z-scores
  vector<lower=0>[mSubj] sd_subj; // Subject-level standard deviation
  matrix[mSubj, nSubj] z_subj; // Subject z-scores
  cholesky_factor_corr[mSubj] L_subj; // cholesky factor of correlation matrix
  
}

transformed parameters {
  
  vector[nLab] r_lab; // lab-level effects
  vector[nYear] r_year; // year-level effects
  matrix[nSubj, mSubj] r_subj; // subject-level effects
  vector[nSubj] r_quiet;
  vector[nSubj] r_snr0;
  vector[nSubj] r_snr2;
  vector[nSubj] r_snr4;
  vector[nSubj] r_snr5;
  vector[nSubj] r_snr6;
  vector[nSubj] r_snr8;
  vector[nSubj] r_snr10;
  vector[nSubj] r_cepcorr;
  vector[nSubj] r_slow;
  vector[nSubj] r_fast;
  vector[nSubj] r_flon;
  vector[nSubj] r_fllow;
  vector[nSubj] r_flhigh;
  vector[nYearSubj] r_yearsubj; // year:subj-level effects
  r_lab = sd_lab[1] * z_lab[1];
  r_year = sd_year[1] * z_year[1];
  r_yearsubj = sd_yearsubj[1] * z_yearsubj[1];
  
  // compute subject-level effects
  r_subj = scale_r_cor(z_subj, sd_subj, L_subj);
  r_quiet = r_subj[ : , 1];
  r_snr0 = r_subj[ : , 2];
  r_snr2 = r_subj[ : , 3];
  r_snr4 = r_subj[ : , 4];
  r_snr5 = r_subj[ : , 5];
  r_snr6 = r_subj[ : , 6];
  r_snr8 = r_subj[ : , 7];
  r_snr10 = r_subj[ : , 8];
  r_cepcorr = r_subj[ : , 9];
  r_slow = r_subj[ : , 10];
  r_fast = r_subj[ : , 11];
  r_flon = r_subj[ : , 12];
  r_fllow = r_subj[ : , 13];
  r_flhigh = r_subj[ : , 14];
  
  vector[N] mu = rep_vector(0.0, N);
  mu += X * b;
  
    for (n in 1 : N) {
      // add more terms to the linear predictor
      mu[n] += r_lab[Lab[n]] + r_year[Year[n]] + r_yearsubj[YearSubj[n]]
               + r_quiet[Subj[n]] + r_snr0[Subj[n]] * Z_SNR0[n]
               + r_snr2[Subj[n]] * Z_SNR2[n] + r_snr4[Subj[n]] * Z_SNR4[n]
               + r_snr5[Subj[n]] * Z_SNR5[n] + r_snr6[Subj[n]] * Z_SNR6[n]
               + r_snr8[Subj[n]] * Z_SNR8[n] + r_snr10[Subj[n]] * Z_SNR10[n]
               + r_cepcorr[Subj[n]] * Z_CepCorr[n] + r_slow[Subj[n]] * Z_Slow[n]
               + r_fast[Subj[n]] * Z_Fast[n] + r_flon[Subj[n]] * Z_FlOn[n] + r_fllow[Subj[n]] * Z_FlLow[n]
               + r_flhigh[Subj[n]] * Z_FlHigh[n];
    }
  
}

model {
  
  // priors n
  real lprior = 0; // prior contributions to the log posterior
  
  lprior += normal_lupdf(b[1:2] | 0, 1);
  lprior += normal_lupdf(b[3:] | 0, 0.5);
  lprior += normal_lupdf(sd_lab | 0, 1);
  lprior += normal_lupdf(sd_year | 0, 1);
  lprior += normal_lupdf(sd_yearsubj | 0, 1);
  lprior += normal_lupdf(sd_subj[1:3] | 1, 1);
  lprior += normal_lupdf(sd_subj[4:] | 0, 1);
  lprior += lkj_corr_cholesky_lupdf(L_subj | 1);
  target += lprior;
  target += std_normal_lupdf(z_lab[1]);
  target += std_normal_lupdf(z_year[1]);
  target += std_normal_lupdf(z_yearsubj[1]);
  target += std_normal_lupdf(to_vector(z_subj));
  
  // likelihood 
    
    target += binomial_logit_lupmf(nCorr | nTrials, mu);
  
}

generated quantities {
  
  vector[N] log_lik;
  
  // Subject-level correlations
  corr_matrix[mSubj] Subj_Cor = multiply_lower_tri_self_transpose(L_subj);
  vector<lower=-1, upper=1>[nSubjCorr] subj_cor;
  
  // Extract upper diagonal of correlation matrix
  for (k in 1 : mSubj) {
    
    for (j in 1 : (k - 1)) {
      subj_cor[choose(k - 1, 2) + j] = Subj_Cor[j, k];
    }
    
  }
  
  // Log-likelihood for LOO
  for (n in 1 : N) {
    
    log_lik[n] = binomial_logit_lpmf(nCorr[n] | nTrials[n], mu[n]);
    
  }
  
}