Discrepancy between TFP and Stan's CholeskyLKJ log prob

[adamhaber]

I'm trying to understand a discrepancy in the log_prob computation between Stan and TFP. I don't know if this is bug, and if it is, where; I'm reporting this here in case anyone thinks this may be of interest.

I generated Cholesky factors of LKJ matrices using the following Stan code:

code <- "
data {
  int K;
}
parameters {
  cholesky_factor_corr[K] L;
}
model {
   target += lkj_corr_cholesky_lpdf(L|2.0);
}

I extracted samples and log_prob:

library(rstan)
data = list(K=5)
fit <- stan(model_code = code, data=data)
e <- as.data.frame(extract(fit))

Finally, I computed the log prob of these matrices using TFP:

m = tfd.CholeskyLKJ(5,2)
b = tfb.CorrelationCholesky()

for datum in data[:k,:-1]:
    L = datum.reshape(5,5).T
    tfp_lps.append(m.log_prob(L).numpy())

When I compare tfp_lps to Stan's lp__, I get a high correlation, but they're not the same. For comparison, running this check with something like a normal distribution passes np.allclose. I've also suspected that this might be related to different handling of log_det_jacobian, but subtracting tfb.CorrelationCholesky().inverse_log_det_jacobian(L,2) didn't help.

Is there something wrong in the way I'm running the comparison, or in the expectation that they would be the same in the first place?

Thanks in advance!

[brianwa84]

Anudhyan/Srinivas have a candidate for the discrepancy.

[adamhaber]

Thanks for addressing this @brianwa84 !

Just wanted to report that I'm still getting the discrepancy. For matrices of dimension 2 I'm getting a correlation of 1 between Stan and TFP's log_probs, but not the same values (difference in mean and standard deviation); For matrices of dimensions 3,4,5 I get very high correlation (>0.98) but not the same values. Perhaps this might be related to differences in the implementation of the bijector?

[brianwa84]

You are testing against TFP nightly? Can you double check tfp.__version__?

[adamhaber]

AFAIK, yes - version is '0.9.0-dev20200110'.

Also, regarding my previous comment - for matrices of size 2, Stan and TFP actually fully agree (after subtracting CorrelationCholesky inv_log_det from TFP's prediction). For dim>2 I'm still getting the discrepancy.

[bloops]

You shouldn't need to subtract CorrelationCholesky inv_log_det. It should match exactly, does it not?

[junpenglao]

If I remembered correctly, Stan might drop some normalizing constant depending on the way you model it, i.e., you might get different result with

target += lkj_corr_cholesky_lpdf(L|2.0);

compare to

L ~ lkj_corr_cholesky(2.0);

[junpenglao]

Yep, see https://mc-stan.org/docs/2_21/functions-reference/cholesky-lkj-correlation-distribution.html

Increment target log probability density with lkj_corr_cholesky_lpdf( L | eta) dropping constant additive terms.

[adamhaber]

I'm pretty sure the tilde notation drops constant additive terms, while the target += notation does not; see here. Unless TFP also drops these terms, it seemed to me that it would be better to use the target += notation for comparison. Would be happy to hear your thoughts on this.

For other distributions we've tried (normal, poisson, exponential, etc) we were able to reproduce Stan's lp__ exactly (np.allclose-exact). Regarding the bijectors - AFAIK, Stan's lp__ contains jacobian corrections, while simply calling CholeskyLKJ.log_prob does not... perhaps the bijector itself is different?

[junpenglao]

I'm pretty sure the tilde notation drops constant additive terms, while the target += notation does not;

Hmmm but that would contradict the Stan doc, which I think it is a better source of truth https://mc-stan.org/docs/2_21/functions-reference/cholesky-lkj-correlation-distribution.html

[junpenglao]

Also - the discrepancy might came from dtype - try rerunning it with dtype=tf.float64

[adamhaber]

I'm not sure I follow - the Stan doc says that if you use L ~ lkj_corr_cholesky(eta), you increment lp__ without adding constant additive terms. My current understanding is that the difference between L ~ lkj_corr_cholesky(eta) and target += lkj_corr_cholesky_lpdf(L|eta) is exactly those additive terms, which is way I chose the latter. What am I missing? :-)