search

pyro-ppl / numpyro

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
https://num.pyro.ai
Apache License 2.0
2.31k stars 246 forks source link

Automatic structured variational inference #1117

Open dirmeier opened 3 years ago

dirmeier commented 3 years ago

Hey all,

would be nice if there was an implementation of Automatic structured variational inference in NumPyro (unless it is already there and I am not seeing it).

A TFP implementation already exists here. If you think this makes sense, I could give it a try.

Cheers, Simon

fehiepsi commented 3 years ago

Hi @dirmeier, FYI there is a AutoStructured class in Pyro which has some similar ideas. I don't know about the scope of the ASVI. Please give it a go if it can be implemented as an auto guide or a separate SVI-style algorithm.

fritzo commented 3 years ago

Hi @dirmeier, we'd welcome contributions of any new automatic guides. For compatibility with Pyro, we'd like to reserve the name AutoStructured to the design in Pyro that @fehiepsi pointed to. Note from that PR https://github.com/pyro-ppl/pyro/pull/2812

Note this [AutoStructured class] does not implement Automatic structured variational inference, a variational family whose stricture is severely limited to dependencies in the model. Nor does this first PR implement automatic suggestion of the guide structure as in Faithful inversion of generative models for effective amortized inference.

In the same direction, @eb8680 has some ideas that could learn the exact optimal Gaussian precision structure based on model dependencies, as he discusses in https://github.com/pyro-ppl/pyro/issues/2813. While I think @eb8680's plan is a good long-term goal and a much more plausibly structured posterior than ASVI, it's probably good for us to have a suite of automatic guides, from the mean field AutoNormal, to the simple AutoLowRankMultivariateNormal, maybe an ASVI, an AutoStructured (which is a little more flexible but less automatic), @eb8680's AutoDependencyTracking guide, and finally the full naive AutoMultivaraiteNormal. WDYT?

dirmeier commented 3 years ago

Hey, @fehiepsi and @fritzo,

Ah, great, didn't see the Pyro implementation. Yeah, completely agree with your concerns about ASVI in general, but haven't used Ambrogioni et al. 21 yet (nor read the paper in depth). Are there any limitations wrt data size for the cascading flows?

it's probably good for us to have a suite of automatic guides

That would be nice!

Anyways, would be happy to help out with the implementations on any of these, since I would like to be able to fit flexible models automatically ,i.e., without having to have much thought about how I construct a surrogate.

Cheers, Simon

gisilvs commented 3 years ago

Hello, I contributed to the TFP implementation of Cascading Flows, and I am very familiar with ASVI. I would be happy to give some suggestions if needed.

esennesh commented 1 year ago

Hi,

I've actually got a relatively short one-file implementation of ASVI as a GuideMessenger subclass. I have a neural amortization extension too (which conditions the guide on data, giving a naive but faithful inverse), though I'd have to move some code around to fully include it with a pull request. Would there be any interest in bringing it over to numpyro and/or primary Pyro?

Cheers.

fehiepsi commented 1 year ago

That would be great, @esennesh! It would be nice to have this guide in both Pyro and NumPyro.

esennesh commented 1 year ago

Hello, I contributed to the TFP implementation of Cascading Flows, and I am very familiar with ASVI. I would be happy to give some suggestions if needed.

Btw @gisilvs , if I could provide some code for ASVI with a naive neural amortization, would you be able to extend it to cascading flows? I don't really understand the neural architecture used for those.

gisilvs commented 1 year ago

Hello, I contributed to the TFP implementation of Cascading Flows, and I am very familiar with ASVI. I would be happy to give some suggestions if needed.

Btw @gisilvs , if I could provide some code for ASVI with a naive neural amortization, would you be able to extend it to cascading flows? I don't really understand the neural architecture used for those.

@esennesh yes, I would be happy to contribute, both to the implementation of Cascading Flows and the follow-up Embedded Model Flows.

rafaelkaufmann commented 1 year ago

+1 to this feature request :)

esennesh commented 1 year ago

+1 to this feature request :)

I've got a feature-branch for it. Just have to clean up the tests I wrote, rebase everything into a clean set of patches, and actually send the pull-request.

howsiyu commented 1 year ago

It would be nice if we can get pyro's StructuredReparam with ASVI too.