Modeling Hierarchical Regression with LKJ Prior

I’m modeling a hierarchical regression model over multiple time series demonstrations. Everything seems to go well until the end where I get warnings regarding the LKJ prior and a recommendation to reparameterize. Any suggestion on what kind of parameterization is needed, would be much appreciated.

data {
  int<lower=1> N; // num demos
  int<lower=1> T; // traj length
  int<lower=1> K; // num basis
  row_vector[T] y[N]; // traj data
  matrix[K, T] psi; // basis feat
  real eps;
}
parameters {
  row_vector[K] mu;
  cholesky_factor_corr[K] Lcorr;
  row_vector<lower=0>[K] sigma;
  row_vector[K] weights[N];
}
model {
  mu ~ normal(0, 10);
  sigma ~ cauchy(0, 10);
  Lcorr ~ lkj_corr_cholesky(1);
  for (n in 1:N) {
      weights[n, :] ~ multi_normal_cholesky(mu, diag_pre_multiply(sigma, Lcorr));
      for (t in 1:T) {
          y[n, t] ~ normal(weights[n, :] * psi[:, t], sqrt(eps));
      }
  }
}
generated quantities {
  matrix[K, K] omega;
  matrix[K, K] cov;
  omega = multiply_lower_tri_self_transpose(Lcorr);
  cov = quad_form_diag(omega, sigma); 
}
WARNING:pystan:Rhat for parameter Lcorr[1,1] is nan!
WARNING:pystan:Rhat for parameter Lcorr[1,2] is nan!
WARNING:pystan:Rhat for parameter Lcorr[1,3] is nan!
WARNING:pystan:Rhat for parameter Lcorr[2,3] is nan!
WARNING:pystan:Rhat for parameter Lcorr[1,4] is nan!
WARNING:pystan:Rhat for parameter Lcorr[2,4] is nan!
WARNING:pystan:Rhat for parameter Lcorr[3,4] is nan!
WARNING:pystan:Rhat for parameter Lcorr[1,5] is nan!
WARNING:pystan:Rhat for parameter Lcorr[2,5] is nan!
WARNING:pystan:Rhat for parameter Lcorr[3,5] is nan!
WARNING:pystan:Rhat for parameter Lcorr[4,5] is nan!
WARNING:pystan:Rhat for parameter Lcorr[1,6] is nan!
WARNING:pystan:Rhat for parameter Lcorr[2,6] is nan!
WARNING:pystan:Rhat for parameter Lcorr[3,6] is nan!
WARNING:pystan:Rhat for parameter Lcorr[4,6] is nan!
WARNING:pystan:Rhat for parameter Lcorr[5,6] is nan!
WARNING:pystan:Rhat for parameter omega[1,1] is nan!
WARNING:pystan:Rhat above 1.1 or below 0.9 indicates that the chains very likely have not mixed
WARNING:pystan:1252 of 10000 iterations ended with a divergence (12.52%).
WARNING:pystan:Try running with adapt_delta larger than 0.8 to remove the divergences.
WARNING:pystan:Chain 3: E-BFMI = 0.19194318452632705
WARNING:pystan:E-BFMI below 0.2 indicates you may need to reparameterize your model

Those warnings do not indicate that there is anything wrong with the LKJ prior, merely that the upper triangle of a Cholesky factor is fixed to zero.

However, there must be something else that is destroying NUTS’ ability to sample from this posterior distribution.

This is weird because some of the entries of Lcorr seem to be sampled “well” while others are suffering.

Lcorr[1,1]     1.0     nan    0.0    1.0    1.0    1.0    1.0    1.0    nan    nan
Lcorr[2,1]    0.22  3.3e-3   0.32  -0.43 3.2e-3   0.24   0.45   0.77   9072    1.0
Lcorr[3,1]    0.19  2.9e-3   0.32  -0.47  -0.03   0.21   0.43   0.75  12487    1.0
Lcorr[1,2]     0.0     nan    0.0    0.0    0.0    0.0    0.0    0.0    nan    nan
Lcorr[2,2]    0.92  9.6e-4    0.1   0.64   0.88   0.96   0.99    1.0  10676    1.0
Lcorr[3,2]    0.39  2.3e-3   0.28  -0.21   0.21   0.42    0.6   0.84  14563    1.0
Lcorr[1,3]     0.0     nan    0.0    0.0    0.0    0.0    0.0    0.0    nan    nan
Lcorr[2,3]     0.0     nan    0.0    0.0    0.0    0.0    0.0    0.0    nan    nan

That isn’t weird. Some of them are random variables and some of them are fixed constants.

Oh thanks! I just noticed that they correspond to the upper triangle matrix values. I hadn’t noticed that before.

This is the relevant warning here. NUTS isn’t able to follow the Hamiltonian accurately. You can try increasing adapt_delta but that usually only works if the divergence rate is lower. This probably requires a non-centered reparameterization.