Stein mixture #1601
Stein mixture #1601
Conversation
| from numpyro.infer.util import _predictive | ||
|
|
||
|
|
||
| class MixtureGuidePredictive: |
There was a problem hiding this comment.
@fehiepsi I believe we can remove the nbatch_dims=1 from Predictive. Let me know what you think.
There was a problem hiding this comment.
It is useful in some cases, e.g. to do predictive with (num_chains, num_samples) or for autoguide posterior samples with multiple sample shapes.
There was a problem hiding this comment.
I feel that this is a bit complicated. I think the Predictive class is pretty flexible. Complicated logic to handle mixture stuff should belong to the guide, not the Predictive (correct me if I'm wrong).
There was a problem hiding this comment.
Hmm, the core is I need the mixture assignments to sample from the model. I specialized the ELBO to SteinLoss account for this.
I suppose the guide is the shared element for both cases (loss and predictive sampling) so hiding the complications there seems more modular. I'll try to lift the guide to a mixture guide (and have it as a SteinVI attribute)--then the lifted guide would work with current Predictive and ELBO. Could also simplify some logic in SteinVI. I'll sketch it and get back.
|
I will update the docs in a separate PR. This PR is already huge (sorry about that c:). |
…es are returned when sites are in guide and model.
|
LGTM, thanks Ola! |
In merging the original Stein mixture paper with ours, we found an inconsistency in the original article. We correct that here and can (probably) connect the current form with optimizing the ELBO (hence a lower bound on the marginal likelihood). Feel free to PM me for details if you are interested.
It was related to #1590. Rebase failed.
TODO
SteinLossMixtureGuidePredictive