Automatic structured variational inference

[dirmeier]

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]

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]

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 pyro-ppl/pyro#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 pyro-ppl/pyro#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]

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]

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]

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]

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

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

[gisilvs]

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]

+1 to this feature request :)

[esennesh]

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

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