/PyTorchDiscreteFlows

Discrete Normalizing Flows implemented in PyTorch

Primary LanguageJupyter Notebook

Acknowledgements

The discrete normalizing flow code is originally taken and modified from: https://github.com/google/edward2/blob/master/edward2/tensorflow/layers/discrete_flows.py and https://github.com/google/edward2/blob/master/edward2/tensorflow/layers/utils.py Which was introduced in the paper: https://arxiv.org/abs/1905.10347

The demo file, MADE, and MLP were modified and taken from: https://github.com/karpathy/pytorch-normalizing-flows

State of Library

To my knowledge as of July 3rd 2020, this is the only functional demo of discrete normalizing flows in PyTorch. The code in edward2 (implemented in TF2 and Keras, lacks any tutorials. Since the release of this repo and because of correspondence with the authors of the original paper, demo code for reproducing Figure 2 using Edward2 has been shared here.

With collaboration from Yashas Annadani and Jan Francu, I have been able to reproduce the paper's Figure 2 discretized mixture of Gaussians with this library.

Use Library

To use this package, clone the repo satisfy the below package requirements, then run Figure2Replication.ipynb.

Requirements: Python 3.0+ PyTorch 1.2.0+ < 1.9.0 (if you are using 1.9.0 you will have issues with the Fast Fourier Transform. See this issue and solution #6 (comment)) Numpy 1.17.2+

Implementation details

NB. Going from Andre Karpathy's notation, flow.reverse() goes from the latent space to the data and flow.forward() goes from the data to the latent space. This is the inverse of some other implementations including the original Tensorflow one. Implements Bipartite and Autoregressive discrete normalizing flows. Also has an implementation of MADE and a simple MLP.

TODOs - Pull requests very welcome!

Replication of Figure 2 Mixture of Gaussians

Figure 2 in the paper looks like this:

PaperFigure2

This library's replication is:

Fig2Reproduction