Implement Rejection Sampling Variational Inference

[fritzo]

Addresses #63

This implements Rejection Sampling Variational Inference (RSVI) of Naesseth et al. (2017).

This implements a new class Rejector as an abstract base class derived from Distribution. The crux to integrating these with Pyro's SVI in TraceGraph_ELBO is a new method Distribution.score_function_term() that can be overridden to compute g_cor in the paper. This method computes a partial score function for a partially reparameterized distribution such that: for an unreparameterized distribution it computes the entire score function, and for a fully reparameterized distribution it computes zero.

Tasks

• integrate with Trace_ELBO
• implement RejectionGamma example, e.g. see @slinderman's notebook
• implement @naesseth's shape augmentation, e.g. see code
• write convergence tests
• write smoke tests
• clean up PR to remove irrelevant testing stuff
• fix unit tests

This was joint work by @martinjankowiak and @fritzo

[fritzo]

cc @rachtsingh @naesseth

[slinderman]

Glad to see this incorporated into pyro, @fritzo and @martinjankowiak! One of the points discussed in the paper is that the score function term tends to be higher variance, and in practice it can often be safely ignored at the cost of a small bias.

You can also ameliorate it to some extent by "shape augmentation," i.e. introducing auxiliary uniform r.v.'s in order to increase the shape parameter and thereby increase the acceptance probability. I did not implement this in the demo notebook, but we used it in our paper. The code for that is in this repo: https://github.com/blei-lab/ars-reparameterization and @naesseth can probably provide more details on its implementation.

[fritzo]

@martinjankowiak This is almost done, there's just one more test to fix.

Could you please review the changes to Rejector? I've simplified and added a log_scale argument that is used to correctly compute the .batch_log_pdf() scale.

[martinjankowiak]

cool awesome. i'll do so in the morning when i have a fresh brain

[martinjankowiak]

maybe explicit names in namedtuple here?

[martinjankowiak]

let's investigate the xfail more

[martinjankowiak]

docstring typo

[fritzo]

@martinjankowiak I believe I've fixed Rejector (sorry for the earlier bug!), but ShapeAugmentedGamma gradients still appear to be wrong. Do you think it's worth merging this as-is and debugging ShapeAugmentedGamma in a follow-up PR that also adds ShapeAugmentedBeta etc?

[martinjankowiak]

lgtm. @fritzo merge?

[fritzo]

Ready to merge, finally.

[rmehta1987]

@fritzo are there any examples of using the accept-reject dirichlet or gamma?

[martinjankowiak]

@rmehta1987 i don't believe there are. but if you changed the Gammas in the guide in examples/sparse_gamma_def.py to ShapeAugmentedGammas (found here: pyro/distributions/testing/rejection_gamma.py) the example should work more or less out of the box (modulo optimization issues)