Allow guides with re-initializable parameters

[ahmadsalim]

The first split of #649 .

It introduces re-initializable guides, which can re-initialize their parameters.
This is useful for Stein inference since the parameters are treated as particles, and we want the particles to be initialized independently instead of all clumped together which would hurt inference.

An alternative would be to have similar strategies which can be called explicitly with a shape and can be used directly in numpyro.param. I do not know which you prefer, so I tried to reuse the existing interface 😄

[fehiepsi]

Thanks for splitting out your PR, @ahmadsalim! IIUC you want two things:

• initializing param randomly, instead of fixing an initial value
• using MCMC for param statements?

I think we can start with a contrib.stein or contrib.svgd and move ReinitGuide and AutoDelta (which depends on ReinitGuide) to there. It is good to add some smoke tests to verify that the implementation works as your intention.

FYI, we just recently move autoguide to the main infer module after many refactoring PRs. In addition, we have cleaned up param stuffs in potential_energy and find_initial_valid_params. I am not sure if we need to introduce it again.

How about the following proposals:

• Implement numpyro.lift to lift param sites to sample sites (this deserves a separate PR - I can implement it if you want). Then you can use the current find_valid_initial_params to get valid initial parameters in the Stein implementation. I guess with that, we also don't have to introduce reinit_param in the current initialization strategy.
• Having a separate PR for init_with_noise, maybe we can start it in contrib module together with ReinitGuide and AutoDelta? I am also a bit curious about how useful it is. In addition, its interface is a bit different from other initial strategy interfaces. Maybe something like init_with_noise(site=None, noise_scale =1.0, init_strategy=init_to_uniform) makes things more consistent.

[ahmadsalim]

Thanks for taking a look!

I do not need MCMC for param statements.
However, the first point is correct, I would like param statements to be randomly initialized.

I would still like them to be categorized as parameters, so that Stein knows that these are the particles to optimize. This is since Stein treats specified parameters as particles, kind of like a generalization of SVI but where we infer multiple values for each parameter using Stein instead of a single one (based on Stein mixture models paper).

[ahmadsalim]

I do not have anything against moving things to contrib and will do that ASAP.

Could you kindly explain how numpyro.lift would work?

[fehiepsi]

I think it will just be a port of http://docs.pyro.ai/en/stable/poutine.html#pyro.poutine.handlers.lift

Do you use those initialization steps in inference? If not, then you can lift to get initial params, then doing inference with your original model as usual.

[ahmadsalim]

Oh, that would be great to have! Then I can drop my other changes to the initialization functions 😄

[fehiepsi]

Sorry for the delay! I'll port pyro.lift today. >>>o<<<

[ahmadsalim]

Awesome, thanks a lot!

I was actually not sure what the conclusion was and started looking into implementing it myself 😄 . But I appreciate you looking into it!

[fehiepsi]

Haha, me too. I thought that you were waiting for it so I just went ahead and implemented. :D

[ahmadsalim]

Wow, this is embarrassing. I missed that the pyro.lift function was integrated.
I was debugging a weird NaN gradient/very slow performance issue with JAX for the lasts months and forgot to continue on this path :)

I will figure out how to re-open the PRs as originally planned, but will close this one.

[ahmadsalim]

@fehiepsi Thanks a lot for pyro.lift 🥇

[fehiepsi]

Sure, please let us know if you need any features else. I hope that NumPyro design should not block much of your work, so hopefully, you only need to play with the algorithm. :)

[ahmadsalim]

Cool! Thanks a lot :)