Add Variational Inference Interface

[Sayam753]

Added Mean Field ADVI along with basic interface. I observe two different ways of adding VI approximations. One by tfd.JointDistributionSequential and other by tfd.MultivariateNormalDiag.

This interface incorporates the Sequential method. It may be possible to use tfd.MultivariateNormalDiag after having a flattened view of parameters. Any feedback at this implementation is welcomed.

• Tests about convergence checks
• Add convergence criteria example in quick start notebook - Will address this in another PR.
• Include samples for untransformed variables as well
• Handle shape issues with ArviZ while sampling
• More user friendly optimizers
• Configure atol for tests
• Configure autobatching - Will address this in another PR.
• Integrate Deterministics callback - Will address this in another PR.

[review-notebook-app]

Check out this pull request on 

Review Jupyter notebook visual diffs & provide feedback on notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

Merging #280 into master will decrease coverage by 0.64%.
The diff coverage is 79.31%.

@@            Coverage Diff             @@
##           master     #280      +/-   ##
==========================================
- Coverage   90.62%   89.97%   -0.65%     
==========================================
  Files          29       32       +3     
  Lines        1919     2035     +116     
==========================================
+ Hits         1739     1831      +92     
- Misses        180      204      +24     

Impacted Files	Coverage Δ
pymc4/distributions/multivariate.py	100.00% <ø> (ø)
pymc4/inference/sampling.py	93.40% <ø> (ø)
pymc4/variational/updates.py	63.63% <63.63%> (ø)
pymc4/variational/approximations.py	80.58% <80.58%> (ø)
pymc4/variational/__init__.py	100.00% <100.00%> (ø)

[Sayam753]

Requesting @ferrine, @twiecki for the review.

[junpenglao]

You should subclass tf.Module in the metaclass, as it capture free parameters (anything that is a tf.Variable or tfp.util.TransformedVariable) automatically.

[Sayam753]

Thanks @junpenglao . I will check how to use tf.Module and integrate here.

[ferrine]

Looks good, I have mostly API comments and suggestions

[twiecki]

@Sayam753 I would plot the elbo on log scale.

[Sayam753]

Yes @twiecki . I am adding the changes. 😃

[Sayam753]

Here are some comparisons drawn for Mean Field ADVI in PyMC4, TFP and PyMC3 (gist). If this is the expected behaviour and for the next steps, shall I go about writing tests for the interface implemented so far? I wish to add Full Rank ADVI in another PR to avoid making this PR large enough to review.

[twiecki]

Pretty weird that the results are so different, might be worth to dig a bit deeper.

[Sayam753]

Yes @twiecki . Indeed the plots were weird. @ferrine suggested to use common random seed and same optimizer. Along with this, doing some minor fixes, I see good results. (gist). I have updated the PR now. It's ready to review.

[twiecki]

Plots look much better now!

…

On Sat, Jun 20, 2020 at 7:57 PM Sayam Kumar ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In pymc4/variational/approximations.py <#280 (comment)>: > + return tf.vectorized_map(lambda samples: logpfn(*samples), q_samples) + + return vectorized_logpfn + + return vectorize_logp_function(logpfn) + + def _build_posterior(self): + raise NotImplementedError + + def flatten_view(self): + """Flattened view of the variational parameters.""" + pass + + def sample(self, n): + """Generate samples from posterior distribution.""" + q_samples = dict(zip(self.unobserved_keys, self.approx.sample(n, seed=self._seed))) Resolved in a1d9679 <a1d9679> — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#280 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFETGGK4HVEQCUD6RI5TZ3RXTZ63ANCNFSM4NWDWMUA> .

[ferrine]

Polising docs and tests is left before shipping. Some work is left beyond this PR (deterministic variables and flattening the model for full rank ADVI)

[ferrine]

Looks good, I'll merge once tests pass