didriknielsen / survae_flows

Code for paper "SurVAE Flows: Surjections to Bridge the Gap between VAEs and Flows"
MIT License
283 stars 34 forks source link

How to add flow priors to VAE #18

Open DylanMannKrzisnik opened 2 years ago

DylanMannKrzisnik commented 2 years ago

Hello,

What would be the simplest way of replacing a Gaussian prior for VAE with a flow-based one?

Thanks

didriknielsen commented 2 years ago

Hi,

I'd say the simplest would be to start from e.g. this example: https://github.com/didriknielsen/survae_flows/blob/master/examples/mnist_vae.py and add some bijective transformations after the VAE on line 50 (the layers are specified from data to latent, i.e. x -> z).

Alternatively, and equivalently, you can replace the base_distribution with a Flow.

Hope this helps. If it's not clear I post an example later.

Didrik

DylanMannKrzisnik commented 2 years ago

Thanks for your response, I had opted for the second option of replacing the base_distribution with a Flow. Much appreciated.