Potential bug in MRVI's get_aggregated_posterior

[rastogiruchir]

The MixtureSameFamily constructed in get_aggregated_posterior computes different log_probs from those I manually compute by iterating over the mixture components. On real datasets, I see that this manual computation shrinks the range of the log_probs calculated under a given sample. Not sure if I'm misunderstanding the intention of get_aggregated_posterior, and I haven't checked in detail why this discrepancy exists.

Here's a small notebook to reproduce: https://colab.research.google.com/drive/1tT8rbLnK5zJqNnZa_Fn5sUpcwNB8Ge90#scrollTo=Zn7z1tEB_zJz.

[ori-kron-wis]

@PierreBoyeau @justjhong

[justjhong]

Hey @rastogiruchir , thanks for raising this issue. I believe you are right, this is in fact a bug. I was able to figure out that if you don't indicate the event dimension of the Normal distribution, the mixture family is basically computing the mixture log probs over each dimension instead of over the full log probs.

In more detail:
Bugged version:
$$\log \Pi_{d}^{D} (\frac{\Sigma_{c}^{C} \mathcal{N}( \mu_{dc}, \sigma_{dc})}{C})$$

Correct version:
$$\log (\frac{\Sigma_{c}^{C} \Pi_{d}^{D} \mathcal{N}(\mu_{dc}, \sigma_{dc})}{N})$$

I was able to reproduce the fix you have with the Multivariate distribution using to_event on the original implementation:

  """Following the code in get_aggregated_posterior and differential_abundance"""
  ap = MixtureSameFamily(
    Categorical(probs=jnp.ones(qu_locs.shape[0]) / qu_locs.shape[0]),
    Normal(qu_locs, qu_scales).to_event(1),
  )
  return ap.log_prob(x).item()

Will make a PR for this. Thanks again Ruchir!

[canergen]

Thanks @rastogiruchir. To highlight Justin's suggested solution fixes it in jax: https://colab.research.google.com/drive/1tT8rbLnK5zJqNnZa_Fn5sUpcwNB8Ge90#scrollTo=Zn7z1tEB_zJz.
In addition, it highlights that MoG in torch were already correct. To note, torch.distributions.Normal requires standard deviation while MultiVariateNormal expects covariance matrices.