Is inverse autoregressive flow supported?

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