How to get an inverse autoregressive flow?

[zhihanyang2022]

Description

In the documentation, the VAE example uses an autoregressive prior rather than an autoregressive posterior conditioned on the input, which is great but a bit unusual. I wanted to try out autoregressive posteriors (e.g., IAF) but couldn't figure out how to do it in zuko.

Implementation

Paper: Improving Variational Inference with Inverse Autoregressive Flow by Kingma et al.

[francois-rozet]

Hi @zhihanyang2022 👋

An IAF is the same as a MAF, but with reversed transformations, such that sampling is fast and log-prob is slow. In Zuko, you can transform any flow to an inverse flow with the following pattern:

maf = zuko.flows.MAF(features=2, transforms=3, hidden_features=(64, 64))  # or NSF, NAF, ...
iaf = zuko.flows.Flow(flow.transform.inv, flow.base)

We use this pattern in the "train from energy" tutorial.

https://zuko.readthedocs.io/stable/tutorials/reverse_kl.html

Please note that to optimize through the sampling, you need to call rsample (or rsample_and_log_prob) instead of sample.

[zhihanyang2022]

Thanks!

Also, for MAF, is mixture of Gaussians (MoG) output distribution available in Zuko?

If not, do you see Zuko supporting this feature in the near future? I'm also happy to contribute but might need some pointers to get started.