Formulating a hierarchical multivariate Gaussian process using the Hilbert space approximation

OK I realized that I made some crucial errors in the revised code I added above. Specifically, w should be in the transformed parameters block instead of the data block, and the indexing for alpha_t in the double loop was incorrect. The corrected code is below:

data {
  int<lower=1> M;  // number of spots
  int<lower=1> N;  // number of gene-spot pairs in long dataframe
  int<lower=1> G;  // number of genes
  int<lower=1> k;  // number of basis functions used to approximate GP
  array[N] int<lower=1, upper=M> spot_id;  // unique ID for each spot
  array[N] int<lower=1, upper=G> gene_id;  // unique ID for each gene
  matrix[M, k] phi;  // matrix of basis functions used to approximate GP
  vector[N] y;  // vector of normalized gene expression used as response variable
}

parameters {
  vector[G] beta0;  // vector of gene-specific intercepts
  matrix[k, G] alpha_t;  // transposed matrix of gene-specific coefficients for each basis function
  real<lower=0> sigma_y;  // observation noise of response variable (normalized, scaled gene expression)
  vector[G] std_log_amplitude;  // vector of logs of gene-specific amplitudes of the approximate GP
  real mu_amplitude;  // mean for the marginal SD
  real<lower=0> sigma_amplitude;  // vector of SDs for the amplitude
  real mu_beta0;  // mean for the gene-specific intercepts
  real<lower=0> sigma_beta0;  // vector of SDs for the gene-specific intercepts
  vector[k] mu_alpha;  // vector of means for the basis function coefficients
  vector<lower=0>[k] sigma_alpha;  // vector of SDs for the basis function coefficients
}

transformed parameters {
  vector[G] amplitude = exp(mu_amplitude + sigma_amplitude * std_log_amplitude);
  vector[G] amplitude_sq = square(amplitude);
  matrix[M, G] phi_alpha = phi * alpha_t;
  vector[N] w;
  for (i in 1:N) {
    w[i] = phi_alpha[spot_id[i], gene_id[i]];
  }
}

model {
  mu_beta0 ~ normal(0, 1);
  sigma_beta0 ~ normal(0, 1);
  mu_alpha ~ normal(0, 1);
  sigma_alpha ~ normal(0, 0.3);
  sigma_y ~ normal(0, 1);
  mu_amplitude ~ normal(0, 1);
  sigma_amplitude ~ normal(0, 0.3);
  beta0 ~ normal(mu_beta0, sigma_beta0);
  std_log_amplitude ~ normal(mu_amplitude, sigma_amplitude);
  for (i in 1:G) {
    for (j in 1:k) {
      alpha_t[j, i] ~ normal(mu_alpha[j], sigma_alpha[j]);
    }
  }
  y ~ normal(beta0[gene_id] + amplitude_sq[gene_id] .* w, sigma_y);
}

However, when I attempt to actually fit the model using the meanfield VI algorithm, the ELBO diverges immediately after the adaptation stage. When I try to use MCMC sampling instead, I get a string of errors saying Exception: normal_lpdf: Scale parameter is 0, but must be positive!. I believe based on this thread this means that, on the log-scale, one or more of the scale parameters named sigma_* is being sampled as a very large negative number, thus ending up very close to zero after it is exponentiated back to the correct scale. This is odd seeing as how the priors are set to not exhibit much variability.

So:

  1. Would utilizing different priors or reparamaterizing help ?
  2. I tried scaling the basis function matrix \Phi prior to passing it into Stan, but this didn’t solve anything.
  3. Did i just mess something up in the Stan code ? It compiles without error, but I’m aware that doesn’t mean there are no problems with it.
  4. Is it possible that the exponentiation and then squaring of the gene-specific amplitude parameter leading to extremely large values ?

Thanks again for all the input, it’s been a huge help !

-Jack