Slow running hierachical model

Modeling
, ,
1

Hi everyone

I have a hierarchical model specified in the following way. It does okay in fitting on low amounts of data, but when I scale it it takes a very long time to fit. I need to be able to fit 300,000 observations. I have looked into vectorisation and reparameterisation, but I am kind of stuck. Can any of you help me? Vectorisation is probably the lowest hanging fruit, but I have a hard time figuring out how to do it, when I have to select the right coefficients for each observation depending on the language, metric, or task. Would constructing a designmatrix instead alleviate my problems?

Any help is appreciated, thanks!

data {
  int<lower=0> N;                       // Number of observations
  int<lower=0> M;                       // Number of groups/models
  int<lower=0> n_language;              // Number of languages
  int<lower=0> n_metric;                // Number of metrics
  int<lower=0> n_task;                  // Number of tasks

  array[N] int<lower=1,upper=M> group;  // Group indicator for each observation
  array[N] int<lower=1,upper=n_language> language; // Language indicator for each observation
  array[N] int<lower=1,upper=n_metric> metric; // Metric indicator for each observation
  array[N] int<lower=1,upper=n_task> task; // Task indicator for each observation

  vector<lower=0,upper=1>[N] y; // Output value for each observation
}

parameters {
  // Hyperpriors
  real mu_alpha;
  real<lower=0> sigma_alpha;

  sum_to_zero_vector[n_language] mu_beta_language;
  real<lower=0> sigma_beta_language;

  sum_to_zero_vector[n_metric] beta_metric;

  sum_to_zero_vector[n_task] mu_beta_task;
  real<lower=0> sigma_beta_task;

  // Group-level parameters
  vector[M] alpha;
  matrix[M, n_language] beta_language;
  matrix[M, n_task] beta_task;

  // Phi parameters
  vector[n_metric] beta_metric_phi;
  vector[n_task] beta_task_phi;
}

transformed parameters {
  vector[N] eta_mu;
  for (i in 1:N) {
    eta_mu[i] = alpha[group[i]]
                + beta_language[group[i], language[i]]
                + beta_metric[metric[i]]
                + beta_task[group[i], task[i]];
  }
  vector[N] mu = inv_logit(eta_mu);

  vector[N] eta_phi;
  for (i in 1:N) {
    eta_phi[i] = beta_metric_phi[metric[i]]
                + beta_task_phi[task[i]];
  }
  vector[N] phi = exp(eta_phi);
}

model {
  beta_metric_phi ~ normal(0, 2);
  beta_task_phi ~ normal(0, 2);

  // Hyperpriors
  mu_alpha ~ normal(0, 2);
  sigma_alpha ~ exponential(1);

  mu_beta_language ~ normal(0, 2);
  sigma_beta_language ~ exponential(1);

  beta_metric ~ normal(0, 2);

  mu_beta_task ~ normal(0, 2);
  sigma_beta_task ~ exponential(1);

  // Group-level priors
  alpha ~ normal(mu_alpha, sigma_alpha);

  for (m in 1:M) {
    beta_language[m] ~ normal(mu_beta_language, sigma_beta_language);
    beta_task[m] ~ normal(mu_beta_task, sigma_beta_task);
  }

  // Likelihood
  vector[N] a = mu .* phi + 1e-9;
  vector[N] b = (1 - mu) .* phi + 1e-9;
  y ~ beta(a, b);
}