Modeling Hierarchical Regression with LKJ Prior

Modeling
1

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
2

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.

3

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
4

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

5

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

6

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.