numpyro icon indicating copy to clipboard operation
numpyro copied to clipboard

Automatic structured variational inference

Open dirmeier opened this issue 3 years ago • 11 comments

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

dirmeier avatar Aug 05 '21 20:08 dirmeier

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.

fehiepsi avatar Aug 06 '21 02:08 fehiepsi

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?

fritzo avatar Aug 06 '21 10:08 fritzo

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

dirmeier avatar Aug 09 '21 07:08 dirmeier

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.

gisilvs avatar Nov 26 '21 10:11 gisilvs

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.

esennesh avatar Feb 13 '23 19:02 esennesh

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

fehiepsi avatar Feb 14 '23 13:02 fehiepsi

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 avatar Feb 17 '23 18:02 esennesh

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.

gisilvs avatar Feb 20 '23 13:02 gisilvs

+1 to this feature request :)

rafaelkaufmann avatar Jul 03 '23 16:07 rafaelkaufmann

+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.

esennesh avatar Jul 03 '23 18:07 esennesh

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

howsiyu avatar Nov 02 '23 07:11 howsiyu