Implementation of DCC inference algorithm

[treigerm]

This is an initial bare bones implementation of the Divide, Conquer, Combine inference algorithm for programs with stochastic support as discussed in #1697. I have also included a simple example to show how to use the interface. This is only a bery basic implementation of the algorithm to keep the size of the PR reasonable. In its current form the algorithm assumes that the branching inside the program is done based on the outcomes of discrete sampling sites which are annotated with infer={"branching": True}.

The algorithm then proceeds as follows:

1. Discover different branches, a.k.a. straight-line programs (SLPs), by sampling from the prior.
2. Run inference on each discovered branch separately.
3. Combine the inference results by weighting each branch proportional to its marginal likelihood estimate.

This is just a draft PR for now to see whether the general approach for the implementation makes sense. If the approach is sensible then I will go and add a more detailed example, some tests and documentation. So the main questions that need to be answered at the moment are:

• Does the general interface for the algorithm seem sensible?
• Did I place the new implementation in the correct location (I located it in the contrib folder for now)?
• Am I using the Numpyro primitives correctly?

[fehiepsi]

Thanks for the contribution, @treigerm! The draft looks great.

Re location: yeah, it makes sense to put the algorithm in contrib. We can move it to infer in the future when the api is solid.
Re primitives: yes, the usage looks correct to me.

[treigerm]

I have now added tests and documentation. Please let me know if you think anything is missing!

After this PR is merged my plan would be to add the SDVI algorithm and after that is added I can write a tutorial about how to use these two inference algorithms and their respective trade-offs.

[fehiepsi]

Is this standard Gaussian a good practical choice for the proposal? Looking at the paper, it seems that the authors used a metropolis-within-gibbs sampler.

[treigerm]

Thank you for spending the time reviewing the PR!

Note that the Gaussian is centered around the MCMC samples (more precisely, each MCMC sample gives rise to a single proposal distribution). As long as the MCMC chain(s) are well-mixed this generally leads to good proposals. This is also what the paper describes (and also what the author's implementation does which I have received upon request).

Looking at the paper, it seems that the authors used a metropolis-within-gibbs sampler.

Actually, the metropoylis-within-gibbs sampler is only used for the local inference tasks. For many models it isn't a very efficient inference algorithm because it only updates one variable at a time. Because the implementation here assumes that the branching is only done based on the outcomes of discrete sampling statements, it can use more efficient algorithms for local inference (like HMC or NUTS).

[fehiepsi]

I see, thanks. How about using sample variance instead of unit variance?

[treigerm]

I have just updated the PR to make the scale in the proposal a parameter that can be set by the user. I agree that unit variance is probably not always desirable. However, I'm not sure whether sample variance is desirable either. The main idea behind the algorithm for estimating the normalization constant (which in more detailed is described in another paper Layered Adaptive Importance Sampling) is that the proposal on top of each sample leads to fairly local proposals. If there are multiple modes in the posterior then using the sample variance could result in lots of proposed samples in the low density regions between the modes.

I'd still be open to adding the option to use the sample variance to the implementation but this is currently complicated by the way the AutoNormal guide is implemented. You would want to compute the sample variance for each individual variable in the program but as far as I can tell there is no way to set variable specific variances in the AutoNormal guide (it's possible to set variable specific means though). So this might be a feature that would be reasonable to add at a later time?

[fehiepsi]

Sorry for the late response! Thank you for the insights - I don't have a strong opinion on whether it's helpful to expose init_scale in AutoNormal. There is another way to substitute sample variance, like substitute(guide, data={"auto_foo_scale": ...}) but we need to be careful at the domain (needs to be unconstrained) of such a foo variable.

[treigerm]

Ah I seem, I guess this would require knowledge about all the variable names in the program though (but that could be extracted automatically)? For now I would lean towards leaving the implementation as is to keep it simple, if that is okay.

[fehiepsi]

LGTM pending linting issues. Thanks for the great contribution, @treigerm!

[treigerm]

Thanks, @fehiepsi ! I'm currently unsure why the tests are failing, it seems to be an import error that it is not able to find the new numpyro.contrib.stochastic_support.dcc module. For me, locally the test are passing if I run XLA_FLAGS="--xla_force_host_platform_device_count=2" pytest -vs test/contrib/stochastic_support/test_dcc.py, so not quite sure where the issue is coming from.

[fehiepsi]

How about adding __init__ file and exposing DCC etc. there? Then you can import

from numpyro.contrib.stochastic_support import DCC

[treigerm]

Ah yes of course! I have added an __init__ file now. I also added the XLA_FLAGS="--xla_force_host_platform_device_count=2" flag to the contrib module tests to make the parallel chain sampling method pass. Tests are passing locally for me.

[fehiepsi]

Could you ignore dcc test here and move dcc to the test chain below instead? thanks

[treigerm]

Done!